commit ec436095dd094d68589d097f82e6cfedc87a51b4 Author: wehub-resource-sync Date: Mon Jul 13 13:30:03 2026 +0800 chore: import upstream snapshot with attribution diff --git a/.github/CODE_OF_CONDUCT.md b/.github/CODE_OF_CONDUCT.md new file mode 100644 index 0000000..b1542d9 --- /dev/null +++ b/.github/CODE_OF_CONDUCT.md @@ -0,0 +1,84 @@ + +# 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 [INSERT CONTACT METHOD]. 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][v2.0]. + +Community Impact Guidelines were inspired by [Mozilla's code of conduct enforcement ladder][Mozilla CoC]. + +For answers to common questions about this code of conduct, see the FAQ at [https://www.contributor-covenant.org/faq][FAQ]. Translations are available at [https://www.contributor-covenant.org/translations][translations]. + +[homepage]: https://www.contributor-covenant.org +[v2.0]: https://www.contributor-covenant.org/version/2/0/code_of_conduct.html +[Mozilla CoC]: https://github.com/mozilla/diversity +[FAQ]: https://www.contributor-covenant.org/faq +[translations]: https://www.contributor-covenant.org/translations diff --git a/.github/CONTRIBUTING.md b/.github/CONTRIBUTING.md new file mode 100644 index 0000000..c16debf --- /dev/null +++ b/.github/CONTRIBUTING.md @@ -0,0 +1,139 @@ +## Before Commit! + +Your commit message must follow Conventional Commits (https://www.conventionalcommits.org/) and your code should be formatted. The Git hooks will do most of the work automatically: + +### Tool Requirements + +You need a recent `clang-format` (>= 18). In a conda environment you can install: + +```shell +conda install -c conda-forge clang-format=18 +``` + +If you previously configured with an older version, remove the build directory and reconfigure: + +```shell +rm -rf kt-kernel/build +``` + +Install `black` for Python formatting: + +```shell +conda install black +``` + +### Install hook: +```shell +bash kt-kernel/scripts/install-git-hooks.sh +#or just cmake the kt-kernel +cmake -S kt-kernel -B kt-kernel/build +``` + +There are manual commands if you need format. + +```shell +cmake -S kt-kernel -B kt-kernel/build +cmake --build kt-kernel/build --target format +``` + +## Developer Note + +Formatting and commit message rules are enforced by Git hooks. After installing `clang-format` and `black`, just commit normally—the hooks will run formatting for you. + +> [!NOTE] +> If formatting modifies files, the commit is aborted after staging those changes. Review them and run `git commit` again. Repeat until no further formatting changes appear. + +--- + +### Conventional Commit Regex (Reference) + +The commit-msg hook enforces this pattern: + +```text +regex='^\[(feat|fix|docs|style|refactor|perf|test|build|ci|chore|revert|wip)\](\([^\)]+\))?(!)?: .+' +``` + +Meaning (English): +* `[type]` required — one of feat|fix|docs|style|refactor|perf|test|build|ci|chore|revert|wip +* Optional scope: `(scope)` — any chars except `)` +* Optional breaking change marker: `!` right after type or scope +* Separator: `: ` (colon + space) +* Subject: free text (at least one character) + +Examples: +```text +[feat]: add adaptive batching +[fix(parser)]: handle empty token list +[docs]!: update API section for breaking rename +``` + +You can bypass locally (not recommended) with: +```shell +git commit --no-verify +``` +## 提交前提醒 + +提交信息必须满足 Conventional Commits 规范 (https://www.conventionalcommits.org/)，代码需要符合格式要求。Git 钩子已经集成了大部分工作： +### 软件要求 + +需要较新的 `clang-format` (>= 18)，在 conda 环境中安装： + +```shell +conda install -c conda-forge clang-format=18 +``` + +如果之前用老版本配置过，请删除构建目录重新配置： + +```shell +rm -rf kt-kernel/build +``` + +安装 `black` 以进行 Python 文件格式化： + +```shell +conda install black +``` +### 安装钩子 +```shell +bash kt-kernel/scripts/install-git-hooks.sh +#or just cmake the kt-kernel +cmake -S kt-kernel -B kt-kernel/build +``` +如果你需要手动格式化： +```shell +cmake -S kt-kernel -B kt-kernel/build +cmake --build kt-kernel/build --target format +``` + +## 开发者说明 + +本仓库通过 Git hooks 自动执行代码格式化与提交信息规范检查。只需安装好 `clang-format` 与 `black` 后正常执行提交即可，钩子会自动格式化。 + +> [!NOTE] +> 如果格式化修改了文件，钩子会终止提交并已暂存这些改动。请查看修改后再次执行 `git commit`，重复直到没有新的格式化变更。 + +### 提交信息正则（参考） + +钩子使用如下正则检查提交信息： +```text +regex='^\[(feat|fix|docs|style|refactor|perf|test|build|ci|chore|revert|wip)\](\([^\)]+\))?(!)?: .+' +``` +含义： +* `[type]` 必填：feat|fix|docs|style|refactor|perf|test|build|ci|chore|revert|wip +* 作用域可选：`(scope)`，不能包含右括号 +* 可选的破坏性标记：`!` +* 分隔符：冒号+空格 `: ` +* 描述：至少一个字符 + +示例： +```text +[feat]: 增加自适应 batch 功能 +[fix(tokenizer)]: 修复空 token 列表处理 +[docs]!: 更新接口文档（存在破坏性修改） +``` + +跳过钩子（不推荐，仅紧急时）： +```shell +git commit --no-verify +``` + diff --git a/.github/ISSUE_TEMPLATE/-bug-.yaml b/.github/ISSUE_TEMPLATE/-bug-.yaml new file mode 100644 index 0000000..060beee --- /dev/null +++ b/.github/ISSUE_TEMPLATE/-bug-.yaml @@ -0,0 +1,55 @@ +name: "\U0001F41B Bug / Help" +description: Create a report to help us improve the ktransformers project +labels: ["pending"] +body: + - type: markdown + attributes: + value: | + Issues included in **[FAQs](https://github.com/kvcache-ai/ktransformers/issues/1608)** or those with **insufficient** information may be closed without a response. + 已经包含在 **[常见问题](https://github.com/kvcache-ai/ktransformers/issues/1608)** 内或提供信息**不完整**的 issues 可能不会被回复。 + + - type: checkboxes + id: reminder + attributes: + label: Reminder + description: | + Please ensure you have read the above rules carefully and searched the existing issues (including FAQs). + 请确保您已经认真阅读了上述规则并且搜索过现有的 issues（包括常见问题）。 + + options: + - label: I have read the above rules and searched the existing issues. + required: true + + - type: textarea + id: system-info + validations: + required: true + attributes: + label: System Info + description: | + Please share your system info with us. You can run the command **lscpu**, ** nvidia-smi ** etc. and copy-paste its output below. + 请提供您的系统信息。您可以在命令行运行 **lscpu**, **nvidia-smi** 等命令，并将其输出复制到该文本框中。 + + placeholder: ktransformers version,sglang version, platform, python version, cpu info, GPU/NPU info ... + + - type: textarea + id: reproduction + validations: + required: true + attributes: + label: Reproduction + description: | + Please provide entry arguments, error messages and stack traces that reproduces the problem. + 请提供入口参数，错误日志以及异常堆栈以便于我们复现问题。 + + value: | + ```text + Put your message here. + ``` + + - type: textarea + id: others + validations: + required: false + attributes: + label: Others \ No newline at end of file diff --git a/.github/ISSUE_TEMPLATE/-feature-.yaml b/.github/ISSUE_TEMPLATE/-feature-.yaml new file mode 100644 index 0000000..9db7f26 --- /dev/null +++ b/.github/ISSUE_TEMPLATE/-feature-.yaml @@ -0,0 +1,41 @@ +name: "\U0001F680 Feature request" +description: Submit a request for a new feature +labels: ["enhancement", "pending"] +body: + - type: markdown + attributes: + value: | + Please do not create issues that are not related to new features under this category. + 请勿在此分类下创建和新特性无关的 issues。 + + - type: checkboxes + id: reminder + attributes: + label: Reminder + description: | + Please ensure you have read the above rules carefully and searched the existing issues. + 请确保您已经认真阅读了上述规则并且搜索过现有的 issues。 + + options: + - label: I have read the above rules and searched the existing issues. + required: true + + - type: textarea + id: description + validations: + required: true + attributes: + label: Description + description: | + A clear and concise description of the feature proposal. + 请详细描述您希望加入的新功能特性。 + + - type: textarea + id: contribution + validations: + required: false + attributes: + label: Pull Request + description: | + Have you already created the relevant PR and submitted the code? + 您是否已经创建了相关 PR 并提交了代码？ \ No newline at end of file diff --git a/.github/ISSUE_TEMPLATE/config.yml b/.github/ISSUE_TEMPLATE/config.yml new file mode 100644 index 0000000..6c53a15 --- /dev/null +++ b/.github/ISSUE_TEMPLATE/config.yml @@ -0,0 +1,5 @@ +blank_issues_enabled: false +contact_links: + - name: 📚 FAQs | 常见问题 + url: https://github.com/kvcache-ai/ktransformers/issues/1608 + about: Reading in advance is recommended | 建议提前阅读 \ No newline at end of file diff --git a/.github/PULL_REQUEST_TEMPLATE.md b/.github/PULL_REQUEST_TEMPLATE.md new file mode 100644 index 0000000..74146bc --- /dev/null +++ b/.github/PULL_REQUEST_TEMPLATE.md @@ -0,0 +1,8 @@ +# What does this PR do? + +Fixes # (issue) + +## Before submitting + +- [ ] Did you read the [contributor guideline](https://github.com/kvcache-ai/ktransformers/blob/main/.github/CONTRIBUTING.md)? +- [ ] Did you write any new necessary tests? \ No newline at end of file diff --git a/.github/SECURITY.md b/.github/SECURITY.md new file mode 100644 index 0000000..25f69d6 --- /dev/null +++ b/.github/SECURITY.md @@ -0,0 +1,7 @@ +# Reporting Security Issues + +To report a security issue, please use the GitHub Security Advisory ["Report a Vulnerability"](https://github.com/kvcache-ai/ktransformers/security/advisories/new) tab. + +We will send a response indicating the next steps in handling your report. After the initial reply to your report, the security team will keep you informed of the progress towards a fix and full announcement, and may ask for additional information or guidance. + +Report security bugs in third-party modules to the person or team maintaining the module. \ No newline at end of file diff --git a/.github/workflows/book-ci.yml b/.github/workflows/book-ci.yml new file mode 100644 index 0000000..61f5f56 --- /dev/null +++ b/.github/workflows/book-ci.yml @@ -0,0 +1,32 @@ +name: Book-CI + +on: + push: + branches: + - main + # - server_support + + pull_request: + branches: + - main + # - server_support +jobs: + test: + name: test + runs-on: ${{ matrix.os }} + strategy: + matrix: + os: [ubuntu-latest, macos-latest, windows-latest] + steps: + - uses: actions/checkout@v4 + - name: Install Rust + run: | + rustup set profile minimal + rustup toolchain install stable + rustup default stable + - name: Setup mdBook + uses: peaceiris/actions-mdbook@v2 + with: + mdbook-version: "latest" + # - name: Run tests + # run: mdbook test \ No newline at end of file diff --git a/.github/workflows/deploy.yml b/.github/workflows/deploy.yml new file mode 100644 index 0000000..dd406df --- /dev/null +++ b/.github/workflows/deploy.yml @@ -0,0 +1,49 @@ +name: Deploy + +on: + push: + branches: + - main + # - server_support + + pull_request: + branches: + - main + # - server_support + +defaults: + run: + shell: bash + +permissions: + contents: write + +jobs: + deploy: + runs-on: ${{ matrix.os }} + strategy: + matrix: + os: [ubuntu-latest, macos-latest, windows-latest] + steps: + - uses: actions/checkout@v4 + - name: Install Rust + run: | + rustup set profile minimal + rustup toolchain install stable + rustup default stable + - name: Setup mdBook + uses: peaceiris/actions-mdbook@v2 + with: + mdbook-version: "latest" + - run: mdbook build + # - name: Copy Assets + # run: | + # chmod +x ci/copy-assets.sh + # ci/copy-assets.sh ${{ matrix.os }} + - name: Deploy + uses: peaceiris/actions-gh-pages@v3 + # or || github.ref == 'refs/heads/server_support' + if: ${{ github.ref == 'refs/heads/main' }} + with: + github_token: ${{ secrets.GITHUB_TOKEN }} + publish_dir: ./book \ No newline at end of file diff --git a/.github/workflows/docker-image.yml b/.github/workflows/docker-image.yml new file mode 100644 index 0000000..6d4f645 --- /dev/null +++ b/.github/workflows/docker-image.yml @@ -0,0 +1,142 @@ +name: DockerHub CI + +on: + release: + types: [published] + workflow_dispatch: + inputs: + push_to_dockerhub: + description: 'Push image to DockerHub? (true/false)' + required: true + default: 'false' + type: boolean + cuda_version: + description: 'CUDA version (e.g., 12.8.1)' + required: false + default: '12.8.1' + type: string + push_simplified_tag: + description: 'Also push simplified tag? (true/false)' + required: false + default: 'true' + type: boolean + ubuntu_mirror: + description: 'Use Tsinghua Ubuntu mirror? (0/1)' + required: false + default: '0' + type: string + + # push: + # branches: + # - main +env: + DOCKERHUB_REPO: ${{ secrets.DOCKERHUB_USERNAME }}/ktransformers +jobs: + test: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + - name: Run tests + run: | + if [ -f docker-compose.test.yml ]; then + docker-compose --file docker-compose.test.yml build + docker-compose --file docker-compose.test.yml run sut + else + docker build . --file docker/Dockerfile + fi + + build-and-push: + needs: test + name: Build and Push Multi-Variant Docker Image + runs-on: ubuntu-latest + + steps: + - name: Checkout repository + uses: actions/checkout@v4 + + - name: Move Docker data directory + run: | + sudo systemctl stop docker + sudo mkdir -p /mnt/docker + sudo rsync -avz /var/lib/docker/ /mnt/docker + sudo rm -rf /var/lib/docker + sudo ln -s /mnt/docker /var/lib/docker + sudo systemctl start docker + + - name: Set up Docker Buildx + uses: docker/setup-buildx-action@v3 + + - name: Login to Docker Hub + uses: docker/login-action@v3 + with: + username: ${{ secrets.DOCKERHUB_USERNAME }} + password: ${{ secrets.DOCKERHUB_TOKEN }} + + - name: Determine build parameters + id: params + run: | + # Determine if we should push + if [ "${{ github.event_name }}" = "release" ]; then + echo "should_push=true" >> $GITHUB_OUTPUT + echo "push_simplified=true" >> $GITHUB_OUTPUT + elif [ "${{ github.event_name }}" = "workflow_dispatch" ]; then + echo "should_push=${{ inputs.push_to_dockerhub }}" >> $GITHUB_OUTPUT + echo "push_simplified=${{ inputs.push_simplified_tag }}" >> $GITHUB_OUTPUT + else + echo "should_push=false" >> $GITHUB_OUTPUT + echo "push_simplified=false" >> $GITHUB_OUTPUT + fi + + # Determine CUDA version + if [ "${{ github.event_name }}" = "workflow_dispatch" ] && [ -n "${{ inputs.cuda_version }}" ]; then + echo "cuda_version=${{ inputs.cuda_version }}" >> $GITHUB_OUTPUT + else + echo "cuda_version=12.8.1" >> $GITHUB_OUTPUT + fi + + # Determine Ubuntu mirror setting + if [ "${{ github.event_name }}" = "workflow_dispatch" ] && [ -n "${{ inputs.ubuntu_mirror }}" ]; then + echo "ubuntu_mirror=${{ inputs.ubuntu_mirror }}" >> $GITHUB_OUTPUT + else + echo "ubuntu_mirror=0" >> $GITHUB_OUTPUT + fi + + - name: Build and push Docker image + run: | + cd docker + + # Build command arguments + BUILD_ARGS=( + --cuda-version "${{ steps.params.outputs.cuda_version }}" + --ubuntu-mirror "${{ steps.params.outputs.ubuntu_mirror }}" + --repository "${{ env.DOCKERHUB_REPO }}" + ) + + # Add simplified tag option if enabled + if [ "${{ steps.params.outputs.push_simplified }}" = "true" ]; then + BUILD_ARGS+=(--also-push-simplified) + fi + + # Add HTTP proxy if available + if [ -n "${{ secrets.HTTP_PROXY }}" ]; then + BUILD_ARGS+=(--http-proxy "${{ secrets.HTTP_PROXY }}") + fi + + # Add HTTPS proxy if available + if [ -n "${{ secrets.HTTPS_PROXY }}" ]; then + BUILD_ARGS+=(--https-proxy "${{ secrets.HTTPS_PROXY }}") + fi + + # Dry run if not pushing + if [ "${{ steps.params.outputs.should_push }}" != "true" ]; then + BUILD_ARGS+=(--dry-run) + fi + + # Execute build script + ./push-to-dockerhub.sh "${BUILD_ARGS[@]}" + + - name: Display image information + if: steps.params.outputs.should_push == 'true' + run: | + echo "::notice title=Docker Image::Image pushed successfully to ${{ env.DOCKERHUB_REPO }}" + echo "Pull command: docker pull ${{ env.DOCKERHUB_REPO }}:v\$(VERSION)-cu\$(CUDA_SHORT)" diff --git a/.github/workflows/kt-kernel-tests.yml b/.github/workflows/kt-kernel-tests.yml new file mode 100644 index 0000000..f4b1e66 --- /dev/null +++ b/.github/workflows/kt-kernel-tests.yml @@ -0,0 +1,104 @@ +name: PR KT-Kernel Test + +on: + pull_request: + branches: + - main + - develop + types: [synchronize, labeled] + workflow_dispatch: + +concurrency: + group: pr-kt-kernel-test-${{ github.ref }} + cancel-in-progress: true + +jobs: + # =============================================== check changes ==================================================== + check-changes: + runs-on: ubuntu-latest + outputs: + kt_kernel: ${{ steps.filter.outputs.kt_kernel }} + steps: + - name: Checkout code + uses: actions/checkout@v4 + + - name: Fail if the PR does not have the 'run-ci' label + if: github.event_name == 'pull_request' && !contains(github.event.pull_request.labels.*.name, 'run-ci') + run: | + echo "This pull request does not have the 'run-ci' label. Failing the workflow." + exit 1 + + - name: Fail if the PR is a draft + if: github.event_name == 'pull_request' && github.event.pull_request.draft == true + run: | + echo "This pull request is a draft. Failing the workflow." + exit 1 + + - name: Detect file changes + id: filter + uses: dorny/paths-filter@v3 + with: + filters: | + kt_kernel: + - "kt-kernel/**" + - ".github/workflows/kt-kernel-tests.yml" + + # =============================================== KT-Kernel tests ==================================================== + per-commit-kt-kernel-cpu: + needs: [check-changes] + if: always() && !failure() && !cancelled() && + (needs.check-changes.outputs.kt_kernel == 'true' || github.event_name == 'workflow_dispatch') + runs-on: kt-cpu + continue-on-error: false + steps: + - name: Cleanup + run: | + sudo rm -rf $GITHUB_WORKSPACE/* || true + + - name: Checkout code + uses: actions/checkout@v4 + with: + submodules: recursive + + - name: Install KT-Kernel + run: | + cd kt-kernel + bash install.sh build + + - name: Run KT-Kernel CPU tests + timeout-minutes: 60 + run: | + cd kt-kernel/test + python3 run_suite.py --hw cpu --suite default + + # =============================================== finish ==================================================== + pr-test-kt-kernel-finish: + needs: [check-changes, per-commit-kt-kernel-cpu] + if: always() + runs-on: ubuntu-latest + steps: + - name: Check all dependent job statuses + run: | + # Convert the 'needs' context to a JSON string + json_needs='${{ toJson(needs) }}' + + # Get a list of all job names from the JSON keys + job_names=$(echo "$json_needs" | jq -r 'keys_unsorted[]') + + for job in $job_names; do + # For each job, extract its result + result=$(echo "$json_needs" | jq -r --arg j "$job" '.[$j].result') + + # Print the job name and its result + echo "$job: $result" + + # Check for failure or cancellation and exit if found + if [[ "$result" == "failure" || "$result" == "cancelled" ]]; then + echo "The above jobs failed." + exit 1 + fi + done + + # If the loop completes, all jobs were successful + echo "All jobs completed successfully" + exit 0 diff --git a/.github/workflows/release-fake-tag.yml b/.github/workflows/release-fake-tag.yml new file mode 100644 index 0000000..475261b --- /dev/null +++ b/.github/workflows/release-fake-tag.yml @@ -0,0 +1,41 @@ +name: Release Fake Tag + +on: + push: + branches: + - main + paths: + - "version.py" + workflow_dispatch: + +permissions: + contents: write + +jobs: + publish: + if: github.repository == 'kvcache-ai/ktransformers' + runs-on: ubuntu-latest + environment: 'prod' + steps: + - name: Checkout repository + uses: actions/checkout@v4 + with: + token: ${{ secrets.GITHUB_TOKEN }} + + - name: Get version + id: get_version + run: | + version=$(cat version.py | grep '__version__' | cut -d'"' -f2) + echo "TAG=v$version" >> $GITHUB_OUTPUT + + - name: Create and push tag + run: | + TAG=${{ steps.get_version.outputs.TAG }} + git config user.name "ktransformers-bot" + git config user.email "ktransformers-bot@users.noreply.github.com" + if git ls-remote --tags --exit-code origin "refs/tags/${TAG}" > /dev/null 2>&1; then + echo "Tag ${TAG} already exists on origin, skipping." + exit 0 + fi + git tag "${TAG}" + git push origin "${TAG}" diff --git a/.github/workflows/release-pypi.yml b/.github/workflows/release-pypi.yml new file mode 100644 index 0000000..6cf8abe --- /dev/null +++ b/.github/workflows/release-pypi.yml @@ -0,0 +1,283 @@ +name: Release to PyPI + +on: + push: + branches: + - main + paths: + - "version.py" + workflow_dispatch: + inputs: + test_pypi: + description: 'Publish to TestPyPI instead of PyPI (for testing)' + required: false + default: 'false' + type: choice + options: + - 'true' + - 'false' + +permissions: + contents: read + +jobs: + # ── sglang-kt (must be on PyPI before users can pip install kt-kernel) ── + build-and-publish-sglang-kt: + name: Build & publish sglang-kt + runs-on: [self-hosted, linux, x64] + if: github.repository == 'kvcache-ai/ktransformers' && github.ref == 'refs/heads/main' + environment: prod + permissions: + id-token: write + contents: read + + steps: + - name: Checkout repository + uses: actions/checkout@v4 + with: + submodules: recursive + + - name: Set up Python + uses: actions/setup-python@v4 + with: + python-version: '3.12' + + - name: Install build tools + run: | + python -m pip install --upgrade pip + pip install build wheel setuptools twine + + - name: Build sglang-kt wheel + working-directory: third_party/sglang/python + run: | + KT_VERSION=$(python3 -c "exec(open('${{ github.workspace }}/version.py').read()); print(__version__)") + export SGLANG_KT_VERSION="$KT_VERSION" + echo "Building sglang-kt v${KT_VERSION} wheel..." + python -m build --wheel -v + ls dist/ | grep -q "sglang_kt" || (echo "ERROR: Wheel name does not contain sglang_kt" && exit 1) + + - name: Publish sglang-kt to PyPI + if: github.event.inputs.test_pypi != 'true' + env: + TWINE_USERNAME: __token__ + TWINE_PASSWORD: ${{ secrets.PYPI_API_TOKEN }} + run: | + python -m twine upload --skip-existing --verbose third_party/sglang/python/dist/*.whl + + - name: Publish sglang-kt to TestPyPI (if requested) + if: github.event.inputs.test_pypi == 'true' + env: + TWINE_USERNAME: __token__ + TWINE_PASSWORD: ${{ secrets.TEST_PYPI_API_TOKEN }} + run: | + python -m twine upload --repository testpypi --skip-existing --verbose third_party/sglang/python/dist/*.whl + + # ── kt-kernel ── + build-kt-kernel: + name: Build kt-kernel (Python ${{ matrix.python-version }}) + runs-on: [self-hosted, linux, x64, gpu] + strategy: + fail-fast: false + matrix: + python-version: ['3.11', '3.12'] + + steps: + - name: Checkout repository + uses: actions/checkout@v4 + with: + submodules: recursive + + - name: Set up Python ${{ matrix.python-version }} + uses: actions/setup-python@v4 + with: + python-version: ${{ matrix.python-version }} + + - name: Verify CUDA availability + run: | + nvidia-smi || (echo "ERROR: GPU not available" && exit 1) + nvcc --version || (echo "ERROR: CUDA toolkit not found" && exit 1) + + - name: Install dependencies + run: | + # System packages (cmake/libhwloc-dev/pkg-config/libnuma-dev) are expected to be + # preinstalled on self-hosted runners. Skip apt-get to avoid sudo dependency. + for pkg in cmake pkg-config; do command -v $pkg >/dev/null || { echo "ERROR: $pkg missing on runner"; exit 1; }; done + python -m pip install --upgrade pip + pip install build wheel setuptools + pip install torch --index-url https://download.pytorch.org/whl/cu128 + + - name: Build kt-kernel wheel + working-directory: kt-kernel + env: + CPUINFER_BUILD_ALL_VARIANTS: '1' + CPUINFER_ENABLE_CPPTRACE: '0' + CPUINFER_USE_CUDA: '1' + CPUINFER_CUDA_ARCHS: '80;86;89;90;120' + CPUINFER_CUDA_STATIC_RUNTIME: '1' + CPUINFER_BUILD_TYPE: 'Release' + CPUINFER_PARALLEL: '4' + CPUINFER_FORCE_REBUILD: '1' + CUDA_HOME: '/usr/local/cuda-12.8' + run: | + echo "Building kt-kernel with:" + echo " - CUDA support (SM 80, 86, 89, 90, 120)" + echo " - CPU multi-variant (AMX, AVX512, AVX2)" + python -m build --wheel -v + + - name: Verify wheel + working-directory: kt-kernel + run: | + echo "Generated wheel:" + ls -lh dist/ + + # Install and test + pip install dist/*.whl + python -c "import kt_kernel; print(f'✓ Version: {kt_kernel.__version__}')" + python -c "import kt_kernel; print(f'✓ CPU variant: {kt_kernel.__cpu_variant__}')" + + # Verify CUDA support + python -c " + from kt_kernel import kt_kernel_ext + cpu_infer = kt_kernel_ext.CPUInfer(4) + methods = dir(cpu_infer) + has_cuda = 'submit_with_cuda_stream' in methods + print(f'✓ CUDA support: {has_cuda}') + " + + # Verify CPU multi-variant support + echo "Checking CPU variants in wheel..." + python -m zipfile -l dist/*.whl | grep "_kt_kernel_ext_" || echo "Warning: No variant .so files found" + python -m zipfile -l dist/*.whl | grep "_kt_kernel_ext_amx.cpython" && echo "✓ AMX variant found" || echo "Note: AMX variant missing" + python -m zipfile -l dist/*.whl | grep "_kt_kernel_ext_avx512" && echo "✓ AVX512 variants found" || echo "Note: AVX512 variants missing" + python -m zipfile -l dist/*.whl | grep "_kt_kernel_ext_avx2.cpython" && echo "✓ AVX2 variant found" || echo "Note: AVX2 variant missing" + + # Verify static linking (should NOT depend on libcudart.so). + # Use $RUNNER_TEMP (honors TMPDIR redirect to /mnt) — /tmp is the + # system disk on self-hosted runners and can be tight. + CHECK_DIR="${RUNNER_TEMP:-/tmp}/check" + rm -rf "$CHECK_DIR" + unzip -q dist/*.whl -d "$CHECK_DIR" + if ldd "$CHECK_DIR"/kt_kernel/*.so 2>/dev/null | grep -q "libcudart.so"; then + echo "ERROR: Dynamic cudart found, should be statically linked" + exit 1 + else + echo "✓ CUDA runtime statically linked" + fi + + - name: Repair wheel for manylinux + working-directory: kt-kernel + run: | + pip install auditwheel patchelf + mkdir -p wheelhouse + for wheel in dist/*.whl; do + auditwheel repair "$wheel" --plat manylinux_2_35_x86_64 --exclude libcuda.so.1 -w wheelhouse/ + done + rm -f dist/*.whl && cp wheelhouse/*.whl dist/ + + - name: Upload artifact + uses: actions/upload-artifact@v4 + with: + name: kt-kernel-wheels-py${{ matrix.python-version }} + path: kt-kernel/dist/*.whl + retention-days: 7 + + publish-pypi: + name: Publish kt-kernel to PyPI + needs: [build-and-publish-sglang-kt, build-kt-kernel] + runs-on: [self-hosted, linux, x64] + if: github.repository == 'kvcache-ai/ktransformers' && github.ref == 'refs/heads/main' + environment: prod + permissions: + id-token: write # For trusted publishing (OIDC) + contents: read + + steps: + - name: Download all wheel artifacts + uses: actions/download-artifact@v4 + with: + path: artifacts/ + + - name: Set up Python + uses: actions/setup-python@v4 + with: + python-version: '3.12' + + - name: Organize wheels into dist/ + run: | + mkdir -p dist/ + find artifacts/ -name "*.whl" -exec cp {} dist/ \; + echo "Wheels to publish:" + ls -lh dist/ + + - name: Get version from wheel + id: get_version + run: | + # Extract version from first wheel filename + wheel_name=$(ls dist/*.whl | head -1 | xargs basename) + # Extract version (format: kt_kernel-X.Y.Z-...) + version=$(echo "$wheel_name" | sed 's/kt_kernel-\([0-9.]*\)-.*/\1/') + echo "VERSION=$version" >> $GITHUB_OUTPUT + echo "Publishing version: $version" + + - name: Install twine + run: | + python -m pip install --upgrade pip + pip install twine + + - name: Publish to TestPyPI (if requested) + if: github.event.inputs.test_pypi == 'true' + env: + TWINE_USERNAME: __token__ + TWINE_PASSWORD: ${{ secrets.TEST_PYPI_API_TOKEN }} + run: | + python -m twine upload \ + --repository testpypi \ + --skip-existing \ + --verbose \ + dist/*.whl + + - name: Publish to PyPI + if: github.event.inputs.test_pypi != 'true' + env: + TWINE_USERNAME: __token__ + TWINE_PASSWORD: ${{ secrets.PYPI_API_TOKEN }} + run: | + python -m twine upload \ + --skip-existing \ + --verbose \ + dist/*.whl + + - name: Create release summary + run: | + echo "## 🎉 kt-kernel v${{ steps.get_version.outputs.VERSION }} Published to PyPI" >> $GITHUB_STEP_SUMMARY + echo "" >> $GITHUB_STEP_SUMMARY + echo "### Installation" >> $GITHUB_STEP_SUMMARY + echo '```bash' >> $GITHUB_STEP_SUMMARY + echo "pip install kt-kernel==${{ steps.get_version.outputs.VERSION }}" >> $GITHUB_STEP_SUMMARY + echo '```' >> $GITHUB_STEP_SUMMARY + echo "" >> $GITHUB_STEP_SUMMARY + echo "### Published Wheels" >> $GITHUB_STEP_SUMMARY + echo "Total: $(ls -1 dist/*.whl | wc -l) wheels (Python 3.10, 3.11, 3.12)" >> $GITHUB_STEP_SUMMARY + echo "" >> $GITHUB_STEP_SUMMARY + echo "### Features" >> $GITHUB_STEP_SUMMARY + echo "" >> $GITHUB_STEP_SUMMARY + echo "**CPU Multi-Variant Support:**" >> $GITHUB_STEP_SUMMARY + echo "- ✅ AMX (Intel Sapphire Rapids+, 2023)" >> $GITHUB_STEP_SUMMARY + echo "- ✅ AVX512 Base/VNNI/VBMI/BF16 (Intel Skylake-X/Ice Lake/Cascade Lake, 2017+)" >> $GITHUB_STEP_SUMMARY + echo "- ✅ AVX2 (Maximum compatibility, 2013+)" >> $GITHUB_STEP_SUMMARY + echo "- 🔧 Runtime CPU detection: Automatically selects optimal variant" >> $GITHUB_STEP_SUMMARY + echo "" >> $GITHUB_STEP_SUMMARY + echo "**CUDA Support:**" >> $GITHUB_STEP_SUMMARY + echo "- ✅ SM 80 (Ampere: A100, RTX 3000 series)" >> $GITHUB_STEP_SUMMARY + echo "- ✅ SM 86 (Ampere: RTX 3060-3090)" >> $GITHUB_STEP_SUMMARY + echo "- ✅ SM 89 (Ada Lovelace: RTX 4000 series)" >> $GITHUB_STEP_SUMMARY + echo "- ✅ SM 90 (Hopper: H100)" >> $GITHUB_STEP_SUMMARY + echo "- 🔧 Static CUDA runtime: Compatible with CUDA 11.8+ and 12.x drivers" >> $GITHUB_STEP_SUMMARY + echo "- 🔧 Works on CPU-only systems (CUDA features disabled gracefully)" >> $GITHUB_STEP_SUMMARY + echo "" >> $GITHUB_STEP_SUMMARY + echo "**Requirements:**" >> $GITHUB_STEP_SUMMARY + echo "- Python 3.10, 3.11, or 3.12" >> $GITHUB_STEP_SUMMARY + echo "- Linux x86-64 (manylinux_2_17 compatible)" >> $GITHUB_STEP_SUMMARY + echo "- For CUDA features: NVIDIA driver with CUDA 11.8+ or 12.x support" >> $GITHUB_STEP_SUMMARY + echo "" >> $GITHUB_STEP_SUMMARY + echo "PyPI link: https://pypi.org/project/kt-kernel/${{ steps.get_version.outputs.VERSION }}/" >> $GITHUB_STEP_SUMMARY diff --git a/.github/workflows/release-sglang-kt.yml b/.github/workflows/release-sglang-kt.yml new file mode 100644 index 0000000..f18221c --- /dev/null +++ b/.github/workflows/release-sglang-kt.yml @@ -0,0 +1,129 @@ +name: Release sglang-kt to PyPI + +on: + push: + branches: + - main + paths: + - "third_party/sglang" + workflow_dispatch: + inputs: + test_pypi: + description: 'Publish to TestPyPI instead of PyPI (for testing)' + required: false + default: 'false' + type: choice + options: + - 'true' + - 'false' + +permissions: + contents: read + +jobs: + build-sglang-kt: + name: Build sglang-kt wheel + runs-on: ubuntu-latest + + steps: + - name: Checkout repository + uses: actions/checkout@v4 + with: + submodules: recursive + + - name: Set up Python + uses: actions/setup-python@v4 + with: + python-version: '3.12' + + - name: Install build tools + run: | + python -m pip install --upgrade pip + pip install build wheel setuptools + + - name: Build sglang-kt wheel + working-directory: third_party/sglang/python + run: | + # Read version from ktransformers version.py + KT_VERSION=$(python3 -c "exec(open('${{ github.workspace }}/version.py').read()); print(__version__)") + export SGLANG_KT_VERSION="$KT_VERSION" + echo "Building sglang-kt v${KT_VERSION} wheel..." + python -m build --wheel -v + + - name: Verify wheel + working-directory: third_party/sglang/python + run: | + echo "Generated wheel:" + ls -lh dist/ + # Verify the wheel has the correct package name + ls dist/ | grep -q "sglang_kt" || (echo "ERROR: Wheel name does not contain sglang_kt" && exit 1) + echo "Wheel name verified." + + - name: Upload artifact + uses: actions/upload-artifact@v4 + with: + name: sglang-kt-wheel + path: third_party/sglang/python/dist/*.whl + retention-days: 7 + + publish-pypi: + name: Publish sglang-kt to PyPI + needs: [build-sglang-kt] + runs-on: ubuntu-latest + if: github.repository == 'kvcache-ai/ktransformers' && github.ref == 'refs/heads/main' + environment: prod + permissions: + id-token: write + contents: read + + steps: + - name: Download wheel artifact + uses: actions/download-artifact@v4 + with: + name: sglang-kt-wheel + path: dist/ + + - name: Display wheels + run: | + echo "Wheels to publish:" + ls -lh dist/ + + - name: Install twine + run: | + python -m pip install --upgrade pip + pip install twine + + - name: Publish to TestPyPI (if requested) + if: github.event.inputs.test_pypi == 'true' + env: + TWINE_USERNAME: __token__ + TWINE_PASSWORD: ${{ secrets.TEST_PYPI_API_TOKEN }} + run: | + python -m twine upload \ + --repository testpypi \ + --skip-existing \ + --verbose \ + dist/*.whl + + - name: Publish to PyPI + if: github.event.inputs.test_pypi != 'true' + env: + TWINE_USERNAME: __token__ + TWINE_PASSWORD: ${{ secrets.PYPI_API_TOKEN }} + run: | + python -m twine upload \ + --skip-existing \ + --verbose \ + dist/*.whl + + - name: Create release summary + run: | + echo "## sglang-kt Published to PyPI" >> $GITHUB_STEP_SUMMARY + echo "" >> $GITHUB_STEP_SUMMARY + echo "### Installation" >> $GITHUB_STEP_SUMMARY + echo '```bash' >> $GITHUB_STEP_SUMMARY + echo "pip install sglang-kt" >> $GITHUB_STEP_SUMMARY + echo '```' >> $GITHUB_STEP_SUMMARY + echo "" >> $GITHUB_STEP_SUMMARY + echo "This is the kvcache-ai fork of SGLang with kt-kernel support." >> $GITHUB_STEP_SUMMARY + echo "PyPI link: https://pypi.org/project/sglang-kt/" >> $GITHUB_STEP_SUMMARY diff --git a/.github/workflows/sync-sglang-submodule.yml b/.github/workflows/sync-sglang-submodule.yml new file mode 100644 index 0000000..7a957c5 --- /dev/null +++ b/.github/workflows/sync-sglang-submodule.yml @@ -0,0 +1,81 @@ +name: Sync sglang submodule + +on: + schedule: + # Run daily at 08:00 UTC + - cron: "0 8 * * *" + workflow_dispatch: + +permissions: + contents: write + pull-requests: write + +jobs: + sync: + name: Check for sglang-kt updates + runs-on: ubuntu-latest + + steps: + - name: Checkout repository + uses: actions/checkout@v4 + with: + submodules: true + fetch-depth: 0 + token: ${{ secrets.GITHUB_TOKEN }} + + - name: Update sglang submodule to latest main + id: update + run: | + OLD_SHA=$(git -C third_party/sglang rev-parse HEAD) + git submodule update --remote third_party/sglang + NEW_SHA=$(git -C third_party/sglang rev-parse HEAD) + + echo "old_sha=$OLD_SHA" >> "$GITHUB_OUTPUT" + echo "new_sha=$NEW_SHA" >> "$GITHUB_OUTPUT" + + if [ "$OLD_SHA" = "$NEW_SHA" ]; then + echo "changed=false" >> "$GITHUB_OUTPUT" + echo "sglang submodule is already up to date ($OLD_SHA)" + else + echo "changed=true" >> "$GITHUB_OUTPUT" + + # Collect commit log between old and new + COMMITS=$(git -C third_party/sglang log --oneline "$OLD_SHA..$NEW_SHA" | head -20) + echo "commits<> "$GITHUB_OUTPUT" + echo "$COMMITS" >> "$GITHUB_OUTPUT" + echo "EOF" >> "$GITHUB_OUTPUT" + + # sglang-kt version = ktransformers version (from version.py) + VERSION=$(python3 -c "exec(open('version.py').read()); print(__version__)" 2>/dev/null || echo "unknown") + echo "version=$VERSION" >> "$GITHUB_OUTPUT" + + echo "sglang submodule updated: $OLD_SHA -> $NEW_SHA (v$VERSION)" + fi + + - name: Create pull request + if: steps.update.outputs.changed == 'true' + uses: peter-evans/create-pull-request@v6 + with: + token: ${{ secrets.GITHUB_TOKEN }} + commit-message: | + [build]: sync sglang submodule to ${{ steps.update.outputs.new_sha }} + branch: auto/sync-sglang + delete-branch: true + title: "[build] Sync sglang-kt submodule (v${{ steps.update.outputs.version }})" + body: | + Automated sync of `third_party/sglang` submodule to latest `main`. + + **Old ref:** `${{ steps.update.outputs.old_sha }}` + **New ref:** `${{ steps.update.outputs.new_sha }}` + **sglang-kt version:** `${{ steps.update.outputs.version }}` + + ### Commits included + ``` + ${{ steps.update.outputs.commits }} + ``` + + --- + *This PR was created automatically by the [sync-sglang-submodule](${{ github.server_url }}/${{ github.repository }}/actions/workflows/sync-sglang-submodule.yml) workflow.* + labels: | + dependencies + automated diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..f44a0ab --- /dev/null +++ b/.gitignore @@ -0,0 +1,33 @@ +__pycache__ +build +.vscode +*.so +*.cache +server.db +logs +node_modules +*.nsys-rep +.vs/ +*pycache* +*build/ +.DS_Store +compile_commands.json +*.egg-info* +*dist/ +ktransformers/server/local_store/ +ktransformers/server_test1.db +*.patch +img/ +tmp*.txt +test.txt +book +ktransformers/tests/chat_txt.txt +mmlu_result* +ktransformers/ktransformers_ext/cuda_musa/ +test_prompt.txt +csrc/demo +build* +CMakeFiles/ +kvc2/ +sched/ +*.png \ No newline at end of file diff --git a/.gitmodules b/.gitmodules new file mode 100644 index 0000000..5b4ca91 --- /dev/null +++ b/.gitmodules @@ -0,0 +1,14 @@ +[submodule "third_party/llama.cpp"] + path = third_party/llama.cpp + url = https://github.com/ggerganov/llama.cpp.git +[submodule "third_party/pybind11"] + path = third_party/pybind11 + url = https://github.com/pybind/pybind11.git +[submodule "third_party/custom_flashinfer"] + path = third_party/custom_flashinfer + url = https://github.com/kvcache-ai/custom_flashinfer.git + branch = fix-precision-mla-merge-main +[submodule "third_party/sglang"] + path = third_party/sglang + url = https://github.com/kvcache-ai/sglang.git + branch = main diff --git a/LICENSE b/LICENSE new file mode 100644 index 0000000..261eeb9 --- /dev/null +++ b/LICENSE @@ -0,0 +1,201 @@ + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright [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/MAINTAINERS.md b/MAINTAINERS.md new file mode 100644 index 0000000..df7ebd8 --- /dev/null +++ b/MAINTAINERS.md @@ -0,0 +1,34 @@ +# Maintainers + +This document lists the current maintainers and outlines their responsibilities. + +## Current Maintainers + +| Name | GitHub | Role | Affiliation | Email | +|------|--------|------|-------------|-------| +| Weiyu Xie | [@ErvinXie](https://github.com/ErvinXie) | Maintainer | [MADSys Lab](https://madsys.cs.tsinghua.edu.cn/) @ Tsinghua University | xwy21@mails.tsinghua.edu.cn | +| Hongtao Chen | [@chenht2022](https://github.com/chenht2022) | Maintainer | [MADSys Lab](https://madsys.cs.tsinghua.edu.cn/) @ Tsinghua University | cht22@mails.tsinghua.edu.cn | +| Jianwei Dong | [@ovowei](https://github.com/ovowei) | Maintainer | [MADSys Lab](https://madsys.cs.tsinghua.edu.cn/) @ Tsinghua University | dongjw24@mails.tsinghua.edu.cn | +| Ziwei Yuan | [@KMSorSMS](https://github.com/KMSorSMS) | Maintainer | [Approaching.AI](http://approaching.ai/) | 2022090910005@std.uestc.edu.cn | +| Qingliang Ou | [@ouqingliang](https://github.com/ouqingliang) | Maintainer | [MADSys Lab](https://madsys.cs.tsinghua.edu.cn/) @ Tsinghua University | oql@bupt.edu.cn | +| Jiaqi Liao | [@SkqLiao](https://github.com/SkqLiao) | Maintainer | [Approaching.AI](http://approaching.ai/) | jiaqi.liao@bit.edu.cn | +| Peilin Li | [@JimmyPeilinLi](https://github.com/JimmyPeilinLi) | Maintainer | [Approaching.AI](http://approaching.ai/) | lipeilin@mail.nwpu.edu.cn | +| Xingxing Hao | [@mrhaoxx](https://github.com/mrhaoxx) | Maintainer | [Approaching.AI](http://approaching.ai/) | mr.haoxx@gmail.com | +| Boxin Zhang | [@Atream](https://github.com/Atream) | Maintainer | [MADSys Lab](https://madsys.cs.tsinghua.edu.cn/) @ Tsinghua University | zhangbx24@mails.tsinghua.edu.cn | +| Jingqi Tang | [@Azure-Tang](https://github.com/Azure-Tang) | Maintainer | [MADSys Lab](https://madsys.cs.tsinghua.edu.cn/) @ Tsinghua University | tangjq25@mails.tsinghua.edu.cn | +| Jiahao Wang | [@qiyuxinlin](https://github.com/qiyuxinlin) | Maintainer | [Approaching.AI](http://approaching.ai/) | 202241050020@hdu.edu.cn | + +## Responsibilities + +Maintainers steward the project and keep it healthy for users and contributors. + +- Review and approve pull requests; ensure changes meet quality, testing, and documentation standards. +- Triage issues, keep labels organized, and respond to questions in a timely manner. +- Uphold the project’s code of conduct and report violations when needed. +- Maintain CI reliability and address regressions promptly. +- Oversee releases and keep compatibility with supported dependency versions. +- Protect project security and follow the security disclosure process. + +## Becoming a Maintainer + +We welcome contributors who show sustained, high-quality contributions and collaborative behavior. If you are interested, please contact an existing maintainer and share your recent contributions and areas of focus. diff --git a/MANIFEST.in b/MANIFEST.in new file mode 100644 index 0000000..2e92268 --- /dev/null +++ b/MANIFEST.in @@ -0,0 +1 @@ +include version.py diff --git a/README.md b/README.md new file mode 100644 index 0000000..819dda0 --- /dev/null +++ b/README.md @@ -0,0 +1,163 @@ +
+

+ + + KTransformers + + + +

+

A Flexible Framework for Experiencing Cutting-edge LLM Inference/Fine-tune Optimizations

+ 🎯 Overview | 🚀 Inference | 🎓 SFT | 🔥 Citation | 🚀 Roadmap(2026Q2) +
+ +## 🎯 Overview + +KTransformers is a research project focused on efficient inference and fine-tuning of large language models through CPU-GPU heterogeneous computing. The project now exposes two user-facing capabilities from the kt-kernel source tree: [Inference](./kt-kernel/README.md) and [SFT](./doc/en/SFT/KTransformers-Fine-Tuning_Quick-Start.md). + +## 🔥 Updates +* **June 21, 2026**: MiniMax-M3 Day0 Support! ([Tutorial](./doc/en/kt-kernel/MiniMax-M3-Tutorial.md)) +* **June 17, 2026**: GLM-5.2 Day0 Support! ([Tutorial](./doc/en/kt-kernel/GLM-5.2-Tutorial.md)) +* **May 6, 2026**: KTransformers at [GOSIM Paris 2026](https://paris2026.gosim.org/zh/schedule/) — "Agentic AI on Edge" track. We'll present KT's inference performance on consumer hardware. +* **May 02, 2026**: DeepSeek-V4-Flash Support! ([Tutorial](./doc/en/DeepSeek-V4-Flash.md)) +* **Apr 30, 2026**: KTransformers v0.6.1 refreshes kt-kernel inference and SFT docs with separate [Inference](./kt-kernel/README.md) and [SFT Quick Start](./doc/en/SFT/KTransformers-Fine-Tuning_Quick-Start.md) entry points. +* **Mar 26, 2026**: Support AVX2-only CPU backend for KT-Kernel inference. ([Tutorial](./doc/en/kt-kernel/AVX2-Tutorial.md)) +* **Feb 13, 2026**: MiniMax-M2.5 Day0 Support! ([Tutorial](./doc/en/MiniMax-M2.5.md)) +* **Feb 12, 2026**: GLM-5 Day0 Support! ([Tutorial](./doc/en/kt-kernel/GLM-5-Tutorial.md)) +* **Jan 27, 2026**: Kimi-K2.5 Day0 Support! ([Tutorial](./doc/en/Kimi-K2.5.md)) ([SFT Tutorial](./doc/en/SFT_Installation_Guide_KimiK2.5.md)) +* **Jan 22, 2026**: Support [CPU-GPU Expert Scheduling](./doc/en/kt-kernel/experts-sched-Tutorial.md), [Native BF16 and FP8 per channel Precision](./doc/en/kt-kernel/Native-Precision-Tutorial.md) and [AutoDL unified fine-tuning and inference](./doc/zh/【云端低价训推】%20KTransformers%2BAutoDL%2BLlamaFactory：随用随租的低成本超大模型「微调%2B推理」一体化流程.pdf) +* **Dec 24, 2025**: Support Native MiniMax-M2.1 inference. ([Tutorial](./doc/en/kt-kernel/MiniMax-M2.1-Tutorial.md)) +* **Dec 22, 2025**: Support RL-DPO fine-tuning with LLaMA-Factory. ([Tutorial](./doc/en/SFT/DPO_tutorial.md)) +* **Dec 5, 2025**: Support Native Kimi-K2-Thinking inference ([Tutorial](./doc/en/kt-kernel/Kimi-K2-Thinking-Native.md)) +* **Nov 6, 2025**: Support Kimi-K2-Thinking inference ([Tutorial](./doc/en/Kimi-K2-Thinking.md)) and fine-tune ([Tutorial](./doc/en/SFT_Installation_Guide_KimiK2.md)) +* **Nov 4, 2025**: KTransformers Fine-Tuning × LLaMA-Factory Integration. ([Tutorial](./doc/en/SFT/KTransformers-Fine-Tuning_User-Guide.md)) +* **Oct 27, 2025**: Support Ascend NPU. ([Tutorial](./doc/zh/DeepseekR1_V3_tutorial_zh_for_Ascend_NPU.md)) +* **Oct 10, 2025**: Integrating into SGLang. ([Roadmap](https://github.com/sgl-project/sglang/issues/11425), [Blog](https://lmsys.org/blog/2025-10-22-KTransformers/)) +* **Sept 11, 2025**: Support Qwen3-Next. ([Tutorial](./doc/en/Qwen3-Next.md)) +* **Sept 05, 2025**: Support Kimi-K2-0905. ([Tutorial](./doc/en/Kimi-K2.md)) +* **July 26, 2025**: Support SmallThinker and GLM4-MoE. ([Tutorial](./doc/en/SmallThinker_and_Glm4moe.md)) +* **July 11, 2025**: Support Kimi-K2. ([Tutorial](./doc/en/Kimi-K2.md)) +* **June 30, 2025**: Support 3-layer (GPU-CPU-Disk) [prefix cache](./doc/en/prefix_cache.md) reuse. +* **May 14, 2025**: Support Intel Arc GPU ([Tutorial](./doc/en/xpu.md)). +* **Apr 29, 2025**: Support AMX-Int8、 AMX-BF16 and Qwen3MoE ([Tutorial](./doc/en/AMX.md)) +* **Apr 9, 2025**: Experimental support for LLaMA 4 models ([Tutorial](./doc/en/llama4.md)). +* **Apr 2, 2025**: Support Multi-concurrency. ([Tutorial](./doc/en/balance-serve.md)). +* **Mar 15, 2025**: Support ROCm on AMD GPU ([Tutorial](./doc/en/ROCm.md)). +* **Mar 5, 2025**: Support unsloth 1.58/2.51 bits weights and [IQ1_S/FP8 hybrid](./doc/en/fp8_kernel.md) weights. Support 139K [Longer Context](./doc/en/DeepseekR1_V3_tutorial.md#v022--v023-longer-context--fp8-kernel) for DeepSeek-V3 and R1 in 24GB VRAM. +* **Feb 25, 2025**: Support [FP8 GPU kernel](./doc/en/fp8_kernel.md) for DeepSeek-V3 and R1; [Longer Context](./doc/en/DeepseekR1_V3_tutorial.md#v022-longer-context). +* **Feb 15, 2025**: Longer Context (from 4K to 8K for 24GB VRAM) & Slightly Faster Speed （+15%, up to 16 Tokens/s), update [docs](./doc/en/DeepseekR1_V3_tutorial.md) and [online books](https://kvcache-ai.github.io/ktransformers/). +* **Feb 10, 2025**: Support Deepseek-R1 and V3 on single (24GB VRAM)/multi gpu and 382G DRAM, up to 3~28x speedup. For detailed show case and reproduction tutorial, see [here](./doc/en/DeepseekR1_V3_tutorial.md). +* **Aug 28, 2024**: Decrease DeepseekV2's required VRAM from 21G to 11G. +* **Aug 15, 2024**: Update detailed [tutorial](doc/en/injection_tutorial.md) for injection and multi-GPU. +* **Aug 14, 2024**: Support llamfile as linear backend. +* **Aug 12, 2024**: Support multiple GPU; Support new model: mixtral 8\*7B and 8\*22B; Support q2k, q3k, q5k dequant on gpu. +* **Aug 9, 2024**: Support windows native. + +--- + +## 📦 Capabilities + +### 🚀 [Inference](./kt-kernel/README.md) - High-Performance kt-kernel Serving + +CPU-optimized kernel operations for heterogeneous LLM inference. + +image + + +**Key Features:** +- **AMX/AVX Acceleration**: Intel AMX and AVX512/AVX2 optimized kernels for INT4/INT8 quantized inference +- **MoE Optimization**: Efficient Mixture-of-Experts inference with NUMA-aware memory management +- **Quantization Support**: CPU-side INT4/INT8 quantized weights, GPU-side GPTQ support +- **Easy Integration**: Clean Python API for SGLang and other frameworks + +**Quick Start:** +```bash +cd kt-kernel +pip install . +``` + +**Use Cases:** + +- CPU-GPU hybrid inference for large MoE models +- Integration with SGLang for production serving +- Heterogeneous expert placement (hot experts on GPU, cold experts on CPU) + +**Performance Examples:** +| Model | Hardware Configuration | Total Throughput | Output Throughput | +|-------|------------------------|------------------|-------------------| +| DeepSeek-R1-0528 (FP8) | 8×L20 GPU + Xeon Gold 6454S | 227.85 tokens/s | 87.58 tokens/s (8-way concurrency) | + +👉 **[Full Documentation →](./kt-kernel/README.md)** + +--- + +### 🎓 [SFT](./doc/en/SFT/KTransformers-Fine-Tuning_Quick-Start.md) - Fine-Tuning with LLaMA-Factory + +KTransformers × LLaMA-Factory integration for ultra-large MoE model fine-tuning. + +![KTransformers SFT](https://raw.githubusercontent.com/kvcache-ai/ktransformers/main/doc/assets/image-20251011010558909.png) + +**Key Features:** +- **Multi-Backend Support**: CPU/GPU hybrid fine-tuning with INT8/INT4 quantization +- **Ultra-Large MoE Support**: Fine-tune models like DeepSeek-V3/R1 on limited GPU memory +- **Faster than ZeRO-Offload**: 6-12x training speedup in benchmarked MoE SFT workloads +- **Lower CPU Memory**: About half the CPU memory of the previous KT SFT path in the benchmarked setup +- **LLaMA-Factory Integration**: Seamless integration with popular fine-tuning framework + +| Model | GPU Memory | Training Speed | Hardware | +|-------|------------|----------------|----------| +| DeepSeek-V3 | ~80GB total | 3.7 it/s | 4x RTX 4090 | +| DeepSeek-R1 | ~80GB total | 3.7 it/s | 4x RTX 4090 | +| Qwen3-30B-A3B | ~24GB total | 8+ it/s | 1x RTX 4090 | + +**Quick Start:** +```bash +cd /path/to/LLaMA-Factory +pip install -e . +pip install -r requirements/ktransformers.txt +CUDA_VISIBLE_DEVICES=0,1,2,3 accelerate launch \ + --config_file examples/ktransformers/accelerate/fsdp2_kt_int8.yaml \ + src/train.py \ + examples/ktransformers/train_lora/qwen3_5moe_lora_sft_kt.yaml +``` + +👉 **[Quick Start →](./doc/en/SFT/KTransformers-Fine-Tuning_Quick-Start.md)** +👉 **[Full Documentation →](./doc/en/SFT/KTransformers-Fine-Tuning_User-Guide.md)** + +--- + +## 🔥 Citation + +If you use KTransformers in your research, please cite our paper: + +```bibtex +@inproceedings{10.1145/3731569.3764843, + title = {KTransformers: Unleashing the Full Potential of CPU/GPU Hybrid Inference for MoE Models}, + author = {Chen, Hongtao and Xie, Weiyu and Zhang, Boxin and Tang, Jingqi and Wang, Jiahao and Dong, Jianwei and Chen, Shaoyuan and Yuan, Ziwei and Lin, Chen and Qiu, Chengyu and Zhu, Yuening and Ou, Qingliang and Liao, Jiaqi and Chen, Xianglin and Ai, Zhiyuan and Wu, Yongwei and Zhang, Mingxing}, + booktitle = {Proceedings of the ACM SIGOPS 31st Symposium on Operating Systems Principles}, + year = {2025} +} +``` + +## 👥 Contributors & Team + +Developed and maintained by: +- [MADSys Lab](https://madsys.cs.tsinghua.edu.cn/) @ Tsinghua University +- [Approaching.AI](http://approaching.ai/) +- [9#AISoft](https://github.com/aisoft9) +- Community contributors + +We welcome contributions! Please feel free to submit issues and pull requests. + +## 💬 Community & Support + +- **GitHub Issues**: [Report bugs or request features](https://github.com/kvcache-ai/ktransformers/issues) +- **WeChat Group**: See [archive/WeChatGroup.png](./archive/WeChatGroup.png) + +## 📦 KT original Code + +The original integrated KTransformers framework has been archived to the [`archive/`](./archive/) directory for reference. The project now organizes the two capabilities above from the kt-kernel source tree for clearer documentation and maintenance. + +For the original documentation with full quick-start guides and examples, see: +- [archive/README.md](./archive/README.md) (English) +- [archive/README_ZH.md](./archive/README_ZH.md) (中文) diff --git a/README.wehub.md b/README.wehub.md new file mode 100644 index 0000000..e9b1ee9 --- /dev/null +++ b/README.wehub.md @@ -0,0 +1,7 @@ +# WeHub 来源说明 + +- 原始项目：`kvcache-ai/ktransformers` +- 原始仓库：https://github.com/kvcache-ai/ktransformers +- 导入方式：上游默认分支的最新快照 +- 原作者、版权和许可证信息以原始仓库及本仓库 LICENSE 为准 +- 本文件仅用于记录来源，不代表 WeHub 是原项目作者 diff --git a/README_ZH.md b/README_ZH.md new file mode 100644 index 0000000..5f397a8 --- /dev/null +++ b/README_ZH.md @@ -0,0 +1,151 @@ +
+

+ + + KTransformers + + + +

+

一个用于体验尖端 LLM 推理/微调优化的灵活框架

+ 🎯 概览 | 🚀 推理 | 🎓 SFT | 🔥 引用 +
+ +## 🎯 概览 + +KTransformers 是一个专注于通过 CPU-GPU 异构计算实现大语言模型高效推理和微调的研究项目。目前两个面向用户的能力都来自 kt-kernel 源码目录：[推理](./kt-kernel/README.md) 和 [SFT](./doc/en/SFT/KTransformers-Fine-Tuning_Quick-Start.md)。 + +## 🔥 更新 + +* **2026 年 4 月 30 日**：KTransformers v0.6.1 更新 kt-kernel 推理和 SFT 文档，提供独立的[推理](./kt-kernel/README.md)和 [SFT Quick Start](./doc/en/SFT/KTransformers-Fine-Tuning_Quick-Start.md)入口。 +* **2025 年 12 月 5 日**：支持原生 Kimi-K2-Thinking 推理（[教程](./doc/en/kt-kernel/Kimi-K2-Thinking-Native.md)） +* **2025 年 11 月 6 日**：支持 Kimi-K2-Thinking 推理（[教程](./doc/en/Kimi-K2-Thinking.md)）和微调（[教程](./doc/en/SFT_Installation_Guide_KimiK2.md)） +* **2025 年 11 月 4 日**：KTransformers 微调 × LLaMA-Factory 集成（[教程](./doc/en/SFT/KTransformers-Fine-Tuning_User-Guide.md)） +* **2025 年 10 月 27 日**：支持昇腾 NPU（[教程](./doc/zh/DeepseekR1_V3_tutorial_zh_for_Ascend_NPU.md)） +* **2025 年 10 月 10 日**：集成到 SGLang（[路线图](https://github.com/sgl-project/sglang/issues/11425)，[博客](https://lmsys.org/blog/2025-10-22-KTransformers/)） +* **2025 年 9 月 11 日**：支持 Qwen3-Next（[教程](./doc/en/Qwen3-Next.md)） +* **2025 年 9 月 5 日**：支持 Kimi-K2-0905（[教程](./doc/en/Kimi-K2.md)） +* **2025 年 7 月 26 日**：支持 SmallThinker 和 GLM4-MoE（[教程](./doc/en/SmallThinker_and_Glm4moe.md)） +* **2025 年 7 月 11 日**：支持 Kimi-K2（[教程](./doc/en/Kimi-K2.md)） +* **2025 年 6 月 30 日**：支持 3 层（GPU-CPU-磁盘）[前缀缓存](./doc/en/prefix_cache.md)复用 +* **2025 年 5 月 14 日**：支持 Intel Arc GPU（[教程](./doc/en/xpu.md)） +* **2025 年 4 月 29 日**：支持 AMX-Int8、AMX-BF16 和 Qwen3MoE（[教程](./doc/en/AMX.md)） +* **2025 年 4 月 9 日**：实验性支持 LLaMA 4 模型（[教程](./doc/en/llama4.md)） +* **2025 年 4 月 2 日**：支持多并发（[教程](./doc/en/balance-serve.md)） +* **2025 年 3 月 15 日**：支持 AMD GPU 上的 ROCm（[教程](./doc/en/ROCm.md)） +* **2025 年 3 月 5 日**：支持 unsloth 1.58/2.51 位权重和 [IQ1_S/FP8 混合](./doc/en/fp8_kernel.md)权重。在 24GB VRAM 中支持 DeepSeek-V3 和 R1 的 139K [更长上下文](./doc/en/DeepseekR1_V3_tutorial.md#v022--v023-longer-context--fp8-kernel) +* **2025 年 2 月 25 日**：为 DeepSeek-V3 和 R1 支持 [FP8 GPU 内核](./doc/en/fp8_kernel.md)；[更长上下文](./doc/en/DeepseekR1_V3_tutorial.md#v022-longer-context) +* **2025 年 2 月 15 日**：更长上下文（24GB VRAM 从 4K 到 8K）& 速度稍快（+15%，最高 16 Tokens/s），更新[文档](./doc/en/DeepseekR1_V3_tutorial.md)和[在线手册](https://kvcache-ai.github.io/ktransformers/) +* **2025 年 2 月 10 日**：支持 Deepseek-R1 和 V3 在单 GPU（24GB VRAM）/多 GPU 和 382GB DRAM 上运行，速度提升高达 3~28 倍。详细案例展示和复现教程请参见[这里](./doc/en/DeepseekR1_V3_tutorial.md) +* **2024 年 8 月 28 日**：将 DeepseekV2 所需的 VRAM 从 21GB 降低到 11GB +* **2024 年 8 月 15 日**：更新了关于注入和多 GPU 的详细[教程](doc/en/injection_tutorial.md) +* **2024 年 8 月 14 日**：支持 llamfile 作为线性后端 +* **2024 年 8 月 12 日**：支持多 GPU；支持新模型：mixtral 8\*7B 和 8\*22B；支持 GPU 上的 q2k、q3k、q5k 去量化 +* **2024 年 8 月 9 日**：支持 Windows 原生环境 + +--- + +## 📦 功能入口 + +### 🚀 [推理](./kt-kernel/README.md) - kt-kernel 高性能推理 + +用于异构 LLM 推理的 CPU 优化内核操作。 + +![image-20251011010558909](./doc/assets/heterogeneous_computing.png) + +**主要特性：** +- **AMX/AVX 加速**：Intel AMX 和 AVX512/AVX2 优化的内核，用于 INT4/INT8 量化推理 +- **MoE 优化**：高效的专家混合推理，具有 NUMA 感知内存管理 +- **量化支持**：CPU 端 INT4/INT8 量化权重，GPU 端 GPTQ 支持 +- **易于集成**：为 SGLang 和其他框架提供简洁的 Python API + +**快速开始：** +```bash +cd kt-kernel +pip install . +``` + +**使用场景：** + +- 大型 MoE 模型的 CPU-GPU 混合推理 +- 与 SGLang 集成用于生产服务 +- 异构专家放置（热专家在 GPU 上，冷专家在 CPU 上） + +**性能示例：** +| 模型 | 硬件配置 | 总吞吐量 | 输出吞吐量 | +|-------|------------------------|------------------|-------------------| +| DeepSeek-R1-0528 (FP8) | 8×L20 GPU + Xeon Gold 6454S | 227.85 tokens/s | 87.58 tokens/s（8 路并发）| + +👉 **[完整文档 →](./kt-kernel/README.md)** + +--- + +### 🎓 [SFT](./doc/en/SFT/KTransformers-Fine-Tuning_Quick-Start.md) - LLaMA-Factory 微调 + +KTransformers × LLaMA-Factory 集成，面向超大 MoE 模型微调。 + +![KTransformers SFT](./doc/assets/image-20251011010558909.png) + +**主要特性:** +- **多后端支持**: CPU/GPU 混合微调，支持 INT8/INT4 量化 +- **超大 MoE 支持**: 在有限 GPU 内存下微调 DeepSeek-V3/R1 等模型 +- **相对 ZeRO-Offload 加速**: 在基准 MoE SFT 任务中训练速度提升 6-12 倍 +- **降低 CPU 内存**: 相比上一版 KT SFT 路径，基准配置下 CPU 内存约降至 1/2 +- **LLaMA-Factory 集成**: 与流行微调框架无缝集成 + +| 模型 | GPU 内存 | 训练速度 | 硬件 | +|-------|------------|----------------|----------| +| DeepSeek-V3 | ~80GB 总计 | 3.7 it/s | 4x RTX 4090 | +| DeepSeek-R1 | ~80GB 总计 | 3.7 it/s | 4x RTX 4090 | +| Qwen3-30B-A3B | ~24GB 总计 | 8+ it/s | 1x RTX 4090 | + +**快速开始:** +```bash +cd /path/to/LLaMA-Factory +pip install -e . +pip install -r requirements/ktransformers.txt +CUDA_VISIBLE_DEVICES=0,1,2,3 accelerate launch \ + --config_file examples/ktransformers/accelerate/fsdp2_kt_int8.yaml \ + src/train.py \ + examples/ktransformers/train_lora/qwen3_5moe_lora_sft_kt.yaml +``` + +👉 **[Quick Start →](./doc/en/SFT/KTransformers-Fine-Tuning_Quick-Start.md)** +👉 **[完整文档 →](./doc/en/SFT/KTransformers-Fine-Tuning_User-Guide.md)** + +--- + +## 🔥 引用 + +如果您在研究中使用了 KTransformers，请引用我们的论文： + +```bibtex +@inproceedings{10.1145/3731569.3764843, + title = {KTransformers: Unleashing the Full Potential of CPU/GPU Hybrid Inference for MoE Models}, + author = {Chen, Hongtao and Xie, Weiyu and Zhang, Boxin and Tang, Jingqi and Wang, Jiahao and Dong, Jianwei and Chen, Shaoyuan and Yuan, Ziwei and Lin, Chen and Qiu, Chengyu and Zhu, Yuening and Ou, Qingliang and Liao, Jiaqi and Chen, Xianglin and Ai, Zhiyuan and Wu, Yongwei and Zhang, Mingxing}, + booktitle = {Proceedings of the ACM SIGOPS 31st Symposium on Operating Systems Principles}, + year = {2025} +} +``` + +## 👥 贡献者与团队 + +由以下团队开发和维护： +- 清华大学 [MADSys 实验室](https://madsys.cs.tsinghua.edu.cn/) +- [Approaching.AI](http://approaching.ai/) +- 社区贡献者 + +我们欢迎贡献！请随时提交问题和拉取请求。 + +## 💬 社区与支持 + +- **GitHub Issues**：[报告问题或请求功能](https://github.com/kvcache-ai/ktransformers/issues) +- **微信群**：请参见 [archive/WeChatGroup.png](./archive/WeChatGroup.png) + +## 📦 KT原仓库 + +原始的集成 KTransformers 框架已归档到 [`archive/`](./archive/) 目录以供参考。该项目现在围绕 kt-kernel 源码树中的上述两个能力入口组织文档和维护。 + +有关原始文档以及完整的快速入门指南和示例，请参见： +- [archive/README.md](./archive/README.md)（英文） +- [archive/README_ZH.md](./archive/README_ZH.md)（中文） diff --git a/archive/.devcontainer/Dockerfile b/archive/.devcontainer/Dockerfile new file mode 100644 index 0000000..e088a94 --- /dev/null +++ b/archive/.devcontainer/Dockerfile @@ -0,0 +1,19 @@ +FROM pytorch/pytorch:2.5.1-cuda12.1-cudnn9-devel as compile_server +WORKDIR /workspace +ENV CUDA_HOME /usr/local/cuda +RUN <> ~/.bashrc && \ + echo "conda activate ktransformers" >> ~/.bashrc + +WORKDIR /ktransformers/ +CMD ["bash"] diff --git a/archive/LICENSE b/archive/LICENSE new file mode 100644 index 0000000..261eeb9 --- /dev/null +++ b/archive/LICENSE @@ -0,0 +1,201 @@ + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright [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/archive/MANIFEST.in b/archive/MANIFEST.in new file mode 100644 index 0000000..4097ce6 --- /dev/null +++ b/archive/MANIFEST.in @@ -0,0 +1,13 @@ +graft third_party +graft ktransformers +graft local_chat.py +graft csrc +include LICENSE README.md +prune ktransformers/website +prune ktransformers/logs +prune ktransformers.egg-info +prune third_party/llama.cpp/models +graft ktransformers/website/dist +global-exclude __pycache__ +include KTransformersOps.*.so +include cpuinfer_ext.*.so diff --git a/archive/Makefile b/archive/Makefile new file mode 100644 index 0000000..74cb3c9 --- /dev/null +++ b/archive/Makefile @@ -0,0 +1,32 @@ +flake_find: + cd ktransformers && flake8 | grep -Eo '[A-Z][0-9]{3}' | sort | uniq| paste -sd ',' - +format: + @cd ktransformers && black . + @black setup.py +dev_install: +# clear build dirs + rm -rf build + rm -rf *.egg-info + rm -rf ktransformers/ktransformers_ext/build + rm -rf ktransformers/ktransformers_ext/cuda/build + rm -rf ktransformers/ktransformers_ext/cuda/dist + rm -rf ktransformers/ktransformers_ext/cuda/*.egg-info + +# install ktransformers + echo "Installing python dependencies from requirements.txt" + pip install -r requirements-local_chat.txt + + echo "Installing ktransformers" + KTRANSFORMERS_FORCE_BUILD=TRUE pip install -e . -v --no-build-isolation + echo "Installation completed successfully" +clean: + rm -rf build + rm -rf *.egg-info + rm -rf ktransformers/ktransformers_ext/build + rm -rf ktransformers/ktransformers_ext/cuda/build + rm -rf ktransformers/ktransformers_ext/cuda/dist + rm -rf ktransformers/ktransformers_ext/cuda/*.egg-info +install_numa: + USE_NUMA=1 make dev_install +install_no_numa: + env -u USE_NUMA make dev_install \ No newline at end of file diff --git a/archive/README.md b/archive/README.md new file mode 100644 index 0000000..25de784 --- /dev/null +++ b/archive/README.md @@ -0,0 +1,136 @@ +
+

+ + KTransformers + +

+

High-Performance CPU-GPU Hybrid Inference for Large Language Models

+
+ +## 🎯 Overview + +KTransformers is a research project focused on efficient inference and fine-tuning of large language models through CPU-GPU heterogeneous computing. The project has evolved into **two core modules**: [kt-kernel](./kt-kernel/) and [kt-sft](./kt-sft/). + +## 🔥 Updates + +* **Nov 6, 2025**: Support Kimi-K2-Thinking inference and fine-tune +* **Nov 4, 2025**: KTransformers Fine-Tuning × LLaMA-Factory Integration +* **Oct 27, 2025**: Support Ascend NPU +* **Oct 10, 2025**: Integrating into SGLang ([Roadmap](https://github.com/sgl-project/sglang/issues/11425), [Blog](https://lmsys.org/blog/2025-10-22-KTransformers/)) +* **Sept 11, 2025**: Support Qwen3-Next +* **Sept 05, 2025**: Support Kimi-K2-0905 +* **July 26, 2025**: Support SmallThinker and GLM4-MoE +* **June 30, 2025**: Support 3-layer (GPU-CPU-Disk) prefix cache reuse +* **May 14, 2025**: Support Intel Arc GPU +* **Apr 29, 2025**: Support AMX-Int8、AMX-BF16 and Qwen3MoE +* **Apr 9, 2025**: Experimental support for LLaMA 4 models +* **Apr 2, 2025**: Support Multi-concurrency +* **Mar 15, 2025**: Support ROCm on AMD GPU +* **Mar 5, 2025**: Support unsloth 1.58/2.51 bits weights and IQ1_S/FP8 hybrid weights; 139K longer context for DeepSeek-V3/R1 +* **Feb 25, 2025**: Support FP8 GPU kernel for DeepSeek-V3 and R1 +* **Feb 10, 2025**: Support Deepseek-R1 and V3, up to 3~28x speedup + +--- + +## 📦 Core Modules + +### 🚀 [kt-kernel](./kt-kernel/) - High-Performance Inference Kernels + +CPU-optimized kernel operations for heterogeneous LLM inference. + +![image-20251011010558909](./doc/assets/heterogeneous_computing.png) + +**Key Features:** +- **AMX/AVX Acceleration**: Intel AMX and AVX512/AVX2 optimized kernels for INT4/INT8 quantized inference +- **MoE Optimization**: Efficient Mixture-of-Experts inference with NUMA-aware memory management +- **Quantization Support**: CPU-side INT4/INT8 quantized weights, GPU-side GPTQ support +- **Easy Integration**: Clean Python API for SGLang and other frameworks + +**Quick Start:** +```bash +cd kt-kernel +pip install . +``` + +**Use Cases:** + +- CPU-GPU hybrid inference for large MoE models +- Integration with SGLang for production serving +- Heterogeneous expert placement (hot experts on GPU, cold experts on CPU) + +**Performance Examples:** +| Model | Hardware Configuration | Total Throughput | Output Throughput | +|-------|------------------------|------------------|-------------------| +| DeepSeek-R1-0528 (FP8) | 8×L20 GPU + Xeon Gold 6454S | 227.85 tokens/s | 87.58 tokens/s (8-way concurrency) | + +👉 **[Full Documentation →](./kt-kernel/README.md)** + +--- + +### 🎓 [kt-sft](./kt-sft/) - Fine-Tuning Framework + +KTransformers × LLaMA-Factory integration for ultra-large MoE model fine-tuning. + +![image-20251011010558909](./doc/assets/image-20251011010558909.png) + +**Key Features:** + +- **Resource Efficient**: Fine-tune 671B DeepSeek-V3 with just **70GB GPU memory** + 1.3TB RAM +- **LoRA Support**: Full LoRA fine-tuning with heterogeneous acceleration +- **LLaMA-Factory Integration**: Seamless integration with popular fine-tuning framework +- **Production Ready**: Chat, batch inference, and metrics evaluation + +**Performance Examples:** + +| Model | Configuration | Throughput | GPU Memory | +|-------|--------------|------------|------------| +| DeepSeek-V3 (671B) | LoRA + AMX | ~40 tokens/s | 70GB (multi-GPU) | +| DeepSeek-V2-Lite (14B) | LoRA + AMX | ~530 tokens/s | 6GB | + +**Quick Start:** +```bash +cd kt-sft +# Install environment following kt-sft/README.md +USE_KT=1 llamafactory-cli train examples/train_lora/deepseek3_lora_sft_kt.yaml +``` + +👉 **[Full Documentation →](./kt-sft/README.md)** + +--- + +## 🔥 Citation + +If you use KTransformers in your research, please cite our paper: + +```bibtex +@inproceedings{10.1145/3731569.3764843, + title = {KTransformers: Unleashing the Full Potential of CPU/GPU Hybrid Inference for MoE Models}, + author = {Chen, Hongtao and Xie, Weiyu and Zhang, Boxin and Tang, Jingqi and Wang, Jiahao and Dong, Jianwei and Chen, Shaoyuan and Yuan, Ziwei and Lin, Chen and Qiu, Chengyu and Zhu, Yuening and Ou, Qingliang and Liao, Jiaqi and Chen, Xianglin and Ai, Zhiyuan and Wu, Yongwei and Zhang, Mingxing}, + booktitle = {Proceedings of the ACM SIGOPS 31st Symposium on Operating Systems Principles}, + year = {2025} +} +``` + +## 👥 Contributors & Team + +Developed and maintained by: +- [MADSys Lab](https://madsys.cs.tsinghua.edu.cn/) @ Tsinghua University +- [Approaching.AI](http://approaching.ai/) +- Community contributors + +We welcome contributions! Please feel free to submit issues and pull requests. + +## 💬 Community & Support + +- **GitHub Issues**: [Report bugs or request features](https://github.com/kvcache-ai/ktransformers/issues) +- **GitHub Discussions**: [Ask questions and share ideas](https://github.com/kvcache-ai/ktransformers/discussions) +- **WeChat Group**: See [archive/WeChatGroup.png](./archive/WeChatGroup.png) + +## 📦 Legacy Code + +The original integrated KTransformers framework has been archived to the [`archive/`](./archive/) directory for reference. The project now focuses on the two core modules above for better modularity and maintainability. + +For the original documentation with full quick-start guides and examples, see: +- [archive/README_LEGACY.md](./archive/README_LEGACY.md) (English) +- [archive/README_ZH_LEGACY.md](./archive/README_ZH_LEGACY.md) (中文) + diff --git a/archive/README_LEGACY.md b/archive/README_LEGACY.md new file mode 100644 index 0000000..92ba915 --- /dev/null +++ b/archive/README_LEGACY.md @@ -0,0 +1,217 @@ +
+ +

+ + + KTransformers + + + +

+

A Flexible Framework for Experiencing Cutting-edge LLM Inference Optimizations

+ 🌟 Show Cases | 🚀 Quick Start | 📃 Tutorial | 🔥 Citation | 💬 Discussion | 🙋 FAQ +
+ +

🎉 Introduction

+KTransformers, pronounced as Quick Transformers, is designed to enhance your 🤗 Transformers experience with advanced kernel optimizations and placement/parallelism strategies. +

+KTransformers is a flexible, Python-centric framework designed with extensibility at its core. +By implementing and injecting an optimized module with a single line of code, users gain access to a Transformers-compatible +interface, RESTful APIs compliant with OpenAI and Ollama, and even a simplified ChatGPT-like web UI. +

+Our vision for KTransformers is to serve as a flexible platform for experimenting with innovative LLM inference optimizations. Please let us know if you need any other features. + +

🔥 Updates

+ +* **Nov 6, 2025**: Support Kimi-K2-Thinking inference ([Tutorial](./doc/en/Kimi-K2-Thinking.md)) and fine-tune ([Tutorial](./doc/en/SFT_Installation_Guide_KimiK2.md)) +* **Nov 4, 2025**: KTransformers Fine-Tuning × LLaMA-Factory Integration. ([Tutorial](./doc/en/KTransformers-Fine-Tuning_User-Guide.md)) +* **Oct 27, 2025**: Support Ascend NPU. ([Tutorial](./doc/zh/DeepseekR1_V3_tutorial_zh_for_Ascend_NPU.md)) +* **Oct 10, 2025**: Integrating into SGLang. ([Roadmap](https://github.com/sgl-project/sglang/issues/11425)) +* **Sept 11, 2025**: Support Qwen3-Next. ([Tutorial](./doc/en/Qwen3-Next.md)) +* **Sept 05, 2025**: Support Kimi-K2-0905. ([Tutorial](./doc/en/Kimi-K2.md)) +* **July 26, 2025**: Support SmallThinker and GLM4-MoE. ([Tutorial](./doc/en/SmallThinker_and_Glm4moe.md)) +* **July 11, 2025**: Support Kimi-K2. ([Tutorial](./doc/en/Kimi-K2.md)) +* **June 30, 2025**: Support 3-layer (GPU-CPU-Disk) [prefix cache](./doc/en/prefix_cache.md) reuse. +* **May 14, 2025**: Support Intel Arc GPU ([Tutorial](./doc/en/xpu.md)). +* **Apr 29, 2025**: Support AMX-Int8、 AMX-BF16 and Qwen3MoE ([Tutorial](./doc/en/AMX.md)) + +https://github.com/user-attachments/assets/fafe8aec-4e22-49a8-8553-59fb5c6b00a2 + +* **Apr 9, 2025**: Experimental support for LLaMA 4 models ([Tutorial](./doc/en/llama4.md)). +* **Apr 2, 2025**: Support Multi-concurrency. ([Tutorial](./doc/en/balance-serve.md)). + +https://github.com/user-attachments/assets/faa3bda2-928b-45a7-b44f-21e12ec84b8a + +* **Mar 15, 2025**: Support ROCm on AMD GPU ([Tutorial](./doc/en/ROCm.md)). +* **Mar 5, 2025**: Support unsloth 1.58/2.51 bits weights and [IQ1_S/FP8 hybrid](./doc/en/fp8_kernel.md) weights. Support 139K [Longer Context](./doc/en/DeepseekR1_V3_tutorial.md#v022--v023-longer-context--fp8-kernel) for DeepSeek-V3 and R1 in 24GB VRAM. +* **Feb 25, 2025**: Support [FP8 GPU kernel](./doc/en/fp8_kernel.md) for DeepSeek-V3 and R1; [Longer Context](./doc/en/DeepseekR1_V3_tutorial.md#v022-longer-context). +* **Feb 15, 2025**: Longer Context (from 4K to 8K for 24GB VRAM) & Slightly Faster Speed （+15%, up to 16 Tokens/s), update [docs](./doc/en/DeepseekR1_V3_tutorial.md) and [online books](https://kvcache-ai.github.io/ktransformers/). +* **Feb 10, 2025**: Support Deepseek-R1 and V3 on single (24GB VRAM)/multi gpu and 382G DRAM, up to 3~28x speedup. For detailed show case and reproduction tutorial, see [here](./doc/en/DeepseekR1_V3_tutorial.md). +* **Aug 28, 2024**: Decrease DeepseekV2's required VRAM from 21G to 11G. +* **Aug 15, 2024**: Update detailed [tutorial](doc/en/injection_tutorial.md) for injection and multi-GPU. +* **Aug 14, 2024**: Support llamfile as linear backend. +* **Aug 12, 2024**: Support multiple GPU; Support new model: mixtral 8\*7B and 8\*22B; Support q2k, q3k, q5k dequant on gpu. +* **Aug 9, 2024**: Support windows native. + + + +

🌟 Show Cases

+ +
+

GPT-4/o1-level Local VSCode Copilot on a Desktop with only 24GB VRAM

+
+ +https://github.com/user-attachments/assets/ebd70bfa-b2c1-4abb-ae3b-296ed38aa285 + +

+ +- **[NEW!!!] Local 671B DeepSeek-Coder-V3/R1:** Running its Q4_K_M version using only 14GB VRAM and 382GB DRAM([Tutorial](./doc/en/DeepseekR1_V3_tutorial.md)). + + - Prefill Speed (tokens/s): + - KTransformers: 54.21 (32 cores) → 74.362 (dual-socket, 2×32 cores) → 255.26 (optimized AMX-based MoE kernel, V0.3 only) → 286.55 (selectively using 6 experts, V0.3 only) + - Compared to 10.31 tokens/s in llama.cpp with 2×32 cores, achieving up to **27.79× speedup**. + - Decode Speed (tokens/s): + - KTransformers: 8.73 (32 cores) → 11.26 (dual-socket, 2×32 cores) → 13.69 (selectively using 6 experts, V0.3 only) + - Compared to 4.51 tokens/s in llama.cpp with 2×32 cores, achieving up to **3.03× speedup**. + - Upcoming Open Source Release: + - AMX optimizations and selective expert activation will be open-sourced in V0.3. + - Currently available only in preview binary distribution, which can be downloaded [here](./doc/en/DeepseekR1_V3_tutorial.md). +- **Local 236B DeepSeek-Coder-V2:** Running its Q4_K_M version using only 21GB VRAM and 136GB DRAM, attainable on a local desktop machine, which scores even better than GPT4-0613 in [BigCodeBench](https://huggingface.co/blog/leaderboard-bigcodebench). + +

+ + DeepSeek-Coder-V2 Score + +

+ +- **Faster Speed:** Achieving 126 tokens/s for 2K prompt prefill and 13.6 tokens/s for generation through MoE offloading and injecting advanced kernels from [Llamafile](https://github.com/Mozilla-Ocho/llamafile/tree/main) and [Marlin](https://github.com/IST-DASLab/marlin). +- **VSCode Integration:** Wrapped into an OpenAI and Ollama compatible API for seamless integration as a backend for [Tabby](https://github.com/TabbyML/tabby) and various other frontends. + +

+ +https://github.com/user-attachments/assets/4c6a8a38-05aa-497d-8eb1-3a5b3918429c + +

+ + + +More advanced features will coming soon, so stay tuned! + +

🚀 Quick Start

+ +Getting started with KTransformers is simple! Follow the steps below to set up and start using it. + +we have already supported vendors: + +- Metax +- Sanechips (ZhuFeng V1.0) +- Intel +- Ascend +- Kunpeng +- AMD + +### 📥 Installation + +To install KTransformers, follow the official [Installation Guide](https://kvcache-ai.github.io/ktransformers/en/install.html). + +

📃 Brief Injection Tutorial

+At the heart of KTransformers is a user-friendly, template-based injection framework. +This allows researchers to easily replace original torch modules with optimized variants. It also simplifies the process of combining multiple optimizations, allowing the exploration of their synergistic effects. + +
+

+ + Inject-Struction + +

+ +Given that vLLM already serves as a great framework for large-scale deployment optimizations, KTransformers is particularly focused on local deployments that are constrained by limited resources. We pay special attention to heterogeneous computing opportunities, such as GPU/CPU offloading of quantized models. For example, we support the efficient Llamafile and Marlin kernels for CPU and GPU, respectively. More details can be found here. + +

Example Usage

+To utilize the provided kernels, users only need to create a YAML-based injection template and add the call to `optimize_and_load_gguf` before using the Transformers model. + +```python +with torch.device("meta"): + model = AutoModelForCausalLM.from_config(config, trust_remote_code=True) +optimize_and_load_gguf(model, optimize_config_path, gguf_path, config) +... +generated = prefill_and_generate(model, tokenizer, input_tensor.cuda(), max_new_tokens=1000) +``` + +In this example, the AutoModel is first initialized on the meta device to avoid occupying any memory resources. Then, `optimize_and_load_gguf` iterates through all sub-modules of the model, matches rules specified in your YAML rule file, and replaces them with advanced modules as specified. + +After injection, the original `generate` interface is available, but we also provide a compatible `prefill_and_generate` method, which enables further optimizations like CUDAGraph to improve generation speed. + +

How to custom your model

+ +A detailed tutorial of the injection and multi-GPU using DeepSeek-V2 as an example is given [here](doc/en/injection_tutorial.md). + +Below is an example of a YAML template for replacing all original Linear modules with Marlin, an advanced 4-bit quantization kernel. + +```yaml +- match: + name: "^model\\.layers\\..*$" # regular expression + class: torch.nn.Linear # only match modules matching name and class simultaneously + replace: + class: ktransformers.operators.linear.KTransformerLinear # optimized Kernel on quantized data types + device: "cpu" # which devices to load this module when initializing + kwargs: + generate_device: "cuda" + generate_linear_type: "QuantizedLinearMarlin" +``` + +Each rule in the YAML file has two parts: `match` and `replace`. The `match` part specifies which module should be replaced, and the `replace` part specifies the module to be injected into the model along with the initialization keywords. + +You can find example rule templates for optimizing DeepSeek-V2 and Qwen2-57B-A14, two SOTA MoE models, in the [ktransformers/optimize/optimize_rules](ktransformers/optimize/optimize_rules) directory. These templates are used to power the `local_chat.py` demo. + +If you are interested in our design principles and the implementation of the injection framework, please refer to the [design document](doc/en/deepseek-v2-injection.md). + +

🔥 Citation

+ +If you use KTransformers for your research, please cite our [paper](https://madsys.cs.tsinghua.edu.cn/publication/ktransformers-unleashing-the-full-potential-of-cpu/gpu-hybrid-inference-for-moe-models/): + +``` +@inproceedings{10.1145/3731569.3764843, +title = {KTransformers: Unleashing the Full Potential of CPU/GPU Hybrid Inference for MoE Models}, +author = {Chen, Hongtao and Xie, Weiyu and Zhang, Boxin and Tang, Jingqi and Wang, Jiahao and Dong, Jianwei and Chen, Shaoyuan and Yuan, Ziwei and Lin, Chen and Qiu, Chengyu and Zhu, Yuening and Ou, Qingliang and Liao, Jiaqi and Chen, Xianglin and Ai, Zhiyuan and Wu, Yongwei and Zhang, Mingxing}, +booktitle = {Proceedings of the ACM SIGOPS 31st Symposium on Operating Systems Principles}, +year = {2025} +} +``` + +

Acknowledgment and Contributors

+ +The development of KTransformers is based on the flexible and versatile framework provided by Transformers. We also benefit from advanced kernels such as GGUF/GGML, Llamafile, Marlin, sglang and flashinfer. We are planning to contribute back to the community by upstreaming our modifications. + +KTransformers is actively maintained and developed by contributors from the MADSys group at Tsinghua University and members from Approaching.AI. We welcome new contributors to join us in making KTransformers faster and easier to use. + +

Discussion

+ +If you have any questions, feel free to open an issue. Alternatively, you can join our WeChat group for further discussion. QR Code: [WeChat Group](WeChatGroup.png) + +

🙋 FAQ

+ +Some common questions are answered in the [FAQ](doc/en/FAQ.md). + diff --git a/archive/README_ZH.md b/archive/README_ZH.md new file mode 100644 index 0000000..4993f86 --- /dev/null +++ b/archive/README_ZH.md @@ -0,0 +1,132 @@ +
+

+ + KTransformers + +

+

高性能 CPU-GPU 异构大语言模型推理

+
+ +## 🎯 项目概述 + +KTransformers 是一个专注于大语言模型高效推理和微调的研究项目，通过 CPU-GPU 异构计算实现资源受限环境下的模型部署。项目已演进为**两个核心模块**：[kt-kernel](./kt-kernel/) 和 [kt-sft](./kt-sft/)。 + +## 🔥 更新 + +* **2025年11月6日**：支持 Kimi-K2-Thinking 推理和微调 +* **2025年11月4日**：KTransformers 微调 × LLaMA-Factory 集成 +* **2025年10月27日**：支持 Ascend NPU +* **2025年10月10日**：集成到 SGLang ([路线图](https://github.com/sgl-project/sglang/issues/11425), [博客](https://lmsys.org/blog/2025-10-22-KTransformers/)) +* **2025年9月11日**：支持 Qwen3-Next +* **2025年9月5日**：支持 Kimi-K2-0905 +* **2025年7月26日**：支持 SmallThinker 和 GLM4-MoE +* **2025年6月30日**：支持 3层（GPU-CPU-磁盘）前缀缓存复用 +* **2025年5月14日**：支持 Intel Arc GPU +* **2025年4月29日**：支持 AMX-Int8、AMX-BF16 和 Qwen3MoE +* **2025年4月9日**：实验性支持 LLaMA 4 模型 +* **2025年4月2日**：支持多并发 +* **2025年3月15日**：支持 AMD GPU 的 ROCm +* **2025年3月5日**：支持 unsloth 1.58/2.51 bits 权重和 IQ1_S/FP8 混合权重；DeepSeek-V3/R1 支持 139K 长上下文 +* **2025年2月25日**：支持 DeepSeek-V3 和 R1 的 FP8 GPU 内核 +* **2025年2月10日**：支持 Deepseek-R1 和 V3，速度提升最高达 3~28 倍 + +--- + +## 📦 核心模块 + +### 🚀 [kt-kernel](./kt-kernel/) - 高性能推理内核 + +面向异构 LLM 推理的 CPU 优化内核操作库。 + +![image-20251011010558909](./doc/assets/heterogeneous_computing.png) + +**核心特性：** +- **AMX/AVX 加速**：Intel AMX 和 AVX512/AVX2 优化内核，支持 INT4/INT8 量化推理 +- **MoE 优化**：高效的专家混合推理，支持 NUMA 感知内存管理 +- **量化支持**：CPU 端 INT4/INT8 量化权重，GPU 端 GPTQ 支持 +- **易于集成**：简洁的 Python API，可集成到 SGLang 等框架 + +**快速开始：** +```bash +cd kt-kernel +pip install . +``` + +**应用场景：** +- 大型 MoE 模型的 CPU-GPU 混合推理 +- 与 SGLang 集成用于生产服务 +- 异构专家放置（热门专家在 GPU，冷门专家在 CPU） + +**性能示例：** +| 模型 | 硬件配置 | 总吞吐量 | 输出吞吐量 | +|------|---------|---------|-----------| +| DeepSeek-R1-0528 (FP8) | 8×L20 GPU + Xeon Gold 6454S | 227.85 tokens/s | 87.58 tokens/s（8路并发）| + +👉 **[完整文档 →](./kt-kernel/README.md)** + +--- + +### 🎓 [kt-sft](./kt-sft/) - 微调框架 + +KTransformers × LLaMA-Factory 集成，支持超大 MoE 模型微调。 + +![image-20251011010558909](./doc/assets/image-20251011010558909.png) + +**核心特性：** +- **资源高效**：仅需 **70GB 显存** + 1.3TB 内存即可微调 671B DeepSeek-V3 +- **LoRA 支持**：完整的 LoRA 微调与异构加速 +- **LLaMA-Factory 集成**：与流行微调框架无缝集成 +- **生产就绪**：支持对话、批量推理和指标评估 + +**性能示例：** +| 模型 | 配置 | 吞吐量 | GPU 显存 | +|------|------|--------|----------| +| DeepSeek-V3 (671B) | LoRA + AMX | ~40 tokens/s | 70GB (多卡) | +| DeepSeek-V2-Lite (14B) | LoRA + AMX | ~530 tokens/s | 6GB | + +**快速开始：** +```bash +cd kt-sft +# 按照 kt-sft/README.md 安装环境 +USE_KT=1 llamafactory-cli train examples/train_lora/deepseek3_lora_sft_kt.yaml +``` + +👉 **[完整文档 →](./kt-sft/README.md)** + +--- + +## 🔥 引用 + +如果您在研究中使用了 KTransformers，请引用我们的论文： + +```bibtex +@inproceedings{10.1145/3731569.3764843, + title = {KTransformers: Unleashing the Full Potential of CPU/GPU Hybrid Inference for MoE Models}, + author = {Chen, Hongtao and Xie, Weiyu and Zhang, Boxin and Tang, Jingqi and Wang, Jiahao and Dong, Jianwei and Chen, Shaoyuan and Yuan, Ziwei and Lin, Chen and Qiu, Chengyu and Zhu, Yuening and Ou, Qingliang and Liao, Jiaqi and Chen, Xianglin and Ai, Zhiyuan and Wu, Yongwei and Zhang, Mingxing}, + booktitle = {Proceedings of the ACM SIGOPS 31st Symposium on Operating Systems Principles}, + year = {2025} +} +``` + +## 👥 贡献者与团队 + +由以下团队开发和维护： +- 清华大学 [MADSys 实验室](https://madsys.cs.tsinghua.edu.cn/) +- [Approaching.AI](http://approaching.ai/) +- 社区贡献者 + +我们欢迎贡献！请随时提交 issues 和 pull requests。 + +## 💬 社区与支持 + +- **GitHub Issues**：[报告 bug 或请求功能](https://github.com/kvcache-ai/ktransformers/issues) +- **GitHub Discussions**：[提问和分享想法](https://github.com/kvcache-ai/ktransformers/discussions) +- **微信群**：查看 [archive/WeChatGroup.png](./archive/WeChatGroup.png) + +## 📦 历史代码 + +原完整的 KTransformers 框架代码已归档至 [`archive/`](./archive/) 目录供参考。项目现专注于上述两个核心模块，以实现更好的模块化和可维护性。 + +关于原始完整文档（包含快速入门指南和示例），请查看： +- [archive/README_LEGACY.md](./archive/README_LEGACY.md) (English) +- [archive/README_ZH_LEGACY.md](./archive/README_ZH_LEGACY.md) (中文) diff --git a/archive/README_ZH_LEGACY.md b/archive/README_ZH_LEGACY.md new file mode 100644 index 0000000..48f28e0 --- /dev/null +++ b/archive/README_ZH_LEGACY.md @@ -0,0 +1,166 @@ +
+ +

+ + + KTransformers + + + +

+

一个用于体验尖端 LLM 推理优化的灵活框架

+ 🌟 案例展示 | 🚀 快速入门 | 📃 教程 | 💬 讨论 | 🙋 常见问题 +
+ +

🎉 介绍

+KTransformers（发音为 Quick Transformers）旨在通过先进的内核优化和放置/并行策略来增强您对 🤗 [Transformers](https://github.com/huggingface/transformers) 的体验。 +

+KTransformers 是一个以 Python 为中心的灵活框架，其核心是可扩展性。通过用一行代码实现并注入优化模块，用户可以获得与 Transformers 兼容的接口、符合 OpenAI 和 Ollama 的 RESTful API，甚至是一个简化的类似 ChatGPT 的 Web 界面。 +

+我们对 KTransformers 的愿景是成为一个用于实验创新 LLM 推理优化的灵活平台。如果您需要任何其他功能，请告诉我们。 + +

🔥 更新

+ +* **2025 年 2 月 15 日**：为DeepSeek-V3/R1支持[FP8 GPU内核](./doc/en/fp8_kernel.md); 支持更长的上下文([教程](./doc/en/DeepseekR1_V3_tutorial.md#v022-longer-context)). +* **2025 年 2 月 15 日**：长上下文(从4K到8K，24GB VRAM) & 稍快的速度(+15%)(最快 16 Tokens/s)，文档请参见 [这里](./doc/en/DeepseekR1_V3_tutorial.md) 和 [在线指南](https://kvcache-ai.github.io/ktransformers/) 。 +* **2025 年 2 月 10 日**：支持 Deepseek-R1 和 V3 在单个（24GB VRAM）/多 GPU 和 382G DRAM 上运行，速度提升高达 3~28 倍。详细教程请参见 [这里](./doc/en/DeepseekR1_V3_tutorial.md)。 +* **2024 年 8 月 28 日**：支持 InternLM2.5-7B-Chat-1M 模型下的 1M 上下文，使用 24GB 的 VRAM 和 150GB 的 DRAM。详细教程请参见 [这里](./doc/en/long_context_tutorial.md)。 +* **2024 年 8 月 28 日**：将 DeepseekV2 所需的 VRAM 从 21G 降低到 11G。 +* **2024 年 8 月 15 日**：更新了详细的 [教程](doc/en/injection_tutorial.md)，介绍注入和多 GPU 的使用。 +* **2024 年 8 月 14 日**：支持 llamfile 作为线性后端。 +* **2024 年 8 月 12 日**：支持多 GPU；支持新模型：mixtral 8\*7B 和 8\*22B；支持 q2k、q3k、q5k 在 GPU 上的去量化。 +* **2024 年 8 月 9 日**：支持 Windows。 + +

🌟 案例展示

+ +
+

在仅 24GB VRAM 的桌面上运行 GPT-4/o1 级别的本地 VSCode Copilot

+
+ +https://github.com/user-attachments/assets/ebd70bfa-b2c1-4abb-ae3b-296ed38aa285 + +

+ +- **[NEW!!!] 本地 671B DeepSeek-Coder-V3/R1**：使用其 Q4_K_M 版本，仅需 14GB VRAM 和 382GB DRAM 即可运行（教程请参见 [这里](./doc/en/DeepseekR1_V3_tutorial.md)）。 + - 预填充速度（tokens/s）： + - KTransformers：54.21（32 核）→ 74.362（双插槽，2×32 核）→ 255.26（优化的 AMX 基 MoE 内核，仅 V0.3）→ 286.55（选择性使用 6 个专家，仅 V0.3） + - 与 llama.cpp 在 2×32 核下相比，达到 **27.79× 速度提升**。 + - 解码速度（tokens/s）： + - KTransformers：8.73（32 核）→ 11.26（双插槽，2×32 核）→ 13.69（选择性使用 6 个专家，仅 V0.3） + - 与 llama.cpp 在 2×32 核下相比，达到 **3.03× 速度提升**。 + - 即将开源发布： + - AMX 优化和选择性专家激活将在 V0.3 中开源。 + - 目前仅在预览二进制分发中可用，可从 [这里](./doc/en/DeepseekR1_V3_tutorial.md) 下载。 + +- **本地 236B DeepSeek-Coder-V2**：使用其 Q4_K_M 版本，仅需 21GB VRAM 和 136GB DRAM 即可运行，甚至在 [BigCodeBench](https://huggingface.co/blog/leaderboard-bigcodebench) 中得分超过 GPT4-0613。 + +

+ + DeepSeek-Coder-V2 Score + +

+ +- **更快的速度**：通过 MoE 卸载和注入来自 [Llamafile](https://github.com/Mozilla-Ocho/llamafile/tree/main) 和 [Marlin](https://github.com/IST-DASLab/marlin) 的高级内核，实现了 2K 提示预填充 126 tokens/s 和生成 13.6 tokens/s 的速度。 +- **VSCode 集成**：封装成符合 OpenAI 和 Ollama 的 API，可无缝集成到 [Tabby](https://github.com/TabbyML/tabby) 和其他前端的后端。 + +

+ +https://github.com/user-attachments/assets/4c6a8a38-05aa-497d-8eb1-3a5b3918429c + +

+ + + + + + +更多高级功能即将推出，敬请期待！ + +

🚀 快速入门

+ + +KTransformers 的入门非常简单！请参考我们的[安装指南]((https://kvcache-ai.github.io/ktransformers/))进行安装。 + +

📃 简要注入教程

+KTransformers 的核心是一个用户友好的、基于模板的注入框架。这使得研究人员可以轻松地将原始 torch 模块替换为优化的变体。它还简化了多种优化的组合过程，允许探索它们的协同效应。 +
+

+ + Inject-Struction + +

+ +鉴于 vLLM 已经是一个用于大规模部署优化的优秀框架，KTransformers 特别关注受资源限制的本地部署。我们特别关注异构计算时机，例如量化模型的 GPU/CPU 卸载。例如，我们支持高效的 Llamafile 和Marlin 内核，分别用于 CPU 和 GPU。 更多详细信息可以在 这里找到。 + + +

示例用法

+要使用提供的内核，用户只需创建一个基于 YAML 的注入模板，并在使用 Transformers 模型之前添加对 `optimize_and_load_gguf` 的调用。 + +```python +with torch.device("meta"): + model = AutoModelForCausalLM.from_config(config, trust_remote_code=True) +optimize_and_load_gguf(model, optimize_config_path, gguf_path, config) +... +generated = prefill_and_generate(model, tokenizer, input_tensor.cuda(), max_new_tokens=1000) +``` + +在这个示例中，首先在 meta 设备上初始化 AutoModel，以避免占用任何内存资源。然后，`optimize_and_load_gguf` 遍历模型的所有子模块，匹配您的 YAML 规则文件中指定的规则，并将它们替换为指定的高级模块。 + +注入后，原始的 `generate` 接口仍然可用，但我们还提供了一个兼容的 `prefill_and_generate` 方法，这使得可以进一步优化，例如使用 CUDAGraph 提高生成速度。 + +

如何自定义您的模型

+ +一个详细的使用 DeepSeek-V2 作为示例的注入和 multi-GPU 教程在 [这里](doc/en/injection_tutorial.md)。 + +以下是一个将所有原始 Linear 模块替换为 Marlin 的 YAML 模板示例，Marlin 是一个高级的 4 位量化内核。 + +```yaml +- match: + name: "^model\\.layers\\..*$" # 正则表达式 + class: torch.nn.Linear # 仅匹配同时符合名称和类的模块 + replace: + class: ktransformers.operators.linear.KTransformerLinear # 量化数据类型的优化内核 + device: "cpu" # 初始化时加载该模块的 device + kwargs: + generate_device: "cuda" + generate_linear_type: "QuantizedLinearMarlin" +``` + +YAML 文件中的每个规则都有两部分：`match` 和 `replace`。`match` 部分指定应替换的模块，`replace` 部分指定要注入到模型中的模块以及初始化关键字。 + +您可以在 [ktransformers/optimize/optimize_rules](ktransformers/optimize/optimize_rules) 目录中找到用于优化 DeepSeek-V2 和 Qwen2-57B-A14 的示例规则模板。这些模板用于为 `local_chat.py` 示例提供支持。 + +如果您对我们的设计原则和注入框架的实现感兴趣，请参考 [设计文档](doc/en/deepseek-v2-injection.md)。 + +

致谢和贡献者

+ +KTransformers 的开发基于 Transformers 提供的灵活和多功能框架。我们还受益于 GGUF/GGML、Llamafile 、 Marlin、sglang和flashinfer 等高级内核。我们计划通过向上游贡献我们的修改来回馈社区。 + +KTransformers 由清华大学 MADSys group 小组的成员以及 Approaching.AI 的成员积极维护和开发。我们欢迎新的贡献者加入我们，使 KTransformers 更快、更易于使用。 + + +

讨论

+ +如果您有任何问题，欢迎随时提出 issue。或者，您可以加入我们的微信群进行进一步讨论。二维码： [微信群](WeChatGroup.png) + +

🙋 常见问题

+ +一些常见问题的答案可以在 [FAQ](doc/en/FAQ.md) 中找到。 diff --git a/archive/SECURITY.md b/archive/SECURITY.md new file mode 100644 index 0000000..034e848 --- /dev/null +++ b/archive/SECURITY.md @@ -0,0 +1,21 @@ +# Security Policy + +## Supported Versions + +Use this section to tell people about which versions of your project are +currently being supported with security updates. + +| Version | Supported | +| ------- | ------------------ | +| 5.1.x | :white_check_mark: | +| 5.0.x | :x: | +| 4.0.x | :white_check_mark: | +| < 4.0 | :x: | + +## Reporting a Vulnerability + +Use this section to tell people how to report a vulnerability. + +Tell them where to go, how often they can expect to get an update on a +reported vulnerability, what to expect if the vulnerability is accepted or +declined, etc. diff --git a/archive/book.toml b/archive/book.toml new file mode 100644 index 0000000..c88d9b7 --- /dev/null +++ b/archive/book.toml @@ -0,0 +1,18 @@ +[book] +authors = ["kvcache-ai"] +language = "zh-CN" +title = "Ktransformers" +src = "doc" + +[output.html] +git-repository-url = "https://github.com/kvcache-ai/ktransformers" +edit-url-template = "https://github.com/kvcache-ai/ktransformers/edit/main/{path}" + +[output.html.playground] +editable = true +copy-js = true +# line-numbers = true + +[output.html.fold] +enable = true +level = 0 \ No newline at end of file diff --git a/archive/config.json b/archive/config.json new file mode 100644 index 0000000..e69de29 diff --git a/archive/csrc/balance_serve/CMakeLists.txt b/archive/csrc/balance_serve/CMakeLists.txt new file mode 100644 index 0000000..48fd6d9 --- /dev/null +++ b/archive/csrc/balance_serve/CMakeLists.txt @@ -0,0 +1,106 @@ +option(KTRANSFORMERS_USE_NPU "ktransformers: use NPU" OFF) +if(KTRANSFORMERS_USE_NPU) + add_definitions(-DKTRANSFORMERS_USE_NPU=1) +endif() + +if(KTRANSFORMERS_USE_NPU) + set(ASCEND_HOME_PATH "$ENV{ASCEND_HOME_PATH}") + message(STATUS "ASCEND_HOME_PATH is ${ASCEND_HOME_PATH}") + include_directories(${ASCEND_HOME_PATH}/include) + + link_directories(${TORCH_INSTALL_PREFIX}/../torch.libs) + # find torch_npu + execute_process( + COMMAND python -c "import torch; import torch_npu; print(torch_npu.__path__[0])" + OUTPUT_VARIABLE TORCH_NPU_PATH + OUTPUT_STRIP_TRAILING_WHITESPACE + ) + message(STATUS "Found PTA at: ${TORCH_NPU_PATH}") + find_library(PTA_LIBRARY torch_npu PATH "${TORCH_NPU_PATH}/lib") +endif() + +cmake_minimum_required(VERSION 3.21) +find_program(GCC_COMPILER NAMES g++-13 g++-12 g++-11 g++ REQUIRED) +set(CMAKE_CXX_COMPILER ${GCC_COMPILER}) + +# 显示选定的编译器 +message(STATUS "Using compiler: ${CMAKE_CXX_COMPILER}") + + +project(balance_serve VERSION 0.1.0) + +set(CMAKE_CXX_STANDARD 20) +set(CMAKE_CXX_FLAGS "-Og -march=native -Wall -Wextra -g -fPIC") +set(CMAKE_BUILD_TYPE "Debug") +# set(CMAKE_CXX_FLAGS "-O3 -march=native -Wall -Wextra -fPIC") +# set(CMAKE_BUILD_TYPE "Release") + + +if(NOT DEFINED _GLIBCXX_USE_CXX11_ABI) + find_package(Python3 REQUIRED COMPONENTS Interpreter) + + execute_process( + COMMAND ${Python3_EXECUTABLE} -c + "import torch; print('1' if torch.compiled_with_cxx11_abi() else '0')" + OUTPUT_VARIABLE ABI_FLAG + OUTPUT_STRIP_TRAILING_WHITESPACE + ) + + set(_GLIBCXX_USE_CXX11_ABI ${ABI_FLAG} CACHE STRING "C++11 ABI setting from PyTorch" FORCE) +endif() + +# 无论是否是自动检测，都传给编译器 +add_compile_definitions(_GLIBCXX_USE_CXX11_ABI=${_GLIBCXX_USE_CXX11_ABI}) + +message(STATUS "_GLIBCXX_USE_CXX11_ABI=${_GLIBCXX_USE_CXX11_ABI}") + +file(GLOB_RECURSE FMT_SOURCES "${CMAKE_CURRENT_SOURCE_DIR}/*.cpp" "${CMAKE_CURRENT_SOURCE_DIR}/*.hpp" "${CMAKE_CURRENT_SOURCE_DIR}/*.h") + +add_custom_target( + format + COMMAND clang-format + -i + -style=file + ${FMT_SOURCES} + COMMENT "Running clang-format on all source files" +) + +set(BUILD_SHARED_LIBS ON) +set(ENABLE_PUSH OFF) +set(ENABLE_COMPRESSION OFF) + +# set(CMAKE_BUILD_TYPE "Release") +set(CMAKE_EXPORT_COMPILE_COMMANDS ON) + +set(THIRD_PARTY_DIR ${CMAKE_CURRENT_SOURCE_DIR}/../../third_party) +set(THIRD_PARTY_BUILD_DIR ${CMAKE_CURRENT_BINARY_DIR}/third_party) +add_subdirectory(${THIRD_PARTY_DIR}/prometheus-cpp ${THIRD_PARTY_BUILD_DIR}/prometheus-cpp EXCLUDE_FROM_ALL) +add_subdirectory(${THIRD_PARTY_DIR}/xxHash/cmake_unofficial ${THIRD_PARTY_BUILD_DIR}/xxHash EXCLUDE_FROM_ALL) +set_target_properties(xxhash PROPERTIES POSITION_INDEPENDENT_CODE ON) + +# add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/third_party/prometheus-cpp ${CMAKE_CURRENT_BINARY_DIR}/third_party/prometheus-cpp) +set(SPDLOG_DIR ${THIRD_PARTY_DIR}/spdlog) +set(FMT_DIR ${THIRD_PARTY_DIR}/fmt) + +set(KVC2_INCLUDE_DIR ${CMAKE_CURRENT_SOURCE_DIR}/kvc2/src) + +include_directories(${THIRD_PARTY_DIR}) + +add_subdirectory(${THIRD_PARTY_DIR}/pybind11 ${THIRD_PARTY_BUILD_DIR}/pybind11) + +execute_process( + COMMAND python3 -c "import torch; print(torch.__path__[0])" + OUTPUT_VARIABLE TORCH_INSTALL_PREFIX + OUTPUT_STRIP_TRAILING_WHITESPACE +) + +message(STATUS "Found PyTorch at: ${TORCH_INSTALL_PREFIX}") + +# set(TORCH_INSTALL_PREFIX "/home/xwy/.conda/envs/kvc/lib/python3.12/site-packages/torch") +find_library(TORCH_PYTHON_LIBRARY torch_python PATH "${TORCH_INSTALL_PREFIX}/lib") +find_package(Torch REQUIRED PATHS "${TORCH_INSTALL_PREFIX}/share/cmake/Torch" NO_DEFAULT_PATH) + +add_subdirectory(kvc2) +add_subdirectory(sched) + +# add_subdirectory(test) diff --git a/archive/csrc/custom_marlin/__init__.py b/archive/csrc/custom_marlin/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/archive/csrc/custom_marlin/binding.cpp b/archive/csrc/custom_marlin/binding.cpp new file mode 100644 index 0000000..184f3e2 --- /dev/null +++ b/archive/csrc/custom_marlin/binding.cpp @@ -0,0 +1,44 @@ +/** + * @Description : + * @Author : Azure-Tang + * @Date : 2024-07-25 13:38:30 + * @Version : 1.0.0 + * @LastEditors : kkk1nak0 + * @LastEditTime : 2024-08-12 03:05:04 + * @Copyright (c) 2024 by KVCache.AI, All Rights Reserved. + **/ + +#include "gptq_marlin/ops.h" +// Python bindings +#include +#include +#include +#include +#include +// namespace py = pybind11; + +PYBIND11_MODULE(vLLMMarlin, m) { + + /*m.def("dequantize_q8_0", &dequantize_q8_0, "Function to dequantize q8_0 + data.", py::arg("data"), py::arg("blk_size"), py::arg("device")); + m.def("dequantize_q6_k", &dequantize_q6_k, "Function to dequantize q6_k + data.", py::arg("data"), py::arg("blk_size"), py::arg("device")); + m.def("dequantize_q5_k", &dequantize_q5_k, "Function to dequantize q5_k + data.", py::arg("data"), py::arg("blk_size"), py::arg("device")); + m.def("dequantize_q4_k", &dequantize_q4_k, "Function to dequantize q4_k + data.", py::arg("data"), py::arg("blk_size"), py::arg("device")); + m.def("dequantize_q3_k", &dequantize_q3_k, "Function to dequantize q3_k + data.", py::arg("data"), py::arg("blk_size"), py::arg("device")); + m.def("dequantize_q2_k", &dequantize_q2_k, "Function to dequantize q2_k + data.", py::arg("data"), py::arg("blk_size"), py::arg("device")); + m.def("dequantize_iq4_xs", &dequantize_iq4_xs, "Function to dequantize + iq4_xs data.", py::arg("data"), py::arg("blk_size"), py::arg("device"));*/ + m.def("gptq_marlin_gemm", &gptq_marlin_gemm, + "Function to perform GEMM using Marlin quantization.", py::arg("a"), + py::arg("b_q_weight"), py::arg("b_scales"), py::arg("g_idx"), + py::arg("perm"), py::arg("workspace"), py::arg("num_bits"), py::arg("size_m_tensor"), + py::arg("size_m"), py::arg("size_n"), py::arg("size_k"), + py::arg("sms"), py::arg("is_k_full")); + m.def("gptq_marlin_repack", &gptq_marlin_repack, + "gptq_marlin repack from GPTQ"); +} \ No newline at end of file diff --git a/archive/csrc/custom_marlin/gptq_marlin/gptq_marlin.cu b/archive/csrc/custom_marlin/gptq_marlin/gptq_marlin.cu new file mode 100644 index 0000000..73ba3dd --- /dev/null +++ b/archive/csrc/custom_marlin/gptq_marlin/gptq_marlin.cu @@ -0,0 +1,2035 @@ +/* + * Modified by Neural Magic + * Copyright (C) Marlin.2024 Elias Frantar + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + + /* + * Adapted from https://github.com/IST-DASLab/marlin + */ + /* + * Adapted from + * https://github.com/vllm-project/vllm/tree/main/csrc/quantization/gptq_marlin + */ +#include "gptq_marlin.cuh" +#include "gptq_marlin_dtypes.cuh" +#include +#define STATIC_ASSERT_SCALAR_TYPE_VALID(scalar_t) \ + static_assert(std::is_same::value || \ + std::is_same::value, \ + "only float16 and bfloat16 is supported"); + +template inline std::string str(T x) { return std::to_string(x); } + +namespace gptq_marlin { + +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800 + + __global__ void permute_cols_kernel(int4 const* __restrict__ a_int4_ptr, + int const* __restrict__ perm_int_ptr, + int4* __restrict__ out_int4_ptr, int size_m, + int size_k, int block_rows) {} + + template shared + // fetch pipeline + const bool has_act_order, // whether act_order is enabled + const int group_blocks = -1 // number of consecutive 16x16 blocks + // with a separate quantization scale + > + __global__ void + Marlin(const int4* __restrict__ A, // fp16 input matrix of shape mxk + const int4* __restrict__ B, // 4bit quantized weight matrix of shape kxn + int4* __restrict__ C, // fp16 output buffer of shape mxn + const int4* __restrict__ scales_ptr, // fp16 quantization scales of shape + // (k/groupsize)xn + const int* __restrict__ g_idx, // int32 group indices of shape k + int num_groups, // number of scale groups per output channel + int prob_m, // batch dimension m + int prob_n, // output dimension n + int prob_k, // reduction dimension k + int* locks // extra global storage for barrier synchronization + ) {} + +} // namespace gptq_marlin + +torch::Tensor gptq_marlin_gemm(torch::Tensor& a, torch::Tensor& b_q_weight, + torch::Tensor& b_scales, torch::Tensor& g_idx, + torch::Tensor& perm, torch::Tensor& workspace, + int64_t num_bits, int64_t size_m, int64_t size_n, + int64_t size_k, bool is_k_full) { + TORCH_CHECK_NOT_IMPLEMENTED(false, + "marlin_gemm(..) requires CUDA_ARCH >= 8.0"); + return torch::empty({ 1, 1 }); +} + +#else + + // m16n8k16 tensor core mma instruction with fp16 inputs and fp32 + // output/accumulation. + template + __device__ inline void mma(const typename ScalarType::FragA& a_frag, + const typename ScalarType::FragB& frag_b, + typename ScalarType::FragC& frag_c) { + const uint32_t* a = reinterpret_cast(&a_frag); + const uint32_t* b = reinterpret_cast(&frag_b); + float* c = reinterpret_cast(&frag_c); + if constexpr (std::is_same::value) { + asm volatile( + "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 " + "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};\n" + : "=f"(c[0]), "=f"(c[1]), "=f"(c[2]), "=f"(c[3]) + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1]), + "f"(c[0]), "f"(c[1]), "f"(c[2]), "f"(c[3])); + } + else if constexpr (std::is_same::value) { + asm volatile( + "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 " + "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};\n" + : "=f"(c[0]), "=f"(c[1]), "=f"(c[2]), "=f"(c[3]) + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1]), + "f"(c[0]), "f"(c[1]), "f"(c[2]), "f"(c[3])); + } + else { + STATIC_ASSERT_SCALAR_TYPE_VALID(scalar_t); + } + } + + // Instruction for loading a full 16x16 matrix fragment of operand A from shared + // memory, directly in tensor core layout. + template + __device__ inline void ldsm4(typename ScalarType::FragA& frag_a, + const void* smem_ptr) { + uint32_t* a = reinterpret_cast(&frag_a); + uint32_t smem = static_cast(__cvta_generic_to_shared(smem_ptr)); + asm volatile( + "ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0,%1,%2,%3}, [%4];\n" + : "=r"(a[0]), "=r"(a[1]), "=r"(a[2]), "=r"(a[3]) + : "r"(smem)); + } + + // Lookup-table based 3-input logical operation; explicitly used for + // dequantization as the compiler does not seem to automatically recognize it in + // all cases. + template __device__ inline int lop3(int a, int b, int c) { + int res; + asm volatile("lop3.b32 %0, %1, %2, %3, %4;\n" + : "=r"(res) + : "r"(a), "r"(b), "r"(c), "n"(lut)); + return res; + } + + // Constructs destination register by taking bytes from 2 sources (based on + // mask) + template + __device__ inline uint32_t prmt(uint32_t a) { + uint32_t res; + asm volatile("prmt.b32 %0, %1, %2, %3;\n" + : "=r"(res) + : "r"(a), "n"(start_byte), "n"(mask)); + return res; + } + + // Efficiently dequantize an int32 value into a full B-fragment of 4 fp16 + // values. We mostly follow the strategy in the link below, with some small + // changes: + // - FP16: + // https://github.com/NVIDIA/FasterTransformer/blob/release/v5.3_tag/src/fastertransformer/cutlass_extensions/include/cutlass_extensions/interleaved_numeric_conversion.h#L215-L287 + // - BF16: + // https://github.com/NVIDIA/FasterTransformer/blob/release/v5.3_tag/src/fastertransformer/cutlass_extensions/include/cutlass_extensions/interleaved_numeric_conversion.h#L327-L385 + template + __device__ inline typename ScalarType::FragB dequant_4bit(int q) { + STATIC_ASSERT_SCALAR_TYPE_VALID(scalar_t); + } + + template <> + __device__ inline typename ScalarType::FragB dequant_4bit(int q) { + const int LO = 0x000f000f; + const int HI = 0x00f000f0; + const int EX = 0x64006400; + // Guarantee that the `(a & b) | c` operations are LOP3s. + int lo = lop3<(0xf0 & 0xcc) | 0xaa>(q, LO, EX); + int hi = lop3<(0xf0 & 0xcc) | 0xaa>(q, HI, EX); + // We want signed int4 outputs, hence we fuse the `-8` symmetric zero point + // directly into `SUB` and `ADD`. + const int SUB = 0x64086408; + const int MUL = 0x2c002c00; + const int ADD = 0xd480d480; + typename ScalarType::FragB frag_b; + frag_b[0] = __hsub2(*reinterpret_cast(&lo), + *reinterpret_cast(&SUB)); + frag_b[1] = __hfma2(*reinterpret_cast(&hi), + *reinterpret_cast(&MUL), + *reinterpret_cast(&ADD)); + return frag_b; + } + + template <> + __device__ inline typename ScalarType::FragB + dequant_4bit(int q) { + static constexpr uint32_t MASK = 0x000f000f; + static constexpr uint32_t EX = 0x43004300; + + // Guarantee that the `(a & b) | c` operations are LOP3s. + + int lo = lop3<(0xf0 & 0xcc) | 0xaa>(q, MASK, EX); + q >>= 4; + int hi = lop3<(0xf0 & 0xcc) | 0xaa>(q, MASK, EX); + + typename ScalarType::FragB frag_b; + static constexpr uint32_t MUL = 0x3F803F80; + static constexpr uint32_t ADD = 0xC308C308; + + frag_b[0] = __hfma2(*reinterpret_cast(&lo), + *reinterpret_cast(&MUL), + *reinterpret_cast(&ADD)); + frag_b[1] = __hfma2(*reinterpret_cast(&hi), + *reinterpret_cast(&MUL), + *reinterpret_cast(&ADD)); + return frag_b; + } + + // Fast Int8ToFp16/Int8ToBf16: Efficiently dequantize 8bit int values to fp16 or + // bf16 Reference: + // - FP16: + // https://github.com/NVIDIA/FasterTransformer/blob/release/v5.3_tag/src/fastertransformer/cutlass_extensions/include/cutlass_extensions/interleaved_numeric_conversion.h#L53-L85 + // - BF16: + // https://github.com/NVIDIA/FasterTransformer/blob/release/v5.3_tag/src/fastertransformer/cutlass_extensions/include/cutlass_extensions/interleaved_numeric_conversion.h#L125-L175 + template + __device__ inline typename ScalarType::FragB dequant_8bit(int q) { + STATIC_ASSERT_SCALAR_TYPE_VALID(scalar_t); + } + + template <> + __device__ inline typename ScalarType::FragB dequant_8bit(int q) { + static constexpr uint32_t mask_for_elt_01 = 0x5250; + static constexpr uint32_t mask_for_elt_23 = 0x5351; + static constexpr uint32_t start_byte_for_fp16 = 0x64646464; + + uint32_t lo = prmt(q); + uint32_t hi = prmt(q); + + static constexpr uint32_t I8s_TO_F16s_MAGIC_NUM = 0x64806480; + + typename ScalarType::FragB frag_b; + frag_b[0] = + __hsub2(*reinterpret_cast(&lo), + *reinterpret_cast(&I8s_TO_F16s_MAGIC_NUM)); + frag_b[1] = + __hsub2(*reinterpret_cast(&hi), + *reinterpret_cast(&I8s_TO_F16s_MAGIC_NUM)); + return frag_b; + } + + template <> + __device__ inline typename ScalarType::FragB + dequant_8bit(int q) { + typename ScalarType::FragB frag_b; + + float fp32_intermediates[4]; + uint32_t* fp32_intermediates_casted = + reinterpret_cast(fp32_intermediates); + + static constexpr uint32_t fp32_base = 0x4B000000; + fp32_intermediates_casted[0] = __byte_perm(q, fp32_base, 0x7650); + fp32_intermediates_casted[1] = __byte_perm(q, fp32_base, 0x7652); + fp32_intermediates_casted[2] = __byte_perm(q, fp32_base, 0x7651); + fp32_intermediates_casted[3] = __byte_perm(q, fp32_base, 0x7653); + + fp32_intermediates[0] -= 8388736.f; + fp32_intermediates[1] -= 8388736.f; + fp32_intermediates[2] -= 8388736.f; + fp32_intermediates[3] -= 8388736.f; + + uint32_t* bf16_result_ptr = reinterpret_cast(&frag_b); + bf16_result_ptr[0] = __byte_perm(fp32_intermediates_casted[0], + fp32_intermediates_casted[1], 0x7632); + bf16_result_ptr[1] = __byte_perm(fp32_intermediates_casted[2], + fp32_intermediates_casted[3], 0x7632); + + return frag_b; + } + + // Multiply dequantized values by the corresponding quantization scale; used + // only for grouped quantization. + template + __device__ inline void scale(typename ScalarType::FragB& frag_b, + typename ScalarType::FragS& frag_s, + int i) { + using scalar_t2 = typename ScalarType::scalar_t2; + scalar_t2 s = ScalarType::num2num2( + reinterpret_cast(&frag_s)[i]); + frag_b[0] = __hmul2(frag_b[0], s); + frag_b[1] = __hmul2(frag_b[1], s); + } + + // Same as above, but for act_order (each K is multiplied individually) + template + __device__ inline void scale4(typename ScalarType::FragB& frag_b, + typename ScalarType::FragS& frag_s_1, + typename ScalarType::FragS& frag_s_2, + typename ScalarType::FragS& frag_s_3, + typename ScalarType::FragS& frag_s_4, + int i) { + using scalar_t2 = typename ScalarType::scalar_t2; + scalar_t2 s_val_1_2; + s_val_1_2.x = reinterpret_cast(&frag_s_1)[i]; + s_val_1_2.y = reinterpret_cast(&frag_s_2)[i]; + + scalar_t2 s_val_3_4; + s_val_3_4.x = reinterpret_cast(&frag_s_3)[i]; + s_val_3_4.y = reinterpret_cast(&frag_s_4)[i]; + + frag_b[0] = __hmul2(frag_b[0], s_val_1_2); + frag_b[1] = __hmul2(frag_b[1], s_val_3_4); + } + + // Given 2 floats multiply by 2 scales (halves) + template + __device__ inline void scale_float(float* c, + typename ScalarType::FragS& s) { + scalar_t* s_ptr = reinterpret_cast(&s); + c[0] = __fmul_rn(c[0], ScalarType::num2float(s_ptr[0])); + c[1] = __fmul_rn(c[1], ScalarType::num2float(s_ptr[1])); + } + + // Wait until barrier reaches `count`, then lock for current threadblock. + __device__ inline void barrier_acquire(int* lock, int count) { + if (threadIdx.x == 0) { + int state = -1; + do + // Guarantee that subsequent writes by this threadblock will be + // visible globally. + asm volatile("ld.global.acquire.gpu.b32 %0, [%1];\n" + : "=r"(state) + : "l"(lock)); + while (state != count); + } + __syncthreads(); + } + + // Release barrier and increment visitation count. + __device__ inline void barrier_release(int* lock, bool reset = false) { + __syncthreads(); + if (threadIdx.x == 0) { + if (reset) { + lock[0] = 0; + return; + } + int val = 1; + // Make sure that all writes since acquiring this barrier are visible + // globally, while releasing the barrier. + asm volatile("fence.acq_rel.gpu;\n"); + asm volatile("red.relaxed.gpu.global.add.s32 [%0], %1;\n" + : + : "l"(lock), "r"(val)); + } + } + + // For a given "a" of size [M,K] performs a permutation of the K columns based + // on the given "perm" indices. + __global__ void permute_cols_kernel(int4 const* __restrict__ a_int4_ptr, + int const* __restrict__ perm_int_ptr, + int4* __restrict__ out_int4_ptr, int size_m, + int size_k, int block_rows) { + int start_row = block_rows * blockIdx.x; + int finish_row = start_row + block_rows; + if (finish_row > size_m) { + finish_row = size_m; + } + int cur_block_rows = finish_row - start_row; + + int row_stride = size_k * sizeof(half) / 16; + + auto permute_row = [&](int row) { + int iters = size_k / default_threads; + int rest = size_k % default_threads; + + int offset = row * row_stride; + + half const* a_row_half = + reinterpret_cast(a_int4_ptr + offset); + half* out_half = reinterpret_cast(out_int4_ptr + offset); + + int base_k = 0; + + for (int i = 0; i < iters; i++) { + int cur_k = base_k + threadIdx.x; + int src_pos = perm_int_ptr[cur_k]; + + out_half[cur_k] = a_row_half[src_pos]; + + base_k += default_threads; + } + + if (rest) { + if (threadIdx.x < rest) { + int cur_k = base_k + threadIdx.x; + int src_pos = perm_int_ptr[cur_k]; + + out_half[cur_k] = a_row_half[src_pos]; + } + } + }; + + for (int i = 0; i < cur_block_rows; i++) { + int cur_row = start_row + i; + if (cur_row < size_m) { + permute_row(cur_row); + } + } + } + + template shared + // fetch pipeline + const bool has_act_order, // whether act_order is enabled + const int group_blocks = -1 // number of consecutive 16x16 blocks + // with a separate quantization scale + > + __device__ void + Marlin(const int4* __restrict__ A, // fp16 input matrix of shape mxk + const int4* __restrict__ B, // 4bit quantized weight matrix of shape kxn + int4* __restrict__ C, // fp16 output buffer of shape mxn + const int4* __restrict__ scales_ptr, // fp16 quantization scales of shape + // (k/groupsize)xn + const int* __restrict__ g_idx, // int32 group indices of shape k + int num_groups, // number of scale groups per output channel + int prob_m, // batch dimension m, should be divisible by (16 * thread_m_blocks) if bigger than that + int prob_n, // output dimension n + int prob_k, // reduction dimension k + int* locks // extra global storage for barrier synchronization + ) { + // Each threadblock processes one "stripe" of the B matrix with (roughly) the + // same size, which might involve multiple column "slices" (of width 16 * + // `thread_n_blocks`). Stripes are defined as shown in the 3x3 matrix 5 SM + // example: + // 0 1 3 + // 0 2 3 + // 1 2 4 + // While this kind of partitioning makes things somewhat more complicated, it + // ensures good utilization of all SMs for many kinds of shape and GPU + // configurations, while requiring as few slow global cross-threadblock + // reductions as possible. + using Dtype = ScalarType; + using scalar_t2 = typename ScalarType::scalar_t2; + using FragA = typename ScalarType::FragA; + using FragB = typename ScalarType::FragB; + using FragC = typename ScalarType::FragC; + using FragS = typename ScalarType::FragS; + + constexpr int pack_factor = 32 / num_bits; + + // int prob_m = *prob_m_ptr; + // const int thread_m_blocks = min(div_ceil(prob_m, 16), template_thread_m_blocks); + // constexpr int thread_m_blocks = template_thread_m_blocks; + + // For larger GEMMs we run multiple batchsize 64 versions in parallel for a + // better partitioning with less reductions + int parallel = 1; + if (prob_m > 16 * thread_m_blocks) { + parallel = prob_m / (16 * thread_m_blocks); + prob_m = 16 * thread_m_blocks; + } + + int k_tiles = prob_k / 16 / thread_k_blocks; + int n_tiles = prob_n / 16 / thread_n_blocks; + int iters = div_ceil(k_tiles * n_tiles * parallel, gridDim.x); + + if constexpr (!has_act_order && group_blocks != -1) { + if (group_blocks >= thread_k_blocks) { + // Ensure that the number of tiles in each stripe is a multiple of the + // groupsize; this avoids an annoying special case where a stripe starts + // in the middle of group. + iters = (group_blocks / thread_k_blocks) * + div_ceil(iters, (group_blocks / thread_k_blocks)); + } + } + + int slice_row = (iters * blockIdx.x) % k_tiles; + int slice_col_par = (iters * blockIdx.x) / k_tiles; + int slice_col = slice_col_par; + int slice_iters; // number of threadblock tiles in the current slice + int slice_count = + 0; // total number of active threadblocks in the current slice + int slice_idx; // index of threadblock in current slice; numbered bottom to + // top + + // We can easily implement parallel problem execution by just remapping + // indices and advancing global pointers + if (slice_col_par >= n_tiles) { + A += (slice_col_par / n_tiles) * 16 * thread_m_blocks * prob_k / 8; + C += (slice_col_par / n_tiles) * 16 * thread_m_blocks * prob_n / 8; + locks += (slice_col_par / n_tiles) * n_tiles; + slice_col = slice_col_par % n_tiles; + } + + // Compute all information about the current slice which is required for + // synchronization. + auto init_slice = [&]() { + slice_iters = + iters * (blockIdx.x + 1) - (k_tiles * slice_col_par + slice_row); + if (slice_iters < 0 || slice_col_par >= n_tiles * parallel) slice_iters = 0; + if (slice_iters == 0) return; + if (slice_row + slice_iters > k_tiles) slice_iters = k_tiles - slice_row; + slice_count = 1; + slice_idx = 0; + int col_first = iters * div_ceil(k_tiles * slice_col_par, iters); + if (col_first <= k_tiles * (slice_col_par + 1)) { + int col_off = col_first - k_tiles * slice_col_par; + slice_count = div_ceil(k_tiles - col_off, iters); + if (col_off > 0) slice_count++; + int delta_first = iters * blockIdx.x - col_first; + if (delta_first < 0 || (col_off == 0 && delta_first == 0)) + slice_idx = slice_count - 1; + else { + slice_idx = slice_count - 1 - delta_first / iters; + if (col_off > 0) slice_idx--; + } + } + if (slice_col == n_tiles) { + A += 16 * thread_m_blocks * prob_k / 8; + C += 16 * thread_m_blocks * prob_n / 8; + locks += n_tiles; + slice_col = 0; + } + }; + init_slice(); + + // A sizes/strides + + // stride of the A matrix in global memory + int a_gl_stride = prob_k / 8; + // stride of an A matrix tile in shared memory + constexpr int a_sh_stride = 16 * thread_k_blocks / 8; + // delta between subsequent A tiles in global memory + constexpr int a_gl_rd_delta_o = 16 * thread_k_blocks / 8; + // between subsequent accesses within a tile + int a_gl_rd_delta_i = a_gl_stride * (threads / a_gl_rd_delta_o); + // between shared memory writes + constexpr int a_sh_wr_delta = a_sh_stride * (threads / a_gl_rd_delta_o); + // between shared memory tile reads + constexpr int a_sh_rd_delta_o = 2 * ((threads / 32) / (thread_n_blocks / 4)); + // within a shared memory tile + constexpr int a_sh_rd_delta_i = a_sh_stride * 16; + // overall size of a tile + constexpr int a_sh_stage = a_sh_stride * (16 * thread_m_blocks); + // number of shared write iterations for a tile + constexpr int a_sh_wr_iters = div_ceil(a_sh_stage, a_sh_wr_delta); + + // B sizes/strides + int b_gl_stride = 16 * prob_n / (pack_factor * 4); + constexpr int b_sh_stride = ((thread_n_blocks * 16) * 16 / pack_factor) / 4; + constexpr int b_thread_vecs = num_bits == 4 ? 1 : 2; + constexpr int b_sh_stride_threads = b_sh_stride / b_thread_vecs; + + int b_gl_rd_delta_o = b_gl_stride * thread_k_blocks; + int b_gl_rd_delta_i = b_gl_stride * (threads / b_sh_stride_threads); + constexpr int b_sh_wr_delta = threads * b_thread_vecs; + constexpr int b_sh_rd_delta = threads * b_thread_vecs; + constexpr int b_sh_stage = b_sh_stride * thread_k_blocks; + constexpr int b_sh_wr_iters = b_sh_stage / b_sh_wr_delta; + + // Scale sizes/strides without act_order + int s_gl_stride = prob_n / 8; + constexpr int s_sh_stride = 16 * thread_n_blocks / 8; + constexpr int s_tb_groups = + !has_act_order && group_blocks != -1 && group_blocks < thread_k_blocks + ? thread_k_blocks / group_blocks + : 1; + constexpr int s_sh_stage = s_tb_groups * s_sh_stride; + int s_gl_rd_delta = s_gl_stride; + + // Scale size/strides with act_order + constexpr int tb_k = 16 * thread_k_blocks; + constexpr int g_idx_stage = has_act_order ? (tb_k * sizeof(int)) / 16 : 0; + // constexpr int act_s_row_stride = 1; + // int act_s_col_stride = act_s_row_stride * num_groups; + int act_s_col_stride = 1; + int act_s_col_warp_stride = act_s_col_stride * 8; + int tb_n_warps = thread_n_blocks / 4; + int act_s_col_tb_stride = act_s_col_warp_stride * tb_n_warps; + + // Global A read index of current thread. + int a_gl_rd = a_gl_stride * (threadIdx.x / a_gl_rd_delta_o) + + (threadIdx.x % a_gl_rd_delta_o); + a_gl_rd += a_gl_rd_delta_o * slice_row; + // Shared write index of current thread. + int a_sh_wr = a_sh_stride * (threadIdx.x / a_gl_rd_delta_o) + + (threadIdx.x % a_gl_rd_delta_o); + // Shared read index. + int a_sh_rd = + a_sh_stride * ((threadIdx.x % 32) % 16) + (threadIdx.x % 32) / 16; + a_sh_rd += 2 * ((threadIdx.x / 32) / (thread_n_blocks / 4)); + + int b_gl_rd = b_gl_stride * (threadIdx.x / b_sh_stride_threads) + + (threadIdx.x % b_sh_stride_threads) * b_thread_vecs; + b_gl_rd += b_sh_stride * slice_col; + b_gl_rd += b_gl_rd_delta_o * slice_row; + int b_sh_wr = threadIdx.x * b_thread_vecs; + int b_sh_rd = threadIdx.x * b_thread_vecs; + + // For act_order + constexpr int k_iter_size = tb_k / b_sh_wr_iters; + int slice_k_start = tb_k * slice_row; + int slice_k_finish = slice_k_start + tb_k * slice_iters; + int slice_k_start_shared_fetch = slice_k_start; + int slice_n_offset = act_s_col_tb_stride * slice_col; + + // No act_order + int s_gl_rd; + if constexpr (!has_act_order) { + if constexpr (group_blocks == -1) { + s_gl_rd = s_sh_stride * slice_col + threadIdx.x; + } + else { + s_gl_rd = s_gl_stride * ((thread_k_blocks * slice_row) / group_blocks) + + s_sh_stride * slice_col + threadIdx.x; + } + } + int s_sh_wr = threadIdx.x; + bool s_sh_wr_pred = threadIdx.x < s_sh_stride; + + // We use a different scale layout for grouped and column-wise quantization as + // we scale a `half2` tile in column-major layout in the former and in + // row-major in the latter case. + int s_sh_rd; + if constexpr (group_blocks != -1) + s_sh_rd = 8 * ((threadIdx.x / 32) % (thread_n_blocks / 4)) + + (threadIdx.x % 32) / 4; + else + s_sh_rd = 8 * ((threadIdx.x / 32) % (thread_n_blocks / 4)) + + (threadIdx.x % 32) % 4; + + // Precompute which thread should not read memory in which iterations; this is + // needed if there are more threads than required for a certain tilesize or + // when the batchsize is not a multiple of 16. + bool a_sh_wr_pred[a_sh_wr_iters]; +#pragma unroll + for (int i = 0; i < a_sh_wr_iters; i++) { + a_sh_wr_pred[i] = a_sh_wr_delta * i + a_sh_wr < a_sh_stride * prob_m; + } + + // To ensure that writing and reading A tiles to/from shared memory, the + // latter in fragment format, is fully bank conflict free, we need to use a + // rather fancy XOR-based layout. The key here is that neither reads nor + // writes of the 16-byte `int4` blocks of 8 consecutive threads involve the + // same shared memory banks. Further, it seems (based on NSight-Compute) that + // each warp must also write a consecutive memory segment? + auto transform_a = [&](int i) { + int row = i / a_gl_rd_delta_o; + return a_gl_rd_delta_o * row + (i % a_gl_rd_delta_o) ^ row; + }; + // Since the computation of this remapping is non-trivial and, due to our main + // loop unrolls, all shared memory accesses are static, we simply precompute + // both transformed reads and writes. + int a_sh_wr_trans[a_sh_wr_iters]; +#pragma unroll + for (int i = 0; i < a_sh_wr_iters; i++) { + a_sh_wr_trans[i] = transform_a(a_sh_wr_delta * i + a_sh_wr); + } + int a_sh_rd_trans[b_sh_wr_iters][thread_m_blocks]; +#pragma unroll + for (int i = 0; i < b_sh_wr_iters; i++) { +#pragma unroll + for (int j = 0; j < thread_m_blocks; j++) + { + a_sh_rd_trans[i][j] = + transform_a(a_sh_rd_delta_o * i + a_sh_rd_delta_i * j + a_sh_rd); + } + } + + // Since B-accesses have non-constant stride they have to be computed at + // runtime; we break dependencies between subsequent accesses with a tile by + // maintining multiple pointers (we have enough registers), a tiny + // optimization. + const int4* B_ptr[b_sh_wr_iters]; +#pragma unroll + for (int i = 0; i < b_sh_wr_iters; i++) + B_ptr[i] = B + b_gl_rd_delta_i * i + b_gl_rd; + + extern __shared__ int4 sh[]; + // Shared memory storage for global fetch pipelines. + int4* sh_a = sh; + int4* sh_b = sh_a + (stages * a_sh_stage); + int4* sh_g_idx = sh_b + (stages * b_sh_stage); + int4* sh_s = sh_g_idx + (stages * g_idx_stage); + + // Register storage for double buffer of shared memory reads. + FragA frag_a[2][thread_m_blocks]; + I4 frag_b_quant[2][b_thread_vecs]; + FragC frag_c[thread_m_blocks][4][2]; + FragS frag_s[2][4]; // No act-order + FragS act_frag_s[2][4][4]; // For act-order + + // Zero accumulators. + auto zero_accums = [&]() { +#pragma unroll + for (int i = 0; i < thread_m_blocks * 4 * 2 * 4; i++) + { + reinterpret_cast(frag_c)[i] = 0; + } + }; + + int sh_first_group_id = -1; + int sh_num_groups = -1; + constexpr int sh_max_num_groups = 32; + + auto fetch_scales_to_shared = [&](bool is_async, int first_group_id, + int last_group_id) { + sh_first_group_id = first_group_id; + sh_num_groups = last_group_id - first_group_id + 1; + + if (sh_num_groups < sh_max_num_groups) { + sh_num_groups = sh_max_num_groups; + } + + if (sh_first_group_id + sh_num_groups > num_groups) { + sh_num_groups = num_groups - sh_first_group_id; + } + + int row_offset = first_group_id * s_gl_stride; + + if (is_async) { + for (int i = 0; i < sh_num_groups; i++) { + if (threadIdx.x < s_sh_stride) { + cp_async4_pred(&sh_s[(i * s_sh_stride) + threadIdx.x], + &scales_ptr[row_offset + (i * s_gl_stride) + + slice_n_offset + threadIdx.x]); + } + } + } + else { + for (int i = 0; i < sh_num_groups; i++) { + if (threadIdx.x < s_sh_stride) { + sh_s[(i * s_sh_stride) + threadIdx.x] = + scales_ptr[row_offset + (i * s_gl_stride) + slice_n_offset + + threadIdx.x]; + } + } + } + }; + // Asynchronously fetch the next A, B and s tile from global to the next + // shared memory pipeline location. + auto fetch_to_shared = [&](int pipe, int a_off, bool pred = true) { + if (pred) { + int4* sh_a_stage = sh_a + a_sh_stage * pipe; +#pragma unroll + for (int i = 0; i < a_sh_wr_iters; i++) { + cp_async4_pred( + &sh_a_stage[a_sh_wr_trans[i]], + &A[a_gl_rd_delta_i * i + a_gl_rd + a_gl_rd_delta_o * a_off], + a_sh_wr_pred[i]); + } + int4* sh_b_stage = sh_b + b_sh_stage * pipe; +#pragma unroll + for (int i = 0; i < b_sh_wr_iters; i++) { +#pragma unroll + for (int j = 0; j < b_thread_vecs; j++) { + cp_async4(&sh_b_stage[b_sh_wr_delta * i + b_sh_wr + j], B_ptr[i] + j); + } + + B_ptr[i] += b_gl_rd_delta_o; + } + + if constexpr (has_act_order) { + // Fetch g_idx thread-block portion + int full_pipe = a_off; + int cur_k = slice_k_start_shared_fetch + tb_k * full_pipe; + if (cur_k < prob_k && cur_k < slice_k_finish) { + int4* sh_g_idx_stage = sh_g_idx + g_idx_stage * pipe; + + int4 const* cur_g_idx_stage_ptr = + reinterpret_cast(&g_idx[cur_k]); + + if (threadIdx.x < g_idx_stage) { + cp_async4_pred(&sh_g_idx_stage[threadIdx.x], + &cur_g_idx_stage_ptr[threadIdx.x]); + } + } + } + else { + if constexpr (group_blocks != -1) { + int4* sh_s_stage = sh_s + s_sh_stage * pipe; + + if constexpr (group_blocks >= thread_k_blocks) { + // Only fetch scales if this tile starts a new group + if (pipe % (group_blocks / thread_k_blocks) == 0) { + if (s_sh_wr_pred) { + cp_async4(&sh_s_stage[s_sh_wr], &scales_ptr[s_gl_rd]); + } + s_gl_rd += s_gl_rd_delta; + } + } + else { + for (int i = 0; i < s_tb_groups; i++) { + if (s_sh_wr_pred) { + cp_async4(&sh_s_stage[i * s_sh_stride + s_sh_wr], + &scales_ptr[s_gl_rd]); + } + s_gl_rd += s_gl_rd_delta; + } + } + } + } + } + // Insert a fence even when we are winding down the pipeline to ensure that + // waiting is also correct at this point. + cp_async_fence(); + }; + + // Wait until the next thread tile has been loaded to shared memory. + auto wait_for_stage = [&]() { + // We only have `stages - 2` active fetches since we are double buffering + // and can only issue the next fetch when it is guaranteed that the previous + // shared memory load is fully complete (as it may otherwise be + // overwritten). + cp_async_wait(); + __syncthreads(); + }; + + // Load the next sub-tile from the current location in the shared memory pipe + // into the current register buffer. + auto fetch_to_registers = [&](int k, int pipe) { + int4* sh_a_stage = sh_a + a_sh_stage * pipe; +#pragma unroll + for (int i = 0; i < thread_m_blocks; i++) + { + ldsm4(frag_a[k % 2][i], + &sh_a_stage[a_sh_rd_trans[k % b_sh_wr_iters][i]]); + } + + int4* sh_b_stage = sh_b + b_sh_stage * pipe; + +#pragma unroll + for (int i = 0; i < b_thread_vecs; i++) { + frag_b_quant[k % 2][i] = *reinterpret_cast( + &sh_b_stage[b_sh_rd_delta * (k % b_sh_wr_iters) + b_sh_rd + i]); + } + }; + + bool is_same_group[stages]; + int same_group_id[stages]; + + auto init_same_group = [&](int pipe) { + if constexpr (!has_act_order) { + is_same_group[pipe] = false; + same_group_id[pipe] = 0; + return; + } + + int4* sh_g_idx_stage = sh_g_idx + g_idx_stage * pipe; + int* sh_g_idx_int_ptr = reinterpret_cast(sh_g_idx_stage); + + int group_id_1 = sh_g_idx_int_ptr[0]; + int group_id_2 = sh_g_idx_int_ptr[tb_k - 1]; + + is_same_group[pipe] = group_id_1 == group_id_2; + same_group_id[pipe] = group_id_1; + }; + + auto fetch_scales_to_registers = [&](int k, int full_pipe) { + int pipe = full_pipe % stages; + + if constexpr (!has_act_order) { + // No act-order case + if constexpr (group_blocks != -1) { + if constexpr (group_blocks >= thread_k_blocks) { + int4* sh_s_stage = + sh_s + s_sh_stage * ((group_blocks / thread_k_blocks) * + (pipe / (group_blocks / thread_k_blocks))); + reinterpret_cast(&frag_s[k % 2])[0] = sh_s_stage[s_sh_rd]; + } + else { + int warp_id = threadIdx.x / 32; + int n_warps = thread_n_blocks / 4; + + int warp_row = warp_id / n_warps; + + int cur_k = warp_row * 16; + cur_k += k_iter_size * (k % b_sh_wr_iters); + + int k_blocks = cur_k / 16; + int cur_group_id = k_blocks / group_blocks; + + int4* sh_s_stage = sh_s + s_sh_stage * pipe; + + reinterpret_cast(&frag_s[k % 2])[0] = + sh_s_stage[s_sh_rd + cur_group_id * s_sh_stride]; + } + } + + return; + } + + // Act-order case + + // Determine K of the "current" thread-block + int cur_k = slice_k_start + tb_k * full_pipe; + if (cur_k >= prob_k || cur_k >= slice_k_finish) { + return; + } + + // Reset (to current thread-block) since we read g_idx portion from the + // shared memory + cur_k = 0; + + // Progress to current iteration + cur_k += k_iter_size * (k % b_sh_wr_iters); + + // Determine "position" inside the thread-block (based on warp and + // thread-id) + int warp_id = threadIdx.x / 32; + int n_warps = + thread_n_blocks / 4; // Each warp processes 4 16-size tiles over N + + int warp_row = warp_id / n_warps; + int warp_col = warp_id % n_warps; + + cur_k += warp_row * 16; + + int th_id = threadIdx.x % 32; + cur_k += (th_id % 4) * 2; // Due to tensor-core layout for fp16 B matrix + + int s_col_shift = + /*slice_n_offset +*/ (act_s_col_warp_stride * warp_col) + + (th_id / 4) * act_s_col_stride; + + if (is_same_group[pipe]) { + if (k % 2 == 0) { + *(reinterpret_cast(&(act_frag_s[k % 2][0][0]))) = + sh_s[(same_group_id[pipe] - sh_first_group_id) * s_sh_stride + + s_col_shift]; + } + else { + *(reinterpret_cast(&(act_frag_s[k % 2][0][0]))) = + *(reinterpret_cast(&(act_frag_s[(k - 1) % 2][0][0]))); + } + + for (int i = 1; i < 4; i++) { + *(reinterpret_cast(&(act_frag_s[k % 2][i][0]))) = + *(reinterpret_cast(&(act_frag_s[k % 2][0][0]))); + } + return; + } + + int4* sh_g_idx_stage = sh_g_idx + g_idx_stage * pipe; + int* sh_g_idx_int_ptr = reinterpret_cast(sh_g_idx_stage); + + constexpr int k_frag_offsets[4] = { 0, 1, 8, + 9 }; // Tensor core offsets per thread + +#pragma unroll + for (int i = 0; i < 4; i++) { + int actual_k = cur_k + k_frag_offsets[i]; + + int group_id = sh_g_idx_int_ptr[actual_k]; + int rel_group_id = group_id - sh_first_group_id; + + *(reinterpret_cast(&(act_frag_s[k % 2][i][0]))) = + sh_s[rel_group_id * s_sh_stride + s_col_shift]; + } + }; + + // Execute the actual tensor core matmul of a sub-tile. + auto matmul = [&](int k) { + // We have the m dimension as the inner loop in order to encourage overlapping + // dequantization and matmul operations. +#pragma unroll + for (int j = 0; j < 4; j++) { + FragB frag_b0; + FragB frag_b1; + if constexpr (num_bits == 4) { + int b_quant = frag_b_quant[k % 2][0][j]; + int b_quant_shift = b_quant >> 8; + + frag_b0 = dequant_4bit(b_quant); + frag_b1 = dequant_4bit(b_quant_shift); + + } + else { + int* frag_b_quant_ptr = reinterpret_cast(frag_b_quant[k % 2]); + int b_quant_0 = frag_b_quant_ptr[j * 2 + 0]; + int b_quant_1 = frag_b_quant_ptr[j * 2 + 1]; + + frag_b0 = dequant_8bit(b_quant_0); + frag_b1 = dequant_8bit(b_quant_1); + } + + // Apply scale to frag_b0 + if constexpr (has_act_order) { + scale4(frag_b0, act_frag_s[k % 2][0][j], + act_frag_s[k % 2][1][j], act_frag_s[k % 2][2][j], + act_frag_s[k % 2][3][j], 0); + } + else { + if constexpr (group_blocks != -1) { + scale(frag_b0, frag_s[k % 2][j], 0); + } + } + + // Apply scale to frag_b1 + if constexpr (has_act_order) { + scale4(frag_b1, act_frag_s[k % 2][0][j], + act_frag_s[k % 2][1][j], act_frag_s[k % 2][2][j], + act_frag_s[k % 2][3][j], 1); + + } + else { + if constexpr (group_blocks != -1) { + scale(frag_b1, frag_s[k % 2][j], 1); + } + } + +#pragma unroll + for (int i = 0; i < thread_m_blocks; i++) { + mma(frag_a[k % 2][i], frag_b0, frag_c[i][j][0]); + mma(frag_a[k % 2][i], frag_b1, frag_c[i][j][1]); + } + } + }; + + // Since we slice across the k dimension of a tile in order to increase the + // number of warps while keeping the n dimension of a tile reasonable, we have + // multiple warps that accumulate their partial sums of the same output + // location; which we have to reduce over in the end. We do in shared memory. + auto thread_block_reduce = [&]() { + constexpr int red_off = threads / b_sh_stride_threads / 2; + if (red_off >= 1) { + int red_idx = threadIdx.x / b_sh_stride_threads; + constexpr int red_sh_stride = b_sh_stride_threads * 4 * 2; + constexpr int red_sh_delta = b_sh_stride_threads; + int red_sh_rd = red_sh_stride * (threadIdx.x / b_sh_stride_threads) + + (threadIdx.x % b_sh_stride_threads); + + // Parallel logarithmic shared memory reduction. We make sure to avoid any + // unnecessary read or write iterations, e.g., for two warps we write only + // once by warp 1 and read only once by warp 0. + +#pragma unroll + for (int m_block = 0; m_block < thread_m_blocks; m_block++) { +#pragma unroll + for (int i = red_off; i > 0; i /= 2) { + if (i <= red_idx && red_idx < 2 * i) { +#pragma unroll + for (int j = 0; j < 4 * 2; j++) { + int red_sh_wr = + red_sh_delta * j + (red_sh_rd - red_sh_stride * i); + if (i < red_off) { + float* c_rd = + reinterpret_cast(&sh[red_sh_delta * j + red_sh_rd]); + float* c_wr = reinterpret_cast(&sh[red_sh_wr]); +#pragma unroll + for (int k = 0; k < 4; k++) + reinterpret_cast(frag_c)[4 * 2 * m_block + j][k] += + c_rd[k] + c_wr[k]; + } + sh[red_sh_wr] = + reinterpret_cast(&frag_c)[4 * 2 * m_block + j]; + } + } + __syncthreads(); + } + if (red_idx == 0) { +#pragma unroll + for (int i = 0; i < 4 * 2; i++) { + float* c_rd = + reinterpret_cast(&sh[red_sh_delta * i + red_sh_rd]); +#pragma unroll + for (int j = 0; j < 4; j++) + reinterpret_cast(frag_c)[4 * 2 * m_block + i][j] += + c_rd[j]; + } + } + __syncthreads(); + } + } + }; + + // Since multiple threadblocks may process parts of the same column slice, we + // finally have to globally reduce over the results. As the striped + // partitioning minimizes the number of such reductions and our outputs are + // usually rather small, we perform this reduction serially in L2 cache. + auto global_reduce = [&](bool first = false, bool last = false) { + // We are very careful here to reduce directly in the output buffer to + // maximize L2 cache utilization in this step. To do this, we write out + // results in FP16 (but still reduce with FP32 compute). + constexpr int active_threads = 32 * thread_n_blocks / 4; + if (threadIdx.x < active_threads) { + int c_gl_stride = prob_n / 8; + int c_gl_wr_delta_o = 8 * c_gl_stride; + int c_gl_wr_delta_i = 4 * (active_threads / 32); + int c_gl_wr = c_gl_stride * ((threadIdx.x % 32) / 4) + + 4 * (threadIdx.x / 32) + threadIdx.x % 4; + c_gl_wr += (2 * thread_n_blocks) * slice_col; + constexpr int c_sh_wr_delta = active_threads; + int c_sh_wr = threadIdx.x; + + int row = (threadIdx.x % 32) / 4; + + if (!first) { + // Interestingly, doing direct global accesses here really seems to mess up + // the compiler and lead to slowdowns, hence we also use async-copies even + // though these fetches are not actually asynchronous. +#pragma unroll + for (int i = 0; i < thread_m_blocks * 4; i++) { + cp_async4_pred( + &sh[c_sh_wr + c_sh_wr_delta * i], + &C[c_gl_wr + c_gl_wr_delta_o * (i / 2) + + c_gl_wr_delta_i * (i % 2)], + i < (thread_m_blocks - 1) * 4 || 8 * (i / 2) + row < prob_m); + } + cp_async_fence(); + cp_async_wait<0>(); + } + +#pragma unroll + for (int i = 0; i < thread_m_blocks * 4; i++) { + if (i < (thread_m_blocks - 1) * 4 || 8 * (i / 2) + row < prob_m) { + if (!first) { + int4 c_red = sh[c_sh_wr + i * c_sh_wr_delta]; +#pragma unroll + for (int j = 0; j < 2 * 4; j++) { + reinterpret_cast( + &frag_c)[4 * 2 * 4 * (i / 4) + 4 * j + (i % 4)] += + Dtype::num2float(reinterpret_cast(&c_red)[j]); + } + } + if (!last) { + int4 c; +#pragma unroll + for (int j = 0; j < 2 * 4; j++) { + reinterpret_cast(&c)[j] = + Dtype::float2num(reinterpret_cast( + &frag_c)[4 * 2 * 4 * (i / 4) + 4 * j + (i % 4)]); + } + C[c_gl_wr + c_gl_wr_delta_o * (i / 2) + c_gl_wr_delta_i * (i % 2)] = + c; + } + } + } + } + }; + + // Write out the reduce final result in the correct layout. We only actually + // reshuffle matrix fragments in this step, the reduction above is performed + // in fragment layout. + auto write_result = [&]() { + int c_gl_stride = prob_n / 8; + constexpr int c_sh_stride = 2 * thread_n_blocks + 1; + int c_gl_wr_delta = c_gl_stride * (threads / (2 * thread_n_blocks)); + constexpr int c_sh_rd_delta = + c_sh_stride * (threads / (2 * thread_n_blocks)); + + int c_gl_wr = c_gl_stride * (threadIdx.x / (2 * thread_n_blocks)) + + (threadIdx.x % (2 * thread_n_blocks)); + c_gl_wr += (2 * thread_n_blocks) * slice_col; + int c_sh_wr = + (4 * c_sh_stride) * ((threadIdx.x % 32) / 4) + (threadIdx.x % 32) % 4; + c_sh_wr += 32 * (threadIdx.x / 32); + int c_sh_rd = c_sh_stride * (threadIdx.x / (2 * thread_n_blocks)) + + (threadIdx.x % (2 * thread_n_blocks)); + + int c_gl_wr_end = c_gl_stride * prob_m; + + // We first reorder in shared memory to guarantee the most efficient final + // global write patterns + auto write = [&](int idx, float c0, float c1, FragS& s) { + scalar_t2 res = + Dtype::nums2num2(Dtype::float2num(c0), Dtype::float2num(c1)); + + // For per-column quantization we finally apply the scale here (only for + // 4-bit) + if constexpr (!has_act_order && group_blocks == -1 && num_bits == 4) { + res = __hmul2(res, s[0]); + } + + ((scalar_t2*)sh)[idx] = res; + }; + + if (threadIdx.x / 32 < thread_n_blocks / 4) { +#pragma unroll + for (int i = 0; i < thread_m_blocks; i++) { +#pragma unroll + for (int j = 0; j < 4; j++) { + int wr = c_sh_wr + 8 * j; + write(wr + (4 * c_sh_stride) * 0 + 0, frag_c[i][j][0][0], + frag_c[i][j][0][1], frag_s[j / 2][2 * (j % 2) + 0]); + write(wr + (4 * c_sh_stride) * 8 + 0, frag_c[i][j][0][2], + frag_c[i][j][0][3], frag_s[j / 2][2 * (j % 2) + 0]); + write(wr + (4 * c_sh_stride) * 0 + 4, frag_c[i][j][1][0], + frag_c[i][j][1][1], frag_s[j / 2][2 * (j % 2) + 1]); + write(wr + (4 * c_sh_stride) * 8 + 4, frag_c[i][j][1][2], + frag_c[i][j][1][3], frag_s[j / 2][2 * (j % 2) + 1]); + } + c_sh_wr += 16 * (4 * c_sh_stride); + } + } + __syncthreads(); + +#pragma unroll + for (int i = 0; + i < div_ceil(16 * thread_m_blocks, threads / (2 * thread_n_blocks)); + i++) { + if (c_gl_wr < c_gl_wr_end) { + C[c_gl_wr] = sh[c_sh_rd]; + c_gl_wr += c_gl_wr_delta; + c_sh_rd += c_sh_rd_delta; + } + } + }; + + // Start global fetch and register load pipelines. + auto start_pipes = [&]() { + +#pragma unroll + for (int i = 0; i < stages - 1; i++) { + if (has_act_order && i == 0) { + int last_g_idx = slice_k_start + stages * tb_k * 2; + if (last_g_idx >= prob_k) { + last_g_idx = prob_k - 1; + } + fetch_scales_to_shared(true, g_idx[slice_k_start], g_idx[last_g_idx]); + } + fetch_to_shared(i, i, i < slice_iters); + } + + zero_accums(); + wait_for_stage(); + init_same_group(0); + fetch_to_registers(0, 0); + fetch_scales_to_registers(0, 0); + a_gl_rd += a_gl_rd_delta_o * (stages - 1); + slice_k_start_shared_fetch += tb_k * (stages - 1); + }; + if (slice_iters) { + start_pipes(); + } + + // Main loop. + while (slice_iters) { + // We unroll over both the global fetch and the register load pipeline to + // ensure all shared memory accesses are static. Note that both pipelines + // have even length meaning that the next iteration will always start at + // index 0. + +#pragma unroll + for (int pipe = 0; pipe < stages;) { +#pragma unroll + for (int k = 0; k < b_sh_wr_iters; k++) { + fetch_to_registers(k + 1, pipe % stages); + fetch_scales_to_registers(k + 1, pipe); + if (k == b_sh_wr_iters - 2) { + fetch_to_shared((pipe + stages - 1) % stages, pipe, + slice_iters >= stages); + pipe++; + wait_for_stage(); + init_same_group(pipe % stages); + } + matmul(k); + } + slice_iters--; + if (slice_iters == 0) { + break; + } + } + + a_gl_rd += a_gl_rd_delta_o * stages; + slice_k_start += tb_k * stages; + slice_k_start_shared_fetch += tb_k * stages; + + if constexpr (has_act_order) { + int first_group_id = g_idx[slice_k_start]; + int last_g_idx = slice_k_start + stages * tb_k * 2; + if (last_g_idx >= prob_k) { + last_g_idx = prob_k - 1; + } + int last_group_id = g_idx[last_g_idx]; + if (last_group_id >= sh_first_group_id + sh_num_groups) { + fetch_scales_to_shared(false, first_group_id, last_group_id); + __syncthreads(); + } + } + + // Process results and, if necessary, proceed to the next column slice. + // While this pattern may not be the most readable, other ways of writing + // the loop seemed to noticeably worse performance after compilation. + if (slice_iters == 0) { + cp_async_wait<0>(); + bool last = slice_idx == slice_count - 1; + // For per-column scales, we only fetch them here in the final step before + // write-out + if constexpr (!has_act_order && group_blocks == -1) { + if constexpr (num_bits == 8) { + if (s_sh_wr_pred) { + cp_async4(&sh_s[s_sh_wr], &scales_ptr[s_gl_rd]); + } + cp_async_fence(); + } + else { + if (last) { + if (s_sh_wr_pred) { + cp_async4(&sh_s[s_sh_wr], &scales_ptr[s_gl_rd]); + } + cp_async_fence(); + } + } + } + + thread_block_reduce(); + if constexpr (!has_act_order && group_blocks == -1) { + if constexpr (num_bits == 8) { + cp_async_wait<0>(); + __syncthreads(); + if (threadIdx.x / 32 < thread_n_blocks / 4) { + reinterpret_cast(&frag_s)[0] = sh_s[s_sh_rd + 0]; + reinterpret_cast(&frag_s)[1] = sh_s[s_sh_rd + 4]; + } + + } + else { + if (last) { + cp_async_wait<0>(); + __syncthreads(); + if (threadIdx.x / 32 < thread_n_blocks / 4) { + reinterpret_cast(&frag_s)[0] = sh_s[s_sh_rd + 0]; + reinterpret_cast(&frag_s)[1] = sh_s[s_sh_rd + 4]; + } + } + } + } + + // For 8-bit channelwise, we apply the scale before the global reduction + // that converts the fp32 results to fp16 (so that we avoid possible + // overflow in fp16) + if constexpr (!has_act_order && group_blocks == -1 && num_bits == 8) { + if (threadIdx.x / 32 < thread_n_blocks / 4) { +#pragma unroll + for (int i = 0; i < thread_m_blocks; i++) { +#pragma unroll + for (int j = 0; j < 4; j++) { + scale_float( + reinterpret_cast(&frag_c[i][j][0][0]), + frag_s[j / 2][2 * (j % 2) + 0]); + scale_float( + reinterpret_cast(&frag_c[i][j][0][2]), + frag_s[j / 2][2 * (j % 2) + 0]); + + scale_float( + reinterpret_cast(&frag_c[i][j][1][0]), + frag_s[j / 2][2 * (j % 2) + 1]); + scale_float( + reinterpret_cast(&frag_c[i][j][1][2]), + frag_s[j / 2][2 * (j % 2) + 1]); + } + } + } + } + + if (slice_count > 1) { // only globally reduce if there is more than one + // block in a slice + barrier_acquire(&locks[slice_col], slice_idx); + global_reduce(slice_idx == 0, last); + barrier_release(&locks[slice_col], last); + } + if (last) // only the last block in a slice actually writes the result + write_result(); + slice_row = 0; + slice_col_par++; + slice_col++; + init_slice(); + if (slice_iters) { + a_gl_rd = a_gl_stride * (threadIdx.x / a_gl_rd_delta_o) + + (threadIdx.x % a_gl_rd_delta_o); +#pragma unroll + for (int i = 0; i < b_sh_wr_iters; i++) + B_ptr[i] += b_sh_stride - b_gl_rd_delta_o * k_tiles; + if (slice_col == 0) { +#pragma unroll + for (int i = 0; i < b_sh_wr_iters; i++) B_ptr[i] -= b_gl_stride; + } + + // Update slice k/n for scales loading + if constexpr (has_act_order) { + slice_k_start = tb_k * slice_row; + slice_k_finish = slice_k_start + tb_k * slice_iters; + slice_k_start_shared_fetch = slice_k_start; + slice_n_offset = act_s_col_tb_stride * slice_col; + + } + else { + s_gl_rd = s_sh_stride * slice_col + threadIdx.x; + } + + start_pipes(); + } + } + } + } + + template shared + // fetch pipeline + const bool has_act_order, // whether act_order is enabled + const int group_blocks = -1 // number of consecutive 16x16 blocks + // with a separate quantization scale + > + __global__ void + Marlin_wrapper(const int4* __restrict__ A, // fp16 input matrix of shape mxk + const int4* __restrict__ B, // 4bit quantized weight matrix of shape kxn + int4* __restrict__ C, // fp16 output buffer of shape mxn + const int4* __restrict__ scales_ptr, // fp16 quantization scales of shape + // (k/groupsize)xn + const int* __restrict__ g_idx, // int32 group indices of shape k + int num_groups, // number of scale groups per output channel + const int* __restrict__ prob_m_ptr, // batch dimension m + int prob_n, // output dimension n + int prob_k, // reduction dimension k + int* locks // extra global storage for barrier synchronization + ) { + int prob_m = *prob_m_ptr; + prob_m = min(prob_m, 1024); + const int thread_m_blocks = min(div_ceil(prob_m, 16), template_thread_m_blocks); + if(prob_m > 16 * thread_m_blocks) + prob_m = (16 * thread_m_blocks) * div_ceil(prob_m, (16 * thread_m_blocks)); + /*if (blockIdx.x == 0 && threadIdx.x == 0) + printf("marlin prob_m %d\n", prob_m);*/ + if (thread_m_blocks == 1) { + Marlin( + A, B, C, scales_ptr, g_idx, num_groups, prob_m, prob_n, + prob_k, locks); + } + else if (thread_m_blocks == 2) { + Marlin( + A, B, C, scales_ptr, g_idx, num_groups, prob_m, prob_n, + prob_k, locks); + } + else if (thread_m_blocks == 3) { + Marlin( + A, B, C, scales_ptr, g_idx, num_groups, prob_m, prob_n, + prob_k, locks); + } + else if (thread_m_blocks == 4) { + Marlin( + A, B, C, scales_ptr, g_idx, num_groups, prob_m, prob_n, + prob_k, locks); + } + } + +#define __CALL_IF(NUM_BITS, THREAD_M_BLOCKS, THREAD_N_BLOCKS, THREAD_K_BLOCKS, \ + HAS_ACT_ORDER, GROUP_BLOCKS, NUM_THREADS) \ + else if (num_bits == NUM_BITS && thread_m_blocks == THREAD_M_BLOCKS && \ + thread_n_blocks == THREAD_N_BLOCKS && \ + thread_k_blocks == THREAD_K_BLOCKS && \ + has_act_order == HAS_ACT_ORDER && group_blocks == GROUP_BLOCKS && \ + num_threads == NUM_THREADS) { \ + cudaFuncSetAttribute( \ + Marlin_wrapper, \ + cudaFuncAttributeMaxDynamicSharedMemorySize, max_shared_mem); \ + Marlin_wrapper<<>>( \ + A_ptr, B_ptr, C_ptr, s_ptr, g_idx_ptr, num_groups, prob_m_ptr, prob_n, \ + prob_k, locks); \ + } + + typedef struct { + int thread_k; + int thread_n; + int num_threads; + } thread_config_t; + + typedef struct { + int max_m_blocks; + thread_config_t tb_cfg; + } exec_config_t; + + thread_config_t small_batch_thread_configs[] = { + // Ordered by priority + + // thread_k, thread_n, num_threads + {128, 128, 256}, + {64, 128, 128}, + {128, 64, 128}, + }; + + thread_config_t large_batch_thread_configs[] = { + // Ordered by priority + + // thread_k, thread_n, num_threads + {64, 256, 256}, + // {128, 128, 256}, + {64, 128, 128}, + {128, 64, 128}, + + }; + + int get_scales_cache_size(thread_config_t const& th_config, int prob_m, + int prob_n, int prob_k, int num_bits, int group_size, + bool has_act_order, bool is_k_full) { + bool cache_scales_chunk = has_act_order && !is_k_full; + + int tb_n = th_config.thread_n; + int tb_k = th_config.thread_k; + + // Get max scale groups per thread-block + int tb_groups; + if (group_size == -1) { + tb_groups = 1; + } + else if (group_size == 0) { + tb_groups = div_ceil(tb_k, 32); // Worst case is 32 group size + } + else { + tb_groups = div_ceil(tb_k, group_size); + } + + if (cache_scales_chunk) { + int load_groups = + tb_groups * pipe_stages * 2; // Chunk size is 2x pipeline over dim K + load_groups = max(load_groups, 32); // We load at least 32 scale groups + return load_groups * tb_n * 2; + + } + else { + int tb_scales = tb_groups * tb_n * 2; + + return tb_scales * pipe_stages; + } + } + + bool is_valid_cache_size(thread_config_t const& th_config, int max_m_blocks, + int prob_m, int prob_n, int prob_k, int num_bits, + int scales_cache_size, int max_shared_mem) { + int pack_factor = 32 / num_bits; + + // Get B size + int tb_k = th_config.thread_k; + int tb_n = th_config.thread_n; + + int b_size = (tb_k * tb_n / pack_factor) * 4; + + // Get A size + int m_blocks = div_ceil(prob_m, 16); + int tb_max_m = 16; + + // zbx: too ugly + // origin + /*while (true) { + if (m_blocks >= max_m_blocks) { + tb_max_m *= max_m_blocks; + break; + } + + max_m_blocks--; + if (max_m_blocks == 0) { + TORCH_CHECK(false, "Unexpected m_blocks = ", m_blocks); + } + }*/ + // refactor + tb_max_m *= std::min(m_blocks, max_m_blocks); + + int a_size = (tb_max_m * tb_k) * 2; + + float pipe_size = (a_size + b_size) * pipe_stages; + + TORCH_CHECK(max_shared_mem / 2 > scales_cache_size); // Sanity + return pipe_size < 0.95f * (max_shared_mem - scales_cache_size); + } + + bool is_valid_config(thread_config_t const& th_config, int max_m_blocks, + int prob_m, int prob_n, int prob_k, int num_bits, + int group_size, bool has_act_order, bool is_k_full, + int max_shared_mem) { + // Sanity + if (th_config.thread_k == -1 || th_config.thread_n == -1 || + th_config.num_threads == -1) { + return false; + } + + // Verify K/N are divisible by thread K/N + if (prob_k % th_config.thread_k != 0 || prob_n % th_config.thread_n != 0) { + return false; + } + + // Verify min for thread K/N + if (th_config.thread_n < min_thread_n || + th_config.thread_k < min_thread_k) { + return false; + } + + // num_threads must be at least 128 (= 4 warps) + if (th_config.num_threads < 128) { + return false; + } + + // Determine cache for scales + int scales_cache_size = + get_scales_cache_size(th_config, prob_m, prob_n, prob_k, num_bits, + group_size, has_act_order, is_k_full); + + // Check that pipeline fits into cache + if (!is_valid_cache_size(th_config, max_m_blocks, prob_m, prob_n, prob_k, + num_bits, scales_cache_size, max_shared_mem)) { + return false; + } + + return true; + } + + exec_config_t determine_thread_config(int prob_m, int prob_n, int prob_k, + int num_bits, int group_size, + bool has_act_order, bool is_k_full, + int max_shared_mem) { + int max_m_blocks = 4; + while (max_m_blocks > 0) { + if (prob_m <= 16) { + for (auto th_config : small_batch_thread_configs) { + if (is_valid_config(th_config, max_m_blocks, prob_m, prob_n, + prob_k, num_bits, group_size, has_act_order, + is_k_full, max_shared_mem)) { + return exec_config_t{ max_m_blocks, th_config }; + } + } + } + else { + for (auto th_config : large_batch_thread_configs) { + if (is_valid_config(th_config, max_m_blocks, prob_m, prob_n, + prob_k, num_bits, group_size, has_act_order, + is_k_full, max_shared_mem)) { + return exec_config_t{ max_m_blocks, th_config }; + } + } + } + + max_m_blocks--; // Process less M blocks per invocation to reduce cache + // usage + } + + return exec_config_t{ 0, {-1, -1, -1} }; + } + +#define CALL_IF(NUM_BITS, N_BLOCKS, K_BLOCKS, NUM_THREADS) \ + __CALL_IF(NUM_BITS, 1, N_BLOCKS, K_BLOCKS, false, 4, NUM_THREADS) \ + __CALL_IF(NUM_BITS, 2, N_BLOCKS, K_BLOCKS, false, 4, NUM_THREADS) \ + __CALL_IF(NUM_BITS, 3, N_BLOCKS, K_BLOCKS, false, 4, NUM_THREADS) \ + __CALL_IF(NUM_BITS, 4, N_BLOCKS, K_BLOCKS, false, 4, NUM_THREADS) \ + __CALL_IF(NUM_BITS, 1, N_BLOCKS, K_BLOCKS, false, 8, NUM_THREADS) \ + __CALL_IF(NUM_BITS, 2, N_BLOCKS, K_BLOCKS, false, 8, NUM_THREADS) \ + __CALL_IF(NUM_BITS, 3, N_BLOCKS, K_BLOCKS, false, 8, NUM_THREADS) \ + __CALL_IF(NUM_BITS, 4, N_BLOCKS, K_BLOCKS, false, 8, NUM_THREADS) + + template + void marlin_mm_f16i4(const void* A, const void* B, void* C, void* s, + void* g_idx, void* perm, void* a_tmp, int* prob_m_ptr, int prob_m, + int prob_n, int prob_k, void* workspace, int num_bits, + bool has_act_order, bool is_k_full, int num_groups, + int group_size, int dev, cudaStream_t stream, int thread_k, + int thread_n, int sms, int max_par) { + TORCH_CHECK(num_bits == 4 || num_bits == 8, + "num_bits must be 4 or 8. Got = ", num_bits); + TORCH_CHECK(prob_m > 0 && prob_n > 0 && prob_k > 0, "Invalid MNK = [", + prob_m, ", ", prob_n, ", ", prob_k, "]"); + + int tot_m = prob_m; + int tot_m_blocks = div_ceil(tot_m, 16); + int pad = 16 * tot_m_blocks - tot_m; + + if (sms == -1) { + cudaDeviceGetAttribute(&sms, cudaDevAttrMultiProcessorCount, dev); + } + + int max_shared_mem = 0; + cudaDeviceGetAttribute(&max_shared_mem, + cudaDevAttrMaxSharedMemoryPerBlockOptin, dev); + TORCH_CHECK(max_shared_mem > 0); + + // Set thread config + exec_config_t exec_cfg; + if (thread_k != -1 && thread_n != -1) { + // User-defined config + exec_cfg = exec_config_t{ + 4, thread_config_t{thread_k, thread_n, default_threads} }; + } + else { + // Auto config + exec_cfg = determine_thread_config(prob_m, prob_n, prob_k, num_bits, + group_size, has_act_order, is_k_full, + max_shared_mem); + } + + TORCH_CHECK( + exec_cfg.max_m_blocks > 0 && + is_valid_config(exec_cfg.tb_cfg, exec_cfg.max_m_blocks, prob_m, + prob_n, prob_k, num_bits, group_size, has_act_order, + is_k_full, max_shared_mem), + "Invalid thread config: max_m_blocks = ", exec_cfg.max_m_blocks, + ", thread_k = ", exec_cfg.tb_cfg.thread_k, + ", thread_n = ", exec_cfg.tb_cfg.thread_n, + ", num_threads = ", exec_cfg.tb_cfg.num_threads, " for MKN = [", prob_m, + ", ", prob_k, ", ", prob_n, "] and num_bits = ", num_bits, + ", group_size = ", group_size, ", has_act_order = ", has_act_order, + ", is_k_full = ", is_k_full, ", max_shared_mem = ", max_shared_mem); + + int num_threads = exec_cfg.tb_cfg.num_threads; + thread_k = exec_cfg.tb_cfg.thread_k; + thread_n = exec_cfg.tb_cfg.thread_n; + + int thread_k_blocks = thread_k / 16; + int thread_n_blocks = thread_n / 16; + + int blocks = sms; + + TORCH_CHECK(prob_n % thread_n == 0, "prob_n = ", prob_n, + " is not divisible by thread_n = ", thread_n); + TORCH_CHECK(prob_k % thread_k == 0, "prob_k = ", prob_k, + " is not divisible by thread_k = ", thread_k); + + int group_blocks = 0; + if (has_act_order) { + if (is_k_full) { + TORCH_CHECK(group_size != -1); + group_blocks = group_size / 16; + TORCH_CHECK(prob_k % group_blocks == 0, "prob_k = ", prob_k, + " is not divisible by group_blocks = ", group_blocks); + } + else { + TORCH_CHECK(group_size == 0); + group_blocks = 0; + } + + } + else { + if (group_size == -1) { + group_blocks = -1; + } + else { + group_blocks = group_size / 16; + TORCH_CHECK(prob_k % group_blocks == 0, "prob_k = ", prob_k, + " is not divisible by group_blocks = ", group_blocks); + } + } + + const int4* A_ptr = (const int4*)A; + const int4* B_ptr = (const int4*)B; + int4* C_ptr = (int4*)C; + const int4* s_ptr = (const int4*)s; + const int* g_idx_ptr = (const int*)g_idx; + const int* perm_ptr = (const int*)perm; + int4* a_tmp_ptr = (int4*)a_tmp; + + int* locks = (int*)workspace; + + if (has_act_order) { + // Permute A columns + int block_rows = div_ceil(prob_m, blocks); + permute_cols_kernel << > > ( + A_ptr, perm_ptr, a_tmp_ptr, prob_m, prob_k, block_rows); + A_ptr = a_tmp_ptr; + } + + // If we have a full K, then we can run the non-act-order version of Marlin + // (since the weight rows are reordered by increasing group ids, and by + // having a full K, we have full original groups) + if (is_k_full) { + has_act_order = false; + } + + // Main loop + for (int i = 0; i < tot_m_blocks; i += exec_cfg.max_m_blocks) { + int thread_m_blocks = tot_m_blocks - i; + prob_m = tot_m - 16 * i; + int par = 1; + if (thread_m_blocks > exec_cfg.max_m_blocks) { + // Note that parallel > 1 currently only works for inputs without + // any padding + par = (16 * thread_m_blocks - pad) / (16 * exec_cfg.max_m_blocks); + if (par > max_par) + par = max_par; + prob_m = (16 * exec_cfg.max_m_blocks) * par; + i += exec_cfg.max_m_blocks * (par - 1); + thread_m_blocks = exec_cfg.max_m_blocks; + } + + // Define kernel configurations +#define undefined_error \ + TORCH_CHECK(false, "Unsupported shapes: MNK = [" + str(prob_m) + ", " + \ + str(prob_n) + ", " + str(prob_k) + "]" + \ + ", has_act_order = " + str(has_act_order) + \ + ", num_groups = " + str(num_groups) + \ + ", group_size = " + str(group_size) + \ + ", thread_m_blocks = " + str(thread_m_blocks) + \ + ", thread_n_blocks = " + str(thread_n_blocks) + \ + ", thread_k_blocks = " + str(thread_k_blocks)); + + /* std::cout << "MNK = [" + str(prob_m) + ", " + \ + str(prob_n) + ", " + str(prob_k) + "]" + \ + ", has_act_order = " + str(has_act_order) + \ + ", num_groups = " + str(num_groups) + \ + ", group_size = " + str(group_size) + \ + ", thread_m_blocks = " + str(thread_m_blocks) + \ + ", thread_n_blocks = " + str(thread_n_blocks) + \ + ", thread_k_blocks = " + str(thread_k_blocks) << std::endl;*/ + + /*if (false) { + } + // CALL_IF(4, 32, 2, 256) + // CALL_IF(4, 16, 4, 256) + __CALL_IF(4, 1, 16, 4, false, 4, 256) + __CALL_IF(4, 2, 16, 4, false, 4, 256) + // CALL_IF(4, 8, 8, 256) + __CALL_IF(4, 1, 8, 8, false, 4, 256) + __CALL_IF(4, 2, 8, 8, false, 4, 256) + // CALL_IF(4, 16, 4, 128) + __CALL_IF(4, 1, 16, 4, false, 4, 128) + __CALL_IF(4, 2, 16, 4, false, 4, 128) + // CALL_IF(4, 8, 8, 128) + __CALL_IF(4, 1, 8, 8, false, 4, 128) + __CALL_IF(4, 2, 8, 8, false, 4, 128) + else {undefined_error}*/ + + if (num_bits == 4 && num_threads == 256) + { + if (false) { + } + CALL_IF(4, 32, 2, 256) + CALL_IF(4, 16, 4, 256) + CALL_IF(4, 8, 8, 256) + else { + undefined_error + } + } + else if (num_bits == 4 && num_threads == 128) + { + if (false) { + } + CALL_IF(4, 8, 4, 128) + CALL_IF(4, 16, 4, 128) + CALL_IF(4, 4, 8, 128) + else { + undefined_error + } + } + // else if (num_bits == 8 && num_threads == 256) + // { + // if (false) { + // } + // CALL_IF(8, 32, 2, 256) + // CALL_IF(8, 16, 4, 256) + // CALL_IF(8, 8, 8, 256) + // else { + // undefined_error + // } + // } + // else if (num_bits == 8 && num_threads == 128) + // { + // if (false) { + // } + // CALL_IF(8, 8, 4, 128) + // CALL_IF(8, 16, 4, 128) + // CALL_IF(8, 4, 8, 128) + // else { + // undefined_error + // } + // } + else { + undefined_error + } + + A_ptr += 16 * thread_m_blocks * (prob_k / 8) * par; + C_ptr += 16 * thread_m_blocks * (prob_n / 8) * par; + } + } + +} // namespace gptq_marlin + +torch::Tensor gptq_marlin_gemm(torch::Tensor& a, torch::Tensor& b_q_weight, + torch::Tensor& b_scales, torch::Tensor& g_idx, + torch::Tensor& perm, torch::Tensor& workspace, + int64_t num_bits, torch::Tensor size_m_tensor, int64_t size_m, int64_t size_n, + int64_t size_k, int sms, bool is_k_full) { + const at::cuda::OptionalCUDAGuard device_guard(device_of(a)); + // Verify num_bits + TORCH_CHECK(num_bits == 4 || num_bits == 8, + "num_bits must be 4 or 8. Got = ", num_bits); + int pack_factor = 32 / num_bits; + + // Verify A + TORCH_CHECK(a.size(0) == size_m, "Shape mismatch: a.size(0) = ", a.size(0), + ", size_m = ", size_m); + TORCH_CHECK(a.size(1) == size_k, "Shape mismatch: a.size(1) = ", a.size(1), + ", size_k = ", size_k); + + // Verify B + TORCH_CHECK(size_k % gptq_marlin::tile_size == 0, "size_k = ", size_k, + " is not divisible by tile_size = ", gptq_marlin::tile_size); + TORCH_CHECK((size_k / gptq_marlin::tile_size) == b_q_weight.size(0), + "Shape mismatch: b_q_weight.size(0) = ", b_q_weight.size(0), + ", size_k = ", size_k, + ", tile_size = ", gptq_marlin::tile_size); + TORCH_CHECK(b_q_weight.size(1) % gptq_marlin::tile_size == 0, + "b_q_weight.size(1) = ", b_q_weight.size(1), + " is not divisible by tile_size = ", gptq_marlin::tile_size); + int actual_size_n = + (b_q_weight.size(1) / gptq_marlin::tile_size) * pack_factor; + TORCH_CHECK(size_n == actual_size_n, "size_n = ", size_n, + ", actual_size_n = ", actual_size_n); + + // Verify device and strides + TORCH_CHECK(a.device().is_cuda(), "A is not on GPU"); + TORCH_CHECK(a.is_contiguous(), "A is not contiguous"); + + TORCH_CHECK(b_q_weight.device().is_cuda(), "b_q_weight is not on GPU"); + TORCH_CHECK(b_q_weight.is_contiguous(), "b_q_weight is not contiguous"); + + TORCH_CHECK(b_scales.device().is_cuda(), "b_scales is not on GPU"); + TORCH_CHECK(b_scales.is_contiguous(), "b_scales is not contiguous"); + + TORCH_CHECK(g_idx.device().is_cuda(), "g_idx is not on GPU"); + TORCH_CHECK(g_idx.is_contiguous(), "g_idx is not contiguous"); + + TORCH_CHECK(perm.device().is_cuda(), "perm is not on GPU"); + TORCH_CHECK(perm.is_contiguous(), "perm is not contiguous"); + + // Alloc buffers + auto options = torch::TensorOptions().dtype(a.dtype()).device(a.device()); + torch::Tensor c = torch::empty({ size_m, size_n }, options); + torch::Tensor a_tmp = torch::empty({ size_m, size_k }, options); + + // thread_k: `k` size of a thread_tile in `weights` (can usually be left as + // auto -1) + int thread_k = -1; + // thread_n: `n` size of a thread_tile in `weights` (can usually be left as + // auto -1) + int thread_n = -1; + // sms: number of SMs to use for the kernel (can usually be left as auto -1) + // int sms = -1; //zbx + + // Verify g_idx and perm + TORCH_CHECK((g_idx.size(0) == 0 && perm.size(0) == 0) || + (g_idx.size(0) == size_k && perm.size(0) == size_k), + "Unexpected g_idx.size(0) = ", g_idx.size(0), + " and perm.size(0) = ", perm.size(0), + ", where size_k = ", size_k); + + // Detect groupsize and act_order + int num_groups = -1; + int group_size = -1; + bool has_act_order = g_idx.size(0) != 0; + + int b_rank = b_scales.sizes().size(); + TORCH_CHECK(b_rank == 2, "b_scales rank = ", b_rank, " is not 2"); + TORCH_CHECK(b_scales.size(1) == size_n, + "b_scales dim 1 = ", b_scales.size(1), + " is not size_n = ", size_n); + num_groups = b_scales.size(0); + + if (has_act_order) { + if (is_k_full) { + TORCH_CHECK(num_groups > 1, + "For act_order, num_groups must be > 1"); + TORCH_CHECK(size_k % num_groups == 0, "size_k = ", size_k, + ", is not divisible by num_groups = ", num_groups); + group_size = size_k / num_groups; + } + else { + group_size = 0; + } + + } + else { + if (num_groups > 1) { + TORCH_CHECK( + size_k % num_groups == 0, "size_k = ", size_k, + ", is not divisible by b_scales.size(0) = ", b_scales.size(0)); + group_size = size_k / num_groups; + } + else { + group_size = -1; + } + } + + // Verify workspace size + TORCH_CHECK( + size_n % gptq_marlin::min_thread_n == 0, "size_n = ", size_n, + ", is not divisible by min_thread_n = ", gptq_marlin::min_thread_n); + int min_workspace_size = + (size_n / gptq_marlin::min_thread_n) * gptq_marlin::max_par; + TORCH_CHECK(workspace.numel() >= min_workspace_size, + "workspace.numel = ", workspace.numel(), + " is below min_workspace_size = ", min_workspace_size); + + int dev = a.get_device(); + if (a.scalar_type() == at::ScalarType::Half) { + gptq_marlin::marlin_mm_f16i4( + a.data_ptr(), b_q_weight.data_ptr(), + c.data_ptr(), b_scales.data_ptr(), + g_idx.data_ptr(), perm.data_ptr(), a_tmp.data_ptr(), + size_m_tensor.data_ptr(), + size_m, size_n, size_k, workspace.data_ptr(), num_bits, + has_act_order, is_k_full, num_groups, group_size, dev, + at::cuda::getCurrentCUDAStream(dev), thread_k, thread_n, sms, + gptq_marlin::max_par); + } + else if (a.scalar_type() == at::ScalarType::BFloat16) { + gptq_marlin::marlin_mm_f16i4( + a.data_ptr(), b_q_weight.data_ptr(), + c.data_ptr(), b_scales.data_ptr(), + g_idx.data_ptr(), perm.data_ptr(), a_tmp.data_ptr(), + size_m_tensor.data_ptr(), + size_m, size_n, size_k, workspace.data_ptr(), num_bits, + has_act_order, is_k_full, num_groups, group_size, dev, + at::cuda::getCurrentCUDAStream(dev), thread_k, thread_n, sms, + gptq_marlin::max_par); + } + else { + TORCH_CHECK(false, + "gpt_marlin_gemm only supports bfloat16 and float16"); + } + + return c; +} + +#endif \ No newline at end of file diff --git a/archive/csrc/custom_marlin/gptq_marlin/gptq_marlin.cuh b/archive/csrc/custom_marlin/gptq_marlin/gptq_marlin.cuh new file mode 100644 index 0000000..5b4b059 --- /dev/null +++ b/archive/csrc/custom_marlin/gptq_marlin/gptq_marlin.cuh @@ -0,0 +1,76 @@ +// Adapted from +// https://github.com/vllm-project/vllm/tree/main/csrc/quantization/gptq_marlin +// Copyrigth 2024 The vLLM team. +// Copyright (c) 2024 by KVCache.AI, All Rights Reserved. +#pragma once + +#include + +#include +#include +#include +#include +#include +#include + +namespace gptq_marlin { + +// 8 warps are a good choice since every SM has 4 schedulers and having more +// than 1 warp per schedule allows some more latency hiding. At the same time, +// we want relatively few warps to have many registers per warp and small tiles. +static constexpr int default_threads = 256; + +static constexpr int pipe_stages = + 4; // 4 pipeline stages fit into shared memory + +static constexpr int min_thread_n = 64; +static constexpr int min_thread_k = 64; + +static constexpr int tile_size = 16; +static constexpr int max_par = 16; + +template struct Vec { + T elems[n]; + __device__ T &operator[](int i) { return elems[i]; } +}; + +using I4 = Vec; + +constexpr int div_ceil(int a, int b) { return (a + b - 1) / b; } + +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800 +// No support for async +#else + +__device__ inline void cp_async4_pred(void *smem_ptr, const void *glob_ptr, + bool pred = true) { + const int BYTES = 16; + uint32_t smem = static_cast(__cvta_generic_to_shared(smem_ptr)); + asm volatile("{\n" + " .reg .pred p;\n" + " setp.ne.b32 p, %0, 0;\n" + " @p cp.async.cg.shared.global [%1], [%2], %3;\n" + "}\n" ::"r"((int)pred), + "r"(smem), "l"(glob_ptr), "n"(BYTES)); +} + +__device__ inline void cp_async4(void *smem_ptr, const void *glob_ptr) { + const int BYTES = 16; + uint32_t smem = static_cast(__cvta_generic_to_shared(smem_ptr)); + asm volatile("{\n" + " cp.async.cg.shared.global [%0], [%1], %2;\n" + "}\n" ::"r"(smem), + "l"(glob_ptr), "n"(BYTES)); +} + +__device__ inline void cp_async_fence() { + asm volatile("cp.async.commit_group;\n" ::); +} + +template __device__ inline void cp_async_wait() { + asm volatile("cp.async.wait_group %0;\n" ::"n"(n)); +} + +#endif + +} // namespace gptq_marlin \ No newline at end of file diff --git a/archive/csrc/custom_marlin/gptq_marlin/gptq_marlin_dtypes.cuh b/archive/csrc/custom_marlin/gptq_marlin/gptq_marlin_dtypes.cuh new file mode 100644 index 0000000..3e8c3ca --- /dev/null +++ b/archive/csrc/custom_marlin/gptq_marlin/gptq_marlin_dtypes.cuh @@ -0,0 +1,77 @@ +// Adapted from +// https://github.com/vllm-project/vllm/tree/main/csrc/quantization/gptq_marlin +// Copyrigth 2024 The vLLM team. +// Copyright (c) 2024 by KVCache.AI, All Rights Reserved. +#ifndef _data_types_cuh +#define _data_types_cuh +#include "gptq_marlin.cuh" +#include +#include + +namespace gptq_marlin { + +template class ScalarType {}; + +template <> class ScalarType { + public: + using scalar_t = half; + using scalar_t2 = half2; + + // Matrix fragments for tensor core instructions; their precise layout is + // documented here: + // https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#matrix-fragments-for-mma-m16n8k16-with-floating-point-type + using FragA = Vec; + using FragB = Vec; + using FragC = Vec; + using FragS = Vec; + + static __device__ float inline num2float(const half x) { + return __half2float(x); + } + + static __device__ half2 inline num2num2(const half x) { + return __half2half2(x); + } + + static __device__ half2 inline nums2num2(const half x1, const half x2) { + return __halves2half2(x1, x2); + } + + static __host__ __device__ half inline float2num(const float x) { + return __float2half(x); + } +}; + +template <> class ScalarType { + public: + using scalar_t = nv_bfloat16; + using scalar_t2 = nv_bfloat162; + + using FragA = Vec; + using FragB = Vec; + using FragC = Vec; + using FragS = Vec; + +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800 + static __device__ float inline num2float(const nv_bfloat16 x) { + return __bfloat162float(x); + } + + static __device__ nv_bfloat162 inline num2num2(const nv_bfloat16 x) { + return __bfloat162bfloat162(x); + } + + static __device__ nv_bfloat162 inline nums2num2(const nv_bfloat16 x1, + const nv_bfloat16 x2) { + return __halves2bfloat162(x1, x2); + } + + static __host__ __device__ nv_bfloat16 inline float2num(const float x) { + return __float2bfloat16(x); + } +#endif +}; + +} // namespace gptq_marlin + +#endif \ No newline at end of file diff --git a/archive/csrc/custom_marlin/gptq_marlin/gptq_marlin_repack.cu b/archive/csrc/custom_marlin/gptq_marlin/gptq_marlin_repack.cu new file mode 100644 index 0000000..4adcbd5 --- /dev/null +++ b/archive/csrc/custom_marlin/gptq_marlin/gptq_marlin_repack.cu @@ -0,0 +1,350 @@ +#include "gptq_marlin.cuh" + +namespace gptq_marlin { + +static constexpr int repack_stages = 8; + +static constexpr int repack_threads = 256; + +static constexpr int tile_k_size = tile_size; +static constexpr int tile_n_size = tile_k_size * 4; + +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800 + +template +__global__ void marlin_repack_kernel( + uint32_t const* __restrict__ b_q_weight_ptr, + uint32_t const* __restrict__ perm_ptr, uint32_t* __restrict__ out_ptr, + int size_k, int size_n) {} + +} // namespace gptq_marlin + +torch::Tensor gptq_marlin_repack(torch::Tensor& b_q_weight, torch::Tensor& perm, + int64_t size_k, int64_t size_n, + int64_t num_bits) { + TORCH_CHECK_NOT_IMPLEMENTED( + false, "marlin_repack_from_gptq(..) requires CUDA_ARCH >= 8.0"); + return torch::empty({1, 1}); +} + +#else + +template +__global__ void marlin_repack_kernel( + uint32_t const* __restrict__ b_q_weight_ptr, + uint32_t const* __restrict__ perm_ptr, uint32_t* __restrict__ out_ptr, + int size_k, int size_n) { + constexpr int pack_factor = 32 / num_bits; + + int k_tiles = size_k / tile_k_size; + int n_tiles = size_n / tile_n_size; + int block_k_tiles = div_ceil(k_tiles, gridDim.x); + + int start_k_tile = blockIdx.x * block_k_tiles; + if (start_k_tile >= k_tiles) { + return; + } + + int finish_k_tile = min(start_k_tile + block_k_tiles, k_tiles); + + // Wait until the next thread tile has been loaded to shared memory. + auto wait_for_stage = [&]() { + // We only have `stages - 2` active fetches since we are double buffering + // and can only issue the next fetch when it is guaranteed that the previous + // shared memory load is fully complete (as it may otherwise be + // overwritten). + cp_async_wait(); + __syncthreads(); + }; + + extern __shared__ int4 sh[]; + + constexpr int perm_size = tile_k_size / 4; + + int4* sh_perm_ptr = sh; + int4* sh_pipe_ptr = sh_perm_ptr; + if constexpr (has_perm) { + sh_pipe_ptr += perm_size; + } + + constexpr int tile_ints = tile_k_size / pack_factor; + + constexpr int stage_n_threads = tile_n_size / 4; + constexpr int stage_k_threads = has_perm ? tile_k_size : tile_ints; + constexpr int stage_size = stage_k_threads * stage_n_threads; + + auto load_perm_to_shared = [&](int k_tile_id) { + int first_k_int4 = (k_tile_id * tile_k_size) / 4; + + int4 const* perm_int4_ptr = reinterpret_cast(perm_ptr); + + if (threadIdx.x < perm_size) { + sh_perm_ptr[threadIdx.x] = perm_int4_ptr[first_k_int4 + threadIdx.x]; + } + __syncthreads(); + }; + + auto fetch_to_shared = [&](int pipe, int k_tile_id, int n_tile_id) { + if (n_tile_id >= n_tiles) { + cp_async_fence(); + return; + } + + int first_n = n_tile_id * tile_n_size; + + int4* sh_ptr = sh_pipe_ptr + stage_size * pipe; + + if constexpr (has_perm) { + if (threadIdx.x < stage_size) { + int k_id = threadIdx.x / stage_n_threads; + int n_id = threadIdx.x % stage_n_threads; + + uint32_t const* sh_perm_int_ptr = + reinterpret_cast(sh_perm_ptr); + + int src_k = sh_perm_int_ptr[k_id]; + int src_k_packed = src_k / pack_factor; + + cp_async4( + &sh_ptr[k_id * stage_n_threads + n_id], + reinterpret_cast(&( + b_q_weight_ptr[src_k_packed * size_n + first_n + (n_id * 4)]))); + } + + } else { + if (threadIdx.x < stage_size) { + int k_id = threadIdx.x / stage_n_threads; + int n_id = threadIdx.x % stage_n_threads; + + int first_k = k_tile_id * tile_k_size; + int first_k_packed = first_k / pack_factor; + + cp_async4(&sh_ptr[k_id * stage_n_threads + n_id], + reinterpret_cast( + &(b_q_weight_ptr[(first_k_packed + k_id) * size_n + + first_n + (n_id * 4)]))); + } + } + + cp_async_fence(); + }; + + auto repack_tile = [&](int pipe, int k_tile_id, int n_tile_id) { + if (n_tile_id >= n_tiles) { + return; + } + + int warp_id = threadIdx.x / 32; + int th_id = threadIdx.x % 32; + + if (warp_id >= 4) { + return; + } + + int tc_col = th_id / 4; + int tc_row = (th_id % 4) * 2; + + constexpr int tc_offsets[4] = {0, 1, 8, 9}; + + int cur_n = warp_id * 16 + tc_col; + + constexpr int sh_stride = 64; + constexpr uint32_t mask = (1 << num_bits) - 1; + + int4* sh_stage_ptr = sh_pipe_ptr + stage_size * pipe; + uint32_t* sh_stage_int_ptr = reinterpret_cast(sh_stage_ptr); + + uint32_t* sh_perm_int_ptr = reinterpret_cast(sh_perm_ptr); + + uint32_t vals[8]; + + if constexpr (has_perm) { + for (int i = 0; i < 4; i++) { + int k_idx = tc_row + tc_offsets[i]; + + uint32_t src_k = sh_perm_int_ptr[k_idx]; + uint32_t src_k_pos = src_k % pack_factor; + + uint32_t b1_val = sh_stage_int_ptr[k_idx * sh_stride + cur_n]; + uint32_t b1_cur_val = (b1_val >> (src_k_pos * num_bits)) & mask; + + uint32_t b2_val = sh_stage_int_ptr[k_idx * sh_stride + cur_n + 8]; + uint32_t b2_cur_val = (b2_val >> (src_k_pos * num_bits)) & mask; + + vals[i] = b1_cur_val; + vals[4 + i] = b2_cur_val; + } + + } else { + uint32_t b1_vals[tile_ints]; + uint32_t b2_vals[tile_ints]; + + #pragma unroll + for (int i = 0; i < tile_ints; i++) { + b1_vals[i] = sh_stage_int_ptr[cur_n + sh_stride * i]; + b2_vals[i] = sh_stage_int_ptr[cur_n + 8 + sh_stride * i]; + } + + #pragma unroll + for (int i = 0; i < 4; i++) { + int cur_elem = tc_row + tc_offsets[i]; + int cur_int = cur_elem / pack_factor; + int cur_pos = cur_elem % pack_factor; + + vals[i] = (b1_vals[cur_int] >> (cur_pos * num_bits)) & mask; + vals[4 + i] = (b2_vals[cur_int] >> (cur_pos * num_bits)) & mask; + } + } + + constexpr int tile_size = tile_k_size * tile_n_size / pack_factor; + int out_offset = (k_tile_id * n_tiles + n_tile_id) * tile_size; + + // Result of: + // https://github.com/NVIDIA/FasterTransformer/blob/main/src/fastertransformer/cutlass_extensions/include/cutlass_extensions/interleaved_numeric_conversion.h + if constexpr (num_bits == 4) { + constexpr int pack_idx[8] = {0, 2, 4, 6, 1, 3, 5, 7}; + + uint32_t res = 0; + #pragma unroll + for (int i = 0; i < 8; i++) { + res |= vals[pack_idx[i]] << (i * 4); + } + + out_ptr[out_offset + th_id * 4 + warp_id] = res; + + } else { + constexpr int pack_idx[4] = {0, 2, 1, 3}; + + uint32_t res1 = 0; + uint32_t res2 = 0; + #pragma unroll + for (int i = 0; i < 4; i++) { + res1 |= vals[pack_idx[i]] << (i * 8); + res2 |= vals[4 + pack_idx[i]] << (i * 8); + } + + out_ptr[out_offset + th_id * 8 + (warp_id * 2) + 0] = res1; + out_ptr[out_offset + th_id * 8 + (warp_id * 2) + 1] = res2; + } + }; + + auto start_pipes = [&](int k_tile_id, int n_tile_id) { + #pragma unroll + for (int pipe = 0; pipe < repack_stages - 1; pipe++) { + fetch_to_shared(pipe, k_tile_id, n_tile_id + pipe); + } + + wait_for_stage(); + }; + #pragma unroll + for (int k_tile_id = start_k_tile; k_tile_id < finish_k_tile; k_tile_id++) { + int n_tile_id = 0; + + if constexpr (has_perm) { + load_perm_to_shared(k_tile_id); + } + + start_pipes(k_tile_id, n_tile_id); + + while (n_tile_id < n_tiles) { + #pragma unroll + for (int pipe = 0; pipe < repack_stages; pipe++) { + fetch_to_shared((pipe + repack_stages - 1) % repack_stages, k_tile_id, + n_tile_id + pipe + repack_stages - 1); + repack_tile(pipe, k_tile_id, n_tile_id + pipe); + wait_for_stage(); + } + n_tile_id += repack_stages; + } + } +} + +} // namespace gptq_marlin + + #define CALL_IF(NUM_BITS, HAS_PERM) \ + else if (num_bits == NUM_BITS && has_perm == HAS_PERM) { \ + cudaFuncSetAttribute( \ + gptq_marlin::marlin_repack_kernel, \ + cudaFuncAttributeMaxDynamicSharedMemorySize, max_shared_mem); \ + gptq_marlin::marlin_repack_kernel \ + <<>>( \ + b_q_weight_ptr, perm_ptr, out_ptr, size_k, size_n); \ + } + +torch::Tensor gptq_marlin_repack(torch::Tensor& b_q_weight, torch::Tensor& perm, + int64_t size_k, int64_t size_n, + int64_t num_bits) { + // Verify compatibility with marlin tile of 16x64 + TORCH_CHECK(size_k % gptq_marlin::tile_k_size == 0, "size_k = ", size_k, + " is not divisible by tile_k_size = ", gptq_marlin::tile_k_size); + TORCH_CHECK(size_n % gptq_marlin::tile_n_size == 0, "size_n = ", size_n, + " is not divisible by tile_n_size = ", gptq_marlin::tile_n_size); + + TORCH_CHECK(num_bits == 4 || num_bits == 8, + "num_bits must be 4 or 8. Got = ", num_bits); + int const pack_factor = 32 / num_bits; + + // Verify B + TORCH_CHECK((size_k / pack_factor) == b_q_weight.size(0), + "Shape mismatch: b_q_weight.size(0) = ", b_q_weight.size(0), + ", size_k = ", size_k, ", pack_factor = ", pack_factor); + TORCH_CHECK(b_q_weight.size(1) == size_n, + "b_q_weight.size(1) = ", b_q_weight.size(1), + " is not size_n = ", size_n); + + // Verify device and strides + TORCH_CHECK(b_q_weight.device().is_cuda(), "b_q_weight is not on GPU"); + TORCH_CHECK(b_q_weight.is_contiguous(), "b_q_weight is not contiguous"); + TORCH_CHECK(b_q_weight.dtype() == at::kInt, "b_q_weight type is not kInt"); + + TORCH_CHECK(perm.device().is_cuda(), "perm is not on GPU"); + TORCH_CHECK(perm.is_contiguous(), "perm is not contiguous"); + TORCH_CHECK(perm.dtype() == at::kInt, "perm type is not at::kInt"); + + // Alloc buffers + const at::cuda::OptionalCUDAGuard device_guard(device_of(b_q_weight)); + auto options = torch::TensorOptions() + .dtype(b_q_weight.dtype()) + .device(b_q_weight.device()); + torch::Tensor out = + torch::empty({size_k / gptq_marlin::tile_size, + size_n * gptq_marlin::tile_size / pack_factor}, + options); + + // Detect if there is act_order + bool has_perm = perm.size(0) != 0; + + // Get ptrs + uint32_t const* b_q_weight_ptr = + reinterpret_cast(b_q_weight.data_ptr()); + uint32_t const* perm_ptr = reinterpret_cast(perm.data_ptr()); + uint32_t* out_ptr = reinterpret_cast(out.data_ptr()); + + // Get dev info + int dev = b_q_weight.get_device(); + cudaStream_t stream = at::cuda::getCurrentCUDAStream(dev); + int blocks; + cudaDeviceGetAttribute(&blocks, cudaDevAttrMultiProcessorCount, dev); + + int max_shared_mem = 0; + cudaDeviceGetAttribute(&max_shared_mem, + cudaDevAttrMaxSharedMemoryPerBlockOptin, dev); + TORCH_CHECK(max_shared_mem > 0); + + if (false) { + } + CALL_IF(4, false) + CALL_IF(4, true) + CALL_IF(8, false) + CALL_IF(8, true) + else { + TORCH_CHECK(false, "Unsupported repack config: num_bits = ", num_bits, + ", has_perm = ", has_perm); + } + + return out; +} + +#endif \ No newline at end of file diff --git a/archive/csrc/custom_marlin/gptq_marlin/ops.h b/archive/csrc/custom_marlin/gptq_marlin/ops.h new file mode 100644 index 0000000..b3327d3 --- /dev/null +++ b/archive/csrc/custom_marlin/gptq_marlin/ops.h @@ -0,0 +1,24 @@ +/** + * @Description : + * @Author : Azure + * @Date : 2024-07-22 09:27:55 + * @Version : 1.0.0 + * @LastEditors : Azure + * @LastEditTime : 2024-07-26 08:35:00 + * @Copyright (c) 2024 by KVCache.AI, All Rights Reserved. + **/ +#pragma once + +#include +#include +#include + +torch::Tensor gptq_marlin_gemm(torch::Tensor &a, torch::Tensor &b_q_weight, + torch::Tensor &b_scales, torch::Tensor &g_idx, + torch::Tensor &perm, torch::Tensor &workspace, + int64_t num_bits, torch::Tensor size_m_tensor, int64_t size_m, int64_t size_n, + int64_t size_k, int sms, bool is_k_full); + +torch::Tensor gptq_marlin_repack(torch::Tensor& b_q_weight, torch::Tensor&perm, + int64_t size_k, int64_t size_n, + int64_t num_bits); \ No newline at end of file diff --git a/archive/csrc/custom_marlin/setup.py b/archive/csrc/custom_marlin/setup.py new file mode 100644 index 0000000..ec710ca --- /dev/null +++ b/archive/csrc/custom_marlin/setup.py @@ -0,0 +1,25 @@ +from setuptools import setup, Extension +from torch.utils import cpp_extension +from torch.utils.cpp_extension import BuildExtension, CUDAExtension +setup( + name='vLLMMarlin', + ext_modules=[ + CUDAExtension( + 'vLLMMarlin', [ + #'custom_gguf/dequant.cu', + 'binding.cpp', + 'gptq_marlin/gptq_marlin.cu', + 'gptq_marlin/gptq_marlin_repack.cu', + ], + extra_compile_args={ + 'cxx': ['-O3'], + 'nvcc': [ + '-O3', + '--use_fast_math', + '-Xcompiler', '-fPIC', + ] + }, + ) + ], + cmdclass={'build_ext': BuildExtension} +) \ No newline at end of file diff --git a/archive/csrc/custom_marlin/test_cuda_graph.py b/archive/csrc/custom_marlin/test_cuda_graph.py new file mode 100644 index 0000000..0024082 --- /dev/null +++ b/archive/csrc/custom_marlin/test_cuda_graph.py @@ -0,0 +1,335 @@ +import csv +import torch +import torch.nn as nn +import vLLMMarlin +torch.set_grad_enabled(False) +from utils.marlin_utils import ( + MarlinWorkspace, + marlin_quantize, + GPTQ_MARLIN_MIN_THREAD_N, + GPTQ_MARLIN_MIN_THREAD_K, + GPTQ_MARLIN_MAX_PARALLEL, +) + +def setup_seed(seed): + torch.manual_seed(seed) + torch.cuda.manual_seed_all(seed) + +setup_seed(20241223) + +torch.set_grad_enabled(False) +torch.set_default_dtype(torch.bfloat16) +global_dtype=torch.bfloat16 +global_device=torch.device("cuda",0) +global_num_cases:int=int(50) +torch.cuda.set_device(0) +torch.backends.cudnn.enabled =True +torch.backends.cudnn.benchmark = True + +max_batch_size = 512 +max_tp = 8 +L2_size = 73728 * 1024 + +def get_usable_mem(): + properties = torch.cuda.get_device_properties(global_device) + #print(f"Total memory: {properties.total_memory / (1024 ** 3):.2f} GB") + allocated_memory = torch.cuda.memory_allocated(global_device) + #print(f"Currently allocated memory: {allocated_memory / (1024 ** 2):.2f} MB") + reserved_memory = torch.cuda.memory_reserved(global_device) + #print(f"Currently reserved memory: {reserved_memory / (1024 ** 2):.2f} MB") + return properties.total_memory - 512 * 1024 ** 2 - allocated_memory# - reserved_memory + +def exp_range(start, stop, step = 2): + now = start + while now <= stop: + yield now + now *= step + +def timing(func, iters, epochs=100): + #warmup + for idx in range(iters): + func(idx) + + torch.cuda.synchronize() + cuda_graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(cuda_graph): + for idx in range(iters): + func(idx) + + for _ in range(2000): + cuda_graph.replay() + + start_event = torch.cuda.Event(enable_timing=True) + end_event = torch.cuda.Event(enable_timing=True) + stream = torch.cuda.Stream() + torch.cuda.synchronize() + #with torch.cuda.stream(stream): + start_event.record() + for _ in range(10): + cuda_graph.replay() + end_event.record() + torch.cuda.synchronize() + elapsed_time_ms0 = start_event.elapsed_time(end_event) + + start_event = torch.cuda.Event(enable_timing=True) + end_event = torch.cuda.Event(enable_timing=True) + torch.cuda.synchronize() + #with torch.cuda.stream(stream): + start_event.record() + for _ in range(epochs+10): + cuda_graph.replay() + end_event.record() + torch.cuda.synchronize() + elapsed_time_ms = start_event.elapsed_time(end_event) - elapsed_time_ms0 + + #print(elapsed_time_ms0, elapsed_time_ms) + return elapsed_time_ms/iters/epochs + +class LinearMarlin(nn.Linear): + marlin_q_w: torch.Tensor + marlin_s: torch.Tensor + g_idx: torch.Tensor + sort_indices: torch.Tensor + has_bias: bool + def __init__( + self, + in_features, + out_features, + bias = False, + device: str = "cuda", + num_bits: int = 4, # 4-bit/8-bit is supported + group_size: int = 64, # -1, 32, 64, 128 + act_order: bool = False, + is_k_full=True, + sms = -1, # sms in GPU + **kwargs, + ): + self.padding = False + assert device.lower() != "cpu", "Marlin quantized linear only supports GPU device" + if in_features%GPTQ_MARLIN_MIN_THREAD_K!=0 or out_features%GPTQ_MARLIN_MIN_THREAD_K!=0: + #print(f"warning!, in_features={in_features} or out_features={out_features} is undivisible by GPTQ_MARLIN_MIN_THREAD_K={GPTQ_MARLIN_MIN_THREAD_K} and GPTQ_MARLIN_MIN_THREAD_N={GPTQ_MARLIN_MIN_THREAD_N}, padding") + self.padding = True + self.orin_in_features = in_features + self.orin_out_features = out_features + in_features = (in_features+GPTQ_MARLIN_MIN_THREAD_K-1)//GPTQ_MARLIN_MIN_THREAD_K*GPTQ_MARLIN_MIN_THREAD_K + out_features = (out_features+GPTQ_MARLIN_MIN_THREAD_N-1)//GPTQ_MARLIN_MIN_THREAD_N*GPTQ_MARLIN_MIN_THREAD_N + #print(f"After padding: in_features={in_features}, out_features={out_features}") + + + super().__init__(in_features, out_features, bias, device) + self.has_bias = bias + self.device = device + self.num_bits = num_bits + self.group_size = group_size + self.act_order = act_order + # TODO: optimize every shape GEMM + + blocks_k, blocks_n = in_features//128, out_features//128 + + self.sms = sms + + self.is_k_full = is_k_full + + self.weight.requires_grad = False + self.weight.t_() + # Pack Marlin linear + #w_ref, marlin_q_w, marlin_s, g_idx, sort_indices, _ = marlin_quantize( + # self.weight, self.num_bits, self.group_size, self.act_order + #) + marlin_q_w = torch.randint(int(-1e9), int(1e9), (in_features//16, out_features*2), device=device, dtype=torch.int) + marlin_s = torch.randn((in_features//64, out_features), device=device) + self.workspace = MarlinWorkspace( + self.out_features, GPTQ_MARLIN_MIN_THREAD_N, GPTQ_MARLIN_MAX_PARALLEL, self.device + ) + self.marlin_q_w = marlin_q_w + self.marlin_s = marlin_s + self.g_idx = torch.empty((0), dtype=torch.int32, device=self.device) + self.sort_indices = torch.empty((0), dtype=torch.int32, device=self.device) + self.k = self.weight.shape[0] + self.n = self.weight.shape[1] + self.weight = None + """ + print(in_features, out_features) + print(marlin_q_w.shape) + print(marlin_q_w.dtype) + print(marlin_s.shape) + print(marlin_s.dtype) + print(self.workspace.scratch.shape) + print(self.workspace.scratch.dtype) + print(self.g_idx.shape) + print(self.g_idx.dtype) + print(self.sort_indices.shape) + print(self.sort_indices.dtype) + #print(w_ref.shape) + #print(w_ref.dtype) + """ + #w_ref = None + + def forward(self, x: torch.Tensor, bsz_tensor: torch.Tensor) -> torch.Tensor: + # Only support input x as BF16 and FP16 + x = x.to(self.device) + orig_shape = list(x.shape) + orig_dtype = x.dtype + x = x.reshape(-1, x.shape[-1]) + if self.padding: + padding_input=torch.empty(x.shape[0], self.in_features, device=x.device, dtype=x.dtype) + padding_input[:,:self.orin_in_features] = x + x = padding_input + marlin_s = self.marlin_s.to(x.dtype) + #print(self.sms * ((orig_shape[0]+63)//64)) + + sms = self.sms + + x = vLLMMarlin.gptq_marlin_gemm( + x, + self.marlin_q_w, + marlin_s, + self.g_idx, + self.sort_indices, + self.workspace.scratch, + self.num_bits, + bsz_tensor, + x.shape[0], + self.n, + x.shape[-1], + sms, + self.is_k_full, + ) + # TODO: don't padding bias + if self.has_bias: + x = x + self.bias + if self.padding: + x = x[:,:self.orin_out_features] + orig_shape[-1] = self.orin_out_features + else: + orig_shape[-1] = self.out_features + return x.reshape(orig_shape).to(orig_dtype) + +def benchLinearMarlin(input_dim, output_dim):#, out_file + print("benchmarking MLP Marlin") + print("-----------------------------------------------------------") + headers = ["batch_size", "tp", "used_time", "bandwidth GB/s", "TFLOPS", "cases", "padding", "sms"] + print(" | ".join(headers) + "\n") + rows = [] + for batch_size in exp_range(1, 64): + for tp in exp_range(1, max_tp): + torch.cuda.empty_cache() + if output_dim % tp != 0: + continue + cur_output_dim = output_dim // tp + modules = [] + inputs = [] + data_size = int(0.53125*input_dim*cur_output_dim) + input_size = int(2*batch_size*input_dim) + output_size = int(2*batch_size*cur_output_dim) + usable_mem = get_usable_mem() - 2 * input_dim * cur_output_dim + min_cases = max(global_num_cases, (2*L2_size) // (data_size+input_size)) + cases = int(min(min_cases, (usable_mem * 0.8) // (data_size+input_size))) + #print(usable_mem, data_size, input_size, cases) + + bsz_tensor = torch.tensor([batch_size], device=global_device, dtype=torch.int32) + + if cases == 0: + row = [f"{batch_size}", "OOM", "OOM", "OOM", "0", "False"] + rows.append(row) + break + for _ in range(cases): + modules.append(LinearMarlin(input_dim, cur_output_dim, sms=56, non_equal_division=False).to(device=global_device).eval()) + inputs.append(torch.randn(batch_size, 1, input_dim, device=global_device)) + + def forward(case_id): + modules[case_id](inputs[case_id], bsz_tensor) + + used_time = timing(forward, iters=cases) + bandwidth = (data_size+input_size+output_size)/used_time/1e6 + flops = 2*batch_size*input_dim*cur_output_dim + tflops = flops/used_time/1e9 + cur_sms = modules[0].sms + row = [f"{batch_size}", f"{tp}", f"{used_time}", f"{bandwidth}", f"{tflops}", f"{cases}", modules[0].padding, cur_sms] + rows.append(row) + print(f"{batch_size}", f"{tp}", f"{used_time}", f"{bandwidth}", f"{tflops}", f"{cases}", modules[0].padding, cur_sms) + + """ + with open(out_file, 'w', newline='') as csvfile: + csvwriter = csv.writer(csvfile) + csvwriter.writerow(headers) + for row in rows: + csvwriter.writerow(row) + """ + + """ + markdown_table = " | ".join(headers) + "\n" + markdown_table += " | ".join(["---"] * len(headers)) + "\n" + for row in rows: + markdown_table += " | ".join(row) + "\n" + + print(markdown_table) + """ + #print("finish write file", out_file) + #print("-------------------------------------------------------------") + +if __name__ == "__main__": + + benchLinearMarlin(5120, 3584) + exit(0) + + max_batch = 1 + cur_batch = 1 + + + marlin_linear = LinearMarlin(5120, 3584) + + input_tensor = torch.randn(max_batch, 1, 5120, device="cuda", dtype=torch.bfloat16) + bsz_tensor = torch.tensor([max_batch], device="cuda", dtype=torch.int32) + + out_truth = marlin_linear(input_tensor, bsz_tensor) + + print(out_truth) + + g = torch.cuda.CUDAGraph() + with torch.cuda.graph(g): + out_buf = marlin_linear(input_tensor, bsz_tensor) + + for i in range(10000): + g.replay() + + #torch.testing.assert_close(out_buf, out_truth, rtol=1e-3, atol=1e-3) + + marlin_linear = LinearMarlin(5120, 3584) + g = torch.cuda.CUDAGraph() + with torch.cuda.graph(g): + out_buf = marlin_linear(input_tensor, bsz_tensor) + + new_input = torch.randn(cur_batch, 1, 5120, device="cuda", dtype=torch.bfloat16) + bsz_tensor.copy_(torch.tensor([cur_batch], device="cuda", dtype=torch.int32)) + + new_out_truth = marlin_linear(new_input, bsz_tensor) + input_tensor[:cur_batch].copy_(new_input) + input_tensor[cur_batch:] = 0 + + g.replay() + + torch.cuda.synchronize() + + def printMinMax(tensor): + abs_tensor = torch.abs(tensor) + + min_val = torch.min(abs_tensor) + max_val = torch.max(abs_tensor) + + min_indices = (abs_tensor == min_val).nonzero(as_tuple=True) + max_indices = (abs_tensor == max_val).nonzero(as_tuple=True) + + print(f"min: {min_val.item()}") + print(f"min idx: {min_indices}") + print(f"max: {max_val.item()}") + print(f"max idx: {max_indices}") + + print(out_buf[:cur_batch].shape) + print(new_out_truth.shape) + + + printMinMax(out_buf[:cur_batch]) + printMinMax(new_out_truth) + + #torch.testing.assert_close(out_buf[:cur_batch, 0, :], new_out_truth[:cur_batch, 0, :], rtol=1e-3, atol=1e-3) diff --git a/archive/csrc/custom_marlin/utils/__init__.py b/archive/csrc/custom_marlin/utils/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/archive/csrc/custom_marlin/utils/format24.py b/archive/csrc/custom_marlin/utils/format24.py new file mode 100644 index 0000000..2434e79 --- /dev/null +++ b/archive/csrc/custom_marlin/utils/format24.py @@ -0,0 +1,308 @@ +# +# Modified by Roberto Lopez Castro (roberto.lopez.castro@udc.es). +# + +import torch + + +# This is PyTorch implementation of main part of reorder_meta() +# function, from tools/util/include/cutlass/util/host_reorder.h file +# of CUTLASS source tree. Furthermore, CUTLASS template for sparse +# GEMM decides upon layout of this matrix, and at the moment for the +# sparse GEMM executed on tensor cores, this is layout described by +# ColumnMajorInterleaved<2> data structure, in +# include/cutlass/layout/matrix.h of CUTLASS source tree. The +# reordering of meta matrix into meta_reordered matrix calculated +# according to these segments of CUTLASS code is re-implemented here. +# Note that this calculation produces offsets for scattering metadata +# matrix elements into reordered metadata matrix elements (or, +# equivalently, for gathering reordered metadata matrix element back +# into metadata matrix elements). +def _calculate_meta_reordering_scatter_offsets(m, meta_ncols, meta_dtype, + device): + dst_rows = torch.arange(0, m, device=device)[:, None].repeat(1, meta_ncols) + dst_cols = torch.arange(0, meta_ncols, device=device).repeat(m, 1) + + # Reorder the rows, then swizzle the 2x2 blocks. + group_x = 64 + group_y = 32 if meta_dtype.itemsize == 2 else 16 + + dst_rows = (dst_rows // group_x * group_x + (dst_rows % 2) * 2 + + (dst_rows % 8) // 4 + ((dst_rows % group_y) % 4) // 2 * 32 + + ((dst_rows % group_x) // 8) * 4) + + topright = ((dst_rows % 2 == 0) & (dst_cols % 2 == 1)).to(torch.int8) + bottomleft = ((dst_rows % 2 == 1) & (dst_cols % 2 == 0)).to(torch.int8) + dst_rows += topright - bottomleft + dst_cols -= topright - bottomleft + + # Assumed that meta tensor is to be stored in CUTLASS + # InterleavedColumnMajor layout, and reverse engineered + # corresponding code to store values into this tensor. + interleave = 2 + cols_maj = dst_cols // interleave + cols_min = dst_cols % interleave + return (cols_maj * m * interleave + dst_rows * interleave + + cols_min).view(-1) + + +# This function converts dense matrix into sparse semi-structured +# representation, producing "compressed" matrix, in the layout used by +# CUTLASS backend, and corresponding metadata matrix. +def sparse_semi_structured_from_dense_cutlass(dense): + if dense.dim() != 2: + raise RuntimeError( + f"Expected 2-dimensional dense tensor, got {dense.dim()}-dimensional tensor" # noqa: E501 + ) + + m, k = dense.shape + device = dense.device + + meta_dtype = torch.int8 + if dense.dtype == torch.int8: + meta_dtype = torch.int32 + elif dense.dtype in [torch.half, torch.bfloat16, torch.float, torch.int32]: + meta_dtype = torch.int16 + else: + raise RuntimeError(f"Invalid datatype {dense.dtype} of dense matrix") + quadbits_per_meta_elem = meta_dtype.itemsize * 8 // 4 + if quadbits_per_meta_elem not in (4, 8): + raise RuntimeError( + "Invalid number of elements per meta element calculated") + + if meta_dtype == torch.int32: + if m % 16 != 0: + raise RuntimeError( + f"Number of rows of dense matrix {m} must be divisible by 16") + else: + if m % 32 != 0: + raise RuntimeError( + f"Number of rows of dense matrix {m} must be divisible by 32") + if k % (4 * quadbits_per_meta_elem) != 0: + raise RuntimeError( + f"Number of columns of dense matrix {k} must be divisible by {4 * quadbits_per_meta_elem}" # noqa: E501 + ) + + if dense.dtype != torch.float: + ksparse = 4 + dense_4 = dense.view(-1, k // ksparse, ksparse) + m0, m1, m2, m3 = (dense_4 != 0).unbind(-1) + else: + ksparse = 2 + dense_2 = dense.view(-1, k // ksparse, ksparse) + m0, m2 = m1, m3 = (dense_2 != 0).unbind(-1) + meta_ncols = k // (ksparse * quadbits_per_meta_elem) + + # Encoding quadruples of True/False values as follows: + # [True, True, False, False] -> 0b0100 + # [True, False, True, False] -> 0b1000 + # [False, True, True, False] -> 0b1001 + # [True, False, False, True ] -> 0b1100 + # [False, True, False, True ] -> 0b1101 + # [False, False, True, True ] -> 0b1110 + # Thus, lower two bits in the encoding are index of the True value + # at the lowest index in the quadruple, and the higher two bits in + # the encoding are index of the other True value in the quadruple. + # In case there are less than two True values, than False value or + # values at some index or indices are considered True for the + # encoding. In case there are more than two True values, then the + # excess True value(s) at some indices are considered False for + # the encoding. The exact encodings used for these cases are as + # follows: + # [False, False, False, False] -> 0b1110 + # [False, False, False, True ] -> 0b1110 + # [False, False, True, False] -> 0b1110 + # [False, True, False, False] -> 0b1001 + # [False, True, True, True ] -> 0b1101 + # [True, False, False, False] -> 0b1000 + # [True, False, True, True ] -> 0b1100 + # [True, True, False, True ] -> 0b0100 + # [True, True, True, False] -> 0b0100 + # [True, True, True, True ] -> 0b0100 + # These particular encodings are chosen, with the help of Espresso + # logic minimizer software, for the purpose of minimization of + # corresponding Boolean functions, that translate non-zero flags + # into encoding bits. Note also possible choices for the first + # and last of these encodings were limited only to (0b0100, + # 0b1110), in order to produce valid encodings for 1:2 sparsity + # case. + + expr0 = m0 & m1 + expr1 = ~m0 & m1 + expr2 = ~m0 & ~m1 + bit0 = expr1 + bit1 = expr2 + bit2 = expr0 | expr2 | m3 + bit3 = expr1 | ~m1 + idxs0 = bit0 | (bit1.to(torch.int64) << 1) + idxs1 = bit2 | (bit3.to(torch.int64) << 1) + + if dense.dtype != torch.float: + sparse0 = dense_4.gather( + -1, idxs0.unsqueeze(-1)) # type: ignore[possibly-undefined] + sparse1 = dense_4.gather(-1, idxs1.unsqueeze(-1)) + sparse = torch.stack((sparse0, sparse1), dim=-1).view(m, k // 2) + else: + sparse = dense_2.gather(-1, + idxs0.unsqueeze(-1) // 2).view( + m, + k // 2) # type: ignore[possibly-undefined] + + meta_4 = idxs0 | (idxs1 << 2) + meta_n = meta_4.view( + (-1, meta_ncols, quadbits_per_meta_elem)).to(meta_dtype) + + if quadbits_per_meta_elem == 4: + meta = (meta_n[:, :, 0] + | (meta_n[:, :, 1] << 4) + | (meta_n[:, :, 2] << 8) + | (meta_n[:, :, 3] << 12)) + elif quadbits_per_meta_elem == 8: + meta = (meta_n[:, :, 0] + | (meta_n[:, :, 1] << 4) + | (meta_n[:, :, 2] << 8) + | (meta_n[:, :, 3] << 12) + | (meta_n[:, :, 4] << 16) + | (meta_n[:, :, 5] << 20) + | (meta_n[:, :, 6] << 24) + | (meta_n[:, :, 7] << 28)) + + # Reorder meta tensor elements. + meta_reordered = meta.new_empty( + (m * meta_ncols, )) # type: ignore[possibly-undefined] + meta_offsets = _calculate_meta_reordering_scatter_offsets( + m, meta_ncols, meta_dtype, device) + meta_reordered.scatter_(0, meta_offsets, meta.view(-1)) + + return (sparse, meta_reordered.view(m, meta_ncols)) + + +# This function performs reverse of the function above - it +# reconstructs dense matrix from a pair of "compressed" matrix, given +# in the layout used by CUTLASS backend, and accompanying metadata +# matrix. +def sparse_semi_structured_to_dense_cutlass(sparse, meta_reordered): + if sparse.dim() != 2: + raise RuntimeError( + f"Expected 2-dimensional sparse tensor, got {sparse.dim()}-dimensional tensor" # noqa: E501 + ) + + m, k = sparse.shape + device = sparse.device + + if meta_reordered.dim() != 2: + raise RuntimeError( + f"Expected 2-dimensional meta tensor, got {meta_reordered.dim()}-dimensional tensor" # noqa: E501 + ) + if meta_reordered.device != device: + raise RuntimeError( + f"Expected meta matrix to be on {device} device, got matrix on {meta_reordered.device} device" # noqa: E501 + ) + + meta_dtype = meta_reordered.dtype + if meta_dtype not in (torch.int16, torch.int32): + raise RuntimeError(f"Invalid datatype {meta_dtype} of meta matrix") + quadbits_per_meta_elem = meta_dtype.itemsize * 8 // 4 + + ksparse = 4 if sparse.dtype != torch.float else 2 + + meta_nrows, meta_ncols = meta_reordered.shape + if meta_nrows != m: + raise RuntimeError( + f"Number of rows of meta matrix {meta_nrows} must be equal to number of columns of spase matrix {m}" # noqa: E501 + ) + if meta_ncols * ksparse * quadbits_per_meta_elem != 2 * k: + raise RuntimeError( + f"Number of columns of sparse matrix {k} different from the {meta_ncols * ksparse * quadbits_per_meta_elem // 2}, " # noqa: E501 + "expected according to the number of columns of meta matrix") + + # Undo meta tensor elements reordering. + meta_offsets = _calculate_meta_reordering_scatter_offsets( + m, meta_ncols, meta_dtype, device) + meta = torch.gather(meta_reordered.view(-1), 0, + meta_offsets).view(m, meta_ncols) + + # Unpack sparse tensor back to original dense tensor, using + # information provided by meta tensor. Note that torch.float + # datatype is handled pretty much the same as + # torch.half/torch.bfloat16, as metadata for a pair of torch.float + # value is encoded as if underlying 8 bytes contain four + # torch.half/torch.bfloat16 values, where either first two or last + # two are zeros. + meta_2 = torch.empty( + (m, meta_ncols, 2 * quadbits_per_meta_elem), + dtype=meta_dtype, + device=device, + ) + if quadbits_per_meta_elem == 4: + meta_2[:, :, 0] = meta & 0b11 + meta_2[:, :, 1] = (meta >> 2) & 0b11 + meta_2[:, :, 2] = (meta >> 4) & 0b11 + meta_2[:, :, 3] = (meta >> 6) & 0b11 + meta_2[:, :, 4] = (meta >> 8) & 0b11 + meta_2[:, :, 5] = (meta >> 10) & 0b11 + meta_2[:, :, 6] = (meta >> 12) & 0b11 + meta_2[:, :, 7] = (meta >> 14) & 0b11 + elif quadbits_per_meta_elem == 8: + meta_2[:, :, 0] = meta & 0b11 + meta_2[:, :, 1] = (meta >> 2) & 0b11 + meta_2[:, :, 2] = (meta >> 4) & 0b11 + meta_2[:, :, 3] = (meta >> 6) & 0b11 + meta_2[:, :, 4] = (meta >> 8) & 0b11 + meta_2[:, :, 5] = (meta >> 10) & 0b11 + meta_2[:, :, 6] = (meta >> 12) & 0b11 + meta_2[:, :, 7] = (meta >> 14) & 0b11 + meta_2[:, :, 8] = (meta >> 16) & 0b11 + meta_2[:, :, 9] = (meta >> 18) & 0b11 + meta_2[:, :, 10] = (meta >> 20) & 0b11 + meta_2[:, :, 11] = (meta >> 22) & 0b11 + meta_2[:, :, 12] = (meta >> 24) & 0b11 + meta_2[:, :, 13] = (meta >> 26) & 0b11 + meta_2[:, :, 14] = (meta >> 28) & 0b11 + meta_2[:, :, 15] = (meta >> 30) & 0b11 + + dense_offsets = meta_2.view(-1) + ( + torch.arange(0, 2 * m * k // ksparse, device=device) * 4).view( + -1, 1).repeat(1, 2).view(-1) + + dense = torch.zeros((m * 2 * k, ), dtype=sparse.dtype, device=device) + if sparse.dtype != torch.float: + # dense.scatter_(0, dense_offsets, sparse.view(-1)) + dense.scatter_(0, dense_offsets, sparse.reshape(-1)) + else: + dense.view(torch.half).scatter_(0, dense_offsets, + sparse.view(torch.half).view(-1)) + + return dense.view(m, 2 * k) + + +def mask_creator(tensor): + """ + Class for creating N:M sparsity masks. + Masks will be created using the N:M ratio, where for every block of + M weights, N will be pruned based on ranked weight value. Each mask + will correspond to the given tensor. + + :param N: The number of weights in a group to keep + :param M: The size of a weight group + """ + N = 2 + M = 4 + + mask = None + # for i, tensor in enumerate(tensors): + if tensor.numel() % M != 0: + raise ValueError( + f"Tensor of size {tensor.shape} can't be evenly divided into " + f"{M} groups") + + num_groups = tensor.numel() // M + + # N:M sparsity for linear layers + tensor_temp = tensor.detach().abs().reshape(num_groups, M) + index = torch.argsort(tensor_temp, dim=1)[:, :int(M - N)] + + w_b = torch.ones(tensor_temp.shape, device=tensor_temp.device) + mask = w_b.scatter_(dim=1, index=index, value=0).reshape(tensor.shape) + + return mask \ No newline at end of file diff --git a/archive/csrc/custom_marlin/utils/marlin_24_perms.py b/archive/csrc/custom_marlin/utils/marlin_24_perms.py new file mode 100644 index 0000000..d79926f --- /dev/null +++ b/archive/csrc/custom_marlin/utils/marlin_24_perms.py @@ -0,0 +1,65 @@ +''' +Date: 2024-11-08 02:46:07 +LastEditors: djw +LastEditTime: 2024-11-08 02:46:41 +''' +"""This file is used for /tests and /benchmarks""" +from typing import Dict, List + +import numpy +import torch + + +# Precompute permutations for Marlin24 weight and scale shuffling # noqa: E501 +# +# Marlin works on [16*2,64] tiles. The goal of the permutations is to reorder the weight data so that it is compatible noqa: # noqa: E501 +# with the tensor-core format that is described here: +# https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#matrix-fragments-for-mma-m16n8k16-with-floating-point-type # noqa: E501 +# +# As a result of this reordering, the vector loads inside the kernel will get the data as it is needed for tensor-core # noqa: E501 +# (without the need to use ldmatrix instructions) # noqa: E501 +def get_perms_24(num_bits: int): + perm_list: List[int] = [] + for i in range(32): + perm1: List[int] = [] + col = i // 4 + col_o = col // 2 + for block in [0, 1]: + for row in [ + 2 * (i % 4), + 2 * (i % 4) + 1, + 2 * (i % 4 + 4), + 2 * (i % 4 + 4) + 1, + ]: + perm1.append(16 * row + col_o * 256 + 8 * (col % 2) + + 4 * block) + for j in range(4): + perm_list.extend([p + 1 * j for p in perm1]) + perm = numpy.array(perm_list) + + if num_bits == 4: + interleave = numpy.array([0, 2, 4, 6, 1, 3, 5, 7]) + elif num_bits == 8: + interleave = numpy.array([0, 2, 1, 3]) + else: + raise ValueError("num_bits must be 4 or 8, got {}".format(num_bits)) + + perm = perm.reshape((-1, len(interleave)))[:, interleave].ravel() + perm = torch.from_numpy(perm) + scale_perm: List[int] = [] + for i in range(8): + scale_perm.extend([i * 8 + j for j in [0, 4, 1, 5, 2, 6, 3, 7]]) + scale_perm_single: List[int] = [] + for i in range(8): + scale_perm_single.extend([8 * i + j for j in [0, 1, 2, 3, 4, 5, 6, 7]]) + return perm, scale_perm, scale_perm_single + + +marlin_24_perm: Dict[int, torch.Tensor] = {} +marlin_24_scale_perm: Dict[int, List[int]] = {} +marlin_24_scale_perm_single: Dict[int, List[int]] = {} +for num_bits in [4, 8]: + perm_24, scale_perm_24, scale_perm_single_24 = get_perms_24(num_bits) + marlin_24_perm[num_bits] = perm_24 + marlin_24_scale_perm[num_bits] = scale_perm_24 + marlin_24_scale_perm_single[num_bits] = scale_perm_single_24 \ No newline at end of file diff --git a/archive/csrc/custom_marlin/utils/marlin_perms.py b/archive/csrc/custom_marlin/utils/marlin_perms.py new file mode 100644 index 0000000..62255ec --- /dev/null +++ b/archive/csrc/custom_marlin/utils/marlin_perms.py @@ -0,0 +1,65 @@ +''' +Date: 2024-11-08 02:46:47 +LastEditors: djw +LastEditTime: 2024-11-08 02:46:55 +''' +"""This file is used for /tests and /benchmarks""" +from typing import Dict, List + +import numpy +import torch + + +# Precompute permutations for Marlin weight and scale shuffling # noqa: E501 +# +# Marlin works on [16,64] tiles. The goal of the permutations is to reorder the weight data so that it is compatible noqa: # noqa: E501 +# with the tensor-core format that is described here: +# https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#matrix-fragments-for-mma-m16n8k16-with-floating-point-type # noqa: E501 +# +# As a result of this reordering, the vector loads inside the kernel will get the data as it is needed for tensor-core # noqa: E501 +# (without the need to use ldmatrix instructions) # noqa: E501 +def get_perms(num_bits: int): + perm_list: List[int] = [] + for i in range(32): + perm1: List[int] = [] + col = i // 4 + for block in [0, 1]: + for row in [ + 2 * (i % 4), + 2 * (i % 4) + 1, + 2 * (i % 4 + 4), + 2 * (i % 4 + 4) + 1, + ]: + perm1.append(16 * row + col + 8 * block) + for j in range(4): + perm_list.extend([p + 256 * j for p in perm1]) + + perm = numpy.array(perm_list) + + if num_bits == 4: + interleave = numpy.array([0, 2, 4, 6, 1, 3, 5, 7]) + elif num_bits == 8: + interleave = numpy.array([0, 2, 1, 3]) + else: + raise Exception("num_bits must be 4 or 8, got {}".format(num_bits)) + + perm = perm.reshape((-1, len(interleave)))[:, interleave].ravel() + perm = torch.from_numpy(perm) + 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 perm, scale_perm, scale_perm_single + + +marlin_perm: Dict[int, torch.Tensor] = {} +marlin_scale_perm: Dict[int, List[int]] = {} +marlin_scale_perm_single: Dict[int, List[int]] = {} +for num_bits in [4, 8]: + perm, scale_perm, scale_perm_single = get_perms(num_bits) + marlin_perm[num_bits] = perm + marlin_scale_perm[num_bits] = scale_perm + marlin_scale_perm_single[num_bits] = scale_perm_single \ No newline at end of file diff --git a/archive/csrc/custom_marlin/utils/marlin_utils.py b/archive/csrc/custom_marlin/utils/marlin_utils.py new file mode 100644 index 0000000..ccecdc3 --- /dev/null +++ b/archive/csrc/custom_marlin/utils/marlin_utils.py @@ -0,0 +1,234 @@ +"""This file is used for /tests and /benchmarks""" +import random + +import numpy +import torch + +from .format24 import ( + mask_creator, sparse_semi_structured_from_dense_cutlass) +from .marlin_24_perms import ( + marlin_24_perm, marlin_24_scale_perm, marlin_24_scale_perm_single) +from .marlin_perms import ( + marlin_perm, marlin_scale_perm, marlin_scale_perm_single) +from .quant_utils import ( + get_pack_factor, quantize_weights, sort_weights, dequantize_weights) + + + +__cuda_arch = torch.cuda.get_device_capability() + +MARLIN_TILE = 16 + +GPTQ_MARLIN_TILE = 16 +GPTQ_MARLIN_MIN_THREAD_N = 64 +GPTQ_MARLIN_MIN_THREAD_K = 128 +GPTQ_MARLIN_MAX_PARALLEL = 16 + +GPTQ_MARLIN_SUPPORTED_NUM_BITS = [4, 8] +GPTQ_MARLIN_SUPPORTED_GROUP_SIZES = [-1, 32, 64, 128] +GPTQ_MARLIN_SUPPORTED_SYM = [True] + +def is_marlin_supported(): + return __cuda_arch[0] >= 8 + + +def marlin_permute_weights(q_w, size_k, size_n, perm, tile=MARLIN_TILE): + assert q_w.shape == (size_k, size_n) + assert size_k % tile == 0, f"size_k = {size_k}, tile = {tile}" + assert size_n % tile == 0, f"size_k = {size_n}, tile = {tile}" + + # Permute weights to 16x64 marlin tiles + q_w = q_w.reshape((size_k // tile, tile, size_n // tile, tile)) + q_w = q_w.permute((0, 2, 1, 3)) + q_w = q_w.reshape((size_k // tile, size_n * tile)) + + q_w = q_w.reshape((-1, perm.numel()))[:, perm].reshape(q_w.shape) + + return q_w + + +def marlin_weights(q_w, size_k, size_n, num_bits, perm): + # Permute + q_w = marlin_permute_weights(q_w, size_k, size_n, perm) + + # Pack + pack_factor = get_pack_factor(num_bits) + orig_device = q_w.device + + q_w = q_w.cpu().numpy().astype(numpy.uint32) + + q_packed = numpy.zeros((q_w.shape[0], q_w.shape[1] // pack_factor), + dtype=numpy.uint32) + for i in range(pack_factor): + q_packed |= q_w[:, i::pack_factor] << num_bits * i + + q_packed = torch.from_numpy(q_packed.astype(numpy.int32)).to(orig_device) + + return q_packed + + +def marlin_permute_scales(s, size_k, size_n, group_size, scale_perm, + scale_perm_single): + 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_quantize( + w: torch.Tensor, + num_bits: int, + group_size: int, + act_order: bool, +): + size_k, size_n = w.shape + + # Normalize group_size + if group_size == -1: + group_size = size_k + assert group_size <= size_k + + # Quantize (and apply act_order if provided) + w_ref, q_w, s, g_idx, rand_perm = quantize_weights(w, num_bits, group_size, + act_order) + + # For act_order, sort the "weights" and "g_idx" so that group ids are + # increasing + sort_indices = torch.empty(0, dtype=torch.int, device=w.device) + if act_order: + q_w, g_idx, sort_indices = sort_weights(q_w, g_idx) + + # Reformat to marlin + marlin_q_w = marlin_weights(q_w, size_k, size_n, num_bits, + marlin_perm[num_bits]) + marlin_s = marlin_permute_scales(s, size_k, size_n, group_size, + marlin_scale_perm[num_bits], + marlin_scale_perm_single[num_bits]) + + # Create result + res_list = [w_ref, marlin_q_w, marlin_s, g_idx, sort_indices, rand_perm] + for i in range(len(res_list)): + res_list[i] = res_list[i].to(w.device) + + return res_list + + +def inject_24(w, size_k, size_n): + assert w.shape == (size_k, size_n) + + mask = mask_creator(w.t()).t().cuda().bool() + + return (mask * w).contiguous(), mask.contiguous() + + +def check_24(w, num_rows_to_sample=50, _verbose=False): + BLOCK_SIZE = 4 + MAX_NON_ZEROS = 2 + + w = w.t().contiguous() + + print("check_24: w.shape = {}".format(w.shape)) + + num_rows, num_cols = w.shape + sampled_row_idxs = random.choices(range(num_rows), k=num_rows_to_sample) + if _verbose: + print(f"Sampled row idxs = {sampled_row_idxs}") + + total_segments = 0 + non_24_segments = 0 + for i in sampled_row_idxs: + for j in range(0, num_cols - BLOCK_SIZE, BLOCK_SIZE): + total_segments += 1 + block = w[i, j:j + BLOCK_SIZE] + num_nonzero = torch.count_nonzero(block) + if num_nonzero > MAX_NON_ZEROS: + print("i = {} j = {} block = {}".format(i, j, block)) + non_24_segments += 1 + + print(f"{non_24_segments} / {total_segments} do not have 2:4 structure.") + + +def compress_quantized_24_weight(q_24, size_k, size_n, num_bits): + assert q_24.shape == (size_k, size_n) + + # Remove zp to normalize over 0 + max_q_val = (1 << num_bits) - 1 + zp = (max_q_val + 1) // 2 + q_24_no_zp = q_24 - zp + + # Compress + q_24_no_zp = q_24_no_zp.t().contiguous() + q_24_no_zp_comp, meta = sparse_semi_structured_from_dense_cutlass( + q_24_no_zp) + q_24_no_zp_comp = q_24_no_zp_comp.t().contiguous() + + # Restore zp + q_24_comp = q_24_no_zp_comp + zp + + # Resize meta to its actual shape (without moving any data) + meta = meta.resize_(meta.shape[1] // 2, meta.shape[0] * 2) + + return q_24_comp, meta + + +def marlin_24_quantize( + w: torch.Tensor, + num_bits: int, + group_size: int, +): + size_k, size_n = w.shape + + # Normalize group_size + if group_size == -1: + group_size = size_k + assert group_size <= size_k + + # Inject 2:4 sparsity + w_24, mask_24 = inject_24(w, size_k, size_n) + + # Quantize + w_24_ref, q_w_24, s, g_idx, rand_perm = quantize_weights(w_24, + num_bits, + group_size, + act_order=False) + + # Compress quantized weight + q_w_24_comp, meta = compress_quantized_24_weight(q_w_24, size_k, size_n, + num_bits) + size_k_comp = size_k // 2 + + # Reformat to marlin + marlin_24_q_w_comp = marlin_weights(q_w_24_comp, size_k_comp, size_n, + num_bits, marlin_24_perm[num_bits]) + marlin_24_s = marlin_permute_scales(s, size_k, size_n, group_size, + marlin_24_scale_perm[num_bits], + marlin_24_scale_perm_single[num_bits]) + + # Create result + res_list = [w_24_ref, marlin_24_q_w_comp, meta, marlin_24_s] + for i in range(len(res_list)): + res_list[i] = res_list[i].to(w.device) + + return res_list + + +def compute_max_diff(output, output_ref): + return torch.mean(torch.abs(output - output_ref)) / torch.mean( + torch.abs(output_ref)) + + +class MarlinWorkspace: + + def __init__(self, out_features, min_thread_n, max_parallel, device): + assert (out_features % min_thread_n == 0), ( + "out_features = {} is undivisible by min_thread_n = {}".format( + out_features, min_thread_n)) + + max_workspace_size = ((out_features // min_thread_n) * max_parallel) + + self.scratch = torch.zeros(max_workspace_size, + dtype=torch.int, + device=device) \ No newline at end of file diff --git a/archive/csrc/custom_marlin/utils/quant_utils.py b/archive/csrc/custom_marlin/utils/quant_utils.py new file mode 100644 index 0000000..077d22f --- /dev/null +++ b/archive/csrc/custom_marlin/utils/quant_utils.py @@ -0,0 +1,195 @@ +"""This file is used for /tests and /benchmarks""" +import numpy +import torch + +SUPPORTED_NUM_BITS = [4, 8] +SUPPORTED_GROUP_SIZES = [-1, 32, 64, 128] + + +def get_pack_factor(num_bits): + assert num_bits in SUPPORTED_NUM_BITS, f"Unsupported num_bits = {num_bits}" + return 32 // num_bits + + +def permute_rows(q_w: torch.Tensor, w_ref: torch.Tensor, group_size: int): + assert q_w.shape == w_ref.shape + + orig_device = q_w.device + k_size, _ = q_w.shape + + g_idx = torch.zeros((k_size, ), dtype=torch.int32) + for i in range(k_size): + g_idx[i] = i // group_size + + # Simulate act_order by doing a random permutation on K + rand_perm = torch.randperm(k_size) + + g_idx = g_idx[rand_perm].contiguous() + q_w = q_w[rand_perm, :].contiguous() + w_ref = w_ref[rand_perm, :].contiguous() + + return ( + w_ref.to(device=orig_device), + q_w.to(device=orig_device), + g_idx.to(device=orig_device), + rand_perm.to(device=orig_device), + ) + + +# Function: Dequantize quantized weights +def dequantize_weights(qweight, qzeros, scales, g_idx, bits=4, group_size=128, device='cuda:0'): + # Create a tensor for bitwise right shift operation + wf = torch.tensor(list(range(0, 32, bits)), dtype=torch.int32, device=device).unsqueeze(0) + + # Apply bitwise right shift and convert qzeros to the appropriate type + zeros = torch.bitwise_right_shift(torch.unsqueeze(qzeros, 2).expand(-1, -1, 32 // bits), wf.unsqueeze(0)).to(torch.int16 if bits == 8 else torch.int8) + torch.bitwise_and(zeros, (2 ** bits) - 1, out=zeros) + + # Reshape the zeros tensor + zeros = zeros + 1 + zeros = zeros.reshape(-1, 1, zeros.shape[1] * zeros.shape[2]) + + # Reshape the scales tensor + scales = scales.reshape(-1, 1, scales.shape[-1]) + + # Similar bitwise right shift operation for qweight and reshape + weight = torch.bitwise_right_shift(torch.unsqueeze(qweight, 1).expand(-1, 32 // bits, -1), wf.unsqueeze(-1)).to(torch.int16 if bits == 8 else torch.int8) + torch.bitwise_and(weight, (2 ** bits) - 1, out=weight) + weight = weight.reshape(-1, group_size, weight.shape[2]) + + # Apply dequantization formula and reshape the final weight + weight = (scales * (weight - zeros)) + weight = weight.reshape(weight.shape[0] * weight.shape[1], weight.shape[2]) + + # Return the transposed weight + return weight.transpose(0, 1) + +def quantize_weights(w: torch.Tensor, num_bits: int, group_size: int, + act_order: bool): + orig_device = w.device + size_k, size_n = w.shape + + assert w.is_floating_point(), "w must be float" + assert num_bits in SUPPORTED_NUM_BITS, f"Unsupported num_bits = {num_bits}" + assert group_size in SUPPORTED_GROUP_SIZES + [ + size_k + ], f"Unsupported groupsize = {group_size}" + + if group_size == -1: + group_size = size_k + assert group_size <= size_k + + max_q_val = 2**num_bits - 1 + half_q_val = (max_q_val + 1) // 2 + + # Reshape to [groupsize, -1] + if group_size < size_k: + w = w.view((-1, group_size, size_n)) + w = w.permute(1, 0, 2) + w = w.reshape((group_size, -1)) + + # Compute scale for each group + s = torch.max(torch.abs(w), 0, keepdim=True)[0] + s *= 2 / max_q_val # 2 => symmetric + + # Quantize + q_w = torch.round(w / s).int() + q_w += half_q_val + q_w = torch.clamp(q_w, 0, max_q_val) + + # Compute ref (dequantized) + w_ref = (q_w - half_q_val).half() * s + + # Restore original shapes + if group_size < size_k: + + def reshape_w(w): + w = w.reshape((group_size, -1, size_n)) + w = w.permute(1, 0, 2) + w = w.reshape((size_k, size_n)).contiguous() + return w + + q_w = reshape_w(q_w) + w_ref = reshape_w(w_ref) + + s = s.reshape((-1, size_n)).contiguous() + + # Apply act_order + g_idx = torch.empty(0, dtype=torch.int, device=w.device) + rand_perm = torch.empty(0, dtype=torch.int, device=w.device) + if act_order: + assert ( + group_size < size_k + ), "For act_order, groupsize = {} must be less than size_k = {}".format( + group_size, size_k) + + w_ref, q_w, g_idx, rand_perm = permute_rows(q_w, w_ref, group_size) + + return ( + w_ref.to(device=orig_device), + q_w.to(device=orig_device), + s.to(device=orig_device), + g_idx.to(device=orig_device), + rand_perm.to(device=orig_device), + ) + + +def sort_weights(q_w: torch.Tensor, g_idx: torch.Tensor): + orig_device = q_w.device + + sort_indices = torch.argsort(g_idx).to( + dtype=torch.int32) # Sort based on g_idx + + g_idx = g_idx[sort_indices].contiguous() + q_w = q_w[sort_indices, :].contiguous() + + return ( + q_w.to(device=orig_device), + g_idx.to(device=orig_device), + sort_indices.to(device=orig_device), + ) + + +def gptq_pack( + q_w: torch.Tensor, + num_bits: int, + size_k: int, + size_n: int, +): + assert q_w.shape == (size_k, size_n) + + pack_factor = get_pack_factor(num_bits) + assert size_k % pack_factor == 0 + + orig_device = q_w.device + + q_w = q_w.cpu().numpy().astype(numpy.uint32) + + q_res = numpy.zeros((size_k // pack_factor, size_n), dtype=numpy.uint32) + + for i in range(pack_factor): + q_res |= q_w[i::pack_factor, :] << num_bits * i + + q_res = torch.from_numpy(q_res.astype(numpy.int32)).to(orig_device) + return q_res + +def gptq_unpack( + q_res: torch.Tensor, + num_bits: int, + size_k: int, + size_n: int, +): + pack_factor = 32 // num_bits + assert size_k % pack_factor == 0 + + orig_device = q_res.device + + q_res = q_res.cpu().numpy() + + q_w = numpy.zeros((size_k, size_n), dtype=numpy.uint32) + + for i in range(pack_factor): + q_w[i::pack_factor, :] = (q_res >> (num_bits * i)) & ((1 << num_bits) - 1) + + q_w = torch.from_numpy(q_w.astype(numpy.int32)).to(orig_device) + return q_w \ No newline at end of file diff --git a/archive/csrc/ktransformers_ext/CMakeLists.txt b/archive/csrc/ktransformers_ext/CMakeLists.txt new file mode 100644 index 0000000..554b5e4 --- /dev/null +++ b/archive/csrc/ktransformers_ext/CMakeLists.txt @@ -0,0 +1,415 @@ +cmake_minimum_required(VERSION 3.16) +project(cpuinfer_ext VERSION 0.1.0) + + +set(CMAKE_CXX_STANDARD 17) + + +set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -O3 -ffast-math -fopenmp") +add_compile_definitions(_GLIBCXX_USE_CXX11_ABI=${_GLIBCXX_USE_CXX11_ABI}) +set(CMAKE_BUILD_TYPE "Release") +# set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -g -ffast-math -fopenmp") +# set(CMAKE_BUILD_TYPE "Debug") +set(CMAKE_EXPORT_COMPILE_COMMANDS ON) + + +include(CheckCXXCompilerFlag) +set(CMAKE_POSITION_INDEPENDENT_CODE ON) + + +option(LLAMA_NATIVE "llama: enable -march=native flag" ON) + +# instruction set specific +if (LLAMA_NATIVE) + set(INS_ENB OFF) +else() + set(INS_ENB ON) +endif() + +option(LLAMA_AVX "llama: enable AVX" OFF) +option(LLAMA_AVX2 "llama: enable AVX2" OFF) +option(LLAMA_AVX512 "llama: enable AVX512" OFF) +option(LLAMA_AVX512_VBMI "llama: enable AVX512-VBMI" OFF) +option(LLAMA_AVX512_VNNI "llama: enable AVX512-VNNI" OFF) +option(LLAMA_AVX512_BF16 "llama: enable AVX512-BF16" OFF) +option(LLAMA_FMA "llama: enable FMA" OFF) +# in MSVC F16C is implied with AVX2/AVX512 +if (NOT MSVC) + option(LLAMA_F16C "llama: enable F16C" OFF) +endif() +option(LLAMA_AVX512_FANCY_SIMD "llama: enable AVX512-VL, AVX512-BW, AVX512-DQ, AVX512-VNNI" OFF) +option(KTRANSFORMERS_USE_CUDA "ktransformers: use CUDA" ON) +option(KTRANSFORMERS_USE_MUSA "ktransformers: use MUSA" OFF) +option(KTRANSFORMERS_USE_ROCM "ktransformers: use ROCM" OFF) +option(KTRANSFORMERS_USE_XPU "ktransformers: use XPU" OFF) +option(KTRANSFORMERS_USE_NPU "ktransformers: use NPU" OFF) + +if(KTRANSFORMERS_USE_NPU) + add_definitions(-DKTRANSFORMERS_USE_NPU=1) +endif() + +# Architecture specific +# TODO: probably these flags need to be tweaked on some architectures +# feel free to update the Makefile for your architecture and send a pull request or issue +message(STATUS "CMAKE_SYSTEM_PROCESSOR: ${CMAKE_SYSTEM_PROCESSOR}") +if (MSVC) + string(TOLOWER "${CMAKE_GENERATOR_PLATFORM}" CMAKE_GENERATOR_PLATFORM_LWR) + message(STATUS "CMAKE_GENERATOR_PLATFORM: ${CMAKE_GENERATOR_PLATFORM}") +else () + set(CMAKE_GENERATOR_PLATFORM_LWR "") +endif () + +if (NOT MSVC) + if (LLAMA_STATIC) + add_link_options(-static) + if (MINGW) + add_link_options(-static-libgcc -static-libstdc++) + endif() + endif() + if (LLAMA_GPROF) + add_compile_options(-pg) + endif() +endif() + +set(ARCH_FLAGS "") + +if (CMAKE_OSX_ARCHITECTURES STREQUAL "arm64" OR CMAKE_GENERATOR_PLATFORM_LWR STREQUAL "arm64" OR + (NOT CMAKE_OSX_ARCHITECTURES AND NOT CMAKE_GENERATOR_PLATFORM_LWR AND + CMAKE_SYSTEM_PROCESSOR MATCHES "^(aarch64|arm.*|ARM64)$")) + message(STATUS "ARM detected") + if (MSVC) + add_compile_definitions(__aarch64__) # MSVC defines _M_ARM64 instead + add_compile_definitions(__ARM_NEON) + add_compile_definitions(__ARM_FEATURE_FMA) + + set(CMAKE_REQUIRED_FLAGS_PREV ${CMAKE_REQUIRED_FLAGS}) + string(JOIN " " CMAKE_REQUIRED_FLAGS ${CMAKE_REQUIRED_FLAGS} "/arch:armv8.2") + check_cxx_source_compiles("#include \nint main() { int8x16_t _a, _b; int32x4_t _s = vdotq_s32(_s, _a, _b); return 0; }" GGML_COMPILER_SUPPORT_DOTPROD) + if (GGML_COMPILER_SUPPORT_DOTPROD) + add_compile_definitions(__ARM_FEATURE_DOTPROD) + endif () + check_cxx_source_compiles("#include \nint main() { float16_t _a; float16x8_t _s = vdupq_n_f16(_a); return 0; }" GGML_COMPILER_SUPPORT_FP16_VECTOR_ARITHMETIC) + if (GGML_COMPILER_SUPPORT_FP16_VECTOR_ARITHMETIC) + add_compile_definitions(__ARM_FEATURE_FP16_VECTOR_ARITHMETIC) + endif () + set(CMAKE_REQUIRED_FLAGS ${CMAKE_REQUIRED_FLAGS_PREV}) + else() + if(KTRANSFORMERS_USE_NPU) + list(APPEND ARCH_FLAGS -march=armv8.2-a+fp16+fp16fml+dotprod -lnuma) + endif() + check_cxx_compiler_flag(-mfp16-format=ieee COMPILER_SUPPORTS_FP16_FORMAT_I3E) + if (NOT "${COMPILER_SUPPORTS_FP16_FORMAT_I3E}" STREQUAL "") + list(APPEND ARCH_FLAGS -mfp16-format=ieee) + endif() + if (${CMAKE_SYSTEM_PROCESSOR} MATCHES "armv6") + # Raspberry Pi 1, Zero + list(APPEND ARCH_FLAGS -mfpu=neon-fp-armv8 -mno-unaligned-access) + endif() + if (${CMAKE_SYSTEM_PROCESSOR} MATCHES "armv7") + if ("${CMAKE_SYSTEM_NAME}" STREQUAL "Android") + # Android armeabi-v7a + list(APPEND ARCH_FLAGS -mfpu=neon-vfpv4 -mno-unaligned-access -funsafe-math-optimizations) + else() + # Raspberry Pi 2 + list(APPEND ARCH_FLAGS -mfpu=neon-fp-armv8 -mno-unaligned-access -funsafe-math-optimizations) + endif() + endif() + if (${CMAKE_SYSTEM_PROCESSOR} MATCHES "armv8") + # Android arm64-v8a + # Raspberry Pi 3, 4, Zero 2 (32-bit) + list(APPEND ARCH_FLAGS -mno-unaligned-access) + endif() + endif() +elseif (CMAKE_OSX_ARCHITECTURES STREQUAL "x86_64" OR CMAKE_GENERATOR_PLATFORM_LWR MATCHES "^(x86_64|i686|amd64|x64|win32)$" OR + (NOT CMAKE_OSX_ARCHITECTURES AND NOT CMAKE_GENERATOR_PLATFORM_LWR AND + CMAKE_SYSTEM_PROCESSOR MATCHES "^(x86_64|i686|AMD64)$")) + message(STATUS "x86 detected") + if(NOT KTRANSFORMERS_USE_NPU) + set(HOST_IS_X86 TRUE) + set(HAS_AVX512 TRUE) + set(__HAS_AMX__ TRUE) + add_compile_definitions(__x86_64__) + # check AVX512 + execute_process( + COMMAND lscpu + OUTPUT_VARIABLE LSCPU_OUTPUT + OUTPUT_STRIP_TRAILING_WHITESPACE + ) + # message(STATUS "LSCPU_OUTPUT: ${LSCPU_OUTPUT}") + + string(FIND "${LSCPU_OUTPUT}" "avx512" COMPILER_SUPPORTS_AVX512F) + + if (COMPILER_SUPPORTS_AVX512F GREATER -1) + message(STATUS "Compiler and CPU support AVX512F (tested by compiling a program)") + add_compile_definitions(__HAS_AVX512F__) + else() + message(STATUS "Compiler and/or CPU do NOT support AVX512F") + set(HAS_AVX512 False) + endif() + + # check AMX + string(FIND "${LSCPU_OUTPUT}" "amx" COMPILER_SUPPORTS_AMX) + + if(COMPILER_SUPPORTS_AMX GREATER -1) + message(STATUS "Compiler supports AMX") + add_compile_definitions(__HAS_AMX__) + else() + message(STATUS "Compiler does NOT support AMX") + endif() + endif() + if (MSVC) + # instruction set detection for MSVC only + if (LLAMA_NATIVE) + include(cmake/FindSIMD.cmake) + endif () + if (LLAMA_AVX512) + list(APPEND ARCH_FLAGS /arch:AVX512) + # MSVC has no compile-time flags enabling specific + # AVX512 extensions, neither it defines the + # macros corresponding to the extensions. + # Do it manually. + if (LLAMA_AVX512_VBMI) + add_compile_definitions($<$:__AVX512VBMI__>) + add_compile_definitions($<$:__AVX512VBMI__>) + endif() + if (LLAMA_AVX512_VNNI) + add_compile_definitions($<$:__AVX512VNNI__>) + add_compile_definitions($<$:__AVX512VNNI__>) + endif() + if (LLAMA_AVX512_FANCY_SIMD) + add_compile_definitions($<$:__AVX512VL__>) + add_compile_definitions($<$:__AVX512VL__>) + add_compile_definitions($<$:__AVX512BW__>) + add_compile_definitions($<$:__AVX512BW__>) + add_compile_definitions($<$:__AVX512DQ__>) + add_compile_definitions($<$:__AVX512DQ__>) + add_compile_definitions($<$:__AVX512VNNI__>) + add_compile_definitions($<$:__AVX512VNNI__>) + endif() + if (LLAMA_AVX512_BF16) + add_compile_definitions($<$:__AVX512BF16__>) + add_compile_definitions($<$:__AVX512BF16__>) + endif() + elseif (LLAMA_AVX2) + list(APPEND ARCH_FLAGS /arch:AVX2) + elseif (LLAMA_AVX) + list(APPEND ARCH_FLAGS /arch:AVX) + endif() + else() + if (LLAMA_NATIVE) + list(APPEND ARCH_FLAGS -mfma -mavx -mavx2) + list(APPEND ARCH_FLAGS -march=native) + endif() + if (LLAMA_F16C) + list(APPEND ARCH_FLAGS -mf16c) + endif() + if (LLAMA_FMA) + list(APPEND ARCH_FLAGS -mfma) + endif() + if (LLAMA_AVX) + list(APPEND ARCH_FLAGS -mavx) + endif() + if (LLAMA_AVX2) + list(APPEND ARCH_FLAGS -mavx2) + endif() + if (LLAMA_AVX512) + list(APPEND ARCH_FLAGS -mavx512f) + list(APPEND ARCH_FLAGS -mavx512bw) + endif() + if (LLAMA_AVX512_VBMI) + list(APPEND ARCH_FLAGS -mavx512vbmi) + endif() + if (LLAMA_AVX512_VNNI) + list(APPEND ARCH_FLAGS -mavx512vnni) + endif() + if (LLAMA_AVX512_FANCY_SIMD) + message(STATUS "AVX512-VL, AVX512-BW, AVX512-DQ, AVX512-VNNI enabled") + list(APPEND ARCH_FLAGS -mavx512vl) + list(APPEND ARCH_FLAGS -mavx512bw) + list(APPEND ARCH_FLAGS -mavx512dq) + list(APPEND ARCH_FLAGS -mavx512vnni) + list(APPEND ARCH_FLAGS -mavx512vpopcntdq) + endif() + if (LLAMA_AVX512_BF16) + list(APPEND ARCH_FLAGS -mavx512bf16) + endif() + endif() +elseif (${CMAKE_SYSTEM_PROCESSOR} MATCHES "ppc64") + message(STATUS "PowerPC detected") + if (${CMAKE_SYSTEM_PROCESSOR} MATCHES "ppc64le") + list(APPEND ARCH_FLAGS -mcpu=powerpc64le) + else() + list(APPEND ARCH_FLAGS -mcpu=native -mtune=native) + #TODO: Add targets for Power8/Power9 (Altivec/VSX) and Power10(MMA) and query for big endian systems (ppc64/le/be) + endif() +else() + message(STATUS "Unknown architecture") +endif() + +# message(STATUS "CUDAToolkit_ROOT:${CUDAToolkit_ROOT}") +# find_package(FindCUDAToolkit REQUIRED) +# if(CUDAToolkit_FOUND) +# message(STATUS "Found CUDA cudart lib at:${CUDAToolkit_LIBRARY_DIR}") +# else() +# message(STATUS "Can't found CUDA lib") +# endif() + +if (NOT EXISTS $ENV{ROCM_PATH}) + if (NOT EXISTS /opt/rocm) + set(ROCM_PATH /usr) + else() + set(ROCM_PATH /opt/rocm) + endif() +else() + set(ROCM_PATH $ENV{ROCM_PATH}) +endif() + +list(APPEND CMAKE_PREFIX_PATH ${ROCM_PATH}) +list(APPEND CMAKE_PREFIX_PATH "${ROCM_PATH}/lib64/cmake") + +if (NOT EXISTS $ENV{MUSA_PATH}) + if (NOT EXISTS /opt/musa) + set(MUSA_PATH /usr/local/musa) + else() + set(MUSA_PATH /opt/musa) + endif() +else() + set(MUSA_PATH $ENV{MUSA_PATH}) +endif() + +list(APPEND CMAKE_MODULE_PATH "${MUSA_PATH}/cmake") + +add_compile_options("$<$:${ARCH_FLAGS}>") +add_compile_options("$<$:${ARCH_FLAGS}>") + +add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/../../third_party/pybind11 ${CMAKE_CURRENT_BINARY_DIR}/third_party/pybind11) +add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/../../third_party/llama.cpp ${CMAKE_CURRENT_BINARY_DIR}/third_party/llama.cpp) + +include_directories(${CMAKE_CURRENT_SOURCE_DIR}/../../third_party) +if (WIN32) + include_directories("$ENV{CUDA_PATH}/include") + add_compile_definitions(KTRANSFORMERS_USE_CUDA=1) +elseif (UNIX) + if (KTRANSFORMERS_USE_ROCM) + find_package(HIP REQUIRED) + if(HIP_FOUND) + include_directories("${HIP_INCLUDE_DIRS}") + add_compile_definitions(KTRANSFORMERS_USE_ROCM=1) + endif() + elseif (KTRANSFORMERS_USE_MUSA) + if (NOT EXISTS $ENV{MUSA_PATH}) + if (NOT EXISTS /opt/musa) + set(MUSA_PATH /usr/local/musa) + else() + set(MUSA_PATH /opt/musa) + endif() + else() + set(MUSA_PATH $ENV{MUSA_PATH}) + endif() + + list(APPEND CMAKE_MODULE_PATH "${MUSA_PATH}/cmake") + + find_package(MUSAToolkit) + if (MUSAToolkit_FOUND) + message(STATUS "MUSA Toolkit found") + add_compile_definitions(KTRANSFORMERS_USE_MUSA=1) + endif() + elseif (KTRANSFORMERS_USE_XPU) + add_compile_definitions(KTRANSFORMERS_USE_XPU=1) + elseif (KTRANSFORMERS_USE_CUDA) + find_package(CUDA REQUIRED) + include_directories("${CUDA_INCLUDE_DIRS}") + include(CheckLanguage) + check_language(CUDA) + if(CMAKE_CUDA_COMPILER) + message(STATUS "CUDA detected") + find_package(CUDAToolkit REQUIRED) + include_directories(${CUDAToolkit_INCLUDE_DIRS}) + endif() + message(STATUS "enabling CUDA") + enable_language(CUDA) + add_compile_definitions(KTRANSFORMERS_USE_CUDA=1) + endif() +endif() + +aux_source_directory(${CMAKE_CURRENT_SOURCE_DIR} SOURCE_DIR1) +aux_source_directory(${CMAKE_CURRENT_SOURCE_DIR}/cpu_backend SOURCE_DIR2) +aux_source_directory(${CMAKE_CURRENT_SOURCE_DIR}/operators/llamafile SOURCE_DIR3) +# aux_source_directory(${CMAKE_CURRENT_SOURCE_DIR}/../../third_party/llamafile SOURCE_DIR4) +file(GLOB LLAMAFILE_SOURCES "${CMAKE_CURRENT_SOURCE_DIR}/../../third_party/llamafile/*.cpp") +list(REMOVE_ITEM LLAMAFILE_SOURCES + "${CMAKE_CURRENT_SOURCE_DIR}/../../third_party/llamafile/sgemm_arm.cpp" + "${CMAKE_CURRENT_SOURCE_DIR}/../../third_party/llamafile/sgemm_x86.cpp" +) +set(SOURCE_DIR4 ${LLAMAFILE_SOURCES}) +aux_source_directory(${CMAKE_CURRENT_SOURCE_DIR}/operators/kvcache SOURCE_DIR5) + +if (HOST_IS_X86 AND HAS_AVX512 AND __HAS_AMX__) + aux_source_directory(${CMAKE_CURRENT_SOURCE_DIR}/operators/amx SOURCE_DIR6) +endif() + + +set(ALL_SOURCES ${SOURCE_DIR1} ${SOURCE_DIR2} ${SOURCE_DIR3} ${SOURCE_DIR4} ${SOURCE_DIR5} ${SOURCE_DIR6}) + +file(GLOB_RECURSE FMT_SOURCES "${CMAKE_CURRENT_SOURCE_DIR}/*.cpp" "${CMAKE_CURRENT_SOURCE_DIR}/*.hpp" "${CMAKE_CURRENT_SOURCE_DIR}/*.h") + +add_custom_target( + format + COMMAND clang-format + -i + -style=file + ${FMT_SOURCES} + COMMENT "Running clang-format on all source files" +) + + +add_library(llamafile STATIC ${SOURCE_DIR4}) + +message(STATUS "CMAKE_CXX_FLAGS: ${CMAKE_CXX_FLAGS}") +message(STATUS "ARCH_FLAGS: ${ARCH_FLAGS}") +pybind11_add_module(${PROJECT_NAME} MODULE ${ALL_SOURCES}) +target_link_libraries(${PROJECT_NAME} PRIVATE llama) + + +if(WIN32) + target_link_libraries(${PROJECT_NAME} PRIVATE "$ENV{CUDA_PATH}/lib/x64/cudart.lib")#CUDA::cudart +elseif(UNIX) + if (KTRANSFORMERS_USE_ROCM) + add_compile_definitions(USE_HIP=1) + target_link_libraries(${PROJECT_NAME} PRIVATE "${ROCM_PATH}/lib/libamdhip64.so") + message(STATUS "Building for HIP") + elseif(KTRANSFORMERS_USE_MUSA) + target_link_libraries(${PROJECT_NAME} PRIVATE MUSA::musart) + elseif(KTRANSFORMERS_USE_XPU) + elseif(KTRANSFORMERS_USE_CUDA AND NOT KTRANSFORMERS_USE_MUSA) + target_link_libraries(${PROJECT_NAME} PRIVATE "${CUDAToolkit_LIBRARY_DIR}/libcudart.so") + endif() +endif() + +# Define the USE_NUMA option +option(USE_NUMA "Disable NUMA support" OFF) + +# Check if the USE_NUMA environment variable is set +if(DEFINED ENV{USE_NUMA}) + set(USE_NUMA ON) +endif() + +if(USE_NUMA) + message(STATUS "NUMA support is enabled") +else() + message(STATUS "NUMA support is disabled") +endif() + +find_library(NUMA_LIBRARY NAMES numa) + +if(NUMA_LIBRARY AND USE_NUMA) + message(STATUS "NUMA library found: ${NUMA_LIBRARY} - enabling NUMA support") + target_link_libraries(${PROJECT_NAME} PRIVATE ${NUMA_LIBRARY}) + target_compile_definitions(${PROJECT_NAME} PRIVATE USE_NUMA) +else() + if(USE_NUMA) + message(FATAL_ERROR "NUMA library not found - maybe sudo apt install libnuma-dev") + else() + message(STATUS "NUMA library not found or user not set USE_NUMA - disabling NUMA support") + endif() +endif() diff --git a/archive/csrc/ktransformers_ext/bench/bench_attention.py b/archive/csrc/ktransformers_ext/bench/bench_attention.py new file mode 100644 index 0000000..1b8b9b8 --- /dev/null +++ b/archive/csrc/ktransformers_ext/bench/bench_attention.py @@ -0,0 +1,178 @@ +#!/usr/bin/env python +# coding=utf-8 +""" +Description : +Author : Jianwei Dong +Date : 2024-08-28 10:32:05 +Version : 1.0.0 +LastEditors : Jianwei Dong +LastEditTime : 2024-08-28 10:32:05 +Copyright (c) 2024 by KVCache.AI, All Rights Reserved. +""" +import os, sys +import time + +sys.path.append(os.path.dirname(__file__) + "/../build") +import cpuinfer_ext +import torch + +layer_num = 10 +kv_head_num = 8 +q_head_num = 32 +head_dim = 128 +block_len = 128 +anchor_num = 1 + +anchor_type = cpuinfer_ext.kvcache.AnchorType.DYNAMIC +kv_type = cpuinfer_ext.kvcache.ggml_type.FP16 +retrieval_type = cpuinfer_ext.kvcache.RetrievalType.LAYER +layer_step: int = 1 +token_step: int = 1 +layer_offset: int = 0 +max_thread_num: int = 64 +max_batch_size: int = 1 +max_block_num: int = 1024 +CPUInfer = cpuinfer_ext.CPUInfer(max_thread_num) + +warm_up_iter = 1000 +test_iter = 10000 + + +def bench_linear(cache_seqlen: int): + with torch.inference_mode(mode=True): + cache_seqlens = torch.tensor([cache_seqlen], dtype=torch.int32, device="cpu") + seqlens_zero = torch.zeros((1,), dtype=torch.int32, device="cpu") + + config = cpuinfer_ext.kvcache.KVCacheConfig( + layer_num, + kv_head_num, + q_head_num, + head_dim, + block_len, + anchor_num, + anchor_type, + kv_type, + retrieval_type, + layer_step, + token_step, + layer_offset, + max_block_num, + max_batch_size, + max_thread_num, + ) + local_kvcache = cpuinfer_ext.kvcache.KVCache(config) + block_table = ( + torch.arange(max_block_num, dtype=torch.int32, device="cpu") + .contiguous() + .view(1, -1) + ) + + for layer_idx in range(layer_num): + k_cache = torch.randn( + (1, cache_seqlen, kv_head_num, head_dim), + dtype=torch.float16, + device="cpu", + ).contiguous() + v_cache = torch.randn( + (1, cache_seqlen, kv_head_num, head_dim), + dtype=torch.float16, + device="cpu", + ).contiguous() + + CPUInfer.submit( + local_kvcache.update_kvcache_fp16( + k_cache.data_ptr(), + v_cache.data_ptr(), + layer_idx, + block_table.data_ptr(), + 1, + max_block_num, + seqlens_zero.data_ptr(), + cache_seqlen, + ) + ) + CPUInfer.sync() + + input = torch.randn( + (1, 1, q_head_num, head_dim), dtype=torch.float16, device="cpu" + ).contiguous() + output = torch.empty( + (1, 1, q_head_num, head_dim), dtype=torch.float16, device="cpu" + ).contiguous() + + # attn_lse: (bsz, q_len, q_head_num) + attn_lse = torch.empty( + (1, 1, q_head_num), dtype=torch.float32, device="cpu" + ).contiguous() + input = input / 100 + + # warm up + for i in range(warm_up_iter): + CPUInfer.submit( + local_kvcache.attn( + input.data_ptr(), + output.data_ptr(), + attn_lse.data_ptr(), + i % layer_num, + 0, + 1, + 1, + max_block_num, + block_table.data_ptr(), + cache_seqlens.data_ptr(), + -1, + -1, + -1, + ) + ) + CPUInfer.sync() + + # test + start = time.perf_counter() + for i in range(test_iter): + CPUInfer.submit( + local_kvcache.attn( + input.data_ptr(), + output.data_ptr(), + attn_lse.data_ptr(), + i % layer_num, + 0, + 1, + 1, + max_block_num, + block_table.data_ptr(), + cache_seqlens.data_ptr(), + -1, + -1, + -1, + ) + ) + CPUInfer.sync() + end = time.perf_counter() + total_time = end - start + print("cache sequence length: ", cache_seqlen) + print("Time(s): ", total_time) + print("Iteration: ", test_iter) + print("Time(us) per iteration: ", total_time / test_iter * 1000000) + print( + "Bandwidth: ", + cache_seqlen + * kv_head_num + * head_dim + * 2 + * 2 + * test_iter + / total_time + / 1000 + / 1000 + / 1000, + "GB/s", + ) + print("") + + +bench_linear(1024) +bench_linear(4096) +bench_linear(16384) +bench_linear(32768) +bench_linear(65536) diff --git a/archive/csrc/ktransformers_ext/bench/bench_attention_torch.py b/archive/csrc/ktransformers_ext/bench/bench_attention_torch.py new file mode 100644 index 0000000..25f20e7 --- /dev/null +++ b/archive/csrc/ktransformers_ext/bench/bench_attention_torch.py @@ -0,0 +1,94 @@ +#!/usr/bin/env python +# coding=utf-8 +""" +Description : +Author : Jianwei Dong +Date : 2024-08-28 10:32:05 +Version : 1.0.0 +LastEditors : Jianwei Dong +LastEditTime : 2024-08-28 10:32:05 +Copyright (c) 2024 by KVCache.AI, All Rights Reserved. +""" +import os, sys +import time + +sys.path.append(os.path.dirname(__file__) + "/../build") +import cpuinfer_ext +import torch + +layer_num = 10 +kv_head_num = 8 +q_head_num = 32 +head_dim = 128 +block_len = 128 +anchor_num = 1 +warm_up_iter = 1000 +test_iter = 10000 + + +def bench_linear(cache_seqlen: int, device): + with torch.inference_mode(mode=True): + + kvcaches = [] + + for layer_idx in range(layer_num): + k_cache = torch.randn( + (1, 32, cache_seqlen, head_dim), + dtype=torch.float16, + device=device, + ).contiguous() + v_cache = torch.randn( + (1, 32, cache_seqlen, head_dim), + dtype=torch.float16, + device=device, + ).contiguous() + + kvcaches.append((k_cache, v_cache)) + + input = torch.randn( + (1, q_head_num, 1, head_dim), dtype=torch.float16, device=device + ).contiguous() + input = input / 100 + + # warm up + for i in range(warm_up_iter): + k_cache = kvcaches[i % layer_num][0] + v_cache = kvcaches[i % layer_num][1] + torch.nn.functional.scaled_dot_product_attention(input, k_cache, v_cache) + + # test + start = time.perf_counter() + for i in range(test_iter): + k_cache = kvcaches[i % layer_num][0] + v_cache = kvcaches[i % layer_num][1] + torch.nn.functional.scaled_dot_product_attention(input, k_cache, v_cache) + end = time.perf_counter() + total_time = end - start + print("cache sequence length: ", cache_seqlen) + print("Time(s): ", total_time) + print("Iteration: ", test_iter) + print("Time(us) per iteration: ", total_time / test_iter * 1000000) + print( + "Bandwidth: ", + cache_seqlen + * q_head_num + * head_dim + * 2 + * 2 + * test_iter + / total_time + / 1000 + / 1000 + / 1000, + "GB/s", + ) + print("") + + +bench_linear(1024, "cpu") +bench_linear(4096, "cpu") +bench_linear(1024, "cuda") +bench_linear(4096, "cuda") +bench_linear(16384, "cuda") +bench_linear(32768, "cuda") +bench_linear(65536, "cuda") diff --git a/archive/csrc/ktransformers_ext/bench/bench_linear.py b/archive/csrc/ktransformers_ext/bench/bench_linear.py new file mode 100644 index 0000000..3189afd --- /dev/null +++ b/archive/csrc/ktransformers_ext/bench/bench_linear.py @@ -0,0 +1,121 @@ +#!/usr/bin/env python +# coding=utf-8 +''' +Description : +Author : chenht2022 +Date : 2024-07-25 10:31:59 +Version : 1.0.0 +LastEditors : chenht2022 +LastEditTime : 2024-08-06 10:35:35 +Copyright (c) 2024 by KVCache.AI, All Rights Reserved. +''' +import os, sys +import time +sys.path.append(os.path.dirname(__file__) + '/../build') +import cpuinfer_ext +import torch + +input_size = 16384 +output_size = 5120 +stride = 16 +group_max_len = 1024 +layer_num = 10 +qlen = 1 +CPUInfer = cpuinfer_ext.CPUInfer(64) +warm_up_iter = 1000 +test_iter = 10000 + +def bench_linear(quant_mode: str): + with torch.inference_mode(mode=True): + + hidden_type = 30 # ggml_type::GGML_TYPE_BF16 + if quant_mode == "fp32": + proj_type = 0 # ggml_type::GGML_TYPE_F32 + bytes_per_elem = 4.000000 + elif quant_mode == "fp16": + proj_type = 1 # ggml_type::GGML_TYPE_F16 + bytes_per_elem = 2.000000 + elif quant_mode == "bf16": + proj_type = 30 # ggml_type::GGML_TYPE_BF16 + bytes_per_elem = 2.000000 + elif quant_mode == "q8_0": + proj_type = 8 # ggml_type::GGML_TYPE_Q8_0 + bytes_per_elem = 1.062500 + elif quant_mode == "q6_k": + proj_type = 14 # ggml_type::GGML_TYPE_Q6_K + bytes_per_elem = 0.820312 + elif quant_mode == "q5_k_m": + proj_type = 13 # ggml_type::GGML_TYPE_Q5_K + bytes_per_elem = 0.687500 + elif quant_mode == "q4_k_m": + proj_type = 12 # ggml_type::GGML_TYPE_Q4_K + bytes_per_elem = 0.562500 + elif quant_mode == "q3_k_m": + proj_type = 11 # ggml_type::GGML_TYPE_Q3_K + bytes_per_elem = 0.429688 + elif quant_mode == "q2_k": + proj_type = 10 # ggml_type::GGML_TYPE_Q2_K + bytes_per_elem = 0.328125 + elif quant_mode == "iq3_xs": + proj_type = 21 # ggml_type::GGML_TYPE_IQ3_S + bytes_per_elem = 0.429688 + elif quant_mode == "iq2_xxs": + proj_type = 16 # ggml_type::GGML_TYPE_IQ2_XXS + bytes_per_elem = 0.257812 + else: + assert(False) + + linears = [] + projs = [] + for _ in range(layer_num): + proj = torch.randn((output_size, input_size), dtype=torch.float32, device = "cuda").to("cpu").contiguous() + config = cpuinfer_ext.linear.LinearConfig(input_size, output_size, stride, group_max_len, proj.data_ptr(), proj_type, hidden_type) + linear = cpuinfer_ext.linear.Linear(config) + projs.append(proj) + linears.append(linear) + input = torch.randn((layer_num, qlen, input_size), dtype=torch.bfloat16, device = "cuda").to("cpu").contiguous() + output = torch.empty((layer_num, qlen, output_size), dtype=torch.bfloat16, device = "cuda").to("cpu").contiguous() + + # warm up + for i in range(warm_up_iter): + CPUInfer.submit( + linears[i % layer_num].forward( + qlen, + input[i % layer_num].data_ptr(), + output[i % layer_num].data_ptr() + ) + ) + CPUInfer.sync() + + # test + start = time.perf_counter() + for i in range(test_iter): + CPUInfer.submit( + linears[i % layer_num].forward( + qlen, + input[i % layer_num].data_ptr(), + output[i % layer_num].data_ptr() + ) + ) + CPUInfer.sync() + end = time.perf_counter() + total_time = end - start + print('Quant mode: ', quant_mode) + print('Time(s): ', total_time) + print('Iteration: ', test_iter) + print('Time(us) per iteration: ', total_time / test_iter * 1000000) + print('Bandwidth: ', input_size * output_size * bytes_per_elem * test_iter / total_time / 1000 / 1000 / 1000, 'GB/s') + print('') + +bench_linear("fp32") +bench_linear("fp16") +bench_linear("bf16") +bench_linear("q8_0") +bench_linear("q6_k") +bench_linear("q5_k_m") +bench_linear("q4_k_m") +bench_linear("q3_k_m") +bench_linear("q2_k") +# Not supported on __x86_64__ +# bench_linear("iq3_xs") +# bench_linear("iq2_xxs") diff --git a/archive/csrc/ktransformers_ext/bench/bench_linear_torch.py b/archive/csrc/ktransformers_ext/bench/bench_linear_torch.py new file mode 100644 index 0000000..72e0e75 --- /dev/null +++ b/archive/csrc/ktransformers_ext/bench/bench_linear_torch.py @@ -0,0 +1,83 @@ +#!/usr/bin/env python +# coding=utf-8 +''' +Description : +Author : chenht2022 +Date : 2024-07-25 10:31:59 +Version : 1.0.0 +LastEditors : chenht2022 +LastEditTime : 2024-07-25 10:32:48 +Copyright (c) 2024 by KVCache.AI, All Rights Reserved. +''' +import os, sys +import time +import torch +import torch.nn.quantized as nnq + +scale, zero_point = 0.1, 0 # Adjust scale and zero_point based on your dataset + +input_size = 16384 +output_size = 5120 +layer_num = 10 +qlen = 1 +warm_up_iter = 1000 +test_iter = 10000 + +def bench_linear(quant_mode: str): + with torch.inference_mode(mode=True): + if quant_mode == "fp32": + proj_type = torch.float32 + bytes_per_elem = 4.000000 + elif quant_mode == "fp16": + proj_type = torch.float16 + bytes_per_elem = 2.000000 + elif quant_mode == "bf16": + proj_type = torch.bfloat16 + bytes_per_elem = 2.000000 + elif quant_mode == "qint8": + proj_type = torch.qint8 + bytes_per_elem = 1.000000 + else: + assert(False) + + projs = [] + for _ in range(layer_num): + proj = torch.randn((output_size, input_size), dtype = torch.float32, device = "cuda").to("cpu").contiguous() + if quant_mode == "qint8": + proj_q = torch.quantize_per_tensor(proj, scale, zero_point, torch.qint8) + quantized_layer = nnq.Linear(input_size, output_size) + quantized_layer.set_weight_bias(proj_q, None) + projs.append(quantized_layer) + else: + projs.append(proj.to(proj_type)) + input = torch.randn((layer_num, qlen, input_size), dtype=torch.bfloat16, device = "cuda").to("cpu").contiguous() + + # warm up + for i in range(warm_up_iter): + if isinstance(projs[i % layer_num], nnq.Linear): + input_q = torch.quantize_per_tensor(input[i % layer_num].to(torch.float32), scale, zero_point, torch.quint8) + t_output = projs[i % layer_num](input_q) + else: + t_output = torch.mm(input[i % layer_num].to(proj_type), projs[i % layer_num].t()) + + # test + start = time.perf_counter() + for i in range(test_iter): + if isinstance(projs[i % layer_num], nnq.Linear): + input_q = torch.quantize_per_tensor(input[i % layer_num].to(torch.float32), scale, zero_point, torch.quint8) + t_output = projs[i % layer_num](input_q) + else: + t_output = torch.mm(input[i % layer_num].to(proj_type), projs[i % layer_num].t()) + end = time.perf_counter() + total_time = end - start + print('Quant mode: ', quant_mode) + print('Time(s): ', total_time) + print('Iteration: ', test_iter) + print('Time(us) per iteration: ', total_time / test_iter * 1000000) + print('Bandwidth: ', input_size * output_size * bytes_per_elem * test_iter / total_time / 1000 / 1000 / 1000, 'GB/s') + print('') + +bench_linear("fp32") +bench_linear("fp16") +bench_linear("bf16") +bench_linear("qint8") diff --git a/archive/csrc/ktransformers_ext/bench/bench_mlp.py b/archive/csrc/ktransformers_ext/bench/bench_mlp.py new file mode 100644 index 0000000..690f9d9 --- /dev/null +++ b/archive/csrc/ktransformers_ext/bench/bench_mlp.py @@ -0,0 +1,150 @@ +#!/usr/bin/env python +# coding=utf-8 +''' +Description : +Author : chenht2022 +Date : 2024-07-16 10:43:18 +Version : 1.0.0 +LastEditors : chenht2022 +LastEditTime : 2024-08-06 10:36:04 +Copyright (c) 2024 by KVCache.AI, All Rights Reserved. +''' +import os, sys +import time +sys.path.append(os.path.dirname(__file__) + '/../build') +import cpuinfer_ext +import torch + +hidden_size = 5120 +intermediate_size = 3072 +stride = 16 +group_max_len = 1024 +layer_num = 10 +qlen = 1 +CPUInfer = cpuinfer_ext.CPUInfer(64) +warm_up_iter = 1000 +test_iter = 10000 + +def bench_mlp(quant_mode: str): + with torch.inference_mode(mode=True): + + hidden_type = 30 # ggml_type::GGML_TYPE_BF16 + if quant_mode == "fp32": + gate_type = 0 # ggml_type::GGML_TYPE_F32 + up_type = 0 # ggml_type::GGML_TYPE_F32 + down_type = 0 # ggml_type::GGML_TYPE_F32 + bytes_per_elem = 4.000000 + elif quant_mode == "fp16": + gate_type = 1 # ggml_type::GGML_TYPE_F16 + up_type = 1 # ggml_type::GGML_TYPE_F16 + down_type = 1 # ggml_type::GGML_TYPE_F16 + bytes_per_elem = 2.000000 + elif quant_mode == "bf16": + gate_type = 30 # ggml_type::GGML_TYPE_BF16 + up_type = 30 # ggml_type::GGML_TYPE_BF16 + down_type = 30 # ggml_type::GGML_TYPE_BF16 + bytes_per_elem = 2.000000 + elif quant_mode == "q8_0": + gate_type = 8 # ggml_type::GGML_TYPE_Q8_0 + up_type = 8 # ggml_type::GGML_TYPE_Q8_0 + down_type = 8 # ggml_type::GGML_TYPE_Q8_0 + bytes_per_elem = 1.062500 + elif quant_mode == "q6_k": + gate_type = 14 # ggml_type::GGML_TYPE_Q6_K + up_type = 14 # ggml_type::GGML_TYPE_Q6_K + down_type = 14 # ggml_type::GGML_TYPE_Q6_K + bytes_per_elem = 0.820312 + elif quant_mode == "q5_k_m": + gate_type = 13 # ggml_type::GGML_TYPE_Q5_K + up_type = 13 # ggml_type::GGML_TYPE_Q5_K + down_type = 14 # ggml_type::GGML_TYPE_Q6_K + bytes_per_elem = 0.731771 + elif quant_mode == "q4_k_m": + gate_type = 12 # ggml_type::GGML_TYPE_Q4_K + up_type = 12 # ggml_type::GGML_TYPE_Q4_K + down_type = 14 # ggml_type::GGML_TYPE_Q6_K + bytes_per_elem = 0.648437 + elif quant_mode == "q3_k_m": + gate_type = 11 # ggml_type::GGML_TYPE_Q3_K + up_type = 11 # ggml_type::GGML_TYPE_Q3_K + down_type = 13 # ggml_type::GGML_TYPE_Q5_K + bytes_per_elem = 0.515625 + elif quant_mode == "q2_k": + gate_type = 10 # ggml_type::GGML_TYPE_Q2_K + up_type = 10 # ggml_type::GGML_TYPE_Q2_K + down_type = 11 # ggml_type::GGML_TYPE_Q3_K + bytes_per_elem = 0.328125 + elif quant_mode == "iq3_xs": + gate_type = 21 # ggml_type::GGML_TYPE_IQ3_S + up_type = 21 # ggml_type::GGML_TYPE_IQ3_S + down_type = 21 # ggml_type::GGML_TYPE_IQ3_S + bytes_per_elem = 0.429688 + elif quant_mode == "iq2_xxs": + gate_type = 16 # ggml_type::GGML_TYPE_IQ2_XXS + up_type = 16 # ggml_type::GGML_TYPE_IQ2_XXS + down_type = 16 # ggml_type::GGML_TYPE_IQ2_XXS + bytes_per_elem = 0.257812 + else: + assert(False) + + + mlps = [] + gate_projs = [] + up_projs = [] + down_projs = [] + for _ in range(layer_num): + gate_proj = torch.randn((intermediate_size, hidden_size), dtype=torch.float32, device = "cuda").to("cpu").contiguous() + up_proj = torch.randn((intermediate_size, hidden_size), dtype=torch.float32, device = "cuda").to("cpu").contiguous() + down_proj = torch.randn((hidden_size, intermediate_size), dtype=torch.float32, device = "cuda").to("cpu").contiguous() + config = cpuinfer_ext.mlp.MLPConfig(hidden_size, intermediate_size, stride, group_max_len, gate_proj.data_ptr(), up_proj.data_ptr(), down_proj.data_ptr(), gate_type, up_type, down_type, hidden_type) + mlp = cpuinfer_ext.mlp.MLP(config) + gate_projs.append(gate_proj) + up_projs.append(up_proj) + down_projs.append(down_proj) + mlps.append(mlp) + input = torch.randn((layer_num, qlen, hidden_size), dtype=torch.bfloat16, device = "cuda").to("cpu").contiguous() + output = torch.empty((layer_num, qlen, hidden_size), dtype=torch.bfloat16, device = "cuda").to("cpu").contiguous() + + # warm up + for i in range(warm_up_iter): + CPUInfer.submit( + mlps[i % layer_num].forward( + qlen, + input[i % layer_num].data_ptr(), + output[i % layer_num].data_ptr() + ) + ) + CPUInfer.sync() + + # test + start = time.perf_counter() + for i in range(test_iter): + CPUInfer.submit( + mlps[i % layer_num].forward( + qlen, + input[i % layer_num].data_ptr(), + output[i % layer_num].data_ptr() + ) + ) + CPUInfer.sync() + end = time.perf_counter() + total_time = end - start + print('Quant mode: ', quant_mode) + print('Time(s): ', total_time) + print('Iteration: ', test_iter) + print('Time(us) per iteration: ', total_time / test_iter * 1000000) + print('Bandwidth: ', hidden_size * intermediate_size * 3 * bytes_per_elem * test_iter / total_time / 1000 / 1000 / 1000, 'GB/s') + print('') + +bench_mlp("fp32") +bench_mlp("fp16") +bench_mlp("bf16") +bench_mlp("q8_0") +bench_mlp("q6_k") +bench_mlp("q5_k_m") +bench_mlp("q4_k_m") +bench_mlp("q3_k_m") +bench_mlp("q2_k") +# Not supported on __x86_64__ +# bench_linear("iq3_xs") +# bench_linear("iq2_xxs") diff --git a/archive/csrc/ktransformers_ext/bench/bench_mlp_torch.py b/archive/csrc/ktransformers_ext/bench/bench_mlp_torch.py new file mode 100644 index 0000000..7b811d8 --- /dev/null +++ b/archive/csrc/ktransformers_ext/bench/bench_mlp_torch.py @@ -0,0 +1,110 @@ +#!/usr/bin/env python +# coding=utf-8 +''' +Description : +Author : chenht2022 +Date : 2024-07-16 10:43:18 +Version : 1.0.0 +LastEditors : chenht2022 +LastEditTime : 2024-07-25 10:32:53 +Copyright (c) 2024 by KVCache.AI, All Rights Reserved. +''' +import os, sys +import time +import torch +import torch.nn.quantized as nnq + +scale, zero_point = 0.1, 0 # Adjust scale and zero_point based on your dataset + +hidden_size = 5120 +intermediate_size = 3072 +layer_num = 10 +qlen = 1 +warm_up_iter = 1000 +test_iter = 10000 + +def act_fn(x): + return x / (1.0 + torch.exp(-x)) + +def mlp_torch(input, gate_proj, up_proj, down_proj): + if isinstance(gate_proj, nnq.Linear): + input_q = torch.quantize_per_tensor(input.to(torch.float32), scale, zero_point, torch.quint8) + gate_buf = gate_proj(input_q) + up_buf = up_proj(input_q) + gate_buf = gate_buf.dequantize() + up_buf = up_buf.dequantize() + intermediate = act_fn(gate_buf) * up_buf + intermediate_q = torch.quantize_per_tensor(intermediate, scale, zero_point, torch.quint8) + expert_output = down_proj(intermediate_q) + ret = expert_output.dequantize() + else: + gate_buf = torch.mm(input.to(gate_proj.dtype), gate_proj.t()) + up_buf = torch.mm(input.to(up_proj.dtype), up_proj.t()) + intermediate = act_fn(gate_buf) * up_buf + ret = torch.mm(intermediate.to(down_proj.dtype), down_proj.t()) + return ret + +def bench_mlp(quant_mode: str): + with torch.inference_mode(mode=True): + if quant_mode == "fp32": + proj_type = torch.float32 + bytes_per_elem = 4.000000 + elif quant_mode == "fp16": + proj_type = torch.float16 + bytes_per_elem = 2.000000 + elif quant_mode == "bf16": + proj_type = torch.bfloat16 + bytes_per_elem = 2.000000 + elif quant_mode == "qint8": + proj_type = torch.qint8 + bytes_per_elem = 1.000000 + else: + assert(False) + + gate_projs = [] + up_projs = [] + down_projs = [] + for _ in range(layer_num): + gate_proj = torch.randn((intermediate_size, hidden_size), dtype=torch.float32, device = "cuda").to("cpu").contiguous() + up_proj = torch.randn((intermediate_size, hidden_size), dtype=torch.float32, device = "cuda").to("cpu").contiguous() + down_proj = torch.randn((hidden_size, intermediate_size), dtype=torch.float32, device = "cuda").to("cpu").contiguous() + if quant_mode == "qint8": + gate_proj_q = torch.quantize_per_tensor(gate_proj, scale, zero_point, torch.qint8) + quantized_gate = nnq.Linear(hidden_size, intermediate_size) + quantized_gate.set_weight_bias(gate_proj_q, None) + up_proj_q = torch.quantize_per_tensor(up_proj, scale, zero_point, torch.qint8) + quantized_up = nnq.Linear(hidden_size, intermediate_size) + quantized_up.set_weight_bias(up_proj_q, None) + down_proj_q = torch.quantize_per_tensor(down_proj, scale, zero_point, torch.qint8) + quantized_down = nnq.Linear(intermediate_size, hidden_size) + quantized_down.set_weight_bias(down_proj_q, None) + gate_projs.append(quantized_gate) + up_projs.append(quantized_up) + down_projs.append(quantized_down) + else: + gate_projs.append(gate_proj.to(proj_type)) + up_projs.append(up_proj.to(proj_type)) + down_projs.append(down_proj.to(proj_type)) + input = torch.randn((layer_num, qlen, hidden_size), dtype=torch.bfloat16, device = "cuda").to("cpu").contiguous() + + # warm up + for i in range(warm_up_iter): + mlp_torch(input[i % layer_num], gate_projs[i % layer_num], up_projs[i % layer_num], down_projs[i % layer_num]) + + # test + start = time.perf_counter() + for i in range(test_iter): + mlp_torch(input[i % layer_num], gate_projs[i % layer_num], up_projs[i % layer_num], down_projs[i % layer_num]) + end = time.perf_counter() + total_time = end - start + print('Quant mode: ', quant_mode) + print('Time(s): ', total_time) + print('Iteration: ', test_iter) + print('Time(us) per iteration: ', total_time / test_iter * 1000000) + print('Bandwidth: ', hidden_size * intermediate_size * 3 * bytes_per_elem * test_iter / total_time / 1000 / 1000 / 1000, 'GB/s') + print('') + +bench_mlp("fp32") +bench_mlp("fp16") +bench_mlp("bf16") +bench_mlp("qint8") diff --git a/archive/csrc/ktransformers_ext/bench/bench_moe.py b/archive/csrc/ktransformers_ext/bench/bench_moe.py new file mode 100644 index 0000000..6d617b7 --- /dev/null +++ b/archive/csrc/ktransformers_ext/bench/bench_moe.py @@ -0,0 +1,160 @@ +#!/usr/bin/env python +# coding=utf-8 +''' +Description : +Author : chenht2022 +Date : 2024-07-25 10:32:05 +Version : 1.0.0 +LastEditors : chenht2022 +LastEditTime : 2024-08-06 10:41:28 +Copyright (c) 2024 by KVCache.AI, All Rights Reserved. +''' +import os, sys +import time +sys.path.append(os.path.dirname(__file__) + '/../build') +import cpuinfer_ext +import torch + +expert_num = 160 +hidden_size = 5120 +intermediate_size = 1536 +stride = 16 +group_min_len = 10 +group_max_len = 1024 +n_routed_experts = 6 +layer_num = 10 +qlen = 1 +CPUInfer = cpuinfer_ext.CPUInfer(64) +warm_up_iter = 1000 +test_iter = 10000 + +def bench_moe(quant_mode: str): + with torch.inference_mode(mode=True): + hidden_type = 30 # ggml_type::GGML_TYPE_BF16 + if quant_mode == "fp32": + gate_type = 0 # ggml_type::GGML_TYPE_F32 + up_type = 0 # ggml_type::GGML_TYPE_F32 + down_type = 0 # ggml_type::GGML_TYPE_F32 + bytes_per_elem = 4.000000 + elif quant_mode == "fp16": + gate_type = 1 # ggml_type::GGML_TYPE_F16 + up_type = 1 # ggml_type::GGML_TYPE_F16 + down_type = 1 # ggml_type::GGML_TYPE_F16 + bytes_per_elem = 2.000000 + elif quant_mode == "bf16": + gate_type = 30 # ggml_type::GGML_TYPE_BF16 + up_type = 30 # ggml_type::GGML_TYPE_BF16 + down_type = 30 # ggml_type::GGML_TYPE_BF16 + bytes_per_elem = 2.000000 + elif quant_mode == "q8_0": + gate_type = 8 # ggml_type::GGML_TYPE_Q8_0 + up_type = 8 # ggml_type::GGML_TYPE_Q8_0 + down_type = 8 # ggml_type::GGML_TYPE_Q8_0 + bytes_per_elem = 1.062500 + elif quant_mode == "q6_k": + gate_type = 14 # ggml_type::GGML_TYPE_Q6_K + up_type = 14 # ggml_type::GGML_TYPE_Q6_K + down_type = 14 # ggml_type::GGML_TYPE_Q6_K + bytes_per_elem = 0.820312 + elif quant_mode == "q5_k_m": + gate_type = 13 # ggml_type::GGML_TYPE_Q5_K + up_type = 13 # ggml_type::GGML_TYPE_Q5_K + down_type = 14 # ggml_type::GGML_TYPE_Q6_K + bytes_per_elem = 0.731771 + elif quant_mode == "q4_k_m": + gate_type = 12 # ggml_type::GGML_TYPE_Q4_K + up_type = 12 # ggml_type::GGML_TYPE_Q4_K + down_type = 14 # ggml_type::GGML_TYPE_Q6_K + bytes_per_elem = 0.648437 + elif quant_mode == "q3_k_m": + gate_type = 11 # ggml_type::GGML_TYPE_Q3_K + up_type = 11 # ggml_type::GGML_TYPE_Q3_K + down_type = 13 # ggml_type::GGML_TYPE_Q5_K + bytes_per_elem = 0.515625 + elif quant_mode == "q2_k": + gate_type = 10 # ggml_type::GGML_TYPE_Q2_K + up_type = 10 # ggml_type::GGML_TYPE_Q2_K + down_type = 11 # ggml_type::GGML_TYPE_Q3_K + bytes_per_elem = 0.328125 + elif quant_mode == "iq3_xs": + gate_type = 21 # ggml_type::GGML_TYPE_IQ3_S + up_type = 21 # ggml_type::GGML_TYPE_IQ3_S + down_type = 21 # ggml_type::GGML_TYPE_IQ3_S + bytes_per_elem = 0.429688 + elif quant_mode == "iq2_xxs": + gate_type = 16 # ggml_type::GGML_TYPE_IQ2_XXS + up_type = 16 # ggml_type::GGML_TYPE_IQ2_XXS + down_type = 16 # ggml_type::GGML_TYPE_IQ2_XXS + bytes_per_elem = 0.257812 + else: + assert(False) + + + moes = [] + gate_projs = [] + up_projs = [] + down_projs = [] + for _ in range(layer_num): + gate_proj = torch.randn((expert_num, intermediate_size, hidden_size), dtype=torch.float32, device = "cuda").to("cpu").contiguous() + up_proj = torch.randn((expert_num, intermediate_size, hidden_size), dtype=torch.float32, device = "cuda").to("cpu").contiguous() + down_proj = torch.randn((expert_num, hidden_size, intermediate_size), dtype=torch.float32, device = "cuda").to("cpu").contiguous() + config = cpuinfer_ext.moe.MOEConfig(expert_num, n_routed_experts, hidden_size, intermediate_size, stride, group_min_len, group_max_len, gate_proj.data_ptr(), up_proj.data_ptr(), down_proj.data_ptr(), gate_type, up_type, down_type, hidden_type) + moe = cpuinfer_ext.moe.MOE(config) + gate_projs.append(gate_proj) + up_projs.append(up_proj) + down_projs.append(down_proj) + moes.append(moe) + expert_ids = torch.stack([torch.stack([torch.randperm(expert_num, dtype=torch.int64, device = "cuda")[:n_routed_experts] for _ in range(qlen)]) for _ in range(layer_num)]).to("cpu").contiguous() + weights = torch.rand((layer_num, qlen, n_routed_experts), dtype=torch.float32, device = "cuda").to("cpu").contiguous() + input = torch.randn((layer_num, qlen, hidden_size), dtype=torch.bfloat16, device = "cuda").to("cpu").contiguous() + output = torch.empty((layer_num, qlen, hidden_size), dtype=torch.bfloat16, device = "cuda").to("cpu").contiguous() + + # warm up + for i in range(warm_up_iter): + CPUInfer.submit( + moes[i % layer_num].forward( + qlen, + n_routed_experts, + expert_ids[i % layer_num].data_ptr(), + weights[i % layer_num].data_ptr(), + input[i % layer_num].data_ptr(), + output[i % layer_num].data_ptr() + ) + ) + CPUInfer.sync() + + # test + start = time.perf_counter() + for i in range(test_iter): + CPUInfer.submit( + moes[i % layer_num].forward( + qlen, + n_routed_experts, + expert_ids[i % layer_num].data_ptr(), + weights[i % layer_num].data_ptr(), + input[i % layer_num].data_ptr(), + output[i % layer_num].data_ptr() + ) + ) + CPUInfer.sync() + end = time.perf_counter() + total_time = end - start + print('Quant mode: ', quant_mode) + print('Time(s): ', total_time) + print('Iteration: ', test_iter) + print('Time(us) per iteration: ', total_time / test_iter * 1000000) + print('Bandwidth: ', hidden_size * intermediate_size * 3 * n_routed_experts * bytes_per_elem * test_iter / total_time / 1000 / 1000 / 1000, 'GB/s') + print('') + +bench_moe("fp32") +bench_moe("fp16") +bench_moe("bf16") +bench_moe("q8_0") +bench_moe("q6_k") +bench_moe("q5_k_m") +bench_moe("q4_k_m") +bench_moe("q3_k_m") +bench_moe("q2_k") +# Not supported on __x86_64__ +# bench_linear("iq3_xs") +# bench_linear("iq2_xxs") diff --git a/archive/csrc/ktransformers_ext/bench/bench_moe_amx.py b/archive/csrc/ktransformers_ext/bench/bench_moe_amx.py new file mode 100644 index 0000000..1101ff8 --- /dev/null +++ b/archive/csrc/ktransformers_ext/bench/bench_moe_amx.py @@ -0,0 +1,107 @@ +#!/usr/bin/env python +# coding=utf-8 +''' +Description : +Author : chenht2022 +Date : 2025-04-25 18:28:12 +Version : 1.0.0 +LastEditors : chenht2022 +LastEditTime : 2025-04-25 18:28:12 +Copyright (c) 2024 by KVCache.AI, All Rights Reserved. +''' +import os, sys +import time +sys.path.append(os.path.dirname(__file__) + '/../build') +import cpuinfer_ext +import torch + +expert_num = 8 +hidden_size = 7168 +intermediate_size = 2048 +max_len = 25600 +n_routed_experts = 8 +layer_num = 10 +qlen = 1024 +CPUInfer = cpuinfer_ext.CPUInfer(65) +warm_up_iter = 100 +test_iter = 100 + +def bench_moe(quant_mode: str): + with torch.inference_mode(mode=True): + if quant_mode == "bf16": + bytes_per_elem = 2.000000 + elif quant_mode == "int8": + bytes_per_elem = 1.000000 + else: + assert(False) + + + moes = [] + gate_projs = [] + up_projs = [] + down_projs = [] + for _ in range(layer_num): + gate_proj = torch.randn((expert_num, intermediate_size, hidden_size), dtype=torch.float32, device = "cuda").to("cpu").contiguous() + up_proj = torch.randn((expert_num, intermediate_size, hidden_size), dtype=torch.float32, device = "cuda").to("cpu").contiguous() + down_proj = torch.randn((expert_num, hidden_size, intermediate_size), dtype=torch.float32, device = "cuda").to("cpu").contiguous() + config = cpuinfer_ext.moe.AMX_MOEConfig(expert_num, n_routed_experts, hidden_size, intermediate_size, max_len, gate_proj.data_ptr(), up_proj.data_ptr(), down_proj.data_ptr()) + if quant_mode == "bf16": + moe = cpuinfer_ext.moe.AMXBF16_MOE(config) + CPUInfer.submit(moe.load_weights()) + CPUInfer.sync() + elif quant_mode == "int8": + moe = cpuinfer_ext.moe.AMXInt8_MOE(config) + CPUInfer.submit(moe.load_weights()) + CPUInfer.sync() + gate_projs.append(gate_proj) + up_projs.append(up_proj) + down_projs.append(down_proj) + moes.append(moe) + expert_ids = torch.stack([torch.stack([torch.randperm(expert_num, dtype=torch.int64, device = "cuda")[:n_routed_experts] for _ in range(qlen)]) for _ in range(layer_num)]).to("cpu").contiguous() + weights = torch.rand((layer_num, qlen, n_routed_experts), dtype=torch.float32, device = "cuda").to("cpu").contiguous() + input = torch.randn((layer_num, qlen, hidden_size), dtype=torch.bfloat16, device = "cuda").to("cpu").contiguous() + output = torch.empty((layer_num, qlen, hidden_size), dtype=torch.bfloat16, device = "cuda").to("cpu").contiguous() + qlen_tensor = torch.tensor([qlen], dtype=torch.int32) + + # warm up + for i in range(warm_up_iter): + CPUInfer.submit( + moes[i % layer_num].forward( + qlen, + n_routed_experts, + expert_ids[i % layer_num].data_ptr(), + weights[i % layer_num].data_ptr(), + input[i % layer_num].data_ptr(), + output[i % layer_num].data_ptr(), + qlen_tensor.data_ptr() + ) + ) + CPUInfer.sync() + + # test + start = time.perf_counter() + for i in range(test_iter): + CPUInfer.submit( + moes[i % layer_num].forward( + qlen, + n_routed_experts, + expert_ids[i % layer_num].data_ptr(), + weights[i % layer_num].data_ptr(), + input[i % layer_num].data_ptr(), + output[i % layer_num].data_ptr(), + qlen_tensor.data_ptr() + ) + ) + CPUInfer.sync() + end = time.perf_counter() + total_time = end - start + print('Quant mode: ', quant_mode) + print('Time(s): ', total_time) + print('Iteration: ', test_iter) + print('Time(us) per iteration: ', total_time / test_iter * 1000000) + print('Bandwidth: ', hidden_size * intermediate_size * 3 * n_routed_experts * bytes_per_elem * test_iter / total_time / 1000 / 1000 / 1000, 'GB/s') + print('Flops: ', hidden_size * intermediate_size * qlen * 3 * n_routed_experts * 2 * test_iter / total_time / 1000 / 1000 / 1000, 'GFLOPS') + print('') + +bench_moe("bf16") +bench_moe("int8") diff --git a/archive/csrc/ktransformers_ext/bench/bench_moe_torch.py b/archive/csrc/ktransformers_ext/bench/bench_moe_torch.py new file mode 100644 index 0000000..1aecf40 --- /dev/null +++ b/archive/csrc/ktransformers_ext/bench/bench_moe_torch.py @@ -0,0 +1,152 @@ +#!/usr/bin/env python +# coding=utf-8 +''' +Description : +Author : chenht2022 +Date : 2024-07-25 10:32:05 +Version : 1.0.0 +LastEditors : chenht2022 +LastEditTime : 2024-07-25 10:32:57 +Copyright (c) 2024 by KVCache.AI, All Rights Reserved. +''' +import os, sys +import time +import torch +import torch.nn.quantized as nnq + +scale, zero_point = 0.1, 0 # Adjust scale and zero_point based on your dataset + +expert_num = 160 +hidden_size = 5120 +intermediate_size = 1536 +n_routed_experts = 6 +layer_num = 10 +qlen = 1 +warm_up_iter = 1000 +test_iter = 10000 + +def act_fn(x): + return x / (1.0 + torch.exp(-x)) + +def mlp_torch(input, gate_proj, up_proj, down_proj): + if isinstance(gate_proj, nnq.Linear): + input_q = torch.quantize_per_tensor(input.to(torch.float32), scale, zero_point, torch.quint8) + gate_buf = gate_proj(input_q) + up_buf = up_proj(input_q) + gate_buf = gate_buf.dequantize() + up_buf = up_buf.dequantize() + intermediate = act_fn(gate_buf) * up_buf + intermediate_q = torch.quantize_per_tensor(intermediate, scale, zero_point, torch.quint8) + expert_output = down_proj(intermediate_q) + ret = expert_output.dequantize() + else: + gate_buf = torch.mm(input.to(gate_proj.dtype), gate_proj.t()) + up_buf = torch.mm(input.to(up_proj.dtype), up_proj.t()) + intermediate = act_fn(gate_buf) * up_buf + ret = torch.mm(intermediate.to(down_proj.dtype), down_proj.t()) + return ret + +def moe_torch(input, expert_ids, weights, gate_proj, up_proj, down_proj): + cnts = expert_ids.new_zeros((expert_ids.shape[0], expert_num)) + cnts.scatter_(1, expert_ids, 1) + tokens_per_expert = cnts.sum(dim=0) + idxs = expert_ids.view(-1).argsort() + sorted_tokens = input[idxs // expert_ids.shape[1]] + + outputs = [] + start_idx = 0 + for i, num_tokens in enumerate(tokens_per_expert): + end_idx = start_idx + num_tokens + if num_tokens == 0: + continue + tokens_for_this_expert = sorted_tokens[start_idx:end_idx] + expert_out = mlp_torch(tokens_for_this_expert, gate_proj[i], up_proj[i], down_proj[i]) + outputs.append(expert_out) + start_idx = end_idx + + outs = torch.cat(outputs, dim=0) if len(outputs) else sorted_tokens.new_empty(0) + + new_x = torch.empty_like(outs) + new_x[idxs] = outs + t_output = ( + new_x.view(*expert_ids.shape, -1) + .type(weights.dtype) + .mul_(weights.unsqueeze(dim=-1)) + .sum(dim=1) + .type(new_x.dtype) + ) + return t_output + +def bench_moe(quant_mode: str): + with torch.inference_mode(mode=True): + if quant_mode == "fp32": + proj_type = torch.float32 + bytes_per_elem = 4.000000 + elif quant_mode == "fp16": + proj_type = torch.float16 + bytes_per_elem = 2.000000 + elif quant_mode == "bf16": + proj_type = torch.bfloat16 + bytes_per_elem = 2.000000 + elif quant_mode == "qint8": + proj_type = torch.qint8 + bytes_per_elem = 1.000000 + else: + assert(False) + + gate_projs = [] + up_projs = [] + down_projs = [] + for _ in range(layer_num): + gate_proj = torch.randn((expert_num, intermediate_size, hidden_size), dtype=torch.float32, device = "cuda").to("cpu").contiguous() + up_proj = torch.randn((expert_num, intermediate_size, hidden_size), dtype=torch.float32, device = "cuda").to("cpu").contiguous() + down_proj = torch.randn((expert_num, hidden_size, intermediate_size), dtype=torch.float32, device = "cuda").to("cpu").contiguous() + if quant_mode == "qint8": + quantized_gate_proj = [] + quantized_up_proj = [] + quantized_down_proj = [] + for i in range(expert_num): + gate_proj_q = torch.quantize_per_tensor(gate_proj[i], scale, zero_point, torch.qint8) + quantized_gate = nnq.Linear(hidden_size, intermediate_size) + quantized_gate.set_weight_bias(gate_proj_q, None) + quantized_gate_proj.append(quantized_gate) + up_proj_q = torch.quantize_per_tensor(up_proj[i], scale, zero_point, torch.qint8) + quantized_up = nnq.Linear(hidden_size, intermediate_size) + quantized_up.set_weight_bias(up_proj_q, None) + quantized_up_proj.append(quantized_up) + down_proj_q = torch.quantize_per_tensor(down_proj[i], scale, zero_point, torch.qint8) + quantized_down = nnq.Linear(intermediate_size, hidden_size) + quantized_down.set_weight_bias(down_proj_q, None) + quantized_down_proj.append(quantized_down) + gate_projs.append(quantized_gate_proj) + up_projs.append(quantized_up_proj) + down_projs.append(quantized_down_proj) + else: + gate_projs.append(gate_proj.to(proj_type)) + up_projs.append(up_proj.to(proj_type)) + down_projs.append(down_proj.to(proj_type)) + expert_ids = torch.stack([torch.stack([torch.randperm(expert_num, dtype=torch.int64, device = "cuda")[:n_routed_experts] for _ in range(qlen)]) for _ in range(layer_num)]).to("cpu").contiguous() + weights = torch.rand((layer_num, qlen, n_routed_experts), dtype=torch.float32, device = "cuda").to("cpu").contiguous() + input = torch.randn((layer_num, qlen, hidden_size), dtype=torch.bfloat16, device = "cuda").to("cpu").contiguous() + + # warm up + for i in range(warm_up_iter): + moe_torch(input[i % layer_num], expert_ids[i % layer_num], weights[i % layer_num], gate_projs[i % layer_num], up_projs[i % layer_num], down_projs[i % layer_num]) + + # test + start = time.perf_counter() + for i in range(test_iter): + moe_torch(input[i % layer_num], expert_ids[i % layer_num], weights[i % layer_num], gate_projs[i % layer_num], up_projs[i % layer_num], down_projs[i % layer_num]) + end = time.perf_counter() + total_time = end - start + print('Quant mode: ', quant_mode) + print('Time(s): ', total_time) + print('Iteration: ', test_iter) + print('Time(us) per iteration: ', total_time / test_iter * 1000000) + print('Bandwidth: ', hidden_size * intermediate_size * 3 * n_routed_experts * bytes_per_elem * test_iter / total_time / 1000 / 1000 / 1000, 'GB/s') + print('') + +bench_moe("fp32") +bench_moe("fp16") +bench_moe("bf16") +bench_moe("qint8") diff --git a/archive/csrc/ktransformers_ext/cmake/FindSIMD.cmake b/archive/csrc/ktransformers_ext/cmake/FindSIMD.cmake new file mode 100644 index 0000000..33377ec --- /dev/null +++ b/archive/csrc/ktransformers_ext/cmake/FindSIMD.cmake @@ -0,0 +1,100 @@ +include(CheckCSourceRuns) + +set(AVX_CODE " + #include + int main() + { + __m256 a; + a = _mm256_set1_ps(0); + return 0; + } +") + +set(AVX512_CODE " + #include + int main() + { + __m512i a = _mm512_set_epi8(0, 0, 0, 0, 0, 0, 0, 0, + 0, 0, 0, 0, 0, 0, 0, 0, + 0, 0, 0, 0, 0, 0, 0, 0, + 0, 0, 0, 0, 0, 0, 0, 0, + 0, 0, 0, 0, 0, 0, 0, 0, + 0, 0, 0, 0, 0, 0, 0, 0, + 0, 0, 0, 0, 0, 0, 0, 0, + 0, 0, 0, 0, 0, 0, 0, 0); + __m512i b = a; + __mmask64 equality_mask = _mm512_cmp_epi8_mask(a, b, _MM_CMPINT_EQ); + return 0; + } +") + +set(AVX2_CODE " + #include + int main() + { + __m256i a = {0}; + a = _mm256_abs_epi16(a); + __m256i x; + _mm256_extract_epi64(x, 0); // we rely on this in our AVX2 code + return 0; + } +") + +set(FMA_CODE " + #include + int main() + { + __m256 acc = _mm256_setzero_ps(); + const __m256 d = _mm256_setzero_ps(); + const __m256 p = _mm256_setzero_ps(); + acc = _mm256_fmadd_ps( d, p, acc ); + return 0; + } +") + +macro(check_sse type flags) + set(__FLAG_I 1) + set(CMAKE_REQUIRED_FLAGS_SAVE ${CMAKE_REQUIRED_FLAGS}) + foreach (__FLAG ${flags}) + if (NOT ${type}_FOUND) + set(CMAKE_REQUIRED_FLAGS ${__FLAG}) + check_c_source_runs("${${type}_CODE}" HAS_${type}_${__FLAG_I}) + if (HAS_${type}_${__FLAG_I}) + set(${type}_FOUND TRUE CACHE BOOL "${type} support") + set(${type}_FLAGS "${__FLAG}" CACHE STRING "${type} flags") + endif() + math(EXPR __FLAG_I "${__FLAG_I}+1") + endif() + endforeach() + set(CMAKE_REQUIRED_FLAGS ${CMAKE_REQUIRED_FLAGS_SAVE}) + + if (NOT ${type}_FOUND) + set(${type}_FOUND FALSE CACHE BOOL "${type} support") + set(${type}_FLAGS "" CACHE STRING "${type} flags") + endif() + + mark_as_advanced(${type}_FOUND ${type}_FLAGS) +endmacro() + +# flags are for MSVC only! +check_sse("AVX" " ;/arch:AVX") +if (NOT ${AVX_FOUND}) + set(LLAMA_AVX OFF) +else() + set(LLAMA_AVX ON) +endif() + +check_sse("AVX2" " ;/arch:AVX2") +check_sse("FMA" " ;/arch:AVX2") +if ((NOT ${AVX2_FOUND}) OR (NOT ${FMA_FOUND})) + set(LLAMA_AVX2 OFF) +else() + set(LLAMA_AVX2 ON) +endif() + +check_sse("AVX512" " ;/arch:AVX512") +if (NOT ${AVX512_FOUND}) + set(LLAMA_AVX512 OFF) +else() + set(LLAMA_AVX512 ON) +endif() diff --git a/archive/csrc/ktransformers_ext/cpu_backend/backend.cpp b/archive/csrc/ktransformers_ext/cpu_backend/backend.cpp new file mode 100644 index 0000000..7478d5c --- /dev/null +++ b/archive/csrc/ktransformers_ext/cpu_backend/backend.cpp @@ -0,0 +1,154 @@ +/** + * @Description : + * @Author : chenht2022 + * @Date : 2024-07-22 02:03:05 + * @Version : 1.0.0 + * @LastEditors : chenht2022 + * @LastEditTime : 2024-07-25 10:33:34 + * @Copyright (c) 2024 by KVCache.AI, All Rights Reserved. + **/ + +#include "backend.h" + +#ifdef USE_NUMA +#include +#include + +thread_local int Backend::numa_node = -1; +#endif + +thread_local int Backend::thread_local_id = -1; + +Backend::Backend(int max_thread_num) { + max_thread_num_ = max_thread_num; + thread_state_.resize(max_thread_num_); + for (int i = 0; i < max_thread_num_; i++) { + thread_state_[i].curr = std::make_unique>(); + thread_state_[i].status = + std::make_unique>(ThreadStatus::WAITING); + } + workers_.resize(max_thread_num_); + for (int i = 1; i < max_thread_num_; i++) { + workers_[i] = std::thread(&Backend::worker_thread, this, i); + } +} + +Backend::~Backend() { + for (int i = 0; i < max_thread_num_; i++) { + thread_state_[i].status->store(ThreadStatus::EXIT, + std::memory_order_release); + } + for (int i = 1; i < max_thread_num_; i++) { + if (workers_[i].joinable()) { + workers_[i].join(); + } + } +} + +int Backend::get_thread_num() { return max_thread_num_; } + +void Backend::do_work_stealing_job(int task_num, + std::function init_func, + std::function compute_func, + std::function finalize_func) { + init_func_ = init_func; + compute_func_ = compute_func; + finalize_func_ = finalize_func; +#ifdef USE_NUMA + // numa node location will be calculated based on the number of threads + thread_num_ = max_thread_num_; +#else + thread_num_ = std::min(max_thread_num_, task_num); +#endif + int base = task_num / thread_num_; + int remain = task_num % thread_num_; + thread_state_[0].end = base + (0 < remain); + + // 为主线程设置 thread_local_id + thread_local_id = 0; + + for (int i = 1; i < thread_num_; i++) { + thread_state_[i].curr->store(thread_state_[i - 1].end, + std::memory_order_relaxed); + thread_state_[i].end = thread_state_[i - 1].end + base + (i < remain); + thread_state_[i].status->store(ThreadStatus::WORKING, + std::memory_order_release); + } + thread_state_[0].curr->store(0, std::memory_order_relaxed); + thread_state_[0].status->store(ThreadStatus::WORKING, + std::memory_order_release); + process_tasks(0); + for (int i = 1; i < thread_num_; i++) { + while (thread_state_[i].status->load(std::memory_order_acquire) == + ThreadStatus::WORKING) { + } + } +} + +void Backend::process_tasks(int thread_id) { + + #ifdef USE_NUMA + if(numa_node == -1){ + numa_node = thread_id * numa_num_configured_nodes() / thread_num_; + struct bitmask* mask = numa_bitmask_alloc(numa_num_configured_nodes()); + numa_bitmask_setbit(mask, numa_node); + numa_bind(mask); + } + #endif + + if (init_func_ != nullptr) { + init_func_(thread_id); + } + while (true) { + int task_id = thread_state_[thread_id].curr->fetch_add( + 1, std::memory_order_acq_rel); + if (task_id >= thread_state_[thread_id].end) { + break; + } + compute_func_(task_id); + } + for (int t_offset = 1; t_offset < thread_num_; t_offset++) { + int t_i = (thread_id + t_offset) % thread_num_; + if (thread_state_[t_i].status->load(std::memory_order_acquire) != + ThreadStatus::WORKING) { + continue; + } + while (true) { + int task_id = thread_state_[t_i].curr->fetch_add( + 1, std::memory_order_acq_rel); + if (task_id >= thread_state_[t_i].end) { + break; + } + compute_func_(task_id); + } + } + if (finalize_func_ != nullptr) { + finalize_func_(thread_id); + } + thread_state_[thread_id].status->store(ThreadStatus::WAITING, + std::memory_order_release); +} + +void Backend::worker_thread(int thread_id) { + auto start = std::chrono::steady_clock::now(); + thread_local_id = thread_id; // 设置线程本地变量 + while (true) { + ThreadStatus status = + thread_state_[thread_id].status->load(std::memory_order_acquire); + if (status == ThreadStatus::WORKING) { + process_tasks(thread_id); + start = std::chrono::steady_clock::now(); + } else if (status == ThreadStatus::WAITING) { + auto now = std::chrono::steady_clock::now(); + auto duration = + std::chrono::duration_cast(now - + start) + .count(); + if (duration > 50) { + std::this_thread::sleep_for(std::chrono::milliseconds(1)); + } + } else if (status == ThreadStatus::EXIT) { + return; + } + } +} \ No newline at end of file diff --git a/archive/csrc/ktransformers_ext/cpu_backend/backend.h b/archive/csrc/ktransformers_ext/cpu_backend/backend.h new file mode 100644 index 0000000..7a95f27 --- /dev/null +++ b/archive/csrc/ktransformers_ext/cpu_backend/backend.h @@ -0,0 +1,58 @@ +/** + * @Description : + * @Author : chenht2022 + * @Date : 2024-07-22 02:03:05 + * @Version : 1.0.0 + * @LastEditors : chenht2022 + * @LastEditTime : 2024-07-25 10:33:38 + * @Copyright (c) 2024 by KVCache.AI, All Rights Reserved. + **/ +#ifndef CPUINFER_BACKEND_H +#define CPUINFER_BACKEND_H + +#include +#include +#include +#include +#include +#include +#include + +enum ThreadStatus { + WORKING, + WAITING, + EXIT, +}; + +struct ThreadState { + std::unique_ptr> status; + std::unique_ptr> curr; + int end; +}; + +class Backend { + public: + Backend(int); + ~Backend(); + int get_thread_num(); + void do_work_stealing_job(int, std::function, + std::function, + std::function); + #ifdef USE_NUMA + static thread_local int numa_node; + #endif + static thread_local int thread_local_id; + + private: + int thread_num_; + int max_thread_num_; + std::vector thread_state_; // [thread_num] + std::function init_func_; + std::function compute_func_; + std::function finalize_func_; + std::vector workers_; + + void process_tasks(int); + void worker_thread(int); +}; +#endif \ No newline at end of file diff --git a/archive/csrc/ktransformers_ext/cpu_backend/cpuinfer.h b/archive/csrc/ktransformers_ext/cpu_backend/cpuinfer.h new file mode 100644 index 0000000..7b1d898 --- /dev/null +++ b/archive/csrc/ktransformers_ext/cpu_backend/cpuinfer.h @@ -0,0 +1,98 @@ +/** + * @Description : + * @Author : chenht2022 + * @Date : 2024-07-16 10:43:18 + * @Version : 1.0.0 + * @LastEditors : chenht2022 + * @LastEditTime : 2024-08-07 09:47:43 + * @Copyright (c) 2024 by KVCache.AI, All Rights Reserved. + **/ + #ifndef CPUINFER_CPUINFER_H + #define CPUINFER_CPUINFER_H + + #include + #include + #include + #include + #include + #include + #include + #include + #ifdef KTRANSFORMERS_USE_CUDA + #include "vendors/cuda.h" + #elif KTRANSFORMERS_USE_MUSA + #include "vendors/musa.h" + #elif KTRANSFORMERS_USE_ROCM + #define __HIP_PLATFORM_AMD__ + #include "vendors/hip.h" + #endif + + #include "backend.h" + #include "task_queue.h" + #include "./vendors/vendor.h" + + #include "llama.cpp/ggml-impl.h" + + class CPUInfer { + public: + CPUInfer(int thread_num) { + backend_ = new Backend(thread_num - 1); + task_queue_ = new TaskQueue(); + for (int i = 0; i < (1 << 16); ++i) { + ggml_table_f32_f16[i] = GGML_COMPUTE_FP16_TO_FP32(i); + } + } + + ~CPUInfer() { + delete backend_; + delete task_queue_; + } + + template + void enqueue(Func f, Obj* obj, Args... args) { + task_queue_->enqueue([=]() { + std::invoke(f, *obj, args..., backend_); + }); + } + + void submit(std::pair params) { + void (*func)(void*) = (void (*)(void*))params.first; + void* args = (void*)params.second; + *((CPUInfer**)args) = this; + func(args); + } + + void sync() { + task_queue_->sync(); + } + + void submit_with_cuda_stream(intptr_t user_cuda_stream, std::pair params) { + #if defined(KTRANSFORMERS_USE_CUDA) || defined(KTRANSFORMERS_USE_MUSA) || defined(KTRANSFORMERS_USE_ROCM) + void (*func)(void*) = (void (*)(void*))params.first; + void* args = (void*)params.second; + *((CPUInfer**)args) = this; + cudaLaunchHostFunc((cudaStream_t)user_cuda_stream, (cudaHostFn_t)func, args); + #else + throw std::runtime_error("submit_with_cuda_stream is not supported on this platforma"); + #endif + } + + static void sync_(void* cpu_infer_ptr) { + CPUInfer* cpuinfer = (CPUInfer*)cpu_infer_ptr; + cpuinfer->sync(); + } + + void sync_with_cuda_stream(intptr_t user_cuda_stream) { + #if defined(KTRANSFORMERS_USE_CUDA) || defined(KTRANSFORMERS_USE_MUSA) || defined(KTRANSFORMERS_USE_ROCM) + cudaLaunchHostFunc((cudaStream_t)user_cuda_stream, (cudaHostFn_t)&sync_, (void*)this); + #else + throw std::runtime_error("sync_with_cuda_stream is not supported on this platforma"); + #endif + } + + public: + Backend* backend_; + TaskQueue* task_queue_; + }; + + #endif \ No newline at end of file diff --git a/archive/csrc/ktransformers_ext/cpu_backend/shared_mem_buffer.cpp b/archive/csrc/ktransformers_ext/cpu_backend/shared_mem_buffer.cpp new file mode 100644 index 0000000..0fe7275 --- /dev/null +++ b/archive/csrc/ktransformers_ext/cpu_backend/shared_mem_buffer.cpp @@ -0,0 +1,56 @@ +/** + * @Description : + * @Author : chenht2022 + * @Date : 2024-08-05 04:49:08 + * @Version : 1.0.0 + * @LastEditors : chenht2022 + * @LastEditTime : 2024-08-05 09:21:29 + * @Copyright (c) 2024 by KVCache.AI, All Rights Reserved. + **/ +#include "shared_mem_buffer.h" +#include + +SharedMemBuffer::SharedMemBuffer() { + buffer_ = nullptr; + size_ = 0; +} + +SharedMemBuffer::~SharedMemBuffer() { + if (buffer_) { + free(buffer_); + } +} + +void SharedMemBuffer::alloc(void* object, std::vector> requests) { + uint64_t size = 0; + for (auto& request : requests) { + size += request.second; + } + if (size > size_) { + if (buffer_) { + free(buffer_); + } + buffer_ = std::aligned_alloc(64, size); + + size_ = size; + for (auto& obj_requests : hist_requests_) { + for (auto& requests : obj_requests.second) { + arrange(requests); + } + } + } + arrange(requests); + hist_requests_[object].push_back(requests); +} + +void SharedMemBuffer::dealloc(void* object) { + hist_requests_.erase(object); +} + +void SharedMemBuffer::arrange(std::vector> requests) { + uint64_t offset = 0; + for (auto& request : requests) { + *(request.first) = (uint8_t*)buffer_ + offset; + offset += request.second; + } +} \ No newline at end of file diff --git a/archive/csrc/ktransformers_ext/cpu_backend/shared_mem_buffer.h b/archive/csrc/ktransformers_ext/cpu_backend/shared_mem_buffer.h new file mode 100644 index 0000000..58e7aa6 --- /dev/null +++ b/archive/csrc/ktransformers_ext/cpu_backend/shared_mem_buffer.h @@ -0,0 +1,37 @@ +/** + * @Description : + * @Author : chenht2022 + * @Date : 2024-08-05 04:49:08 + * @Version : 1.0.0 + * @LastEditors : chenht2022 + * @LastEditTime : 2024-08-05 06:36:41 + * @Copyright (c) 2024 by KVCache.AI, All Rights Reserved. + **/ + + #ifndef CPUINFER_SHAREDMEMBUFFER_H + #define CPUINFER_SHAREDMEMBUFFER_H + + #include + #include + #include + #include + + class SharedMemBuffer { + public: + SharedMemBuffer(); + ~SharedMemBuffer(); + + void alloc(void* object, std::vector> requests); + void dealloc(void* object); + + private: + void* buffer_; + uint64_t size_; + std::map>>> hist_requests_; + + void arrange(std::vector> requests); + }; + + static SharedMemBuffer shared_mem_buffer; + + #endif \ No newline at end of file diff --git a/archive/csrc/ktransformers_ext/cpu_backend/task_queue.cpp b/archive/csrc/ktransformers_ext/cpu_backend/task_queue.cpp new file mode 100644 index 0000000..297a4ea --- /dev/null +++ b/archive/csrc/ktransformers_ext/cpu_backend/task_queue.cpp @@ -0,0 +1,67 @@ +/** + * @Description : + * @Author : chenht2022 + * @Date : 2024-07-17 12:25:51 + * @Version : 1.0.0 + * @LastEditors : chenht2022 + * @LastEditTime : 2024-10-09 11:08:10 + * @Copyright (c) 2024 by KVCache.AI, All Rights Reserved. + **/ +#include "task_queue.h" + +TaskQueue::TaskQueue() { + worker = std::thread(&TaskQueue::processTasks, this); + sync_flag.store(true, std::memory_order_seq_cst); + exit_flag.store(false, std::memory_order_seq_cst); +} + +TaskQueue::~TaskQueue() { + { + mutex.lock(); + exit_flag.store(true, std::memory_order_seq_cst); + mutex.unlock(); + } + cv.notify_all(); + if (worker.joinable()) { + worker.join(); + } +} + +void TaskQueue::enqueue(std::function task) { + { + mutex.lock(); + tasks.push(task); + sync_flag.store(false, std::memory_order_seq_cst); + mutex.unlock(); + } + cv.notify_one(); +} + +void TaskQueue::sync() { + while (!sync_flag.load(std::memory_order_seq_cst)) + ; +} + +void TaskQueue::processTasks() { + while (true) { + std::function task; + { + mutex.lock(); + cv.wait(mutex, [this]() { return !tasks.empty() || exit_flag.load(std::memory_order_seq_cst); }); + if (exit_flag.load(std::memory_order_seq_cst) && tasks.empty()) { + return; + } + task = tasks.front(); + tasks.pop(); + mutex.unlock(); + } + task(); + { + mutex.lock(); + if (tasks.empty()) { + sync_flag.store(true, std::memory_order_seq_cst); + } + mutex.unlock(); + } + } +} \ No newline at end of file diff --git a/archive/csrc/ktransformers_ext/cpu_backend/task_queue.h b/archive/csrc/ktransformers_ext/cpu_backend/task_queue.h new file mode 100644 index 0000000..4f9f112 --- /dev/null +++ b/archive/csrc/ktransformers_ext/cpu_backend/task_queue.h @@ -0,0 +1,138 @@ +/** + * @Description : + * @Author : chenht2022 + * @Date : 2024-07-16 10:43:18 + * @Version : 1.0.0 + * @LastEditors : chenht + * @LastEditTime : 2024-10-09 11:08:07 + * @Copyright (c) 2024 by KVCache.AI, All Rights Reserved. + **/ +#ifndef CPUINFER_TASKQUEUE_H +#define CPUINFER_TASKQUEUE_H + +#include +#include +#include +#include +#include +#include +#include +#ifdef _WIN32 +#include +#endif + +class custom_mutex { + private: +#ifdef _WIN32 + CRITICAL_SECTION cs; +#else + std::mutex mtx; +#endif + + public: + custom_mutex() { +#ifdef _WIN32 + InitializeCriticalSection(&cs); +#else + // No initialization required for std::mutex +#endif + } + + ~custom_mutex() { +#ifdef _WIN32 + DeleteCriticalSection(&cs); +#endif + } + + void lock() { +#ifdef _WIN32 + EnterCriticalSection(&cs); +#else + mtx.lock(); +#endif + } + + void unlock() { +#ifdef _WIN32 + LeaveCriticalSection(&cs); +#else + mtx.unlock(); +#endif + } + +#ifdef _WIN32 + CRITICAL_SECTION* get_handle() { + return &cs; + } +#else + std::mutex* get_handle() { + return &mtx; + } +#endif +}; + +class custom_condition_variable { + private: +#ifdef _WIN32 + CONDITION_VARIABLE cond_var; +#else + std::condition_variable cond_var; +#endif + + public: + custom_condition_variable() { +#ifdef _WIN32 + InitializeConditionVariable(&cond_var); +#endif + } + + template + void wait(custom_mutex& mutex, Predicate pred) { +#ifdef _WIN32 + while (!pred()) { + SleepConditionVariableCS(&cond_var, mutex.get_handle(), INFINITE); + } +#else + std::unique_lock lock(*mutex.get_handle(), std::adopt_lock); + cond_var.wait(lock, pred); + lock.release(); +#endif + } + + void notify_one() { +#ifdef _WIN32 + WakeConditionVariable(&cond_var); +#else + cond_var.notify_one(); +#endif + } + + void notify_all() { +#ifdef _WIN32 + WakeAllConditionVariable(&cond_var); +#else + cond_var.notify_all(); +#endif + } +}; + +class TaskQueue { + public: + TaskQueue(); + ~TaskQueue(); + + void enqueue(std::function); + + void sync(); + + private: + void processTasks(); + + std::queue> tasks; + custom_mutex mutex; + custom_condition_variable cv; + std::thread worker; + std::atomic sync_flag; + std::atomic exit_flag; +}; +#endif \ No newline at end of file diff --git a/archive/csrc/ktransformers_ext/cpu_backend/vendors/README.md b/archive/csrc/ktransformers_ext/cpu_backend/vendors/README.md new file mode 100644 index 0000000..d179f66 --- /dev/null +++ b/archive/csrc/ktransformers_ext/cpu_backend/vendors/README.md @@ -0,0 +1,3 @@ +## TODO + +This directory can be removed after updating the version of `llama.cpp`. \ No newline at end of file diff --git a/archive/csrc/ktransformers_ext/cpu_backend/vendors/cuda.h b/archive/csrc/ktransformers_ext/cpu_backend/vendors/cuda.h new file mode 100644 index 0000000..1746b07 --- /dev/null +++ b/archive/csrc/ktransformers_ext/cpu_backend/vendors/cuda.h @@ -0,0 +1,15 @@ +#pragma once + +#include +#include +#include +#include +#include + +#if CUDART_VERSION < 11020 +#define CU_DEVICE_ATTRIBUTE_VIRTUAL_MEMORY_MANAGEMENT_SUPPORTED CU_DEVICE_ATTRIBUTE_VIRTUAL_ADDRESS_MANAGEMENT_SUPPORTED +#define CUBLAS_TF32_TENSOR_OP_MATH CUBLAS_TENSOR_OP_MATH +#define CUBLAS_COMPUTE_16F CUDA_R_16F +#define CUBLAS_COMPUTE_32F CUDA_R_32F +#define cublasComputeType_t cudaDataType_t +#endif // CUDART_VERSION < 11020 diff --git a/archive/csrc/ktransformers_ext/cpu_backend/vendors/hip.h b/archive/csrc/ktransformers_ext/cpu_backend/vendors/hip.h new file mode 100644 index 0000000..abbc1e8 --- /dev/null +++ b/archive/csrc/ktransformers_ext/cpu_backend/vendors/hip.h @@ -0,0 +1,172 @@ +#pragma once + +#define HIP_ENABLE_WARP_SYNC_BUILTINS 1 +#include +#include +#include +#include +#ifdef __HIP_PLATFORM_AMD__ +// for rocblas_initialize() +#include "rocblas/rocblas.h" +#endif // __HIP_PLATFORM_AMD__ + +#define CUBLAS_COMPUTE_16F HIPBLAS_R_16F +#define CUBLAS_COMPUTE_32F HIPBLAS_R_32F +#define CUBLAS_COMPUTE_32F_FAST_16F HIPBLAS_R_32F +#define CUBLAS_GEMM_DEFAULT HIPBLAS_GEMM_DEFAULT +#define CUBLAS_GEMM_DEFAULT_TENSOR_OP HIPBLAS_GEMM_DEFAULT +#define CUBLAS_OP_N HIPBLAS_OP_N +#define CUBLAS_OP_T HIPBLAS_OP_T +#define CUBLAS_STATUS_SUCCESS HIPBLAS_STATUS_SUCCESS +#define CUBLAS_TF32_TENSOR_OP_MATH 0 +#define CUDA_R_16F HIPBLAS_R_16F +#define CUDA_R_32F HIPBLAS_R_32F +#define CU_DEVICE_ATTRIBUTE_VIRTUAL_MEMORY_MANAGEMENT_SUPPORTED hipDeviceAttributeVirtualMemoryManagementSupported +#define CU_MEM_ALLOC_GRANULARITY_RECOMMENDED hipMemAllocationGranularityRecommended +#define CU_MEM_ALLOCATION_TYPE_PINNED hipMemAllocationTypePinned +#define CU_MEM_LOCATION_TYPE_DEVICE hipMemLocationTypeDevice +#define CU_MEM_ACCESS_FLAGS_PROT_READWRITE hipMemAccessFlagsProtReadWrite +#define CU_CHECK(fn) {hipError_t err = fn; if(err != hipSuccess) { GGML_ABORT("HipVMM Failure: %s\n", hipGetErrorString(err)); }} +#define __shfl_sync(mask, var, laneMask, width) __shfl(var, laneMask, width) +#define __shfl_xor_sync(mask, var, laneMask, width) __shfl_xor(var, laneMask, width) +#define cublasComputeType_t hipblasDatatype_t //deprecated, new hipblasComputeType_t not in 5.6 +#define cublasCreate hipblasCreate +#define cublasDestroy hipblasDestroy +#define cublasGemmEx hipblasGemmEx +#define cublasGemmBatchedEx hipblasGemmBatchedEx +#define cublasGemmStridedBatchedEx hipblasGemmStridedBatchedEx +#define cublasHandle_t hipblasHandle_t +#define cublasSetMathMode(handle, mode) CUBLAS_STATUS_SUCCESS +#define cublasSetStream hipblasSetStream +#define cublasSgemm hipblasSgemm +#define cublasStatus_t hipblasStatus_t +#define cublasOperation_t hipblasOperation_t +#define cudaDataType_t hipblasDatatype_t //deprecated, new hipblasDatatype not in 5.6 +#define cudaDeviceCanAccessPeer hipDeviceCanAccessPeer +#define cudaDeviceDisablePeerAccess hipDeviceDisablePeerAccess +#define cudaDeviceEnablePeerAccess hipDeviceEnablePeerAccess +#define cudaDeviceProp hipDeviceProp_t +#define cudaDeviceSynchronize hipDeviceSynchronize +#define cudaError_t hipError_t +#define cudaErrorPeerAccessAlreadyEnabled hipErrorPeerAccessAlreadyEnabled +#define cudaErrorPeerAccessNotEnabled hipErrorPeerAccessNotEnabled +#define cudaEventCreateWithFlags hipEventCreateWithFlags +#define cudaEventDisableTiming hipEventDisableTiming +#define cudaEventRecord hipEventRecord +#define cudaEventSynchronize hipEventSynchronize +#define cudaEvent_t hipEvent_t +#define cudaEventDestroy hipEventDestroy +#define cudaFree hipFree +#define cudaFreeHost hipHostFree +#define cudaGetDevice hipGetDevice +#define cudaGetDeviceCount hipGetDeviceCount +#define cudaGetDeviceProperties hipGetDeviceProperties +#define cudaGetErrorString hipGetErrorString +#define cudaGetLastError hipGetLastError +#define cudaHostRegister hipHostRegister +#define cudaHostRegisterPortable hipHostRegisterPortable +#define cudaHostRegisterReadOnly hipHostRegisterReadOnly +#define cudaHostUnregister hipHostUnregister +#define cudaLaunchHostFunc hipLaunchHostFunc +#define cudaMalloc hipMalloc +#define cudaMallocHost(ptr, size) hipHostMalloc(ptr, size, hipHostMallocDefault) +#define cudaMemcpy hipMemcpy +#define cudaMemcpyAsync hipMemcpyAsync +#define cudaMemcpyPeerAsync hipMemcpyPeerAsync +#define cudaMemcpy2DAsync hipMemcpy2DAsync +#define cudaMemcpyDeviceToDevice hipMemcpyDeviceToDevice +#define cudaMemcpyDeviceToHost hipMemcpyDeviceToHost +#define cudaMemcpyHostToDevice hipMemcpyHostToDevice +#define cudaMemcpyKind hipMemcpyKind +#define cudaMemset hipMemset +#define cudaMemsetAsync hipMemsetAsync +#define cudaMemGetInfo hipMemGetInfo +#define cudaOccupancyMaxPotentialBlockSize hipOccupancyMaxPotentialBlockSize +#define cudaSetDevice hipSetDevice +#define cuDeviceGet hipDeviceGet +#define CUdevice hipDevice_t +#define CUdeviceptr hipDeviceptr_t +#define cuMemUnmap hipMemUnmap +#define CUmemAccessDesc hipMemAccessDesc +#define cuMemAddressFree hipMemAddressFree +#define cuMemRelease hipMemRelease +#define CUmemGenericAllocationHandle hipMemGenericAllocationHandle_t +#define cuMemCreate hipMemCreate +#define cuMemAddressReserve hipMemAddressReserve +#define cuMemMap hipMemMap +#define cuMemSetAccess hipMemSetAccess +#define cuMemGetAllocationGranularity hipMemGetAllocationGranularity +#define CUmemAllocationProp hipMemAllocationProp +#define cuDeviceGetAttribute hipDeviceGetAttribute +#define cudaStreamCreateWithFlags hipStreamCreateWithFlags +#define cudaStreamDestroy hipStreamDestroy +#define cudaStreamFireAndForget hipStreamFireAndForget +#define cudaStreamNonBlocking hipStreamNonBlocking +#define cudaStreamPerThread hipStreamPerThread +#define cudaStreamSynchronize hipStreamSynchronize +#define cudaStreamWaitEvent(stream, event, flags) hipStreamWaitEvent(stream, event, flags) +#define cudaGraphExec_t hipGraphExec_t +#define cudaGraphNode_t hipGraphNode_t +#define cudaKernelNodeParams hipKernelNodeParams +#define cudaKernelNodeParams hipKernelNodeParams +#define cudaGraphExecDestroy hipGraphExecDestroy +#define cudaGraphLaunch hipGraphLaunch +#define cudaErrorGraphExecUpdateFailure hipErrorGraphExecUpdateFailure +#define cudaGraphExecUpdateResultInfo hipGraphExecUpdateResult +#define cudaGraphNodeType hipGraphNodeType +#define cudaGraphNodeTypeKernel hipGraphNodeTypeKernel +#define cudaGraphInstantiate hipGraphInstantiate +#define cudaStreamEndCapture hipStreamEndCapture +#define cudaGraphDestroy hipGraphDestroy +#define cudaGraphKernelNodeSetParams hipGraphKernelNodeSetParams +#define cudaErrorInvalidDeviceFunction hipErrorInvalidDeviceFunction +#define cudaGraphKernelNodeGetParams hipGraphKernelNodeGetParams +#define cudaGraphNodeGetType hipGraphNodeGetType +#define cudaGraphGetNodes hipGraphGetNodes +#define cudaGraphExecUpdate hipGraphExecUpdate +#define cudaStreamCaptureModeRelaxed hipStreamCaptureModeRelaxed +#define cudaStreamBeginCapture hipStreamBeginCapture +#define cudaGraph_t hipGraph_t +#define cudaStream_t hipStream_t +#define cudaSuccess hipSuccess +#define cudaHostFn_t hipHostFn_t +#define __trap() do { abort(); __builtin_unreachable(); } while(0) +#define CUBLAS_STATUS_SUCCESS HIPBLAS_STATUS_SUCCESS +#define CUBLAS_STATUS_NOT_INITIALIZED HIPBLAS_STATUS_NOT_INITIALIZED +#define CUBLAS_STATUS_ALLOC_FAILED HIPBLAS_STATUS_ALLOC_FAILED +#define CUBLAS_STATUS_INVALID_VALUE HIPBLAS_STATUS_INVALID_VALUE +#define CUBLAS_STATUS_ARCH_MISMATCH HIPBLAS_STATUS_ARCH_MISMATCH +#define CUBLAS_STATUS_MAPPING_ERROR HIPBLAS_STATUS_MAPPING_ERROR +#define CUBLAS_STATUS_EXECUTION_FAILED HIPBLAS_STATUS_EXECUTION_FAILED +#define CUBLAS_STATUS_INTERNAL_ERROR HIPBLAS_STATUS_INTERNAL_ERROR +#define CUBLAS_STATUS_NOT_SUPPORTED HIPBLAS_STATUS_NOT_SUPPORTED + +#define __CUDA_ARCH__ 1300 + +#if defined(__gfx803__) || defined(__gfx900__) || defined(__gfx906__) +#define GCN +#endif + +#if defined(__gfx908__) || defined(__gfx90a__) || defined(__gfx942__) +#define CDNA +#endif + +#if defined(__gfx1100__) || defined(__gfx1101__) || defined(__gfx1102__) || defined(__gfx1103__) || \ + defined(__gfx1150__) || defined(__gfx1151__) +#define RDNA3 +#endif + +#if defined(__gfx1030__) || defined(__gfx1031__) || defined(__gfx1032__) || defined(__gfx1033__) || \ + defined(__gfx1034__) || defined(__gfx1035__) || defined(__gfx1036__) || defined(__gfx1037__) +#define RDNA2 +#endif + +#if defined(__gfx1010__) || defined(__gfx1012__) +#define RDNA1 +#endif + +#ifndef __has_builtin + #define __has_builtin(x) 0 +#endif + +typedef hip_bfloat16 nv_bfloat16; diff --git a/archive/csrc/ktransformers_ext/cpu_backend/vendors/musa.h b/archive/csrc/ktransformers_ext/cpu_backend/vendors/musa.h new file mode 100644 index 0000000..6cc1b69 --- /dev/null +++ b/archive/csrc/ktransformers_ext/cpu_backend/vendors/musa.h @@ -0,0 +1,137 @@ +#pragma once + +#include +#include +#include +#include +#include +#define CUBLAS_COMPUTE_16F CUDA_R_16F +#define CUBLAS_COMPUTE_32F CUDA_R_32F +#define CUBLAS_COMPUTE_32F_FAST_16F MUBLAS_COMPUTE_32F_FAST_16F +#define CUBLAS_GEMM_DEFAULT MUBLAS_GEMM_DEFAULT +#define CUBLAS_GEMM_DEFAULT_TENSOR_OP MUBLAS_GEMM_DEFAULT +#define CUBLAS_OP_N MUBLAS_OP_N +#define CUBLAS_OP_T MUBLAS_OP_T +#define CUBLAS_STATUS_SUCCESS MUBLAS_STATUS_SUCCESS +#define CUBLAS_TF32_TENSOR_OP_MATH MUBLAS_MATH_MODE_DEFAULT +#define CUDA_R_16F MUSA_R_16F +#define CUDA_R_32F MUSA_R_32F +#define cublasComputeType_t cudaDataType_t +#define cublasCreate mublasCreate +#define cublasDestroy mublasDestroy +#define cublasGemmEx mublasGemmEx +#define cublasGemmBatchedEx mublasGemmBatchedEx +#define cublasGemmStridedBatchedEx mublasGemmStridedBatchedEx +#define cublasHandle_t mublasHandle_t +#define cublasSetMathMode mublasSetMathMode +#define cublasSetStream mublasSetStream +#define cublasSgemm mublasSgemm +#define cublasStatus_t mublasStatus_t +#define cublasOperation_t mublasOperation_t +#define cublasGetStatusString mublasStatus_to_string +#define cudaDataType_t musaDataType_t +#define cudaDeviceCanAccessPeer musaDeviceCanAccessPeer +#define cudaDeviceDisablePeerAccess musaDeviceDisablePeerAccess +#define cudaDeviceEnablePeerAccess musaDeviceEnablePeerAccess +#define cudaDeviceProp musaDeviceProp +#define cudaDeviceSynchronize musaDeviceSynchronize +#define cudaError_t musaError_t +#define cudaErrorPeerAccessAlreadyEnabled musaErrorPeerAccessAlreadyEnabled +#define cudaErrorPeerAccessNotEnabled musaErrorPeerAccessNotEnabled +#define cudaEventCreateWithFlags musaEventCreateWithFlags +#define cudaEventDisableTiming musaEventDisableTiming +#define cudaEventRecord musaEventRecord +#define cudaEventSynchronize musaEventSynchronize +#define cudaEvent_t musaEvent_t +#define cudaEventDestroy musaEventDestroy +#define cudaFree musaFree +#define cudaFreeHost musaFreeHost +#define cudaGetDevice musaGetDevice +#define cudaGetDeviceCount musaGetDeviceCount +#define cudaGetDeviceProperties musaGetDeviceProperties +#define cudaGetErrorString musaGetErrorString +#define cudaGetLastError musaGetLastError +#define cudaHostRegister musaHostRegister +#define cudaHostRegisterPortable musaHostRegisterPortable +#define cudaHostRegisterReadOnly musaHostRegisterReadOnly +#define cudaHostUnregister musaHostUnregister +#define cudaLaunchHostFunc musaLaunchHostFunc +#define cudaMalloc musaMalloc +#define cudaMallocHost musaMallocHost +#define cudaMallocManaged musaMallocManaged +#define cudaMemcpy musaMemcpy +#define cudaMemcpyAsync musaMemcpyAsync +#define cudaMemcpyPeerAsync musaMemcpyPeerAsync +#define cudaMemcpy2DAsync musaMemcpy2DAsync +#define cudaMemcpyDeviceToDevice musaMemcpyDeviceToDevice +#define cudaMemcpyDeviceToHost musaMemcpyDeviceToHost +#define cudaMemcpyHostToDevice musaMemcpyHostToDevice +#define cudaMemcpyKind musaMemcpyKind +#define cudaMemset musaMemset +#define cudaMemsetAsync musaMemsetAsync +#define cudaMemGetInfo musaMemGetInfo +#define cudaOccupancyMaxPotentialBlockSize musaOccupancyMaxPotentialBlockSize +#define cudaSetDevice musaSetDevice +#define cudaStreamCreateWithFlags musaStreamCreateWithFlags +#define cudaStreamDestroy musaStreamDestroy +#define cudaStreamFireAndForget musaStreamFireAndForget +#define cudaStreamNonBlocking musaStreamNonBlocking +#define cudaStreamPerThread musaStreamPerThread +#define cudaStreamSynchronize musaStreamSynchronize +#define cudaStreamWaitEvent musaStreamWaitEvent +#define cudaStream_t musaStream_t +#define cudaSuccess musaSuccess + +// Additional mappings for MUSA virtual memory pool +#define CU_DEVICE_ATTRIBUTE_VIRTUAL_MEMORY_MANAGEMENT_SUPPORTED MU_DEVICE_ATTRIBUTE_VIRTUAL_ADDRESS_MANAGEMENT_SUPPORTED +#define CU_MEM_ACCESS_FLAGS_PROT_READWRITE MU_MEM_ACCESS_FLAGS_PROT_READWRITE +#define CU_MEM_ALLOC_GRANULARITY_RECOMMENDED MU_MEM_ALLOC_GRANULARITY_RECOMMENDED +#define CU_MEM_ALLOCATION_TYPE_PINNED MU_MEM_ALLOCATION_TYPE_PINNED +#define CU_MEM_LOCATION_TYPE_DEVICE MU_MEM_LOCATION_TYPE_DEVICE +#define CUdevice MUdevice +#define CUdeviceptr MUdeviceptr +#define CUmemAccessDesc MUmemAccessDesc +#define CUmemAllocationProp MUmemAllocationProp +#define CUmemGenericAllocationHandle MUmemGenericAllocationHandle +#define cuDeviceGet muDeviceGet +#define cuDeviceGetAttribute muDeviceGetAttribute +#define cuMemAddressFree muMemAddressFree +#define cuMemAddressReserve muMemAddressReserve +#define cuMemCreate muMemCreate +#define cuMemGetAllocationGranularity muMemGetAllocationGranularity +#define cuMemMap muMemMap +#define cuMemRelease muMemRelease +#define cuMemSetAccess muMemSetAccess +#define cuMemUnmap muMemUnmap +#define cudaFuncAttributeMaxDynamicSharedMemorySize musaFuncAttributeMaxDynamicSharedMemorySize +#define cudaFuncSetAttribute musaFuncSetAttribute +#define cudaMemcpy3DPeerParms musaMemcpy3DPeerParms +#define make_cudaExtent make_musaExtent +#define make_cudaPitchedPtr make_musaPitchedPtr + +// Additional mappings for MUSA graphs +#define CUDA_SUCCESS MUSA_SUCCESS +#define CUresult MUresult +#define cuGetErrorString muGetErrorString +#define cudaErrorGraphExecUpdateFailure musaErrorGraphExecUpdateFailure +#define cudaErrorInvalidDeviceFunction musaErrorInvalidDeviceFunction +#define cudaGraphDestroy musaGraphDestroy +#define cudaGraphExecDestroy musaGraphExecDestroy +#define cudaGraphExec_t musaGraphExec_t +#define cudaGraphExecUpdate musaGraphExecUpdate +#define cudaGraphExecUpdateResultInfo musaGraphExecUpdateResult +#define cudaGraphGetNodes musaGraphGetNodes +#define cudaGraphInstantiate musaGraphInstantiate +#define cudaGraphKernelNodeGetParams musaGraphKernelNodeGetParams +#define cudaGraphKernelNodeSetParams musaGraphKernelNodeSetParams +#define cudaGraphLaunch musaGraphLaunch +#define cudaGraphNodeGetType musaGraphNodeGetType +#define cudaGraphNode_t musaGraphNode_t +#define cudaGraphNodeType musaGraphNodeType +#define cudaGraphNodeTypeKernel musaGraphNodeTypeKernel +#define cudaGraph_t musaGraph_t +#define cudaKernelNodeParams musaKernelNodeParams +#define cudaStreamCaptureModeRelaxed musaStreamCaptureModeRelaxed +#define cudaStreamEndCapture musaStreamEndCapture + +typedef mt_bfloat16 nv_bfloat16; diff --git a/archive/csrc/ktransformers_ext/cpu_backend/vendors/vendor.h b/archive/csrc/ktransformers_ext/cpu_backend/vendors/vendor.h new file mode 100644 index 0000000..8470438 --- /dev/null +++ b/archive/csrc/ktransformers_ext/cpu_backend/vendors/vendor.h @@ -0,0 +1,13 @@ +#ifndef CPUINFER_VENDOR_VENDOR_H +#define CPUINFER_VENDOR_VENDOR_H + +#ifdef USE_CUDA +#include "cuda.h" +#elif USE_HIP +#define __HIP_PLATFORM_AMD__ +#include "hip.h" +#elif USE_MUSA +#include "musa.h" +#endif + +#endif // CPUINFER_VENDOR_VENDOR_H \ No newline at end of file diff --git a/archive/csrc/ktransformers_ext/cuda/binding.cpp b/archive/csrc/ktransformers_ext/cuda/binding.cpp new file mode 100644 index 0000000..4aa00c6 --- /dev/null +++ b/archive/csrc/ktransformers_ext/cuda/binding.cpp @@ -0,0 +1,71 @@ +/** + * @Description : + * @Author : Azure-Tang, Boxin Zhang + * @Date : 2024-07-25 13:38:30 + * @Version : 0.2.2 + * @Copyright (c) 2024 by KVCache.AI, All Rights Reserved. +**/ + +#include "custom_gguf/ops.h" +#ifdef KTRANSFORMERS_USE_CUDA +#include "gptq_marlin/ops.h" +#endif +// Python bindings +#include +#include +#include +#include +#include +// namespace py = pybind11; + +PYBIND11_MODULE(KTransformersOps, m) { + + m.def("dequantize_q8_0", [](const intptr_t data, int num_bytes, int blk_size, const int ele_per_blk, torch::Device device, py::object target_dtype) { + torch::Dtype dtype = torch::python::detail::py_object_to_dtype(target_dtype); + return dequantize_q8_0((int8_t*)data, num_bytes, blk_size, ele_per_blk, device, dtype); + }, "Function to dequantize q8_0 data.", + py::arg("data"), py::arg("num_bytes"), py::arg("blk_size"), py::arg("ele_per_blk"), py::arg("device"), py::arg("target_dtype")); + + m.def("dequantize_q6_k", [](const intptr_t data, int num_bytes, int blk_size, const int ele_per_blk, torch::Device device, py::object target_dtype) { + torch::Dtype dtype = torch::python::detail::py_object_to_dtype(target_dtype); + return dequantize_q6_k((int8_t*)data, num_bytes, blk_size, ele_per_blk, device, dtype); + }, "Function to dequantize q6_k data.", + py::arg("data"), py::arg("num_bytes"), py::arg("blk_size"), py::arg("ele_per_blk"), py::arg("device"), py::arg("target_dtype")); + + m.def("dequantize_q5_k", [](const intptr_t data, int num_bytes, int blk_size, const int ele_per_blk, torch::Device device, py::object target_dtype) { + torch::Dtype dtype = torch::python::detail::py_object_to_dtype(target_dtype); + return dequantize_q5_k((int8_t*)data, num_bytes, blk_size, ele_per_blk, device, dtype); + }, "Function to dequantize q5_k data.", + py::arg("data"), py::arg("num_bytes"), py::arg("blk_size"), py::arg("ele_per_blk"), py::arg("device"), py::arg("target_dtype")); + + m.def("dequantize_q4_k", [](const intptr_t data, int num_bytes, int blk_size, const int ele_per_blk, torch::Device device, py::object target_dtype) { + torch::Dtype dtype = torch::python::detail::py_object_to_dtype(target_dtype); + return dequantize_q4_k((int8_t*)data, num_bytes, blk_size, ele_per_blk, device, dtype); + }, "Function to dequantize q4_k data.", + py::arg("data"), py::arg("num_bytes"), py::arg("blk_size"), py::arg("ele_per_blk"), py::arg("device"), py::arg("target_dtype")); + + m.def("dequantize_q3_k", [](const intptr_t data, int num_bytes, int blk_size, const int ele_per_blk, torch::Device device, py::object target_dtype) { + torch::Dtype dtype = torch::python::detail::py_object_to_dtype(target_dtype); + return dequantize_q3_k((int8_t*)data, num_bytes, blk_size, ele_per_blk, device, dtype); + }, "Function to dequantize q3_k data.", + py::arg("data"), py::arg("num_bytes"), py::arg("blk_size"), py::arg("ele_per_blk"), py::arg("device"), py::arg("target_dtype")); + + m.def("dequantize_q2_k", [](const intptr_t data, int num_bytes, int blk_size, const int ele_per_blk, torch::Device device, py::object target_dtype) { + torch::Dtype dtype = torch::python::detail::py_object_to_dtype(target_dtype); + return dequantize_q2_k((int8_t*)data, num_bytes, blk_size, ele_per_blk, device, dtype); + }, "Function to dequantize q2_k data.", + py::arg("data"), py::arg("num_bytes"), py::arg("blk_size"), py::arg("ele_per_blk"), py::arg("device"), py::arg("target_dtype")); + + m.def("dequantize_iq4_xs", [](const intptr_t data, int num_bytes, int blk_size, const int ele_per_blk, torch::Device device, py::object target_dtype) { + torch::Dtype dtype = torch::python::detail::py_object_to_dtype(target_dtype); + return dequantize_iq4_xs((int8_t*)data, num_bytes, blk_size, ele_per_blk, device, dtype); + }, "Function to dequantize iq4_xs data.", + py::arg("data"), py::arg("num_bytes"), py::arg("blk_size"), py::arg("ele_per_blk"), py::arg("device"), py::arg("target_dtype")); + +#ifdef KTRANSFORMERS_USE_CUDA + m.def("gptq_marlin_gemm", &gptq_marlin_gemm, "Function to perform GEMM using Marlin quantization.", + py::arg("a"), py::arg("b_q_weight"), py::arg("b_scales"), py::arg("g_idx"), + py::arg("perm"), py::arg("workspace"), py::arg("num_bits"), py::arg("size_m"), + py::arg("size_n"), py::arg("size_k"), py::arg("is_k_full")); +#endif +} \ No newline at end of file diff --git a/archive/csrc/ktransformers_ext/cuda/custom_gguf/dequant.cu b/archive/csrc/ktransformers_ext/cuda/custom_gguf/dequant.cu new file mode 100644 index 0000000..8d80160 --- /dev/null +++ b/archive/csrc/ktransformers_ext/cuda/custom_gguf/dequant.cu @@ -0,0 +1,882 @@ +/* + * @Description : + * @Author : Azure-Tang, Boxin Zhang + * @Date : 2024-07-25 13:38:30 + * @Version : 0.2.2 + * Adapted from https://github.com/ggerganov/ggml/blob/fca1caafea7de9fbd7efc733b9818f9cf2da3050/src/ggml-quants.c + * Copyright (c) 2023-2024 The ggml authors + * Copyright (c) 2024 by KVCache.AI, All Rights Reserved. + */ +#include +#include +#include +#include +#include +#include +#include +#include + +#ifdef __HIP_PLATFORM_AMD__ +typedef __hip_bfloat16 nv_bfloat16; +#endif + +__global__ void dequantize_q8_0_fp32_kernel(const int8_t* data, float* output, const int blk_size, const int ele_per_blk, const int num_blocks) { + long long global_idx = blockIdx.x * blockDim.x + threadIdx.x; + for (long long block_id = global_idx; block_id < num_blocks; block_id += blockDim.x * gridDim.x){ + float* __restrict__ output_blk = (float*)(output + block_id * ele_per_blk); + const int8_t* cur_block = data + block_id * blk_size; + float scale = __half2float(*((half*)cur_block)); + cur_block += 2; + for (int i = 0; i < ele_per_blk; i++){ + output_blk[i] = scale * cur_block[i]; + } + } +} + +__global__ void dequantize_q8_0_fp16_kernel(const int8_t* data, __half* output, const int blk_size, const int ele_per_blk, const int num_blocks) { + long long global_idx = blockIdx.x * blockDim.x + threadIdx.x; + for (long long block_id = global_idx; block_id < num_blocks; block_id += blockDim.x * gridDim.x) { + __half* __restrict__ output_blk = (__half*)(output + block_id * ele_per_blk); + const int8_t* cur_block = data + block_id * blk_size; + float scale = __half2float(*((half*)cur_block)); + cur_block += 2; + for (int i = 0; i < ele_per_blk; i++) { + output_blk[i] = __float2half(scale * cur_block[i]); + } + } +} + +__global__ void dequantize_q8_0_bf16_kernel(const int8_t* data, nv_bfloat16* output, const int blk_size, const int ele_per_blk, const int num_blocks) { + long long global_idx = blockIdx.x * blockDim.x + threadIdx.x; + for (long long block_id = global_idx; block_id < num_blocks; block_id += blockDim.x * gridDim.x) { + nv_bfloat16* __restrict__ output_blk = (nv_bfloat16*)(output + block_id * ele_per_blk); + const int8_t* cur_block = data + block_id * blk_size; + float scale = __half2float(*((half*)cur_block)); + cur_block += 2; + for (int i = 0; i < ele_per_blk; i++) { + output_blk[i] = __float2bfloat16(scale * cur_block[i]); + } + } +} + +// __device__ void get_scale_min_k4(int j, const uint8_t * __restrict__ q, uint8_t * __restrict__ d, uint8_t * __restrict__ m) { +__device__ void get_scale_min_k4(int j, const uint8_t * q, uint8_t * __restrict__ d, uint8_t * __restrict__ m) { + if (j < 4) { + *d = q[j] & 63; *m = q[j + 4] & 63; + } else { + *d = (q[j+4] & 0xF) | ((q[j-4] >> 6) << 4); + *m = (q[j+4] >> 4) | ((q[j-0] >> 6) << 4); + } +} + +__global__ void dequantize_q2_k_fp32_kernel(const int8_t* data, float* output, const int blk_size, const int ele_per_blk, const int num_blocks) { + long long global_idx = blockIdx.x * blockDim.x + threadIdx.x; + for (long long block_id=global_idx; block_id(data + block_id * blk_size + 80))); + const float min = __half2float(*(reinterpret_cast(data + block_id * blk_size + 82))); + + const uint8_t * __restrict__ q = (uint8_t*)(data + block_id * blk_size + 16); + + int is = 0; + float dl, ml; + + for (int n = 0; n < 256; n += 128) { + int shift = 0; + for (int j = 0; j < 4; ++j) { + uint8_t* scales = (uint8_t*)(data + block_id * blk_size + (is++)); + uint8_t sc = *scales; + dl = d * (sc & 0xF); ml = min * (sc >> 4); + for (int l = 0; l < 16; ++l) *output_blk++ = dl * ((int8_t)((q[l] >> shift) & 3)) - ml; + + scales = (uint8_t*)(data + block_id * blk_size + (is++)); + sc = *scales; + + dl = d * (sc & 0xF); ml = min * (sc >> 4); + for (int l = 0; l < 16; ++l) *output_blk++ = dl * ((int8_t)((q[l+16] >> shift) & 3)) - ml; + + shift += 2; + } + q += 32; + } + } +} + +__global__ void dequantize_q2_k_fp16_kernel(const int8_t* data, __half* output, const int blk_size, const int ele_per_blk, const int num_blocks) { + long long global_idx = blockIdx.x * blockDim.x + threadIdx.x; + for (long long block_id=global_idx; block_id(data + block_id * blk_size + 80))); + const float min = __half2float(*(reinterpret_cast(data + block_id * blk_size + 82))); + + const uint8_t * __restrict__ q = (uint8_t*)(data + block_id * blk_size + 16); + + int is = 0; + float dl, ml; + + for (int n = 0; n < 256; n += 128) { + int shift = 0; + for (int j = 0; j < 4; ++j) { + uint8_t* scales = (uint8_t*)(data + block_id * blk_size + (is++)); + uint8_t sc = *scales; + dl = d * (sc & 0xF); ml = min * (sc >> 4); + for (int l = 0; l < 16; ++l) *output_blk++ = __float2half(dl * ((int8_t)((q[l] >> shift) & 3)) - ml); + + scales = (uint8_t*)(data + block_id * blk_size + (is++)); + sc = *scales; + + dl = d * (sc & 0xF); ml = min * (sc >> 4); + for (int l = 0; l < 16; ++l) *output_blk++ = __float2half(dl * ((int8_t)((q[l+16] >> shift) & 3)) - ml); + + shift += 2; + } + q += 32; + } + } +} + +__global__ void dequantize_q2_k_bf16_kernel(const int8_t* data, nv_bfloat16* output, const int blk_size, const int ele_per_blk, const int num_blocks) { + long long global_idx = blockIdx.x * blockDim.x + threadIdx.x; + for (long long block_id=global_idx; block_id(data + block_id * blk_size + 80))); + const float min = __half2float(*(reinterpret_cast(data + block_id * blk_size + 82))); + + const uint8_t * __restrict__ q = (uint8_t*)(data + block_id * blk_size + 16); + + int is = 0; + float dl, ml; + + for (int n = 0; n < 256; n += 128) { + int shift = 0; + for (int j = 0; j < 4; ++j) { + uint8_t* scales = (uint8_t*)(data + block_id * blk_size + (is++)); + uint8_t sc = *scales; + dl = d * (sc & 0xF); ml = min * (sc >> 4); + for (int l = 0; l < 16; ++l) *output_blk++ = __float2bfloat16(dl * ((int8_t)((q[l] >> shift) & 3)) - ml); + + scales = (uint8_t*)(data + block_id * blk_size + (is++)); + sc = *scales; + + dl = d * (sc & 0xF); ml = min * (sc >> 4); + for (int l = 0; l < 16; ++l) *output_blk++ = __float2bfloat16(dl * ((int8_t)((q[l+16] >> shift) & 3)) - ml); + + shift += 2; + } + q += 32; + } + } +} + +__global__ void dequantize_q3_k_fp32_kernel(const int8_t* data, float* output, const int blk_size, const int ele_per_blk, const int num_blocks) { + + long long global_idx = blockIdx.x * blockDim.x + threadIdx.x; + const uint32_t kmask1 = 0x03030303; + const uint32_t kmask2 = 0x0f0f0f0f; + for (long long block_id=global_idx; block_id(data + block_id * blk_size + 108))); + + const uint8_t * __restrict__ q = (uint8_t*)(data + block_id * blk_size + 32); + const uint8_t * __restrict__ hm = (uint8_t*)(data + block_id * blk_size + 0); + uint8_t m = 1; + + + uint8_t* block_scales = (uint8_t*)(data + block_id * blk_size + 96); + + for (int i = 0; i < 3; i++) { + aux[i] = 0; + for (int j = 0; j < 4; j++) { + aux[i] |= ((uint32_t)block_scales[i * 4 + j]) << (j * 8); + } + } + + uint32_t tmp = aux[2]; + aux[2] = ((aux[0] >> 4) & kmask2) | (((tmp >> 4) & kmask1) << 4); + aux[3] = ((aux[1] >> 4) & kmask2) | (((tmp >> 6) & kmask1) << 4); + aux[0] = (aux[0] & kmask2) | (((tmp >> 0) & kmask1) << 4); + aux[1] = (aux[1] & kmask2) | (((tmp >> 2) & kmask1) << 4); + + int is = 0; + float dl; + for (int n = 0; n < 256; n += 128) { + int shift = 0; + for (int j = 0; j < 4; ++j) { + + dl = d_all * (scales[is++] - 32); + for (int l = 0; l < 16; ++l) { + *output_blk++ = dl * ((int8_t)((q[l+ 0] >> shift) & 3) - ((hm[l+ 0] & m) ? 0 : 4)); + } + + dl = d_all * (scales[is++] - 32); + for (int l = 0; l < 16; ++l) { + *output_blk++ = dl * ((int8_t)((q[l+16] >> shift) & 3) - ((hm[l+16] & m) ? 0 : 4)); + } + + shift += 2; + m <<= 1; + } + q += 32; + } + } +} + +__global__ void dequantize_q3_k_fp16_kernel(const int8_t* data, __half* output, const int blk_size, const int ele_per_blk, const int num_blocks) { + + long long global_idx = blockIdx.x * blockDim.x + threadIdx.x; + const uint32_t kmask1 = 0x03030303; + const uint32_t kmask2 = 0x0f0f0f0f; + for (long long block_id=global_idx; block_id(data + block_id * blk_size + 108))); + + const uint8_t * __restrict__ q = (uint8_t*)(data + block_id * blk_size + 32); + const uint8_t * __restrict__ hm = (uint8_t*)(data + block_id * blk_size + 0); + uint8_t m = 1; + + + uint8_t* block_scales = (uint8_t*)(data + block_id * blk_size + 96); + + for (int i = 0; i < 3; i++) { + aux[i] = 0; + for (int j = 0; j < 4; j++) { + aux[i] |= ((uint32_t)block_scales[i * 4 + j]) << (j * 8); + } + } + + uint32_t tmp = aux[2]; + aux[2] = ((aux[0] >> 4) & kmask2) | (((tmp >> 4) & kmask1) << 4); + aux[3] = ((aux[1] >> 4) & kmask2) | (((tmp >> 6) & kmask1) << 4); + aux[0] = (aux[0] & kmask2) | (((tmp >> 0) & kmask1) << 4); + aux[1] = (aux[1] & kmask2) | (((tmp >> 2) & kmask1) << 4); + + int is = 0; + float dl; + for (int n = 0; n < 256; n += 128) { + int shift = 0; + for (int j = 0; j < 4; ++j) { + + dl = d_all * (scales[is++] - 32); + for (int l = 0; l < 16; ++l) { + *output_blk++ = __float2half(dl * ((int8_t)((q[l+ 0] >> shift) & 3) - ((hm[l+ 0] & m) ? 0 : 4))); + } + + dl = d_all * (scales[is++] - 32); + for (int l = 0; l < 16; ++l) { + *output_blk++ = __float2half(dl * ((int8_t)((q[l+16] >> shift) & 3) - ((hm[l+16] & m) ? 0 : 4))); + } + + shift += 2; + m <<= 1; + } + q += 32; + } + } +} + +__global__ void dequantize_q3_k_bf16_kernel(const int8_t* data, nv_bfloat16* output, const int blk_size, const int ele_per_blk, const int num_blocks) { + + long long global_idx = blockIdx.x * blockDim.x + threadIdx.x; + const uint32_t kmask1 = 0x03030303; + const uint32_t kmask2 = 0x0f0f0f0f; + for (long long block_id=global_idx; block_id(data + block_id * blk_size + 108))); + + const uint8_t * __restrict__ q = (uint8_t*)(data + block_id * blk_size + 32); + const uint8_t * __restrict__ hm = (uint8_t*)(data + block_id * blk_size + 0); + uint8_t m = 1; + + + uint8_t* block_scales = (uint8_t*)(data + block_id * blk_size + 96); + + for (int i = 0; i < 3; i++) { + aux[i] = 0; + for (int j = 0; j < 4; j++) { + aux[i] |= ((uint32_t)block_scales[i * 4 + j]) << (j * 8); + } + } + + uint32_t tmp = aux[2]; + aux[2] = ((aux[0] >> 4) & kmask2) | (((tmp >> 4) & kmask1) << 4); + aux[3] = ((aux[1] >> 4) & kmask2) | (((tmp >> 6) & kmask1) << 4); + aux[0] = (aux[0] & kmask2) | (((tmp >> 0) & kmask1) << 4); + aux[1] = (aux[1] & kmask2) | (((tmp >> 2) & kmask1) << 4); + + int is = 0; + float dl; + for (int n = 0; n < 256; n += 128) { + int shift = 0; + for (int j = 0; j < 4; ++j) { + + dl = d_all * (scales[is++] - 32); + for (int l = 0; l < 16; ++l) { + *output_blk++ = __float2bfloat16(dl * ((int8_t)((q[l+ 0] >> shift) & 3) - ((hm[l+ 0] & m) ? 0 : 4))); + } + + dl = d_all * (scales[is++] - 32); + for (int l = 0; l < 16; ++l) { + *output_blk++ = __float2bfloat16(dl * ((int8_t)((q[l+16] >> shift) & 3) - ((hm[l+16] & m) ? 0 : 4))); + } + + shift += 2; + m <<= 1; + } + q += 32; + } + } +} + + +__global__ void dequantize_q4_k_fp32_kernel(const int8_t* data, float* output, const int blk_size, const int ele_per_blk, const int num_blocks) { + long long global_idx = blockIdx.x * blockDim.x + threadIdx.x; + for (long long block_id=global_idx; block_id(data + block_id * 144 + 0))); + const float min = __half2float(*(reinterpret_cast(data + block_id * 144 + 2))); + int is = 0; + uint8_t sc, m; + for (int j = 0; j < ele_per_blk; j += 64) { + uint8_t* scales = (uint8_t*)(data + block_id * 144 + 4); + get_scale_min_k4(is + 0, scales, &sc, &m); + const float d1 = d * sc; const float m1 = min * m; + get_scale_min_k4(is + 1, scales, &sc, &m); + const float d2 = d * sc; const float m2 = min * m; + for (int l = 0; l < 32; ++l) *output_blk++ = d1 * (q[l] & 0xF) - m1; + for (int l = 0; l < 32; ++l) *output_blk++ = d2 * (q[l] >> 4) - m2; + q += 32; is += 2; + } + } +} + +__global__ void dequantize_q4_k_fp16_kernel(const int8_t* data, __half* output, const int blk_size, const int ele_per_blk, const int num_blocks) { + long long global_idx = blockIdx.x * blockDim.x + threadIdx.x; + for (long long block_id=global_idx; block_id(data + block_id * 144 + 0))); + const float min = __half2float(*(reinterpret_cast(data + block_id * 144 + 2))); + int is = 0; + uint8_t sc, m; + for (int j = 0; j < ele_per_blk; j += 64) { + uint8_t* scales = (uint8_t*)(data + block_id * 144 + 4); + get_scale_min_k4(is + 0, scales, &sc, &m); + const float d1 = d * sc; const float m1 = min * m; + get_scale_min_k4(is + 1, scales, &sc, &m); + const float d2 = d * sc; const float m2 = min * m; + for (int l = 0; l < 32; ++l) *output_blk++ = __float2half(d1 * (q[l] & 0xF) - m1); + for (int l = 0; l < 32; ++l) *output_blk++ = __float2half(d2 * (q[l] >> 4) - m2); + q += 32; is += 2; + } + } +} + +__global__ void dequantize_q4_k_bf16_kernel(const int8_t* data, nv_bfloat16* output, const int blk_size, const int ele_per_blk, const int num_blocks) { + long long global_idx = blockIdx.x * blockDim.x + threadIdx.x; + for (long long block_id=global_idx; block_id(data + block_id * 144 + 0))); + const float min = __half2float(*(reinterpret_cast(data + block_id * 144 + 2))); + int is = 0; + uint8_t sc, m; + for (int j = 0; j < ele_per_blk; j += 64) { + uint8_t* scales = (uint8_t*)(data + block_id * 144 + 4); + get_scale_min_k4(is + 0, scales, &sc, &m); + const float d1 = d * sc; const float m1 = min * m; + get_scale_min_k4(is + 1, scales, &sc, &m); + const float d2 = d * sc; const float m2 = min * m; + for (int l = 0; l < 32; ++l) *output_blk++ = __float2bfloat16(d1 * (q[l] & 0xF) - m1); + for (int l = 0; l < 32; ++l) *output_blk++ = __float2bfloat16(d2 * (q[l] >> 4) - m2); + q += 32; is += 2; + } + } +} + +__global__ void dequantize_q5_k_fp32_kernel(const int8_t* data, float* output, const int blk_size, const int ele_per_blk, const int num_blocks) { + long long global_idx = blockIdx.x * blockDim.x + threadIdx.x; + for (long long block_id = global_idx; block_id < num_blocks; block_id += blockDim.x * gridDim.x){ + float* __restrict__ output_blk = (float*)(output + block_id * ele_per_blk); + + const float d = __half2float(*(reinterpret_cast(data + block_id * blk_size + 0))); + const float min = __half2float(*(reinterpret_cast(data + block_id * blk_size + 2))); + + const uint8_t * __restrict__ qh = (uint8_t*)(data + block_id * blk_size + 16); + const uint8_t * __restrict__ ql = (uint8_t*)(data + block_id * blk_size + 48); + + int is = 0; + uint8_t sc, m; + uint8_t u1 = 1, u2 = 2; + uint8_t* scales = (uint8_t*)(data + block_id * blk_size + 4); + + for (int j = 0; j < 256; j += 64) { + get_scale_min_k4(is + 0, scales, &sc, &m); + const float d1 = d * sc; const float m1 = min * m; + get_scale_min_k4(is + 1, scales, &sc, &m); + const float d2 = d * sc; const float m2 = min * m; + for (int l = 0; l < 32; ++l) *output_blk++ = d1 * ((ql[l] & 0xF) + (qh[l] & u1 ? 16 : 0)) - m1; + for (int l = 0; l < 32; ++l) *output_blk++ = d2 * ((ql[l] >> 4) + (qh[l] & u2 ? 16 : 0)) - m2; + ql += 32; is += 2; + u1 <<= 2; u2 <<= 2; + } + } +} + +__global__ void dequantize_q5_k_fp16_kernel(const int8_t* data, __half* output, const int blk_size, const int ele_per_blk, const int num_blocks) { + long long global_idx = blockIdx.x * blockDim.x + threadIdx.x; + for (long long block_id = global_idx; block_id < num_blocks; block_id += blockDim.x * gridDim.x){ + __half* __restrict__ output_blk = (__half*)(output + block_id * ele_per_blk); + + const float d = __half2float(*(reinterpret_cast(data + block_id * blk_size + 0))); + const float min = __half2float(*(reinterpret_cast(data + block_id * blk_size + 2))); + + const uint8_t * __restrict__ qh = (uint8_t*)(data + block_id * blk_size + 16); + const uint8_t * __restrict__ ql = (uint8_t*)(data + block_id * blk_size + 48); + + int is = 0; + uint8_t sc, m; + uint8_t u1 = 1, u2 = 2; + uint8_t* scales = (uint8_t*)(data + block_id * blk_size + 4); + + for (int j = 0; j < 256; j += 64) { + get_scale_min_k4(is + 0, scales, &sc, &m); + const float d1 = d * sc; const float m1 = min * m; + get_scale_min_k4(is + 1, scales, &sc, &m); + const float d2 = d * sc; const float m2 = min * m; + for (int l = 0; l < 32; ++l) *output_blk++ = __float2half(d1 * ((ql[l] & 0xF) + (qh[l] & u1 ? 16 : 0)) - m1); + for (int l = 0; l < 32; ++l) *output_blk++ = __float2half(d2 * ((ql[l] >> 4) + (qh[l] & u2 ? 16 : 0)) - m2); + ql += 32; is += 2; + u1 <<= 2; u2 <<= 2; + } + } +} + +__global__ void dequantize_q5_k_bf16_kernel(const int8_t* data, nv_bfloat16* output, const int blk_size, const int ele_per_blk, const int num_blocks) { + long long global_idx = blockIdx.x * blockDim.x + threadIdx.x; + for (long long block_id = global_idx; block_id < num_blocks; block_id += blockDim.x * gridDim.x){ + nv_bfloat16* __restrict__ output_blk = (nv_bfloat16*)(output + block_id * ele_per_blk); + + const float d = __half2float(*(reinterpret_cast(data + block_id * blk_size + 0))); + const float min = __half2float(*(reinterpret_cast(data + block_id * blk_size + 2))); + + const uint8_t * __restrict__ qh = (uint8_t*)(data + block_id * blk_size + 16); + const uint8_t * __restrict__ ql = (uint8_t*)(data + block_id * blk_size + 48); + + int is = 0; + uint8_t sc, m; + uint8_t u1 = 1, u2 = 2; + uint8_t* scales = (uint8_t*)(data + block_id * blk_size + 4); + + for (int j = 0; j < 256; j += 64) { + get_scale_min_k4(is + 0, scales, &sc, &m); + const float d1 = d * sc; const float m1 = min * m; + get_scale_min_k4(is + 1, scales, &sc, &m); + const float d2 = d * sc; const float m2 = min * m; + for (int l = 0; l < 32; ++l) *output_blk++ = __float2bfloat16(d1 * ((ql[l] & 0xF) + (qh[l] & u1 ? 16 : 0)) - m1); + for (int l = 0; l < 32; ++l) *output_blk++ = __float2bfloat16(d2 * ((ql[l] >> 4) + (qh[l] & u2 ? 16 : 0)) - m2); + ql += 32; is += 2; + u1 <<= 2; u2 <<= 2; + } + } +} + +__global__ void dequantize_q6_k_fp32_kernel(const int8_t* data, float* output, const int blk_size, const int ele_per_blk, const int num_blocks) { + long long global_idx = blockIdx.x * blockDim.x + threadIdx.x; + for (long long block_id=global_idx; block_id(data + block_id * blk_size + 208))); + + const uint8_t * __restrict__ ql = (uint8_t*)(data + block_id * blk_size); + const uint8_t * __restrict__ qh = (uint8_t*)(data + block_id * blk_size + 128); + const int8_t * __restrict__ sc = (int8_t*)(data + block_id * blk_size + 192); + + + for (int n = 0; n < ele_per_blk; n += 128) { + for (int l = 0; l < 32; ++l) { + int is = l/16; + const int8_t q1 = (int8_t)((ql[l + 0] & 0xF) | (((qh[l] >> 0) & 3) << 4)) - 32; + const int8_t q2 = (int8_t)((ql[l + 32] & 0xF) | (((qh[l] >> 2) & 3) << 4)) - 32; + const int8_t q3 = (int8_t)((ql[l + 0] >> 4) | (((qh[l] >> 4) & 3) << 4)) - 32; + const int8_t q4 = (int8_t)((ql[l + 32] >> 4) | (((qh[l] >> 6) & 3) << 4)) - 32; + output_blk[l + 0] = d * sc[is + 0] * q1; + output_blk[l + 32] = d * sc[is + 2] * q2; + output_blk[l + 64] = d * sc[is + 4] * q3; + output_blk[l + 96] = d * sc[is + 6] * q4; + } + output_blk += 128; + ql += 64; + qh += 32; + sc += 8; + } + } +} + +__global__ void dequantize_q6_k_fp16_kernel(const int8_t* data, __half* output, const int blk_size, const int ele_per_blk, const int num_blocks) { + long long global_idx = blockIdx.x * blockDim.x + threadIdx.x; + for (long long block_id=global_idx; block_id(data + block_id * blk_size + 208))); + + const uint8_t * __restrict__ ql = (uint8_t*)(data + block_id * blk_size); + const uint8_t * __restrict__ qh = (uint8_t*)(data + block_id * blk_size + 128); + const int8_t * __restrict__ sc = (int8_t*)(data + block_id * blk_size + 192); + + + for (int n = 0; n < ele_per_blk; n += 128) { + for (int l = 0; l < 32; ++l) { + int is = l/16; + const int8_t q1 = (int8_t)((ql[l + 0] & 0xF) | (((qh[l] >> 0) & 3) << 4)) - 32; + const int8_t q2 = (int8_t)((ql[l + 32] & 0xF) | (((qh[l] >> 2) & 3) << 4)) - 32; + const int8_t q3 = (int8_t)((ql[l + 0] >> 4) | (((qh[l] >> 4) & 3) << 4)) - 32; + const int8_t q4 = (int8_t)((ql[l + 32] >> 4) | (((qh[l] >> 6) & 3) << 4)) - 32; + output_blk[l + 0] = __float2half(d * sc[is + 0] * q1); + output_blk[l + 32] = __float2half(d * sc[is + 2] * q2); + output_blk[l + 64] = __float2half(d * sc[is + 4] * q3); + output_blk[l + 96] = __float2half(d * sc[is + 6] * q4); + } + output_blk += 128; + ql += 64; + qh += 32; + sc += 8; + } + } +} + +__global__ void dequantize_q6_k_bf16_kernel(const int8_t* data, nv_bfloat16* output, const int blk_size, const int ele_per_blk, const int num_blocks) { + long long global_idx = blockIdx.x * blockDim.x + threadIdx.x; + for (long long block_id=global_idx; block_id(data + block_id * blk_size + 208))); + + const uint8_t * __restrict__ ql = (uint8_t*)(data + block_id * blk_size); + const uint8_t * __restrict__ qh = (uint8_t*)(data + block_id * blk_size + 128); + const int8_t * __restrict__ sc = (int8_t*)(data + block_id * blk_size + 192); + + + for (int n = 0; n < ele_per_blk; n += 128) { + for (int l = 0; l < 32; ++l) { + int is = l/16; + const int8_t q1 = (int8_t)((ql[l + 0] & 0xF) | (((qh[l] >> 0) & 3) << 4)) - 32; + const int8_t q2 = (int8_t)((ql[l + 32] & 0xF) | (((qh[l] >> 2) & 3) << 4)) - 32; + const int8_t q3 = (int8_t)((ql[l + 0] >> 4) | (((qh[l] >> 4) & 3) << 4)) - 32; + const int8_t q4 = (int8_t)((ql[l + 32] >> 4) | (((qh[l] >> 6) & 3) << 4)) - 32; + output_blk[l + 0] = __float2bfloat16(d * sc[is + 0] * q1); + output_blk[l + 32] = __float2bfloat16(d * sc[is + 2] * q2); + output_blk[l + 64] = __float2bfloat16(d * sc[is + 4] * q3); + output_blk[l + 96] = __float2bfloat16(d * sc[is + 6] * q4); + } + output_blk += 128; + ql += 64; + qh += 32; + sc += 8; + } + } +} + +static constexpr __device__ int8_t kvalues_iq4nl[16] = {-127, -104, -83, -65, -49, -35, -22, -10, 1, 13, 25, 38, 53, 69, 89, 113}; + +__global__ void dequantize_iq4_xs_fp32_kernel(const int8_t* data, float* output, const int blk_size, const int ele_per_blk, const int num_blocks) { + long long global_idx = blockIdx.x * blockDim.x + threadIdx.x; + for (long long block_id=global_idx; block_id(data + block_id * blk_size))); + const uint16_t scales_h = *(reinterpret_cast(data + block_id * blk_size + 2)); + const uint8_t* scales_l = (uint8_t*)(data + block_id * blk_size + 2 + 2); + const uint8_t* qs = (uint8_t*)(data + block_id * blk_size + 2 + 2 + 4); + + for (int ib = 0; ib < 8; ++ib) { + const int ls = ((scales_l[ib / 2] >> 4 * (ib % 2)) & 0xf) | (((scales_h >> 2 * ib) & 3) << 4); + const float dl = d * (ls - 32); + for (int j = 0; j < 16; ++j) { + output_blk[j + 0] = dl * kvalues_iq4nl[qs[j] & 0xf]; + output_blk[j + 16] = dl * kvalues_iq4nl[qs[j] >> 4]; + } + output_blk += 32; + qs += 16; + } + } +} + +__global__ void dequantize_iq4_xs_fp16_kernel(const int8_t* data, __half* output, const int blk_size, const int ele_per_blk, const int num_blocks) { + long long global_idx = blockIdx.x * blockDim.x + threadIdx.x; + for (long long block_id=global_idx; block_id(data + block_id * blk_size))); + const uint16_t scales_h = *(reinterpret_cast(data + block_id * blk_size + 2)); + const uint8_t* scales_l = (uint8_t*)(data + block_id * blk_size + 2 + 2); + const uint8_t* qs = (uint8_t*)(data + block_id * blk_size + 2 + 2 + 4); + + for (int ib = 0; ib < 8; ++ib) { + const int ls = ((scales_l[ib / 2] >> 4 * (ib % 2)) & 0xf) | (((scales_h >> 2 * ib) & 3) << 4); + const float dl = d * (ls - 32); + for (int j = 0; j < 16; ++j) { + output_blk[j + 0] = __float2half(dl * kvalues_iq4nl[qs[j] & 0xf]); + output_blk[j + 16] = __float2half(dl * kvalues_iq4nl[qs[j] >> 4]); + } + output_blk += 32; + qs += 16; + } + } +} + +__global__ void dequantize_iq4_xs_bf16_kernel(const int8_t* data, nv_bfloat16* output, const int blk_size, const int ele_per_blk, const int num_blocks) { + long long global_idx = blockIdx.x * blockDim.x + threadIdx.x; + for (long long block_id=global_idx; block_id(data + block_id * blk_size))); + const uint16_t scales_h = *(reinterpret_cast(data + block_id * blk_size + 2)); + const uint8_t* scales_l = (uint8_t*)(data + block_id * blk_size + 2 + 2); + const uint8_t* qs = (uint8_t*)(data + block_id * blk_size + 2 + 2 + 4); + + for (int ib = 0; ib < 8; ++ib) { + const int ls = ((scales_l[ib / 2] >> 4 * (ib % 2)) & 0xf) | (((scales_h >> 2 * ib) & 3) << 4); + const float dl = d * (ls - 32); + for (int j = 0; j < 16; ++j) { + output_blk[j + 0] = __float2bfloat16(dl * kvalues_iq4nl[qs[j] & 0xf]); + output_blk[j + 16] = __float2bfloat16(dl * kvalues_iq4nl[qs[j] >> 4]); + } + output_blk += 32; + qs += 16; + } + } +} + +torch::Tensor dequantize_q8_0(const int8_t* data, const int num_bytes, const int blk_size, const int ele_per_blk, const torch::Device device, const torch::Dtype target_dtype) { + int num_blocks = num_bytes / blk_size; + const at::cuda::OptionalCUDAGuard device_guard(device); + + auto options = torch::TensorOptions().dtype(torch::kInt8).device(device).memory_format(torch::MemoryFormat::Contiguous); + auto data_gpu = torch::empty({ num_bytes }, options); + + cudaMemcpy(data_gpu.data_ptr(), data, num_bytes, cudaMemcpyHostToDevice); + //data_gpu.copy_(data, false); + + // Create output tensor + auto output = torch::zeros({ num_blocks, 32 }, torch::dtype(target_dtype).device(device)); + + switch (target_dtype) { + case torch::kFloat16: + dequantize_q8_0_fp16_kernel<<<512, 256>>>(data_gpu.data_ptr(), (__half*)output.data_ptr(), blk_size, ele_per_blk, num_blocks); + break; + case torch::kBFloat16: + dequantize_q8_0_bf16_kernel<<<512, 256>>>(data_gpu.data_ptr(), (nv_bfloat16*)output.data_ptr(), blk_size, ele_per_blk, num_blocks); + break; + case torch::kFloat32: + dequantize_q8_0_fp32_kernel<<<512, 256>>>(data_gpu.data_ptr(), output.data_ptr(), blk_size, ele_per_blk, num_blocks); + break; + default: + printf("target type not support\n"); + exit(0); + } + + cudaDeviceSynchronize(); + return output; +} + + +torch::Tensor dequantize_q6_k(const int8_t* data, const int num_bytes, const int blk_size, const int ele_per_blk, const torch::Device device, const torch::Dtype target_dtype) { + // data.numel%blk_size should be 0, else raise err + int num_blocks = num_bytes / blk_size; + + const at::cuda::OptionalCUDAGuard device_guard(device); + auto options = torch::TensorOptions().dtype(torch::kInt8).device(device).memory_format(torch::MemoryFormat::Contiguous); + auto data_gpu = torch::empty({num_bytes}, options); + + cudaMemcpy(data_gpu.data_ptr(), data, num_bytes, cudaMemcpyHostToDevice); + //data_gpu.copy_(data, false); + + // Create output tensor + auto output = torch::zeros({num_blocks, 256}, torch::dtype(target_dtype).device(device)); + + switch (target_dtype) { + case torch::kFloat16: + dequantize_q6_k_fp16_kernel<<<512, 256>>>(data_gpu.data_ptr(), (__half*)output.data_ptr(), blk_size, ele_per_blk, num_blocks); + break; + case torch::kBFloat16: + dequantize_q6_k_bf16_kernel<<<512, 256>>>(data_gpu.data_ptr(), (nv_bfloat16*)output.data_ptr(), blk_size, ele_per_blk, num_blocks); + break; + case torch::kFloat32: + dequantize_q6_k_fp32_kernel<<<512, 256>>>(data_gpu.data_ptr(), output.data_ptr(), blk_size, ele_per_blk, num_blocks); + break; + default: + printf("target type not support\n"); + exit(0); + } + cudaDeviceSynchronize(); + return output; +} + +torch::Tensor dequantize_q5_k(const int8_t* data, const int num_bytes, const int blk_size, const int ele_per_blk, const torch::Device device, const torch::Dtype target_dtype) { + int num_blocks = num_bytes / blk_size; + const at::cuda::OptionalCUDAGuard device_guard(device); + + auto options = torch::TensorOptions().dtype(torch::kInt8).device(device).memory_format(torch::MemoryFormat::Contiguous); + auto data_gpu = torch::empty({num_bytes}, options); + + cudaMemcpy(data_gpu.data_ptr(), data, num_bytes, cudaMemcpyHostToDevice); + //data_gpu.copy_(data, false); + + // Create output tensor + auto output = torch::zeros({num_blocks, 256}, torch::dtype(target_dtype).device(device)); + + switch (target_dtype) { + case torch::kFloat16: + dequantize_q5_k_fp16_kernel<<<512, 256>>>(data_gpu.data_ptr(), (__half*)output.data_ptr(), blk_size, ele_per_blk, num_blocks); + break; + case torch::kBFloat16: + dequantize_q5_k_bf16_kernel<<<512, 256>>>(data_gpu.data_ptr(), (nv_bfloat16*)output.data_ptr(), blk_size, ele_per_blk, num_blocks); + break; + case torch::kFloat32: + dequantize_q5_k_fp32_kernel<<<512, 256>>>(data_gpu.data_ptr(), output.data_ptr(), blk_size, ele_per_blk, num_blocks); + break; + default: + printf("target type not support\n"); + exit(0); + } + cudaDeviceSynchronize(); + return output; +} + +torch::Tensor dequantize_q4_k(const int8_t* data, const int num_bytes, const int blk_size, const int ele_per_blk, const torch::Device device, const torch::Dtype target_dtype) { + // data.numel%blk_size should be 0, else raise err + int num_blocks = num_bytes / blk_size; + const at::cuda::OptionalCUDAGuard device_guard(device); + + auto options = torch::TensorOptions().dtype(torch::kInt8).device(device).memory_format(torch::MemoryFormat::Contiguous); + auto data_gpu = torch::empty({num_bytes}, options); + + cudaMemcpy(data_gpu.data_ptr(), data, num_bytes, cudaMemcpyHostToDevice); + //data_gpu.copy_(data, false); + + // Create output tensor + auto output = torch::zeros({num_blocks, 256}, torch::dtype(target_dtype).device(device)); + + switch (target_dtype) { + case torch::kFloat16: + dequantize_q4_k_fp16_kernel<<<512, 256>>>(data_gpu.data_ptr(), (__half*)output.data_ptr(), blk_size, ele_per_blk, num_blocks); + break; + case torch::kBFloat16: + dequantize_q4_k_bf16_kernel<<<512, 256>>>(data_gpu.data_ptr(), (nv_bfloat16*)output.data_ptr(), blk_size, ele_per_blk, num_blocks); + break; + case torch::kFloat32: + dequantize_q4_k_fp32_kernel<<<512, 256>>>(data_gpu.data_ptr(), output.data_ptr(), blk_size, ele_per_blk, num_blocks); + break; + default: + printf("target type not support\n"); + exit(0); + } + cudaDeviceSynchronize(); + return output; +} + +torch::Tensor dequantize_q3_k(const int8_t* data, const int num_bytes, const int blk_size, const int ele_per_blk, const torch::Device device, const torch::Dtype target_dtype) { + int num_blocks = num_bytes / blk_size; + const at::cuda::OptionalCUDAGuard device_guard(device); + + auto options = torch::TensorOptions().dtype(torch::kInt8).device(device).memory_format(torch::MemoryFormat::Contiguous); + auto data_gpu = torch::empty({num_bytes}, options); + + cudaMemcpy(data_gpu.data_ptr(), data, num_bytes, cudaMemcpyHostToDevice); + //data_gpu.copy_(data, false); + + // Create output tensor + auto output = torch::zeros({num_blocks, 256}, torch::dtype(target_dtype).device(device)); + + switch (target_dtype) { + case torch::kFloat16: + dequantize_q3_k_fp16_kernel<<<512, 256>>>(data_gpu.data_ptr(), (__half*)output.data_ptr(), blk_size, ele_per_blk, num_blocks); + break; + case torch::kBFloat16: + dequantize_q3_k_bf16_kernel<<<512, 256>>>(data_gpu.data_ptr(), (nv_bfloat16*)output.data_ptr(), blk_size, ele_per_blk, num_blocks); + break; + case torch::kFloat32: + dequantize_q3_k_fp32_kernel<<<512, 256>>>(data_gpu.data_ptr(), output.data_ptr(), blk_size, ele_per_blk, num_blocks); + break; + default: + printf("target type not support\n"); + exit(0); + } + cudaDeviceSynchronize(); + return output; +} + +torch::Tensor dequantize_q2_k(const int8_t* data, const int num_bytes, const int blk_size, const int ele_per_blk, const torch::Device device, const torch::Dtype target_dtype) { + int num_blocks = num_bytes / blk_size; + const at::cuda::OptionalCUDAGuard device_guard(device); + + auto options = torch::TensorOptions().dtype(torch::kInt8).device(device).memory_format(torch::MemoryFormat::Contiguous); + auto data_gpu = torch::empty({num_bytes}, options); + + cudaMemcpy(data_gpu.data_ptr(), data, num_bytes, cudaMemcpyHostToDevice); + //data_gpu.copy_(data, false); + + // Create output tensor + auto output = torch::zeros({num_blocks, 256}, torch::dtype(target_dtype).device(device)); + + switch (target_dtype) { + case torch::kFloat16: + dequantize_q2_k_fp16_kernel<<<512, 256>>>(data_gpu.data_ptr(), (__half*)output.data_ptr(), blk_size, ele_per_blk, num_blocks); + break; + case torch::kBFloat16: + dequantize_q2_k_bf16_kernel<<<512, 256>>>(data_gpu.data_ptr(), (nv_bfloat16*)output.data_ptr(), blk_size, ele_per_blk, num_blocks); + break; + case torch::kFloat32: + dequantize_q2_k_fp32_kernel<<<512, 256>>>(data_gpu.data_ptr(), output.data_ptr(), blk_size, ele_per_blk, num_blocks); + break; + default: + printf("target type not support\n"); + exit(0); + } + cudaDeviceSynchronize(); + return output; +} + +torch::Tensor dequantize_iq4_xs(const int8_t* data, const int num_bytes, const int blk_size, const int ele_per_blk, const torch::Device device, const torch::Dtype target_dtype) { + int num_blocks = num_bytes / blk_size; + const at::cuda::OptionalCUDAGuard device_guard(device); + + auto options = torch::TensorOptions().dtype(torch::kInt8).device(device).memory_format(torch::MemoryFormat::Contiguous); + auto data_gpu = torch::empty({num_bytes}, options); + + cudaMemcpy(data_gpu.data_ptr(), data, num_bytes, cudaMemcpyHostToDevice); + //data_gpu.copy_(data, false); + + // Create output tensor + auto output = torch::zeros({num_blocks, 256}, torch::dtype(target_dtype).device(device)); + + switch (target_dtype) { + case torch::kFloat16: + dequantize_iq4_xs_fp16_kernel<<<512, 256>>>(data_gpu.data_ptr(), (__half*)output.data_ptr(), blk_size, ele_per_blk, num_blocks); + break; + case torch::kBFloat16: + dequantize_iq4_xs_bf16_kernel<<<512, 256>>>(data_gpu.data_ptr(), (nv_bfloat16*)output.data_ptr(), blk_size, ele_per_blk, num_blocks); + break; + case torch::kFloat32: + dequantize_iq4_xs_fp32_kernel<<<512, 256>>>(data_gpu.data_ptr(), output.data_ptr(), blk_size, ele_per_blk, num_blocks); + break; + default: + printf("target type not support\n"); + exit(0); + } + cudaDeviceSynchronize(); + return output; +} diff --git a/archive/csrc/ktransformers_ext/cuda/custom_gguf/ops.h b/archive/csrc/ktransformers_ext/cuda/custom_gguf/ops.h new file mode 100644 index 0000000..fe9161a --- /dev/null +++ b/archive/csrc/ktransformers_ext/cuda/custom_gguf/ops.h @@ -0,0 +1,22 @@ +/** + * @Description : + * @Author : Azure-Tang + * @Date : 2024-07-22 09:27:55 + * @Version : 1.0.0 + * @LastEditors : kkk1nak0 + * @LastEditTime : 2024-08-12 03:48:46 + * @Copyright (c) 2024 by KVCache.AI, All Rights Reserved. +**/ +#pragma once + +#include +#include +#include + +torch::Tensor dequantize_q8_0(const int8_t* data, const int num_bytes, const int blk_size, const int ele_per_blk, const torch::Device device, const torch::Dtype target_dtype); +torch::Tensor dequantize_q6_k(const int8_t* data, const int num_bytes, const int blk_size, const int ele_per_blk, const torch::Device device, const torch::Dtype target_dtype); +torch::Tensor dequantize_q5_k(const int8_t* data, const int num_bytes, const int blk_size, const int ele_per_blk, const torch::Device device, const torch::Dtype target_dtype); +torch::Tensor dequantize_q4_k(const int8_t* data, const int num_bytes, const int blk_size, const int ele_per_blk, const torch::Device device, const torch::Dtype target_dtype); +torch::Tensor dequantize_q3_k(const int8_t* data, const int num_bytes, const int blk_size, const int ele_per_blk, const torch::Device device, const torch::Dtype target_dtype); +torch::Tensor dequantize_q2_k(const int8_t* data, const int num_bytes, const int blk_size, const int ele_per_blk, const torch::Device device, const torch::Dtype target_dtype); +torch::Tensor dequantize_iq4_xs(const int8_t* data, const int num_bytes, const int blk_size, const int ele_per_blk, const torch::Device device, const torch::Dtype target_dtype); \ No newline at end of file diff --git a/archive/csrc/ktransformers_ext/cuda/gptq_marlin/gptq_marlin.cu b/archive/csrc/ktransformers_ext/cuda/gptq_marlin/gptq_marlin.cu new file mode 100644 index 0000000..87f4581 --- /dev/null +++ b/archive/csrc/ktransformers_ext/cuda/gptq_marlin/gptq_marlin.cu @@ -0,0 +1,1907 @@ +/* + * Modified by Neural Magic + * Copyright (C) Marlin.2024 Elias Frantar + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +/* + * Adapted from https://github.com/IST-DASLab/marlin + */ +/* + * Adapted from https://github.com/vllm-project/vllm/tree/main/csrc/quantization/gptq_marlin + */ +#include "gptq_marlin.cuh" +#include "gptq_marlin_dtypes.cuh" +#include +#define STATIC_ASSERT_SCALAR_TYPE_VALID(scalar_t) \ + static_assert(std::is_same::value || \ + std::is_same::value, \ + "only float16 and bfloat16 is supported"); + +template +inline std::string str(T x) { + return std::to_string(x); +} + +namespace gptq_marlin { + +#if (defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800) || defined(__HIP_PLATFORM_AMD__) + +__global__ void permute_cols_kernel(int4 const* __restrict__ a_int4_ptr, + int const* __restrict__ perm_int_ptr, + int4* __restrict__ out_int4_ptr, int size_m, + int size_k, int block_rows) {} + +template shared + // fetch pipeline + const bool has_act_order, // whether act_order is enabled + const int group_blocks = -1 // number of consecutive 16x16 blocks + // with a separate quantization scale + > +__global__ void Marlin( + const int4* __restrict__ A, // fp16 input matrix of shape mxk + const int4* __restrict__ B, // 4bit quantized weight matrix of shape kxn + int4* __restrict__ C, // fp16 output buffer of shape mxn + const int4* __restrict__ scales_ptr, // fp16 quantization scales of shape + // (k/groupsize)xn + const int* __restrict__ g_idx, // int32 group indices of shape k + int num_groups, // number of scale groups per output channel + int prob_m, // batch dimension m + int prob_n, // output dimension n + int prob_k, // reduction dimension k + int* locks // extra global storage for barrier synchronization +) {} + +} // namespace gptq_marlin + +torch::Tensor gptq_marlin_gemm(torch::Tensor& a, torch::Tensor& b_q_weight, + torch::Tensor& b_scales, torch::Tensor& g_idx, + torch::Tensor& perm, torch::Tensor& workspace, + int64_t num_bits, int64_t size_m, int64_t size_n, + int64_t size_k, bool is_k_full) { + TORCH_CHECK_NOT_IMPLEMENTED(false, + "marlin_gemm(..) requires CUDA_ARCH >= 8.0"); + return torch::empty({1, 1}); +} + +#else + +// m16n8k16 tensor core mma instruction with fp16 inputs and fp32 +// output/accumulation. +template +__device__ inline void mma(const typename ScalarType::FragA& a_frag, + const typename ScalarType::FragB& frag_b, + typename ScalarType::FragC& frag_c) { + const uint32_t* a = reinterpret_cast(&a_frag); + const uint32_t* b = reinterpret_cast(&frag_b); + float* c = reinterpret_cast(&frag_c); + if constexpr (std::is_same::value) { + asm volatile( + "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 " + "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};\n" + : "=f"(c[0]), "=f"(c[1]), "=f"(c[2]), "=f"(c[3]) + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1]), + "f"(c[0]), "f"(c[1]), "f"(c[2]), "f"(c[3])); + } else if constexpr (std::is_same::value) { + asm volatile( + "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 " + "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};\n" + : "=f"(c[0]), "=f"(c[1]), "=f"(c[2]), "=f"(c[3]) + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1]), + "f"(c[0]), "f"(c[1]), "f"(c[2]), "f"(c[3])); + } else { + STATIC_ASSERT_SCALAR_TYPE_VALID(scalar_t); + } +} + +// Instruction for loading a full 16x16 matrix fragment of operand A from shared +// memory, directly in tensor core layout. +template +__device__ inline void ldsm4(typename ScalarType::FragA& frag_a, + const void* smem_ptr) { + uint32_t* a = reinterpret_cast(&frag_a); + uint32_t smem = static_cast(__cvta_generic_to_shared(smem_ptr)); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0,%1,%2,%3}, [%4];\n" + : "=r"(a[0]), "=r"(a[1]), "=r"(a[2]), "=r"(a[3]) + : "r"(smem)); +} + +// Lookup-table based 3-input logical operation; explicitly used for +// dequantization as the compiler does not seem to automatically recognize it in +// all cases. +template +__device__ inline int lop3(int a, int b, int c) { + int res; + asm volatile("lop3.b32 %0, %1, %2, %3, %4;\n" + : "=r"(res) + : "r"(a), "r"(b), "r"(c), "n"(lut)); + return res; +} + +// Constructs destination register by taking bytes from 2 sources (based on +// mask) +template +__device__ inline uint32_t prmt(uint32_t a) { + uint32_t res; + asm volatile("prmt.b32 %0, %1, %2, %3;\n" + : "=r"(res) + : "r"(a), "n"(start_byte), "n"(mask)); + return res; +} + +// Efficiently dequantize an int32 value into a full B-fragment of 4 fp16 +// values. We mostly follow the strategy in the link below, with some small +// changes: +// - FP16: +// https://github.com/NVIDIA/FasterTransformer/blob/release/v5.3_tag/src/fastertransformer/cutlass_extensions/include/cutlass_extensions/interleaved_numeric_conversion.h#L215-L287 +// - BF16: +// https://github.com/NVIDIA/FasterTransformer/blob/release/v5.3_tag/src/fastertransformer/cutlass_extensions/include/cutlass_extensions/interleaved_numeric_conversion.h#L327-L385 +template +__device__ inline typename ScalarType::FragB dequant_4bit(int q) { + STATIC_ASSERT_SCALAR_TYPE_VALID(scalar_t); +} + +template <> +__device__ inline typename ScalarType::FragB dequant_4bit(int q) { + const int LO = 0x000f000f; + const int HI = 0x00f000f0; + const int EX = 0x64006400; + // Guarantee that the `(a & b) | c` operations are LOP3s. + int lo = lop3<(0xf0 & 0xcc) | 0xaa>(q, LO, EX); + int hi = lop3<(0xf0 & 0xcc) | 0xaa>(q, HI, EX); + // We want signed int4 outputs, hence we fuse the `-8` symmetric zero point + // directly into `SUB` and `ADD`. + const int SUB = 0x64086408; + const int MUL = 0x2c002c00; + const int ADD = 0xd480d480; + typename ScalarType::FragB frag_b; + frag_b[0] = __hsub2(*reinterpret_cast(&lo), + *reinterpret_cast(&SUB)); + frag_b[1] = __hfma2(*reinterpret_cast(&hi), + *reinterpret_cast(&MUL), + *reinterpret_cast(&ADD)); + return frag_b; +} + +template <> +__device__ inline typename ScalarType::FragB +dequant_4bit(int q) { + static constexpr uint32_t MASK = 0x000f000f; + static constexpr uint32_t EX = 0x43004300; + + // Guarantee that the `(a & b) | c` operations are LOP3s. + + int lo = lop3<(0xf0 & 0xcc) | 0xaa>(q, MASK, EX); + q >>= 4; + int hi = lop3<(0xf0 & 0xcc) | 0xaa>(q, MASK, EX); + + typename ScalarType::FragB frag_b; + static constexpr uint32_t MUL = 0x3F803F80; + static constexpr uint32_t ADD = 0xC308C308; + + frag_b[0] = __hfma2(*reinterpret_cast(&lo), + *reinterpret_cast(&MUL), + *reinterpret_cast(&ADD)); + frag_b[1] = __hfma2(*reinterpret_cast(&hi), + *reinterpret_cast(&MUL), + *reinterpret_cast(&ADD)); + return frag_b; +} + +// Fast Int8ToFp16/Int8ToBf16: Efficiently dequantize 8bit int values to fp16 or +// bf16 Reference: +// - FP16: +// https://github.com/NVIDIA/FasterTransformer/blob/release/v5.3_tag/src/fastertransformer/cutlass_extensions/include/cutlass_extensions/interleaved_numeric_conversion.h#L53-L85 +// - BF16: +// https://github.com/NVIDIA/FasterTransformer/blob/release/v5.3_tag/src/fastertransformer/cutlass_extensions/include/cutlass_extensions/interleaved_numeric_conversion.h#L125-L175 +template +__device__ inline typename ScalarType::FragB dequant_8bit(int q) { + STATIC_ASSERT_SCALAR_TYPE_VALID(scalar_t); +} + +template <> +__device__ inline typename ScalarType::FragB dequant_8bit(int q) { + static constexpr uint32_t mask_for_elt_01 = 0x5250; + static constexpr uint32_t mask_for_elt_23 = 0x5351; + static constexpr uint32_t start_byte_for_fp16 = 0x64646464; + + uint32_t lo = prmt(q); + uint32_t hi = prmt(q); + + static constexpr uint32_t I8s_TO_F16s_MAGIC_NUM = 0x64806480; + + typename ScalarType::FragB frag_b; + frag_b[0] = __hsub2(*reinterpret_cast(&lo), + *reinterpret_cast(&I8s_TO_F16s_MAGIC_NUM)); + frag_b[1] = __hsub2(*reinterpret_cast(&hi), + *reinterpret_cast(&I8s_TO_F16s_MAGIC_NUM)); + return frag_b; +} + +template <> +__device__ inline typename ScalarType::FragB +dequant_8bit(int q) { + typename ScalarType::FragB frag_b; + + float fp32_intermediates[4]; + uint32_t* fp32_intermediates_casted = + reinterpret_cast(fp32_intermediates); + + static constexpr uint32_t fp32_base = 0x4B000000; + fp32_intermediates_casted[0] = __byte_perm(q, fp32_base, 0x7650); + fp32_intermediates_casted[1] = __byte_perm(q, fp32_base, 0x7652); + fp32_intermediates_casted[2] = __byte_perm(q, fp32_base, 0x7651); + fp32_intermediates_casted[3] = __byte_perm(q, fp32_base, 0x7653); + + fp32_intermediates[0] -= 8388736.f; + fp32_intermediates[1] -= 8388736.f; + fp32_intermediates[2] -= 8388736.f; + fp32_intermediates[3] -= 8388736.f; + + uint32_t* bf16_result_ptr = reinterpret_cast(&frag_b); + bf16_result_ptr[0] = __byte_perm(fp32_intermediates_casted[0], + fp32_intermediates_casted[1], 0x7632); + bf16_result_ptr[1] = __byte_perm(fp32_intermediates_casted[2], + fp32_intermediates_casted[3], 0x7632); + + return frag_b; +} + +// Multiply dequantized values by the corresponding quantization scale; used +// only for grouped quantization. +template +__device__ inline void scale(typename ScalarType::FragB& frag_b, + typename ScalarType::FragS& frag_s, + int i) { + using scalar_t2 = typename ScalarType::scalar_t2; + scalar_t2 s = + ScalarType::num2num2(reinterpret_cast(&frag_s)[i]); + frag_b[0] = __hmul2(frag_b[0], s); + frag_b[1] = __hmul2(frag_b[1], s); +} + +// Same as above, but for act_order (each K is multiplied individually) +template +__device__ inline void scale4(typename ScalarType::FragB& frag_b, + typename ScalarType::FragS& frag_s_1, + typename ScalarType::FragS& frag_s_2, + typename ScalarType::FragS& frag_s_3, + typename ScalarType::FragS& frag_s_4, + int i) { + using scalar_t2 = typename ScalarType::scalar_t2; + scalar_t2 s_val_1_2; + s_val_1_2.x = reinterpret_cast(&frag_s_1)[i]; + s_val_1_2.y = reinterpret_cast(&frag_s_2)[i]; + + scalar_t2 s_val_3_4; + s_val_3_4.x = reinterpret_cast(&frag_s_3)[i]; + s_val_3_4.y = reinterpret_cast(&frag_s_4)[i]; + + frag_b[0] = __hmul2(frag_b[0], s_val_1_2); + frag_b[1] = __hmul2(frag_b[1], s_val_3_4); +} + +// Given 2 floats multiply by 2 scales (halves) +template +__device__ inline void scale_float(float* c, + typename ScalarType::FragS& s) { + scalar_t* s_ptr = reinterpret_cast(&s); + c[0] = __fmul_rn(c[0], ScalarType::num2float(s_ptr[0])); + c[1] = __fmul_rn(c[1], ScalarType::num2float(s_ptr[1])); +} + +// Wait until barrier reaches `count`, then lock for current threadblock. +__device__ inline void barrier_acquire(int* lock, int count) { + if (threadIdx.x == 0) { + int state = -1; + do + // Guarantee that subsequent writes by this threadblock will be visible + // globally. + asm volatile("ld.global.acquire.gpu.b32 %0, [%1];\n" + : "=r"(state) + : "l"(lock)); + while (state != count); + } + __syncthreads(); +} + +// Release barrier and increment visitation count. +__device__ inline void barrier_release(int* lock, bool reset = false) { + __syncthreads(); + if (threadIdx.x == 0) { + if (reset) { + lock[0] = 0; + return; + } + int val = 1; + // Make sure that all writes since acquiring this barrier are visible + // globally, while releasing the barrier. + asm volatile("fence.acq_rel.gpu;\n"); + asm volatile("red.relaxed.gpu.global.add.s32 [%0], %1;\n" + : + : "l"(lock), "r"(val)); + } +} + +// For a given "a" of size [M,K] performs a permutation of the K columns based +// on the given "perm" indices. +__global__ void permute_cols_kernel(int4 const* __restrict__ a_int4_ptr, + int const* __restrict__ perm_int_ptr, + int4* __restrict__ out_int4_ptr, int size_m, + int size_k, int block_rows) { + int start_row = block_rows * blockIdx.x; + int finish_row = start_row + block_rows; + if (finish_row > size_m) { + finish_row = size_m; + } + int cur_block_rows = finish_row - start_row; + + int row_stride = size_k * sizeof(half) / 16; + + auto permute_row = [&](int row) { + int iters = size_k / default_threads; + int rest = size_k % default_threads; + + int offset = row * row_stride; + + half const* a_row_half = reinterpret_cast(a_int4_ptr + offset); + half* out_half = reinterpret_cast(out_int4_ptr + offset); + + int base_k = 0; + + for (int i = 0; i < iters; i++) { + int cur_k = base_k + threadIdx.x; + int src_pos = perm_int_ptr[cur_k]; + + out_half[cur_k] = a_row_half[src_pos]; + + base_k += default_threads; + } + + if (rest) { + if (threadIdx.x < rest) { + int cur_k = base_k + threadIdx.x; + int src_pos = perm_int_ptr[cur_k]; + + out_half[cur_k] = a_row_half[src_pos]; + } + } + }; + + for (int i = 0; i < cur_block_rows; i++) { + int cur_row = start_row + i; + if (cur_row < size_m) { + permute_row(cur_row); + } + } +} + +template shared + // fetch pipeline + const bool has_act_order, // whether act_order is enabled + const int group_blocks = -1 // number of consecutive 16x16 blocks + // with a separate quantization scale + > +__global__ void Marlin( + const int4* __restrict__ A, // fp16 input matrix of shape mxk + const int4* __restrict__ B, // 4bit quantized weight matrix of shape kxn + int4* __restrict__ C, // fp16 output buffer of shape mxn + const int4* __restrict__ scales_ptr, // fp16 quantization scales of shape + // (k/groupsize)xn + const int* __restrict__ g_idx, // int32 group indices of shape k + int num_groups, // number of scale groups per output channel + int prob_m, // batch dimension m + int prob_n, // output dimension n + int prob_k, // reduction dimension k + int* locks // extra global storage for barrier synchronization +) { + // Each threadblock processes one "stripe" of the B matrix with (roughly) the + // same size, which might involve multiple column "slices" (of width 16 * + // `thread_n_blocks`). Stripes are defined as shown in the 3x3 matrix 5 SM + // example: + // 0 1 3 + // 0 2 3 + // 1 2 4 + // While this kind of partitioning makes things somewhat more complicated, it + // ensures good utilization of all SMs for many kinds of shape and GPU + // configurations, while requiring as few slow global cross-threadblock + // reductions as possible. + using Dtype = ScalarType; + using scalar_t2 = typename ScalarType::scalar_t2; + using FragA = typename ScalarType::FragA; + using FragB = typename ScalarType::FragB; + using FragC = typename ScalarType::FragC; + using FragS = typename ScalarType::FragS; + + constexpr int pack_factor = 32 / num_bits; + + // For larger GEMMs we run multiple batchsize 64 versions in parallel for a + // better partitioning with less reductions + int parallel = 1; + if (prob_m > 16 * thread_m_blocks) { + parallel = prob_m / (16 * thread_m_blocks); + prob_m = 16 * thread_m_blocks; + } + + int k_tiles = prob_k / 16 / thread_k_blocks; + int n_tiles = prob_n / 16 / thread_n_blocks; + int iters = div_ceil(k_tiles * n_tiles * parallel, gridDim.x); + + if constexpr (!has_act_order && group_blocks != -1) { + if (group_blocks >= thread_k_blocks) { + // Ensure that the number of tiles in each stripe is a multiple of the + // groupsize; this avoids an annoying special case where a stripe starts + // in the middle of group. + iters = (group_blocks / thread_k_blocks) * + div_ceil(iters, (group_blocks / thread_k_blocks)); + } + } + + int slice_row = (iters * blockIdx.x) % k_tiles; + int slice_col_par = (iters * blockIdx.x) / k_tiles; + int slice_col = slice_col_par; + int slice_iters; // number of threadblock tiles in the current slice + int slice_count = + 0; // total number of active threadblocks in the current slice + int slice_idx; // index of threadblock in current slice; numbered bottom to + // top + + // We can easily implement parallel problem execution by just remapping + // indices and advancing global pointers + if (slice_col_par >= n_tiles) { + A += (slice_col_par / n_tiles) * 16 * thread_m_blocks * prob_k / 8; + C += (slice_col_par / n_tiles) * 16 * thread_m_blocks * prob_n / 8; + locks += (slice_col_par / n_tiles) * n_tiles; + slice_col = slice_col_par % n_tiles; + } + + // Compute all information about the current slice which is required for + // synchronization. + auto init_slice = [&]() { + slice_iters = + iters * (blockIdx.x + 1) - (k_tiles * slice_col_par + slice_row); + if (slice_iters < 0 || slice_col_par >= n_tiles * parallel) slice_iters = 0; + if (slice_iters == 0) return; + if (slice_row + slice_iters > k_tiles) slice_iters = k_tiles - slice_row; + slice_count = 1; + slice_idx = 0; + int col_first = iters * div_ceil(k_tiles * slice_col_par, iters); + if (col_first <= k_tiles * (slice_col_par + 1)) { + int col_off = col_first - k_tiles * slice_col_par; + slice_count = div_ceil(k_tiles - col_off, iters); + if (col_off > 0) slice_count++; + int delta_first = iters * blockIdx.x - col_first; + if (delta_first < 0 || (col_off == 0 && delta_first == 0)) + slice_idx = slice_count - 1; + else { + slice_idx = slice_count - 1 - delta_first / iters; + if (col_off > 0) slice_idx--; + } + } + if (slice_col == n_tiles) { + A += 16 * thread_m_blocks * prob_k / 8; + C += 16 * thread_m_blocks * prob_n / 8; + locks += n_tiles; + slice_col = 0; + } + }; + init_slice(); + + // A sizes/strides + + // stride of the A matrix in global memory + int a_gl_stride = prob_k / 8; + // stride of an A matrix tile in shared memory + constexpr int a_sh_stride = 16 * thread_k_blocks / 8; + // delta between subsequent A tiles in global memory + constexpr int a_gl_rd_delta_o = 16 * thread_k_blocks / 8; + // between subsequent accesses within a tile + int a_gl_rd_delta_i = a_gl_stride * (threads / a_gl_rd_delta_o); + // between shared memory writes + constexpr int a_sh_wr_delta = a_sh_stride * (threads / a_gl_rd_delta_o); + // between shared memory tile reads + constexpr int a_sh_rd_delta_o = 2 * ((threads / 32) / (thread_n_blocks / 4)); + // within a shared memory tile + constexpr int a_sh_rd_delta_i = a_sh_stride * 16; + // overall size of a tile + constexpr int a_sh_stage = a_sh_stride * (16 * thread_m_blocks); + // number of shared write iterations for a tile + constexpr int a_sh_wr_iters = div_ceil(a_sh_stage, a_sh_wr_delta); + + // B sizes/strides + int b_gl_stride = 16 * prob_n / (pack_factor * 4); + constexpr int b_sh_stride = ((thread_n_blocks * 16) * 16 / pack_factor) / 4; + constexpr int b_thread_vecs = num_bits == 4 ? 1 : 2; + constexpr int b_sh_stride_threads = b_sh_stride / b_thread_vecs; + + int b_gl_rd_delta_o = b_gl_stride * thread_k_blocks; + int b_gl_rd_delta_i = b_gl_stride * (threads / b_sh_stride_threads); + constexpr int b_sh_wr_delta = threads * b_thread_vecs; + constexpr int b_sh_rd_delta = threads * b_thread_vecs; + constexpr int b_sh_stage = b_sh_stride * thread_k_blocks; + constexpr int b_sh_wr_iters = b_sh_stage / b_sh_wr_delta; + + // Scale sizes/strides without act_order + int s_gl_stride = prob_n / 8; + constexpr int s_sh_stride = 16 * thread_n_blocks / 8; + constexpr int s_tb_groups = + !has_act_order && group_blocks != -1 && group_blocks < thread_k_blocks + ? thread_k_blocks / group_blocks + : 1; + constexpr int s_sh_stage = s_tb_groups * s_sh_stride; + int s_gl_rd_delta = s_gl_stride; + + // Scale size/strides with act_order + constexpr int tb_k = 16 * thread_k_blocks; + constexpr int g_idx_stage = has_act_order ? (tb_k * sizeof(int)) / 16 : 0; + // constexpr int act_s_row_stride = 1; + // int act_s_col_stride = act_s_row_stride * num_groups; + int act_s_col_stride = 1; + int act_s_col_warp_stride = act_s_col_stride * 8; + int tb_n_warps = thread_n_blocks / 4; + int act_s_col_tb_stride = act_s_col_warp_stride * tb_n_warps; + + // Global A read index of current thread. + int a_gl_rd = a_gl_stride * (threadIdx.x / a_gl_rd_delta_o) + + (threadIdx.x % a_gl_rd_delta_o); + a_gl_rd += a_gl_rd_delta_o * slice_row; + // Shared write index of current thread. + int a_sh_wr = a_sh_stride * (threadIdx.x / a_gl_rd_delta_o) + + (threadIdx.x % a_gl_rd_delta_o); + // Shared read index. + int a_sh_rd = + a_sh_stride * ((threadIdx.x % 32) % 16) + (threadIdx.x % 32) / 16; + a_sh_rd += 2 * ((threadIdx.x / 32) / (thread_n_blocks / 4)); + + int b_gl_rd = b_gl_stride * (threadIdx.x / b_sh_stride_threads) + + (threadIdx.x % b_sh_stride_threads) * b_thread_vecs; + b_gl_rd += b_sh_stride * slice_col; + b_gl_rd += b_gl_rd_delta_o * slice_row; + int b_sh_wr = threadIdx.x * b_thread_vecs; + int b_sh_rd = threadIdx.x * b_thread_vecs; + + // For act_order + constexpr int k_iter_size = tb_k / b_sh_wr_iters; + int slice_k_start = tb_k * slice_row; + int slice_k_finish = slice_k_start + tb_k * slice_iters; + int slice_k_start_shared_fetch = slice_k_start; + int slice_n_offset = act_s_col_tb_stride * slice_col; + + // No act_order + int s_gl_rd; + if constexpr (!has_act_order) { + if constexpr (group_blocks == -1) { + s_gl_rd = s_sh_stride * slice_col + threadIdx.x; + } else { + s_gl_rd = s_gl_stride * ((thread_k_blocks * slice_row) / group_blocks) + + s_sh_stride * slice_col + threadIdx.x; + } + } + int s_sh_wr = threadIdx.x; + bool s_sh_wr_pred = threadIdx.x < s_sh_stride; + + // We use a different scale layout for grouped and column-wise quantization as + // we scale a `half2` tile in column-major layout in the former and in + // row-major in the latter case. + int s_sh_rd; + if constexpr (group_blocks != -1) + s_sh_rd = 8 * ((threadIdx.x / 32) % (thread_n_blocks / 4)) + + (threadIdx.x % 32) / 4; + else + s_sh_rd = 8 * ((threadIdx.x / 32) % (thread_n_blocks / 4)) + + (threadIdx.x % 32) % 4; + + // Precompute which thread should not read memory in which iterations; this is + // needed if there are more threads than required for a certain tilesize or + // when the batchsize is not a multiple of 16. + bool a_sh_wr_pred[a_sh_wr_iters]; + #pragma unroll + for (int i = 0; i < a_sh_wr_iters; i++) + a_sh_wr_pred[i] = a_sh_wr_delta * i + a_sh_wr < a_sh_stride * prob_m; + + // To ensure that writing and reading A tiles to/from shared memory, the + // latter in fragment format, is fully bank conflict free, we need to use a + // rather fancy XOR-based layout. The key here is that neither reads nor + // writes of the 16-byte `int4` blocks of 8 consecutive threads involve the + // same shared memory banks. Further, it seems (based on NSight-Compute) that + // each warp must also write a consecutive memory segment? + auto transform_a = [&](int i) { + int row = i / a_gl_rd_delta_o; + return a_gl_rd_delta_o * row + (i % a_gl_rd_delta_o) ^ row; + }; + // Since the computation of this remapping is non-trivial and, due to our main + // loop unrolls, all shared memory accesses are static, we simply precompute + // both transformed reads and writes. + int a_sh_wr_trans[a_sh_wr_iters]; + #pragma unroll + for (int i = 0; i < a_sh_wr_iters; i++) + a_sh_wr_trans[i] = transform_a(a_sh_wr_delta * i + a_sh_wr); + int a_sh_rd_trans[b_sh_wr_iters][thread_m_blocks]; + #pragma unroll + for (int i = 0; i < b_sh_wr_iters; i++) { + #pragma unroll + for (int j = 0; j < thread_m_blocks; j++) + a_sh_rd_trans[i][j] = + transform_a(a_sh_rd_delta_o * i + a_sh_rd_delta_i * j + a_sh_rd); + } + + // Since B-accesses have non-constant stride they have to be computed at + // runtime; we break dependencies between subsequent accesses with a tile by + // maintining multiple pointers (we have enough registers), a tiny + // optimization. + const int4* B_ptr[b_sh_wr_iters]; + #pragma unroll + for (int i = 0; i < b_sh_wr_iters; i++) + B_ptr[i] = B + b_gl_rd_delta_i * i + b_gl_rd; + + extern __shared__ int4 sh[]; + // Shared memory storage for global fetch pipelines. + int4* sh_a = sh; + int4* sh_b = sh_a + (stages * a_sh_stage); + int4* sh_g_idx = sh_b + (stages * b_sh_stage); + int4* sh_s = sh_g_idx + (stages * g_idx_stage); + + // Register storage for double buffer of shared memory reads. + FragA frag_a[2][thread_m_blocks]; + I4 frag_b_quant[2][b_thread_vecs]; + FragC frag_c[thread_m_blocks][4][2]; + FragS frag_s[2][4]; // No act-order + FragS act_frag_s[2][4][4]; // For act-order + + // Zero accumulators. + auto zero_accums = [&]() { + #pragma unroll + for (int i = 0; i < thread_m_blocks * 4 * 2 * 4; i++) + reinterpret_cast(frag_c)[i] = 0; + }; + + int sh_first_group_id = -1; + int sh_num_groups = -1; + constexpr int sh_max_num_groups = 32; + + auto fetch_scales_to_shared = [&](bool is_async, int first_group_id, + int last_group_id) { + sh_first_group_id = first_group_id; + sh_num_groups = last_group_id - first_group_id + 1; + + if (sh_num_groups < sh_max_num_groups) { + sh_num_groups = sh_max_num_groups; + } + + if (sh_first_group_id + sh_num_groups > num_groups) { + sh_num_groups = num_groups - sh_first_group_id; + } + + int row_offset = first_group_id * s_gl_stride; + + if (is_async) { + for (int i = 0; i < sh_num_groups; i++) { + if (threadIdx.x < s_sh_stride) { + cp_async4_pred(&sh_s[(i * s_sh_stride) + threadIdx.x], + &scales_ptr[row_offset + (i * s_gl_stride) + + slice_n_offset + threadIdx.x]); + } + } + } else { + for (int i = 0; i < sh_num_groups; i++) { + if (threadIdx.x < s_sh_stride) { + sh_s[(i * s_sh_stride) + threadIdx.x] = + scales_ptr[row_offset + (i * s_gl_stride) + slice_n_offset + + threadIdx.x]; + } + } + } + }; + // Asynchronously fetch the next A, B and s tile from global to the next + // shared memory pipeline location. + auto fetch_to_shared = [&](int pipe, int a_off, bool pred = true) { + if (pred) { + int4* sh_a_stage = sh_a + a_sh_stage * pipe; + #pragma unroll + for (int i = 0; i < a_sh_wr_iters; i++) { + cp_async4_pred( + &sh_a_stage[a_sh_wr_trans[i]], + &A[a_gl_rd_delta_i * i + a_gl_rd + a_gl_rd_delta_o * a_off], + a_sh_wr_pred[i]); + } + int4* sh_b_stage = sh_b + b_sh_stage * pipe; + #pragma unroll + for (int i = 0; i < b_sh_wr_iters; i++) { + #pragma unroll + for (int j = 0; j < b_thread_vecs; j++) { + cp_async4(&sh_b_stage[b_sh_wr_delta * i + b_sh_wr + j], B_ptr[i] + j); + } + + B_ptr[i] += b_gl_rd_delta_o; + } + + if constexpr (has_act_order) { + // Fetch g_idx thread-block portion + int full_pipe = a_off; + int cur_k = slice_k_start_shared_fetch + tb_k * full_pipe; + if (cur_k < prob_k && cur_k < slice_k_finish) { + int4* sh_g_idx_stage = sh_g_idx + g_idx_stage * pipe; + + int4 const* cur_g_idx_stage_ptr = + reinterpret_cast(&g_idx[cur_k]); + + if (threadIdx.x < g_idx_stage) { + cp_async4_pred(&sh_g_idx_stage[threadIdx.x], + &cur_g_idx_stage_ptr[threadIdx.x]); + } + } + } else { + if constexpr (group_blocks != -1) { + int4* sh_s_stage = sh_s + s_sh_stage * pipe; + + if constexpr (group_blocks >= thread_k_blocks) { + // Only fetch scales if this tile starts a new group + if (pipe % (group_blocks / thread_k_blocks) == 0) { + if (s_sh_wr_pred) { + cp_async4(&sh_s_stage[s_sh_wr], &scales_ptr[s_gl_rd]); + } + s_gl_rd += s_gl_rd_delta; + } + } else { + for (int i = 0; i < s_tb_groups; i++) { + if (s_sh_wr_pred) { + cp_async4(&sh_s_stage[i * s_sh_stride + s_sh_wr], + &scales_ptr[s_gl_rd]); + } + s_gl_rd += s_gl_rd_delta; + } + } + } + } + } + // Insert a fence even when we are winding down the pipeline to ensure that + // waiting is also correct at this point. + cp_async_fence(); + }; + + // Wait until the next thread tile has been loaded to shared memory. + auto wait_for_stage = [&]() { + // We only have `stages - 2` active fetches since we are double buffering + // and can only issue the next fetch when it is guaranteed that the previous + // shared memory load is fully complete (as it may otherwise be + // overwritten). + cp_async_wait(); + __syncthreads(); + }; + + // Load the next sub-tile from the current location in the shared memory pipe + // into the current register buffer. + auto fetch_to_registers = [&](int k, int pipe) { + int4* sh_a_stage = sh_a + a_sh_stage * pipe; + #pragma unroll + for (int i = 0; i < thread_m_blocks; i++) + ldsm4(frag_a[k % 2][i], + &sh_a_stage[a_sh_rd_trans[k % b_sh_wr_iters][i]]); + int4* sh_b_stage = sh_b + b_sh_stage * pipe; + + #pragma unroll + for (int i = 0; i < b_thread_vecs; i++) { + frag_b_quant[k % 2][i] = *reinterpret_cast( + &sh_b_stage[b_sh_rd_delta * (k % b_sh_wr_iters) + b_sh_rd + i]); + } + }; + + bool is_same_group[stages]; + int same_group_id[stages]; + + auto init_same_group = [&](int pipe) { + if constexpr (!has_act_order) { + is_same_group[pipe] = false; + same_group_id[pipe] = 0; + return; + } + + int4* sh_g_idx_stage = sh_g_idx + g_idx_stage * pipe; + int* sh_g_idx_int_ptr = reinterpret_cast(sh_g_idx_stage); + + int group_id_1 = sh_g_idx_int_ptr[0]; + int group_id_2 = sh_g_idx_int_ptr[tb_k - 1]; + + is_same_group[pipe] = group_id_1 == group_id_2; + same_group_id[pipe] = group_id_1; + }; + + auto fetch_scales_to_registers = [&](int k, int full_pipe) { + int pipe = full_pipe % stages; + + if constexpr (!has_act_order) { + // No act-order case + if constexpr (group_blocks != -1) { + if constexpr (group_blocks >= thread_k_blocks) { + int4* sh_s_stage = + sh_s + s_sh_stage * ((group_blocks / thread_k_blocks) * + (pipe / (group_blocks / thread_k_blocks))); + reinterpret_cast(&frag_s[k % 2])[0] = sh_s_stage[s_sh_rd]; + } else { + int warp_id = threadIdx.x / 32; + int n_warps = thread_n_blocks / 4; + + int warp_row = warp_id / n_warps; + + int cur_k = warp_row * 16; + cur_k += k_iter_size * (k % b_sh_wr_iters); + + int k_blocks = cur_k / 16; + int cur_group_id = k_blocks / group_blocks; + + int4* sh_s_stage = sh_s + s_sh_stage * pipe; + + reinterpret_cast(&frag_s[k % 2])[0] = + sh_s_stage[s_sh_rd + cur_group_id * s_sh_stride]; + } + } + + return; + } + + // Act-order case + + // Determine K of the "current" thread-block + int cur_k = slice_k_start + tb_k * full_pipe; + if (cur_k >= prob_k || cur_k >= slice_k_finish) { + return; + } + + // Reset (to current thread-block) since we read g_idx portion from the + // shared memory + cur_k = 0; + + // Progress to current iteration + cur_k += k_iter_size * (k % b_sh_wr_iters); + + // Determine "position" inside the thread-block (based on warp and + // thread-id) + int warp_id = threadIdx.x / 32; + int n_warps = + thread_n_blocks / 4; // Each warp processes 4 16-size tiles over N + + int warp_row = warp_id / n_warps; + int warp_col = warp_id % n_warps; + + cur_k += warp_row * 16; + + int th_id = threadIdx.x % 32; + cur_k += (th_id % 4) * 2; // Due to tensor-core layout for fp16 B matrix + + int s_col_shift = + /*slice_n_offset +*/ (act_s_col_warp_stride * warp_col) + + (th_id / 4) * act_s_col_stride; + + if (is_same_group[pipe]) { + if (k % 2 == 0) { + *(reinterpret_cast(&(act_frag_s[k % 2][0][0]))) = + sh_s[(same_group_id[pipe] - sh_first_group_id) * s_sh_stride + + s_col_shift]; + } else { + *(reinterpret_cast(&(act_frag_s[k % 2][0][0]))) = + *(reinterpret_cast(&(act_frag_s[(k - 1) % 2][0][0]))); + } + + for (int i = 1; i < 4; i++) { + *(reinterpret_cast(&(act_frag_s[k % 2][i][0]))) = + *(reinterpret_cast(&(act_frag_s[k % 2][0][0]))); + } + return; + } + + int4* sh_g_idx_stage = sh_g_idx + g_idx_stage * pipe; + int* sh_g_idx_int_ptr = reinterpret_cast(sh_g_idx_stage); + + constexpr int k_frag_offsets[4] = {0, 1, 8, + 9}; // Tensor core offsets per thread + + #pragma unroll + for (int i = 0; i < 4; i++) { + int actual_k = cur_k + k_frag_offsets[i]; + + int group_id = sh_g_idx_int_ptr[actual_k]; + int rel_group_id = group_id - sh_first_group_id; + + *(reinterpret_cast(&(act_frag_s[k % 2][i][0]))) = + sh_s[rel_group_id * s_sh_stride + s_col_shift]; + } + }; + + // Execute the actual tensor core matmul of a sub-tile. + auto matmul = [&](int k) { + // We have the m dimension as the inner loop in order to encourage overlapping + // dequantization and matmul operations. + #pragma unroll + for (int j = 0; j < 4; j++) { + FragB frag_b0; + FragB frag_b1; + if constexpr (num_bits == 4) { + int b_quant = frag_b_quant[k % 2][0][j]; + int b_quant_shift = b_quant >> 8; + + frag_b0 = dequant_4bit(b_quant); + frag_b1 = dequant_4bit(b_quant_shift); + + } else { + int* frag_b_quant_ptr = reinterpret_cast(frag_b_quant[k % 2]); + int b_quant_0 = frag_b_quant_ptr[j * 2 + 0]; + int b_quant_1 = frag_b_quant_ptr[j * 2 + 1]; + + frag_b0 = dequant_8bit(b_quant_0); + frag_b1 = dequant_8bit(b_quant_1); + } + + // Apply scale to frag_b0 + if constexpr (has_act_order) { + scale4(frag_b0, act_frag_s[k % 2][0][j], + act_frag_s[k % 2][1][j], act_frag_s[k % 2][2][j], + act_frag_s[k % 2][3][j], 0); + } else { + if constexpr (group_blocks != -1) { + scale(frag_b0, frag_s[k % 2][j], 0); + } + } + + // Apply scale to frag_b1 + if constexpr (has_act_order) { + scale4(frag_b1, act_frag_s[k % 2][0][j], + act_frag_s[k % 2][1][j], act_frag_s[k % 2][2][j], + act_frag_s[k % 2][3][j], 1); + + } else { + if constexpr (group_blocks != -1) { + scale(frag_b1, frag_s[k % 2][j], 1); + } + } + + #pragma unroll + for (int i = 0; i < thread_m_blocks; i++) { + mma(frag_a[k % 2][i], frag_b0, frag_c[i][j][0]); + mma(frag_a[k % 2][i], frag_b1, frag_c[i][j][1]); + } + } + }; + + // Since we slice across the k dimension of a tile in order to increase the + // number of warps while keeping the n dimension of a tile reasonable, we have + // multiple warps that accumulate their partial sums of the same output + // location; which we have to reduce over in the end. We do in shared memory. + auto thread_block_reduce = [&]() { + constexpr int red_off = threads / b_sh_stride_threads / 2; + if (red_off >= 1) { + int red_idx = threadIdx.x / b_sh_stride_threads; + constexpr int red_sh_stride = b_sh_stride_threads * 4 * 2; + constexpr int red_sh_delta = b_sh_stride_threads; + int red_sh_rd = red_sh_stride * (threadIdx.x / b_sh_stride_threads) + + (threadIdx.x % b_sh_stride_threads); + + // Parallel logarithmic shared memory reduction. We make sure to avoid any + // unnecessary read or write iterations, e.g., for two warps we write only + // once by warp 1 and read only once by warp 0. + + #pragma unroll + for (int m_block = 0; m_block < thread_m_blocks; m_block++) { + #pragma unroll + for (int i = red_off; i > 0; i /= 2) { + if (i <= red_idx && red_idx < 2 * i) { + #pragma unroll + for (int j = 0; j < 4 * 2; j++) { + int red_sh_wr = + red_sh_delta * j + (red_sh_rd - red_sh_stride * i); + if (i < red_off) { + float* c_rd = + reinterpret_cast(&sh[red_sh_delta * j + red_sh_rd]); + float* c_wr = reinterpret_cast(&sh[red_sh_wr]); + #pragma unroll + for (int k = 0; k < 4; k++) + reinterpret_cast(frag_c)[4 * 2 * m_block + j][k] += + c_rd[k] + c_wr[k]; + } + sh[red_sh_wr] = + reinterpret_cast(&frag_c)[4 * 2 * m_block + j]; + } + } + __syncthreads(); + } + if (red_idx == 0) { + #pragma unroll + for (int i = 0; i < 4 * 2; i++) { + float* c_rd = + reinterpret_cast(&sh[red_sh_delta * i + red_sh_rd]); + #pragma unroll + for (int j = 0; j < 4; j++) + reinterpret_cast(frag_c)[4 * 2 * m_block + i][j] += + c_rd[j]; + } + } + __syncthreads(); + } + } + }; + + // Since multiple threadblocks may process parts of the same column slice, we + // finally have to globally reduce over the results. As the striped + // partitioning minimizes the number of such reductions and our outputs are + // usually rather small, we perform this reduction serially in L2 cache. + auto global_reduce = [&](bool first = false, bool last = false) { + // We are very careful here to reduce directly in the output buffer to + // maximize L2 cache utilization in this step. To do this, we write out + // results in FP16 (but still reduce with FP32 compute). + constexpr int active_threads = 32 * thread_n_blocks / 4; + if (threadIdx.x < active_threads) { + int c_gl_stride = prob_n / 8; + int c_gl_wr_delta_o = 8 * c_gl_stride; + int c_gl_wr_delta_i = 4 * (active_threads / 32); + int c_gl_wr = c_gl_stride * ((threadIdx.x % 32) / 4) + + 4 * (threadIdx.x / 32) + threadIdx.x % 4; + c_gl_wr += (2 * thread_n_blocks) * slice_col; + constexpr int c_sh_wr_delta = active_threads; + int c_sh_wr = threadIdx.x; + + int row = (threadIdx.x % 32) / 4; + + if (!first) { + // Interestingly, doing direct global accesses here really seems to mess up + // the compiler and lead to slowdowns, hence we also use async-copies even + // though these fetches are not actually asynchronous. + #pragma unroll + for (int i = 0; i < thread_m_blocks * 4; i++) { + cp_async4_pred( + &sh[c_sh_wr + c_sh_wr_delta * i], + &C[c_gl_wr + c_gl_wr_delta_o * (i / 2) + + c_gl_wr_delta_i * (i % 2)], + i < (thread_m_blocks - 1) * 4 || 8 * (i / 2) + row < prob_m); + } + cp_async_fence(); + cp_async_wait<0>(); + } + + #pragma unroll + for (int i = 0; i < thread_m_blocks * 4; i++) { + if (i < (thread_m_blocks - 1) * 4 || 8 * (i / 2) + row < prob_m) { + if (!first) { + int4 c_red = sh[c_sh_wr + i * c_sh_wr_delta]; + #pragma unroll + for (int j = 0; j < 2 * 4; j++) { + reinterpret_cast( + &frag_c)[4 * 2 * 4 * (i / 4) + 4 * j + (i % 4)] += + Dtype::num2float(reinterpret_cast(&c_red)[j]); + } + } + if (!last) { + int4 c; + #pragma unroll + for (int j = 0; j < 2 * 4; j++) { + reinterpret_cast(&c)[j] = + Dtype::float2num(reinterpret_cast( + &frag_c)[4 * 2 * 4 * (i / 4) + 4 * j + (i % 4)]); + } + C[c_gl_wr + c_gl_wr_delta_o * (i / 2) + c_gl_wr_delta_i * (i % 2)] = + c; + } + } + } + } + }; + + // Write out the reduce final result in the correct layout. We only actually + // reshuffle matrix fragments in this step, the reduction above is performed + // in fragment layout. + auto write_result = [&]() { + int c_gl_stride = prob_n / 8; + constexpr int c_sh_stride = 2 * thread_n_blocks + 1; + int c_gl_wr_delta = c_gl_stride * (threads / (2 * thread_n_blocks)); + constexpr int c_sh_rd_delta = + c_sh_stride * (threads / (2 * thread_n_blocks)); + + int c_gl_wr = c_gl_stride * (threadIdx.x / (2 * thread_n_blocks)) + + (threadIdx.x % (2 * thread_n_blocks)); + c_gl_wr += (2 * thread_n_blocks) * slice_col; + int c_sh_wr = + (4 * c_sh_stride) * ((threadIdx.x % 32) / 4) + (threadIdx.x % 32) % 4; + c_sh_wr += 32 * (threadIdx.x / 32); + int c_sh_rd = c_sh_stride * (threadIdx.x / (2 * thread_n_blocks)) + + (threadIdx.x % (2 * thread_n_blocks)); + + int c_gl_wr_end = c_gl_stride * prob_m; + + // We first reorder in shared memory to guarantee the most efficient final + // global write patterns + auto write = [&](int idx, float c0, float c1, FragS& s) { + scalar_t2 res = + Dtype::nums2num2(Dtype::float2num(c0), Dtype::float2num(c1)); + + // For per-column quantization we finally apply the scale here (only for + // 4-bit) + if constexpr (!has_act_order && group_blocks == -1 && num_bits == 4) { + res = __hmul2(res, s[0]); + } + + ((scalar_t2*)sh)[idx] = res; + }; + + if (threadIdx.x / 32 < thread_n_blocks / 4) { + #pragma unroll + for (int i = 0; i < thread_m_blocks; i++) { + #pragma unroll + for (int j = 0; j < 4; j++) { + int wr = c_sh_wr + 8 * j; + write(wr + (4 * c_sh_stride) * 0 + 0, frag_c[i][j][0][0], + frag_c[i][j][0][1], frag_s[j / 2][2 * (j % 2) + 0]); + write(wr + (4 * c_sh_stride) * 8 + 0, frag_c[i][j][0][2], + frag_c[i][j][0][3], frag_s[j / 2][2 * (j % 2) + 0]); + write(wr + (4 * c_sh_stride) * 0 + 4, frag_c[i][j][1][0], + frag_c[i][j][1][1], frag_s[j / 2][2 * (j % 2) + 1]); + write(wr + (4 * c_sh_stride) * 8 + 4, frag_c[i][j][1][2], + frag_c[i][j][1][3], frag_s[j / 2][2 * (j % 2) + 1]); + } + c_sh_wr += 16 * (4 * c_sh_stride); + } + } + __syncthreads(); + + #pragma unroll + for (int i = 0; + i < div_ceil(16 * thread_m_blocks, threads / (2 * thread_n_blocks)); + i++) { + if (c_gl_wr < c_gl_wr_end) { + C[c_gl_wr] = sh[c_sh_rd]; + c_gl_wr += c_gl_wr_delta; + c_sh_rd += c_sh_rd_delta; + } + } + }; + + // Start global fetch and register load pipelines. + auto start_pipes = [&]() { + + #pragma unroll + for (int i = 0; i < stages - 1; i++) { + if (has_act_order && i == 0) { + int last_g_idx = slice_k_start + stages * tb_k * 2; + if (last_g_idx >= prob_k) { + last_g_idx = prob_k - 1; + } + fetch_scales_to_shared(true, g_idx[slice_k_start], g_idx[last_g_idx]); + } + fetch_to_shared(i, i, i < slice_iters); + } + + zero_accums(); + wait_for_stage(); + init_same_group(0); + fetch_to_registers(0, 0); + fetch_scales_to_registers(0, 0); + a_gl_rd += a_gl_rd_delta_o * (stages - 1); + slice_k_start_shared_fetch += tb_k * (stages - 1); + }; + if (slice_iters) { + start_pipes(); + } + + // Main loop. + while (slice_iters) { + // We unroll over both the global fetch and the register load pipeline to + // ensure all shared memory accesses are static. Note that both pipelines + // have even length meaning that the next iteration will always start at + // index 0. + + #pragma unroll + for (int pipe = 0; pipe < stages;) { + #pragma unroll + for (int k = 0; k < b_sh_wr_iters; k++) { + fetch_to_registers(k + 1, pipe % stages); + fetch_scales_to_registers(k + 1, pipe); + if (k == b_sh_wr_iters - 2) { + fetch_to_shared((pipe + stages - 1) % stages, pipe, + slice_iters >= stages); + pipe++; + wait_for_stage(); + init_same_group(pipe % stages); + } + matmul(k); + } + slice_iters--; + if (slice_iters == 0) { + break; + } + } + + a_gl_rd += a_gl_rd_delta_o * stages; + slice_k_start += tb_k * stages; + slice_k_start_shared_fetch += tb_k * stages; + + if constexpr (has_act_order) { + int first_group_id = g_idx[slice_k_start]; + int last_g_idx = slice_k_start + stages * tb_k * 2; + if (last_g_idx >= prob_k) { + last_g_idx = prob_k - 1; + } + int last_group_id = g_idx[last_g_idx]; + if (last_group_id >= sh_first_group_id + sh_num_groups) { + fetch_scales_to_shared(false, first_group_id, last_group_id); + __syncthreads(); + } + } + + // Process results and, if necessary, proceed to the next column slice. + // While this pattern may not be the most readable, other ways of writing + // the loop seemed to noticeably worse performance after compilation. + if (slice_iters == 0) { + cp_async_wait<0>(); + bool last = slice_idx == slice_count - 1; + // For per-column scales, we only fetch them here in the final step before + // write-out + if constexpr (!has_act_order && group_blocks == -1) { + if constexpr (num_bits == 8) { + if (s_sh_wr_pred) { + cp_async4(&sh_s[s_sh_wr], &scales_ptr[s_gl_rd]); + } + cp_async_fence(); + } else { + if (last) { + if (s_sh_wr_pred) { + cp_async4(&sh_s[s_sh_wr], &scales_ptr[s_gl_rd]); + } + cp_async_fence(); + } + } + } + + thread_block_reduce(); + if constexpr (!has_act_order && group_blocks == -1) { + if constexpr (num_bits == 8) { + cp_async_wait<0>(); + __syncthreads(); + if (threadIdx.x / 32 < thread_n_blocks / 4) { + reinterpret_cast(&frag_s)[0] = sh_s[s_sh_rd + 0]; + reinterpret_cast(&frag_s)[1] = sh_s[s_sh_rd + 4]; + } + + } else { + if (last) { + cp_async_wait<0>(); + __syncthreads(); + if (threadIdx.x / 32 < thread_n_blocks / 4) { + reinterpret_cast(&frag_s)[0] = sh_s[s_sh_rd + 0]; + reinterpret_cast(&frag_s)[1] = sh_s[s_sh_rd + 4]; + } + } + } + } + + // For 8-bit channelwise, we apply the scale before the global reduction + // that converts the fp32 results to fp16 (so that we avoid possible + // overflow in fp16) + if constexpr (!has_act_order && group_blocks == -1 && num_bits == 8) { + if (threadIdx.x / 32 < thread_n_blocks / 4) { + #pragma unroll + for (int i = 0; i < thread_m_blocks; i++) { + #pragma unroll + for (int j = 0; j < 4; j++) { + scale_float( + reinterpret_cast(&frag_c[i][j][0][0]), + frag_s[j / 2][2 * (j % 2) + 0]); + scale_float( + reinterpret_cast(&frag_c[i][j][0][2]), + frag_s[j / 2][2 * (j % 2) + 0]); + + scale_float( + reinterpret_cast(&frag_c[i][j][1][0]), + frag_s[j / 2][2 * (j % 2) + 1]); + scale_float( + reinterpret_cast(&frag_c[i][j][1][2]), + frag_s[j / 2][2 * (j % 2) + 1]); + } + } + } + } + + if (slice_count > 1) { // only globally reduce if there is more than one + // block in a slice + barrier_acquire(&locks[slice_col], slice_idx); + global_reduce(slice_idx == 0, last); + barrier_release(&locks[slice_col], last); + } + if (last) // only the last block in a slice actually writes the result + write_result(); + slice_row = 0; + slice_col_par++; + slice_col++; + init_slice(); + if (slice_iters) { + a_gl_rd = a_gl_stride * (threadIdx.x / a_gl_rd_delta_o) + + (threadIdx.x % a_gl_rd_delta_o); + #pragma unroll + for (int i = 0; i < b_sh_wr_iters; i++) + B_ptr[i] += b_sh_stride - b_gl_rd_delta_o * k_tiles; + if (slice_col == 0) { + #pragma unroll + for (int i = 0; i < b_sh_wr_iters; i++) B_ptr[i] -= b_gl_stride; + } + + // Update slice k/n for scales loading + if constexpr (has_act_order) { + slice_k_start = tb_k * slice_row; + slice_k_finish = slice_k_start + tb_k * slice_iters; + slice_k_start_shared_fetch = slice_k_start; + slice_n_offset = act_s_col_tb_stride * slice_col; + + } else { + s_gl_rd = s_sh_stride * slice_col + threadIdx.x; + } + + start_pipes(); + } + } + } +} + + #define __CALL_IF(NUM_BITS, THREAD_M_BLOCKS, THREAD_N_BLOCKS, \ + THREAD_K_BLOCKS, HAS_ACT_ORDER, GROUP_BLOCKS, NUM_THREADS) \ + else if (num_bits == NUM_BITS && thread_m_blocks == THREAD_M_BLOCKS && \ + thread_n_blocks == THREAD_N_BLOCKS && \ + thread_k_blocks == THREAD_K_BLOCKS && \ + has_act_order == HAS_ACT_ORDER && group_blocks == GROUP_BLOCKS && \ + num_threads == NUM_THREADS) { \ + cudaFuncSetAttribute( \ + Marlin, \ + cudaFuncAttributeMaxDynamicSharedMemorySize, max_shared_mem); \ + Marlin<<>>( \ + A_ptr, B_ptr, C_ptr, s_ptr, g_idx_ptr, num_groups, prob_m, prob_n, \ + prob_k, locks); \ + } + +typedef struct { + int thread_k; + int thread_n; + int num_threads; +} thread_config_t; + +typedef struct { + int max_m_blocks; + thread_config_t tb_cfg; +} exec_config_t; + +thread_config_t small_batch_thread_configs[] = { + // Ordered by priority + + // thread_k, thread_n, num_threads + {128, 128, 256}, + {64, 128, 128}, + {128, 64, 128}, +}; + +thread_config_t large_batch_thread_configs[] = { + // Ordered by priority + + // thread_k, thread_n, num_threads + {64, 256, 256}, + {64, 128, 128}, + {128, 64, 128}, + +}; + +int get_scales_cache_size(thread_config_t const& th_config, int prob_m, + int prob_n, int prob_k, int num_bits, int group_size, + bool has_act_order, bool is_k_full) { + bool cache_scales_chunk = has_act_order && !is_k_full; + + int tb_n = th_config.thread_n; + int tb_k = th_config.thread_k; + + // Get max scale groups per thread-block + int tb_groups; + if (group_size == -1) { + tb_groups = 1; + } else if (group_size == 0) { + tb_groups = div_ceil(tb_k, 32); // Worst case is 32 group size + } else { + tb_groups = div_ceil(tb_k, group_size); + } + + if (cache_scales_chunk) { + int load_groups = + tb_groups * pipe_stages * 2; // Chunk size is 2x pipeline over dim K + load_groups = max(load_groups, 32); // We load at least 32 scale groups + return load_groups * tb_n * 2; + + } else { + int tb_scales = tb_groups * tb_n * 2; + + return tb_scales * pipe_stages; + } +} + +bool is_valid_cache_size(thread_config_t const& th_config, int max_m_blocks, + int prob_m, int prob_n, int prob_k, int num_bits, + int scales_cache_size, int max_shared_mem) { + int pack_factor = 32 / num_bits; + + // Get B size + int tb_k = th_config.thread_k; + int tb_n = th_config.thread_n; + + int b_size = (tb_k * tb_n / pack_factor) * 4; + + // Get A size + int m_blocks = div_ceil(prob_m, 16); + int tb_max_m = 16; + + while (true) { + if (m_blocks >= max_m_blocks) { + tb_max_m *= max_m_blocks; + break; + } + + max_m_blocks--; + if (max_m_blocks == 0) { + TORCH_CHECK(false, "Unexpected m_blocks = ", m_blocks); + } + } + + int a_size = (tb_max_m * tb_k) * 2; + + float pipe_size = (a_size + b_size) * pipe_stages; + + TORCH_CHECK(max_shared_mem / 2 > scales_cache_size); // Sanity + + return pipe_size < 0.95f * (max_shared_mem - scales_cache_size); +} + +bool is_valid_config(thread_config_t const& th_config, int max_m_blocks, + int prob_m, int prob_n, int prob_k, int num_bits, + int group_size, bool has_act_order, bool is_k_full, + int max_shared_mem) { + // Sanity + if (th_config.thread_k == -1 || th_config.thread_n == -1 || + th_config.num_threads == -1) { + return false; + } + + // Verify K/N are divisible by thread K/N + if (prob_k % th_config.thread_k != 0 || prob_n % th_config.thread_n != 0) { + return false; + } + + // Verify min for thread K/N + if (th_config.thread_n < min_thread_n || th_config.thread_k < min_thread_k) { + return false; + } + + // num_threads must be at least 128 (= 4 warps) + if (th_config.num_threads < 128) { + return false; + } + + // Determine cache for scales + int scales_cache_size = + get_scales_cache_size(th_config, prob_m, prob_n, prob_k, num_bits, + group_size, has_act_order, is_k_full); + + // Check that pipeline fits into cache + if (!is_valid_cache_size(th_config, max_m_blocks, prob_m, prob_n, prob_k, + num_bits, scales_cache_size, max_shared_mem)) { + return false; + } + + return true; +} + +exec_config_t determine_thread_config(int prob_m, int prob_n, int prob_k, + int num_bits, int group_size, + bool has_act_order, bool is_k_full, + int max_shared_mem) { + int max_m_blocks = 4; + while (max_m_blocks > 0) { + if (prob_m <= 16) { + for (auto th_config : small_batch_thread_configs) { + if (is_valid_config(th_config, max_m_blocks, prob_m, prob_n, prob_k, + num_bits, group_size, has_act_order, is_k_full, + max_shared_mem)) { + return exec_config_t{max_m_blocks, th_config}; + } + } + } else { + for (auto th_config : large_batch_thread_configs) { + if (is_valid_config(th_config, max_m_blocks, prob_m, prob_n, prob_k, + num_bits, group_size, has_act_order, is_k_full, + max_shared_mem)) { + return exec_config_t{max_m_blocks, th_config}; + } + } + } + + max_m_blocks--; // Process less M blocks per invocation to reduce cache + // usage + } + + return exec_config_t{0, {-1, -1, -1}}; +} + + #define CALL_IF(NUM_BITS, N_BLOCKS, K_BLOCKS, NUM_THREADS) \ + __CALL_IF(NUM_BITS, 1, N_BLOCKS, K_BLOCKS, true, 0, NUM_THREADS) \ + __CALL_IF(NUM_BITS, 2, N_BLOCKS, K_BLOCKS, true, 0, NUM_THREADS) \ + __CALL_IF(NUM_BITS, 3, N_BLOCKS, K_BLOCKS, true, 0, NUM_THREADS) \ + __CALL_IF(NUM_BITS, 4, N_BLOCKS, K_BLOCKS, true, 0, NUM_THREADS) \ + \ + __CALL_IF(NUM_BITS, 1, N_BLOCKS, K_BLOCKS, false, -1, NUM_THREADS) \ + __CALL_IF(NUM_BITS, 1, N_BLOCKS, K_BLOCKS, false, 2, NUM_THREADS) \ + __CALL_IF(NUM_BITS, 1, N_BLOCKS, K_BLOCKS, false, 4, NUM_THREADS) \ + __CALL_IF(NUM_BITS, 1, N_BLOCKS, K_BLOCKS, false, 8, NUM_THREADS) \ + \ + __CALL_IF(NUM_BITS, 2, N_BLOCKS, K_BLOCKS, false, -1, NUM_THREADS) \ + __CALL_IF(NUM_BITS, 2, N_BLOCKS, K_BLOCKS, false, 2, NUM_THREADS) \ + __CALL_IF(NUM_BITS, 2, N_BLOCKS, K_BLOCKS, false, 4, NUM_THREADS) \ + __CALL_IF(NUM_BITS, 2, N_BLOCKS, K_BLOCKS, false, 8, NUM_THREADS) \ + \ + __CALL_IF(NUM_BITS, 3, N_BLOCKS, K_BLOCKS, false, -1, NUM_THREADS) \ + __CALL_IF(NUM_BITS, 3, N_BLOCKS, K_BLOCKS, false, 2, NUM_THREADS) \ + __CALL_IF(NUM_BITS, 3, N_BLOCKS, K_BLOCKS, false, 4, NUM_THREADS) \ + __CALL_IF(NUM_BITS, 3, N_BLOCKS, K_BLOCKS, false, 8, NUM_THREADS) \ + \ + __CALL_IF(NUM_BITS, 4, N_BLOCKS, K_BLOCKS, false, -1, NUM_THREADS) \ + __CALL_IF(NUM_BITS, 4, N_BLOCKS, K_BLOCKS, false, 2, NUM_THREADS) \ + __CALL_IF(NUM_BITS, 4, N_BLOCKS, K_BLOCKS, false, 4, NUM_THREADS) \ + __CALL_IF(NUM_BITS, 4, N_BLOCKS, K_BLOCKS, false, 8, NUM_THREADS) + +template +void marlin_mm_f16i4(const void* A, const void* B, void* C, void* s, + void* g_idx, void* perm, void* a_tmp, int prob_m, + int prob_n, int prob_k, void* workspace, int num_bits, + bool has_act_order, bool is_k_full, int num_groups, + int group_size, int dev, cudaStream_t stream, int thread_k, + int thread_n, int sms, int max_par) { + TORCH_CHECK(num_bits == 4 || num_bits == 8, + "num_bits must be 4 or 8. Got = ", num_bits); + TORCH_CHECK(prob_m > 0 && prob_n > 0 && prob_k > 0, "Invalid MNK = [", prob_m, + ", ", prob_n, ", ", prob_k, "]"); + + int tot_m = prob_m; + int tot_m_blocks = div_ceil(tot_m, 16); + int pad = 16 * tot_m_blocks - tot_m; + + if (sms == -1) { + cudaDeviceGetAttribute(&sms, cudaDevAttrMultiProcessorCount, dev); + } + + int max_shared_mem = 0; + cudaDeviceGetAttribute(&max_shared_mem, + cudaDevAttrMaxSharedMemoryPerBlockOptin, dev); + TORCH_CHECK(max_shared_mem > 0); + + // Set thread config + exec_config_t exec_cfg; + if (thread_k != -1 && thread_n != -1) { + // User-defined config + exec_cfg = + exec_config_t{4, thread_config_t{thread_k, thread_n, default_threads}}; + } else { + // Auto config + exec_cfg = + determine_thread_config(prob_m, prob_n, prob_k, num_bits, group_size, + has_act_order, is_k_full, max_shared_mem); + } + + TORCH_CHECK(exec_cfg.max_m_blocks > 0 && + is_valid_config(exec_cfg.tb_cfg, exec_cfg.max_m_blocks, + prob_m, prob_n, prob_k, num_bits, group_size, + has_act_order, is_k_full, max_shared_mem), + "Invalid thread config: max_m_blocks = ", exec_cfg.max_m_blocks, + ", thread_k = ", exec_cfg.tb_cfg.thread_k, + ", thread_n = ", exec_cfg.tb_cfg.thread_n, + ", num_threads = ", exec_cfg.tb_cfg.num_threads, " for MKN = [", + prob_m, ", ", prob_k, ", ", prob_n, "] and num_bits = ", num_bits, + ", group_size = ", group_size, + ", has_act_order = ", has_act_order, ", is_k_full = ", is_k_full, + ", max_shared_mem = ", max_shared_mem); + + int num_threads = exec_cfg.tb_cfg.num_threads; + thread_k = exec_cfg.tb_cfg.thread_k; + thread_n = exec_cfg.tb_cfg.thread_n; + + int thread_k_blocks = thread_k / 16; + int thread_n_blocks = thread_n / 16; + + int blocks = sms; + + TORCH_CHECK(prob_n % thread_n == 0, "prob_n = ", prob_n, + " is not divisible by thread_n = ", thread_n); + TORCH_CHECK(prob_k % thread_k == 0, "prob_k = ", prob_k, + " is not divisible by thread_k = ", thread_k); + + int group_blocks = 0; + if (has_act_order) { + if (is_k_full) { + TORCH_CHECK(group_size != -1); + group_blocks = group_size / 16; + TORCH_CHECK(prob_k % group_blocks == 0, "prob_k = ", prob_k, + " is not divisible by group_blocks = ", group_blocks); + } else { + TORCH_CHECK(group_size == 0); + group_blocks = 0; + } + + } else { + if (group_size == -1) { + group_blocks = -1; + } else { + group_blocks = group_size / 16; + TORCH_CHECK(prob_k % group_blocks == 0, "prob_k = ", prob_k, + " is not divisible by group_blocks = ", group_blocks); + } + } + + const int4* A_ptr = (const int4*)A; + const int4* B_ptr = (const int4*)B; + int4* C_ptr = (int4*)C; + const int4* s_ptr = (const int4*)s; + const int* g_idx_ptr = (const int*)g_idx; + const int* perm_ptr = (const int*)perm; + int4* a_tmp_ptr = (int4*)a_tmp; + + int* locks = (int*)workspace; + + if (has_act_order) { + // Permute A columns + int block_rows = div_ceil(prob_m, blocks); + permute_cols_kernel<<>>( + A_ptr, perm_ptr, a_tmp_ptr, prob_m, prob_k, block_rows); + A_ptr = a_tmp_ptr; + } + + // If we have a full K, then we can run the non-act-order version of Marlin + // (since the weight rows are reordered by increasing group ids, and by having + // a full K, we have full original groups) + if (is_k_full) { + has_act_order = false; + } + + // Main loop + for (int i = 0; i < tot_m_blocks; i += exec_cfg.max_m_blocks) { + int thread_m_blocks = tot_m_blocks - i; + prob_m = tot_m - 16 * i; + int par = 1; + if (thread_m_blocks > exec_cfg.max_m_blocks) { + // Note that parallel > 1 currently only works for inputs without any + // padding + par = (16 * thread_m_blocks - pad) / (16 * exec_cfg.max_m_blocks); + if (par > max_par) par = max_par; + prob_m = (16 * exec_cfg.max_m_blocks) * par; + i += exec_cfg.max_m_blocks * (par - 1); + thread_m_blocks = exec_cfg.max_m_blocks; + } + + + + // Define kernel configurations +#define undefined_error TORCH_CHECK(false, "Unsupported shapes: MNK = [" + str(prob_m) + ", " + \ + str(prob_n) + ", " + str(prob_k) + "]" + \ + ", has_act_order = " + str(has_act_order) + \ + ", num_groups = " + str(num_groups) + \ + ", group_size = " + str(group_size) + \ + ", thread_m_blocks = " + str(thread_m_blocks) + \ + ", thread_n_blocks = " + str(thread_n_blocks) + \ + ", thread_k_blocks = " + str(thread_k_blocks)); + + + if (num_bits == 4 && num_threads == 256) + { + if (false) { + } + CALL_IF(4, 32, 2, 256) + CALL_IF(4, 16, 4, 256) + CALL_IF(4, 8, 8, 256) + else { + undefined_error + } + } + else if (num_bits == 4 && num_threads == 128) + { + if (false) { + } + CALL_IF(4, 8, 4, 128) + CALL_IF(4, 4, 8, 128) + else { + undefined_error + } + } + else if (num_bits == 8 && num_threads == 256) + { + if (false) { + } + CALL_IF(8, 32, 2, 256) + CALL_IF(8, 16, 4, 256) + CALL_IF(8, 8, 8, 256) + else { + undefined_error + } + } + else if (num_bits == 8 && num_threads == 128) + { + if (false) { + } + CALL_IF(8, 8, 4, 128) + CALL_IF(8, 4, 8, 128) + else { + undefined_error + } + } + else { + undefined_error + } + + A_ptr += 16 * thread_m_blocks * (prob_k / 8) * par; + C_ptr += 16 * thread_m_blocks * (prob_n / 8) * par; + } +} + +} // namespace gptq_marlin + +torch::Tensor gptq_marlin_gemm(torch::Tensor& a, torch::Tensor& b_q_weight, + torch::Tensor& b_scales, torch::Tensor& g_idx, + torch::Tensor& perm, torch::Tensor& workspace, + int64_t num_bits, int64_t size_m, int64_t size_n, + int64_t size_k, bool is_k_full) { + const at::cuda::OptionalCUDAGuard device_guard(device_of(a)); + // Verify num_bits + TORCH_CHECK(num_bits == 4 || num_bits == 8, + "num_bits must be 4 or 8. Got = ", num_bits); + int pack_factor = 32 / num_bits; + + // Verify A + TORCH_CHECK(a.size(0) == size_m, "Shape mismatch: a.size(0) = ", a.size(0), + ", size_m = ", size_m); + TORCH_CHECK(a.size(1) == size_k, "Shape mismatch: a.size(1) = ", a.size(1), + ", size_k = ", size_k); + + // Verify B + TORCH_CHECK(size_k % gptq_marlin::tile_size == 0, "size_k = ", size_k, + " is not divisible by tile_size = ", gptq_marlin::tile_size); + TORCH_CHECK((size_k / gptq_marlin::tile_size) == b_q_weight.size(0), + "Shape mismatch: b_q_weight.size(0) = ", b_q_weight.size(0), + ", size_k = ", size_k, ", tile_size = ", gptq_marlin::tile_size); + TORCH_CHECK(b_q_weight.size(1) % gptq_marlin::tile_size == 0, + "b_q_weight.size(1) = ", b_q_weight.size(1), + " is not divisible by tile_size = ", gptq_marlin::tile_size); + int actual_size_n = + (b_q_weight.size(1) / gptq_marlin::tile_size) * pack_factor; + TORCH_CHECK(size_n == actual_size_n, "size_n = ", size_n, + ", actual_size_n = ", actual_size_n); + + // Verify device and strides + TORCH_CHECK(a.device().is_cuda(), "A is not on GPU"); + TORCH_CHECK(a.is_contiguous(), "A is not contiguous"); + + TORCH_CHECK(b_q_weight.device().is_cuda(), "b_q_weight is not on GPU"); + TORCH_CHECK(b_q_weight.is_contiguous(), "b_q_weight is not contiguous"); + + TORCH_CHECK(b_scales.device().is_cuda(), "b_scales is not on GPU"); + TORCH_CHECK(b_scales.is_contiguous(), "b_scales is not contiguous"); + + TORCH_CHECK(g_idx.device().is_cuda(), "g_idx is not on GPU"); + TORCH_CHECK(g_idx.is_contiguous(), "g_idx is not contiguous"); + + TORCH_CHECK(perm.device().is_cuda(), "perm is not on GPU"); + TORCH_CHECK(perm.is_contiguous(), "perm is not contiguous"); + + // Alloc buffers + auto options = torch::TensorOptions().dtype(a.dtype()).device(a.device()); + torch::Tensor c = torch::empty({size_m, size_n}, options); + torch::Tensor a_tmp = torch::empty({size_m, size_k}, options); + + // thread_k: `k` size of a thread_tile in `weights` (can usually be left as + // auto -1) + int thread_k = -1; + // thread_n: `n` size of a thread_tile in `weights` (can usually be left as + // auto -1) + int thread_n = -1; + // sms: number of SMs to use for the kernel (can usually be left as auto -1) + int sms = -1; + + // Verify g_idx and perm + TORCH_CHECK((g_idx.size(0) == 0 && perm.size(0) == 0) || + (g_idx.size(0) == size_k && perm.size(0) == size_k), + "Unexpected g_idx.size(0) = ", g_idx.size(0), + " and perm.size(0) = ", perm.size(0), + ", where size_k = ", size_k); + + // Detect groupsize and act_order + int num_groups = -1; + int group_size = -1; + bool has_act_order = g_idx.size(0) != 0; + + int b_rank = b_scales.sizes().size(); + TORCH_CHECK(b_rank == 2, "b_scales rank = ", b_rank, " is not 2"); + TORCH_CHECK(b_scales.size(1) == size_n, "b_scales dim 1 = ", b_scales.size(1), + " is not size_n = ", size_n); + num_groups = b_scales.size(0); + + if (has_act_order) { + if (is_k_full) { + TORCH_CHECK(num_groups > 1, "For act_order, num_groups must be > 1"); + TORCH_CHECK(size_k % num_groups == 0, "size_k = ", size_k, + ", is not divisible by num_groups = ", num_groups); + group_size = size_k / num_groups; + } else { + group_size = 0; + } + + } else { + if (num_groups > 1) { + TORCH_CHECK( + size_k % num_groups == 0, "size_k = ", size_k, + ", is not divisible by b_scales.size(0) = ", b_scales.size(0)); + group_size = size_k / num_groups; + } else { + group_size = -1; + } + } + + // Verify workspace size + TORCH_CHECK( + size_n % gptq_marlin::min_thread_n == 0, "size_n = ", size_n, + ", is not divisible by min_thread_n = ", gptq_marlin::min_thread_n); + int min_workspace_size = + (size_n / gptq_marlin::min_thread_n) * gptq_marlin::max_par; + TORCH_CHECK(workspace.numel() >= min_workspace_size, + "workspace.numel = ", workspace.numel(), + " is below min_workspace_size = ", min_workspace_size); + + int dev = a.get_device(); + if (a.scalar_type() == at::ScalarType::Half) { + gptq_marlin::marlin_mm_f16i4( + a.data_ptr(), b_q_weight.data_ptr(), c.data_ptr(), + b_scales.data_ptr(), g_idx.data_ptr(), perm.data_ptr(), + a_tmp.data_ptr(), size_m, size_n, size_k, + workspace.data_ptr(), num_bits, has_act_order, is_k_full, num_groups, + group_size, dev, at::cuda::getCurrentCUDAStream(dev), thread_k, + thread_n, sms, gptq_marlin::max_par); + } else if (a.scalar_type() == at::ScalarType::BFloat16) { + gptq_marlin::marlin_mm_f16i4( + a.data_ptr(), b_q_weight.data_ptr(), + c.data_ptr(), b_scales.data_ptr(), + g_idx.data_ptr(), perm.data_ptr(), a_tmp.data_ptr(), + size_m, size_n, size_k, workspace.data_ptr(), num_bits, has_act_order, + is_k_full, num_groups, group_size, dev, + at::cuda::getCurrentCUDAStream(dev), thread_k, thread_n, sms, + gptq_marlin::max_par); + } else { + TORCH_CHECK(false, "gpt_marlin_gemm only supports bfloat16 and float16"); + } + + return c; +} + +#endif diff --git a/archive/csrc/ktransformers_ext/cuda/gptq_marlin/gptq_marlin.cuh b/archive/csrc/ktransformers_ext/cuda/gptq_marlin/gptq_marlin.cuh new file mode 100644 index 0000000..ccf9cfd --- /dev/null +++ b/archive/csrc/ktransformers_ext/cuda/gptq_marlin/gptq_marlin.cuh @@ -0,0 +1,80 @@ +// Adapted from +// https://github.com/vllm-project/vllm/tree/main/csrc/quantization/gptq_marlin +// Copyrigth 2024 The vLLM team. +// Copyright (c) 2024 by KVCache.AI, All Rights Reserved. +#pragma once + +#include + +#include +#include +#include +#include +#include +#include + +namespace gptq_marlin { + +// 8 warps are a good choice since every SM has 4 schedulers and having more +// than 1 warp per schedule allows some more latency hiding. At the same time, +// we want relatively few warps to have many registers per warp and small tiles. +static constexpr int default_threads = 256; + +static constexpr int pipe_stages = + 4; // 4 pipeline stages fit into shared memory + +static constexpr int min_thread_n = 64; +static constexpr int min_thread_k = 64; + +static constexpr int tile_size = 16; +static constexpr int max_par = 16; + +template +struct Vec { + T elems[n]; + __device__ T& operator[](int i) { return elems[i]; } +}; + +using I4 = Vec; + +constexpr int div_ceil(int a, int b) { return (a + b - 1) / b; } + +#if (defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800) || defined (__HIP_PLATFORM_AMD__) +// No support for async +#else + +__device__ inline void cp_async4_pred(void* smem_ptr, const void* glob_ptr, + bool pred = true) { + const int BYTES = 16; + uint32_t smem = static_cast(__cvta_generic_to_shared(smem_ptr)); + asm volatile( + "{\n" + " .reg .pred p;\n" + " setp.ne.b32 p, %0, 0;\n" + " @p cp.async.cg.shared.global [%1], [%2], %3;\n" + "}\n" ::"r"((int)pred), + "r"(smem), "l"(glob_ptr), "n"(BYTES)); +} + +__device__ inline void cp_async4(void* smem_ptr, const void* glob_ptr) { + const int BYTES = 16; + uint32_t smem = static_cast(__cvta_generic_to_shared(smem_ptr)); + asm volatile( + "{\n" + " cp.async.cg.shared.global [%0], [%1], %2;\n" + "}\n" ::"r"(smem), + "l"(glob_ptr), "n"(BYTES)); +} + +__device__ inline void cp_async_fence() { + asm volatile("cp.async.commit_group;\n" ::); +} + +template +__device__ inline void cp_async_wait() { + asm volatile("cp.async.wait_group %0;\n" ::"n"(n)); +} + +#endif + +} // namespace gptq_marlin diff --git a/archive/csrc/ktransformers_ext/cuda/gptq_marlin/gptq_marlin_dtypes.cuh b/archive/csrc/ktransformers_ext/cuda/gptq_marlin/gptq_marlin_dtypes.cuh new file mode 100644 index 0000000..80f6ea4 --- /dev/null +++ b/archive/csrc/ktransformers_ext/cuda/gptq_marlin/gptq_marlin_dtypes.cuh @@ -0,0 +1,85 @@ +// Adapted from +// https://github.com/vllm-project/vllm/tree/main/csrc/quantization/gptq_marlin +// Copyrigth 2024 The vLLM team. +// Copyright (c) 2024 by KVCache.AI, All Rights Reserved. +#ifndef _data_types_cuh +#define _data_types_cuh +#include "gptq_marlin.cuh" +#include +#include + +#ifdef __HIP_PLATFORM_AMD__ +typedef __hip_bfloat16 nv_bfloat16; +typedef __hip_bfloat162 nv_bfloat162; +#endif + +namespace gptq_marlin { + +template +class ScalarType {}; + +template <> +class ScalarType { + public: + using scalar_t = half; + using scalar_t2 = half2; + + // Matrix fragments for tensor core instructions; their precise layout is + // documented here: + // https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#matrix-fragments-for-mma-m16n8k16-with-floating-point-type + using FragA = Vec; + using FragB = Vec; + using FragC = Vec; + using FragS = Vec; + + static __device__ float inline num2float(const half x) { + return __half2float(x); + } + + static __device__ half2 inline num2num2(const half x) { + return __half2half2(x); + } + + static __device__ half2 inline nums2num2(const half x1, const half x2) { + return __halves2half2(x1, x2); + } + + static __host__ __device__ half inline float2num(const float x) { + return __float2half(x); + } +}; + +template <> +class ScalarType { + public: + using scalar_t = nv_bfloat16; + using scalar_t2 = nv_bfloat162; + + using FragA = Vec; + using FragB = Vec; + using FragC = Vec; + using FragS = Vec; + +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800 + static __device__ float inline num2float(const nv_bfloat16 x) { + return __bfloat162float(x); + } + + static __device__ nv_bfloat162 inline num2num2(const nv_bfloat16 x) { + return __bfloat162bfloat162(x); + } + + static __device__ nv_bfloat162 inline nums2num2(const nv_bfloat16 x1, + const nv_bfloat16 x2) { + return __halves2bfloat162(x1, x2); + } + + static __host__ __device__ nv_bfloat16 inline float2num(const float x) { + return __float2bfloat16(x); + } +#endif +}; + +} // namespace gptq_marlin + +#endif diff --git a/archive/csrc/ktransformers_ext/cuda/gptq_marlin/ops.h b/archive/csrc/ktransformers_ext/cuda/gptq_marlin/ops.h new file mode 100644 index 0000000..47cefea --- /dev/null +++ b/archive/csrc/ktransformers_ext/cuda/gptq_marlin/ops.h @@ -0,0 +1,24 @@ +/** + * @Description : + * @Author : Azure + * @Date : 2024-07-22 09:27:55 + * @Version : 1.0.0 + * @LastEditors : Azure + * @LastEditTime : 2024-07-26 08:35:00 + * @Copyright (c) 2024 by KVCache.AI, All Rights Reserved. +**/ +#pragma once + +#include +#include +#include + +torch::Tensor gptq_marlin_gemm(torch::Tensor& a, torch::Tensor& b_q_weight, + torch::Tensor& b_scales, torch::Tensor& g_idx, + torch::Tensor& perm, torch::Tensor& workspace, + int64_t num_bits, int64_t size_m, int64_t size_n, + int64_t size_k, bool is_k_full); + +// torch::Tensor gptq_marlin_repack(torch::Tensor& b_q_weight, torch::Tensor& perm, +// int64_t size_k, int64_t size_n, +// int64_t num_bits); \ No newline at end of file diff --git a/archive/csrc/ktransformers_ext/cuda/setup.py b/archive/csrc/ktransformers_ext/cuda/setup.py new file mode 100644 index 0000000..156bb0e --- /dev/null +++ b/archive/csrc/ktransformers_ext/cuda/setup.py @@ -0,0 +1,26 @@ + +from setuptools import setup, Extension +from torch.utils import cpp_extension +from torch.utils.cpp_extension import BuildExtension, CUDAExtension +setup( + name='KTransformersOps', + ext_modules=[ + CUDAExtension( + 'KTransformersOps', [ + 'custom_gguf/dequant.cu', + 'binding.cpp', + 'gptq_marlin/gptq_marlin.cu', + # 'gptq_marlin_repack.cu', + ], + extra_compile_args={ + 'cxx': ['-O3'], + 'nvcc': [ + '-O3', + '--use_fast_math', + '-Xcompiler', '-fPIC', + ] + }, + ) + ], + cmdclass={'build_ext': BuildExtension} +) \ No newline at end of file diff --git a/archive/csrc/ktransformers_ext/cuda/test_dequant.py b/archive/csrc/ktransformers_ext/cuda/test_dequant.py new file mode 100644 index 0000000..c39d6c7 --- /dev/null +++ b/archive/csrc/ktransformers_ext/cuda/test_dequant.py @@ -0,0 +1,16 @@ +import os +import sys +sys.path.insert(0,"/home/zbx/ktransformers") +from ktransformers.util.custom_loader import GGUFLoader +import torch + +gguf_loader_1 = GGUFLoader("/mnt/data/model/DeepseekV3-q4km-gguf") +gguf_loader_2 = GGUFLoader("/mnt/data/chenht/model/gguf_for_ktransformers/DeepSeek-V3-bf16/") + +torch.set_default_dtype(torch.bfloat16) + +tensor_1 = gguf_loader_1.load_gguf_tensor("blk.0.attn_kv_a_mqa.weight", "cuda") +tensor_2 = gguf_loader_2.load_gguf_tensor("blk.0.attn_kv_a_mqa.weight", "cuda") + +print(tensor_1[0, -64:]) +print(tensor_2[0, -64:]) \ No newline at end of file diff --git a/archive/csrc/ktransformers_ext/examples/test_attention.py b/archive/csrc/ktransformers_ext/examples/test_attention.py new file mode 100644 index 0000000..5627a0e --- /dev/null +++ b/archive/csrc/ktransformers_ext/examples/test_attention.py @@ -0,0 +1,142 @@ +#!/usr/bin/env python +# coding=utf-8 +""" +Description : +Author : Jianwei Dong +Date : 2024-08-28 10:32:05 +Version : 1.0.0 +LastEditors : chenht2022 +LastEditTime : 2024-08-28 10:32:05 +Copyright (c) 2024 by KVCache.AI, All Rights Reserved. +""" +import os, sys +import time + +sys.path.append(os.path.dirname(__file__) + "/../build") +import cpuinfer_ext +from flash_attn import flash_attn_with_kvcache +import torch + +layer_num = 10 +kv_head_num = 8 +q_head_num = 32 +head_dim = 128 +block_len = 128 +anchor_num = 1 +cache_seqlen = 8192 +cache_seqlens = torch.tensor([cache_seqlen], dtype=torch.int32, device="cpu") +seqlens_zero = torch.zeros((1,), dtype=torch.int32, device="cpu") +anchor_type = cpuinfer_ext.kvcache.AnchorType.DYNAMIC +kv_type = cpuinfer_ext.kvcache.ggml_type.FP16 +retrieval_type = cpuinfer_ext.kvcache.RetrievalType.LAYER +layer_step: int = 1 +token_step: int = 1 +layer_offset: int = 0 +max_thread_num: int = 2 +max_batch_size: int = 1 +max_block_num: int = 512 +CPUInfer = cpuinfer_ext.CPUInfer(max_thread_num) +validation_iter = 100 + +with torch.inference_mode(mode=True): + config = cpuinfer_ext.kvcache.KVCacheConfig( + layer_num, + kv_head_num, + q_head_num, + head_dim, + block_len, + anchor_num, + anchor_type, + kv_type, + retrieval_type, + layer_step, + token_step, + layer_offset, + max_block_num, + max_batch_size, + max_thread_num, + ) + local_kvcache = cpuinfer_ext.kvcache.KVCache(config) + + kvcaches = [] + block_table = ( + torch.arange(max_block_num, dtype=torch.int32, device="cpu") + .contiguous() + .view(1, -1) + ) + + for layer_idx in range(layer_num): + k_cache = torch.randn( + (1, cache_seqlen, kv_head_num, head_dim), dtype=torch.float16, device="cpu" + ).contiguous() + v_cache = torch.randn( + (1, cache_seqlen, kv_head_num, head_dim), dtype=torch.float16, device="cpu" + ).contiguous() + + CPUInfer.submit( + local_kvcache.update_kvcache_fp16( + k_cache.data_ptr(), + v_cache.data_ptr(), + layer_idx, + block_table.data_ptr(), + 1, + max_block_num, + seqlens_zero.data_ptr(), + cache_seqlen, + ) + ) + CPUInfer.sync() + + kvcaches.append((k_cache.to("cuda"), v_cache.to("cuda"))) + + # validation + for i in range(validation_iter): + + k_cache = kvcaches[i % layer_num][0] + v_cache = kvcaches[i % layer_num][1] + input = torch.randn( + (1, 1, q_head_num, head_dim), dtype=torch.float16, device="cpu" + ).contiguous() + output = torch.empty( + (1, 1, q_head_num, head_dim), dtype=torch.float16, device="cpu" + ).contiguous() + + # attn_lse: (bsz, q_len, q_head_num) + attn_lse = torch.empty( + (1, 1, q_head_num), dtype=torch.float32, device="cpu" + ).contiguous() + input = input / 100 + + CPUInfer.submit( + local_kvcache.attn( + input.data_ptr(), + output.data_ptr(), + attn_lse.data_ptr(), + i % layer_num, + 0, + 1, + 1, + max_block_num, + block_table.data_ptr(), + cache_seqlens.data_ptr(), + -1, + -1, + -1, + ) + ) + CPUInfer.sync() + # print("cpuinfer output", output) + + t_output = flash_attn_with_kvcache( + q=input.to("cuda"), + k_cache=k_cache, + v_cache=v_cache, + cache_seqlens=cache_seqlens.to("cuda"), + ) + # print("torch output", t_output) + + diff = torch.mean(torch.abs(output.to("cuda") - t_output)) / torch.mean( + torch.abs(t_output) + ) + print("diff = ", diff) + assert diff < 0.001 diff --git a/archive/csrc/ktransformers_ext/examples/test_linear.py b/archive/csrc/ktransformers_ext/examples/test_linear.py new file mode 100644 index 0000000..7a331db --- /dev/null +++ b/archive/csrc/ktransformers_ext/examples/test_linear.py @@ -0,0 +1,62 @@ +#!/usr/bin/env python +# coding=utf-8 +''' +Description : +Author : chenht2022 +Date : 2024-07-25 10:32:05 +Version : 1.0.0 +LastEditors : chenht2022 +LastEditTime : 2024-08-06 10:36:59 +Copyright (c) 2024 by KVCache.AI, All Rights Reserved. +''' +import os, sys +import time +sys.path.append(os.path.dirname(__file__) + '/../build') +import cpuinfer_ext +import torch + +input_size = 16384 +output_size = 5120 +stride = 32 +group_max_len = 1024 +proj_type = 1 # ggml_type::GGML_TYPE_F16 +hidden_type = 1 # ggml_type::GGML_TYPE_F16 +qlen = 30 +layer_num = 10 +CPUInfer = cpuinfer_ext.CPUInfer(48) +validation_iter = 100 + +with torch.inference_mode(mode=True): + linears = [] + projs = [] + for _ in range(layer_num): + proj = torch.randn((output_size, input_size), dtype=torch.float16, device = "cuda").to("cpu").contiguous() + config = cpuinfer_ext.linear.LinearConfig(input_size, output_size, stride, group_max_len, proj.data_ptr(), proj_type, hidden_type) + linear = cpuinfer_ext.linear.Linear(config) + projs.append(proj) + linears.append(linear) + + # validation + for i in range(validation_iter): + linear = linears[i % layer_num] + input = torch.randn((qlen, input_size), dtype=torch.float16).contiguous() + output = torch.empty((qlen, output_size), dtype=torch.float16).contiguous() + input = input / 100 + + CPUInfer.submit( + linear.forward( + qlen, + input.data_ptr(), + output.data_ptr() + ) + ) + CPUInfer.sync() + # print('cpuinfer output', output) + + proj = projs[i%layer_num] + t_output = torch.mm(input, proj.t()) + # print('torch output', t_output) + + diff = torch.mean(torch.abs(output - t_output)) / torch.mean(torch.abs(t_output)) + print('diff = ', diff) + assert(diff < 0.001) diff --git a/archive/csrc/ktransformers_ext/examples/test_mlp.py b/archive/csrc/ktransformers_ext/examples/test_mlp.py new file mode 100644 index 0000000..9805e72 --- /dev/null +++ b/archive/csrc/ktransformers_ext/examples/test_mlp.py @@ -0,0 +1,82 @@ +#!/usr/bin/env python +# coding=utf-8 +''' +Description : +Author : chenht2022 +Date : 2024-07-25 10:32:05 +Version : 1.0.0 +LastEditors : chenht2022 +LastEditTime : 2024-08-06 10:37:28 +Copyright (c) 2024 by KVCache.AI, All Rights Reserved. +''' +import os, sys +import time +sys.path.append(os.path.dirname(__file__) + '/../build') +import cpuinfer_ext +import torch + +hidden_size = 5120 +intermediate_size = 3072 +stride = 32 +group_max_len = 1024 +gate_type = 1 # ggml_type::GGML_TYPE_F16 +up_type = 1 # ggml_type::GGML_TYPE_F16 +down_type = 1 # ggml_type::GGML_TYPE_F16 +hidden_type = 1 # ggml_type::GGML_TYPE_F16 +qlen = 30 +layer_num = 10 +CPUInfer = cpuinfer_ext.CPUInfer(48) +validation_iter = 100 + +def act_fn(x): + return x / (1.0 + torch.exp(-x)) + +def mlp_torch(input, gate_proj, up_proj, down_proj): + gate_buf = torch.mm(input, gate_proj.t()) + up_buf = torch.mm(input, up_proj.t()) + intermediate = act_fn(gate_buf) * up_buf + ret = torch.mm(intermediate, down_proj.t()) + return ret + +with torch.inference_mode(mode=True): + mlps = [] + gate_projs = [] + up_projs = [] + down_projs = [] + for _ in range(layer_num): + gate_proj = torch.randn((intermediate_size, hidden_size), dtype=torch.float16, device = "cuda").to("cpu").contiguous() + up_proj = torch.randn((intermediate_size, hidden_size), dtype=torch.float16, device = "cuda").to("cpu").contiguous() + down_proj = torch.randn((hidden_size, intermediate_size), dtype=torch.float16, device = "cuda").to("cpu").contiguous() + config = cpuinfer_ext.mlp.MLPConfig(hidden_size, intermediate_size, stride, group_max_len, gate_proj.data_ptr(), up_proj.data_ptr(), down_proj.data_ptr(), gate_type, up_type, down_type, hidden_type) + mlp = cpuinfer_ext.mlp.MLP(config) + gate_projs.append(gate_proj) + up_projs.append(up_proj) + down_projs.append(down_proj) + mlps.append(mlp) + + # validation + for i in range(validation_iter): + mlp = mlps[i % layer_num] + input = torch.randn((qlen, hidden_size), dtype=torch.float16).contiguous() + output = torch.empty((qlen, hidden_size), dtype=torch.float16).contiguous() + input = input / 100 + + CPUInfer.submit( + mlp.forward( + qlen, + input.data_ptr(), + output.data_ptr() + ) + ) + CPUInfer.sync() + # print('cpuinfer output', output) + + gate_proj = gate_projs[i%layer_num] + up_proj = up_projs[i%layer_num] + down_proj = down_projs[i%layer_num] + t_output = mlp_torch(input, gate_proj, up_proj, down_proj) + # print('torch output', t_output) + + diff = torch.mean(torch.abs(output - t_output)) / torch.mean(torch.abs(t_output)) + print('diff = ', diff) + assert(diff < 0.001) diff --git a/archive/csrc/ktransformers_ext/examples/test_moe.py b/archive/csrc/ktransformers_ext/examples/test_moe.py new file mode 100644 index 0000000..3fa4dbd --- /dev/null +++ b/archive/csrc/ktransformers_ext/examples/test_moe.py @@ -0,0 +1,121 @@ +#!/usr/bin/env python +# coding=utf-8 +''' +Description : +Author : chenht2022 +Date : 2024-07-25 10:32:05 +Version : 1.0.0 +LastEditors : chenht2022 +LastEditTime : 2024-08-06 10:38:05 +Copyright (c) 2024 by KVCache.AI, All Rights Reserved. +''' +import os, sys +import time +sys.path.append(os.path.dirname(__file__) + '/../build') +import cpuinfer_ext +import torch + +expert_num = 160 +hidden_size = 5120 +intermediate_size = 1536 +stride = 32 +group_min_len = 10 +group_max_len = 1024 +gate_type = 1 # ggml_type::GGML_TYPE_F16 +up_type = 1 # ggml_type::GGML_TYPE_F16 +down_type = 1 # ggml_type::GGML_TYPE_F16 +hidden_type = 1 # ggml_type::GGML_TYPE_F16 +n_routed_experts = 6 +qlen = 30 +layer_num = 10 +CPUInfer = cpuinfer_ext.CPUInfer(48) +validation_iter = 100 + +def act_fn(x): + return x / (1.0 + torch.exp(-x)) + +def mlp_torch(input, gate_proj, up_proj, down_proj): + gate_buf = torch.mm(input, gate_proj.t()) + up_buf = torch.mm(input, up_proj.t()) + intermediate = act_fn(gate_buf) * up_buf + ret = torch.mm(intermediate, down_proj.t()) + return ret + +def moe_torch(input, expert_ids, weights, gate_proj, up_proj, down_proj): + cnts = expert_ids.new_zeros((expert_ids.shape[0], expert_num)) + cnts.scatter_(1, expert_ids, 1) + tokens_per_expert = cnts.sum(dim=0) + idxs = expert_ids.view(-1).argsort() + sorted_tokens = input[idxs // expert_ids.shape[1]] + + outputs = [] + start_idx = 0 + for i, num_tokens in enumerate(tokens_per_expert): + end_idx = start_idx + num_tokens + if num_tokens == 0: + continue + tokens_for_this_expert = sorted_tokens[start_idx:end_idx] + expert_out = mlp_torch(tokens_for_this_expert, gate_proj[i], up_proj[i], down_proj[i]) + outputs.append(expert_out) + start_idx = end_idx + + outs = torch.cat(outputs, dim=0) if len(outputs) else sorted_tokens.new_empty(0) + + new_x = torch.empty_like(outs) + new_x[idxs] = outs + t_output = ( + new_x.view(*expert_ids.shape, -1) + .type(weights.dtype) + .mul_(weights.unsqueeze(dim=-1)) + .sum(dim=1) + .type(new_x.dtype) + ) + return t_output + +with torch.inference_mode(mode=True): + moes = [] + gate_projs = [] + up_projs = [] + down_projs = [] + for _ in range(layer_num): + gate_proj = torch.randn((expert_num, intermediate_size, hidden_size), dtype=torch.float16, device = "cuda").to("cpu").contiguous() + up_proj = torch.randn((expert_num, intermediate_size, hidden_size), dtype=torch.float16, device = "cuda").to("cpu").contiguous() + down_proj = torch.randn((expert_num, hidden_size, intermediate_size), dtype=torch.float16, device = "cuda").to("cpu").contiguous() + config = cpuinfer_ext.moe.MOEConfig(expert_num, n_routed_experts, hidden_size, intermediate_size, stride, group_min_len, group_max_len, gate_proj.data_ptr(), up_proj.data_ptr(), down_proj.data_ptr(), gate_type, up_type, down_type, hidden_type) + moe = cpuinfer_ext.moe.MOE(config) + gate_projs.append(gate_proj) + up_projs.append(up_proj) + down_projs.append(down_proj) + moes.append(moe) + + # validation + for i in range(validation_iter): + expert_ids = torch.stack([torch.randperm(expert_num)[:n_routed_experts] for _ in range(qlen)]).contiguous() + weights = torch.rand((qlen, n_routed_experts), dtype=torch.float32).contiguous() + input = torch.randn((qlen, hidden_size), dtype=torch.float16).contiguous() + output = torch.empty((qlen, hidden_size), dtype=torch.float16).contiguous() + input = input / 100 + + moe = moes[i % layer_num] + CPUInfer.submit( + moe.forward( + qlen, + n_routed_experts, + expert_ids.data_ptr(), + weights.data_ptr(), + input.data_ptr(), + output.data_ptr() + ) + ) + CPUInfer.sync() + # print('cpuinfer output', output) + + gate_proj = gate_projs[i%layer_num] + up_proj = up_projs[i%layer_num] + down_proj = down_projs[i%layer_num] + t_output = moe_torch(input, expert_ids, weights, gate_proj, up_proj, down_proj) + # print('torch output', t_output) + + diff = torch.mean(torch.abs(output - t_output)) / torch.mean(torch.abs(t_output)) + print('diff = ', diff) + assert(diff < 0.001) diff --git a/archive/csrc/ktransformers_ext/ext_bindings.cpp b/archive/csrc/ktransformers_ext/ext_bindings.cpp new file mode 100644 index 0000000..a2dd733 --- /dev/null +++ b/archive/csrc/ktransformers_ext/ext_bindings.cpp @@ -0,0 +1,790 @@ +/** + * @Description : + * @Author : chenht2022, Jianwei Dong + * @Date : 2024-07-22 02:03:22 + * @Version : 1.0.0 + * @LastEditors : Jianwei Dong + * @LastEditTime : 2024-08-26 22:47:06 + * @Copyright (c) 2024 by KVCache.AI, All Rights Reserved. + **/ +// Python bindings +#include "cpu_backend/cpuinfer.h" +#if !defined(KTRANSFORMERS_USE_ROCM) && !defined(KTRANSFORMERS_USE_XPU) && !defined(KTRANSFORMERS_USE_NPU) +#include "device_launch_parameters.h" +#endif +#include "llamafile/flags.h" +#include "operators/kvcache/kvcache.h" +#include "operators/llamafile/linear.h" +#include "operators/llamafile/mlp.h" +#include "operators/llamafile/moe.h" + +#if defined(__x86_64__) && defined(__HAS_AVX512F__) && defined(__HAS_AMX__) +#include "operators/amx/moe.hpp" +#endif + +#include "pybind11/functional.h" +#include "pybind11/operators.h" +#include "pybind11/pybind11.h" +#include "pybind11/stl.h" +#include +#include +#include + +namespace py = pybind11; +using namespace pybind11::literals; + +// Binding functions for the KVCache class +class KVCacheBindings { + public: + class AttnBindings { + public: + struct Args { + CPUInfer *cpuinfer; + KVCache *kv_cache; + const ggml_fp16_t *q_in; + ggml_fp16_t *output; + float *attn_lse; + int layer_idx; + int generate_token_idx; + int q_len; + int batch_size; + int max_block_num; + int *block_table; + int *cache_seqlens; + int pick_block_num; + int init_block_num; + int local_block_num; + }; + static void inner(void *args) { + Args *args_ = (Args *)args; + args_->cpuinfer->enqueue( + &KVCache::attn, args_->kv_cache, args_->q_in, args_->output, + args_->attn_lse, args_->layer_idx, args_->generate_token_idx, + args_->q_len, args_->batch_size, args_->max_block_num, + args_->block_table, args_->cache_seqlens, args_->pick_block_num, + args_->init_block_num, args_->local_block_num); + } + static std::pair + cpuinfer_interface(KVCache &kv_cache, intptr_t q_in, intptr_t output, + intptr_t attn_lse, int layer_idx, + int generate_token_idx, int q_len, int batch_size, + int max_block_num, intptr_t block_table, + intptr_t cache_seqlens, int pick_block_num, + int init_block_num, int local_block_num) { + Args *args = new Args{nullptr, + &kv_cache, + (const ggml_fp16_t *)q_in, + (ggml_fp16_t *)output, + (float *)attn_lse, + layer_idx, + generate_token_idx, + q_len, + batch_size, + max_block_num, + (int *)block_table, + (int *)cache_seqlens, + pick_block_num, + init_block_num, + local_block_num}; + return std::make_pair((intptr_t)&inner, (intptr_t)args); + } + }; + + class GetAllKVCacheOneLayerBindings { + public: + struct Args { + CPUInfer *cpuinfer; + KVCache *kv_cache; + int layer_id; + ggml_fp16_t *k_in; + ggml_fp16_t *v_in; + }; + static void inner(void *args) { + Args *args_ = (Args *)args; + args_->cpuinfer->enqueue(&KVCache::get_all_kvcache_one_layer, + args_->kv_cache, args_->layer_id, + args_->k_in, args_->v_in); + } + static std::pair + cpuinfer_interface(KVCache &kv_cache, intptr_t k_in, intptr_t v_in, + int layer_id) { + Args *args = new Args{nullptr, &kv_cache, layer_id, + (ggml_fp16_t *)k_in, (ggml_fp16_t *)v_in}; + return std::make_pair((intptr_t)&inner, (intptr_t)args); + } + }; + + class GetAndUpdateKVCacheFp16Bindings { + public: + struct Args { + CPUInfer *cpuinfer; + KVCache *kv_cache; + ggml_fp16_t *k_in; + ggml_fp16_t *v_in; + int layer_id; + int *block_table; + int batch_size; + int max_block_num; + int *cache_seqlens; + int q_len; + }; + static void inner(void *args) { + Args *args_ = (Args *)args; + args_->cpuinfer->enqueue(&KVCache::get_and_update_kvcache_fp16, + args_->kv_cache, args_->k_in, args_->v_in, + args_->layer_id, args_->block_table, + args_->batch_size, args_->max_block_num, + args_->cache_seqlens, args_->q_len); + } + static std::pair + cpuinfer_interface(KVCache &kv_cache, intptr_t k_in, intptr_t v_in, + int layer_id, intptr_t block_table, int batch_size, + int max_block_num, intptr_t cache_seqlens, + int q_len) { + Args *args = new Args{nullptr, + &kv_cache, + (ggml_fp16_t *)k_in, + (ggml_fp16_t *)v_in, + layer_id, + (int *)block_table, + batch_size, + max_block_num, + (int *)cache_seqlens, + q_len}; + return std::make_pair((intptr_t)&inner, (intptr_t)args); + } + }; + class GetKVCacheFp16Bindings { + public: + struct Args { + CPUInfer *cpuinfer; + KVCache *kv_cache; + ggml_fp16_t *k_in; + ggml_fp16_t *v_in; + int layer_id; + int *block_table; + int batch_size; + int max_block_num; + int *cache_seqlens; + }; + static void inner(void *args) { + Args *args_ = (Args *)args; + args_->cpuinfer->enqueue( + &KVCache::get_kvcache_fp16, args_->kv_cache, args_->k_in, + args_->v_in, args_->layer_id, args_->block_table, + args_->batch_size, args_->max_block_num, args_->cache_seqlens); + } + static std::pair + cpuinfer_interface(KVCache &kv_cache, intptr_t k_in, intptr_t v_in, + int layer_id, intptr_t block_table, int batch_size, + int max_block_num, intptr_t cache_seqlens) { + Args *args = new Args{nullptr, + &kv_cache, + (ggml_fp16_t *)k_in, + (ggml_fp16_t *)v_in, + layer_id, + (int *)block_table, + batch_size, + max_block_num, + (int *)cache_seqlens}; + return std::make_pair((intptr_t)&inner, (intptr_t)args); + } + }; + + class UpdateKVCacheFp16Bindings { + public: + struct Args { + CPUInfer *cpuinfer; + KVCache *kv_cache; + ggml_fp16_t *k_in; + ggml_fp16_t *v_in; + int layer_id; + int *block_table; + int batch_size; + int max_block_num; + int *cache_seqlens; + int q_len; + }; + static void inner(void *args) { + Args *args_ = (Args *)args; + args_->cpuinfer->enqueue(&KVCache::update_kvcache_fp16, + args_->kv_cache, args_->k_in, args_->v_in, + args_->layer_id, args_->block_table, + args_->batch_size, args_->max_block_num, + args_->cache_seqlens, args_->q_len); + } + static std::pair + cpuinfer_interface(KVCache &kv_cache, intptr_t k_in, intptr_t v_in, + int layer_id, intptr_t block_table, int batch_size, + int max_block_num, intptr_t cache_seqlens, + int q_len) { + Args *args = new Args{nullptr, + &kv_cache, + (ggml_fp16_t *)k_in, + (ggml_fp16_t *)v_in, + layer_id, + (int *)block_table, + batch_size, + max_block_num, + (int *)cache_seqlens, + q_len}; + return std::make_pair((intptr_t)&inner, (intptr_t)args); + } + }; + + class UpdateImportanceBindings { + public: + struct Args { + CPUInfer *cpuinfer; + KVCache *kv_cache; + const ggml_fp16_t *importance; + int layer_id; + int *block_table; + int batch_size; + int max_block_num; + int *offset; + int width; + }; + static void inner(void *args) { + Args *args_ = (Args *)args; + args_->cpuinfer->enqueue( + &KVCache::update_importance, args_->kv_cache, args_->importance, + args_->layer_id, args_->block_table, args_->batch_size, + args_->max_block_num, args_->offset, args_->width); + } + static std::pair + cpuinfer_interface(KVCache &kv_cache, intptr_t importance, int layer_id, + intptr_t block_table, int batch_size, + int max_block_num, intptr_t offset, int width) { + Args *args = new Args{nullptr, + &kv_cache, + (const ggml_fp16_t *)importance, + layer_id, + (int *)block_table, + batch_size, + max_block_num, + (int *)offset, + width}; + return std::make_pair((intptr_t)&inner, (intptr_t)args); + } + }; + + class AttnWithKVCacheBindings { + public: + struct Args { + CPUInfer *cpuinfer; + KVCache *kv_cache; + const ggml_fp16_t *q_in; + const ggml_fp16_t *k_in; + const ggml_fp16_t *v_in; + ggml_fp16_t *output; + float *attn_lse; + int layer_idx; + int generate_token_idx; + int q_len; + int batch_size; + int max_block_num; + int *block_table; + int *cache_seqlens; + int topk; + int local; + }; + static void inner(void *args) { + Args *args_ = (Args *)args; + args_->cpuinfer->enqueue( + &KVCache::attn_with_kvcache, args_->kv_cache, args_->q_in, + args_->k_in, args_->v_in, args_->output, args_->attn_lse, + args_->layer_idx, args_->generate_token_idx, args_->q_len, + args_->batch_size, args_->max_block_num, args_->block_table, + args_->cache_seqlens, args_->topk, args_->local); + } + static std::pair + cpuinfer_interface(KVCache &kv_cache, intptr_t q_in, intptr_t k_in, + intptr_t v_in, intptr_t output, intptr_t attn_lse, + int layer_idx, int generate_token_idx, int q_len, + int batch_size, int max_block_num, + intptr_t block_table, intptr_t cache_seqlens, + int topk, int local) { + Args *args = new Args{nullptr, + &kv_cache, + (const ggml_fp16_t *)q_in, + (const ggml_fp16_t *)k_in, + (const ggml_fp16_t *)v_in, + (ggml_fp16_t *)output, + (float *)attn_lse, + layer_idx, + generate_token_idx, + q_len, + batch_size, + max_block_num, + (int *)block_table, + (int *)cache_seqlens, + topk, + local}; + return std::make_pair((intptr_t)&inner, (intptr_t)args); + } + }; + + class ClearImportanceAllLayersBindings { + public: + struct Args { + CPUInfer *cpuinfer; + KVCache *kv_cache; + int *block_table; + int *cache_seqlens; + int batch_size; + int max_block_num; + }; + static void inner(void *args) { + Args *args_ = (Args *)args; + args_->cpuinfer->enqueue(&KVCache::clear_importance_all_layers, + args_->kv_cache, args_->block_table, + args_->cache_seqlens, args_->batch_size, + args_->max_block_num); + } + static std::pair + cpuinfer_interface(KVCache &kv_cache, intptr_t block_table, + intptr_t cache_seqlens, int batch_size, + int max_block_num) { + Args *args = new Args{nullptr, + &kv_cache, + (int *)block_table, + (int *)cache_seqlens, + batch_size, + max_block_num}; + return std::make_pair((intptr_t)&inner, (intptr_t)args); + } + }; + + class CalcAnchorAllLayersBindinds { + public: + struct Args { + CPUInfer *cpuinfer; + KVCache *kv_cache; + int *block_table; + int *cache_seqlens; + int batch_size; + int max_block_num; + }; + static void inner(void *args) { + Args *args_ = (Args *)args; + args_->cpuinfer->enqueue(&KVCache::calc_anchor_all_layers, + args_->kv_cache, args_->block_table, + args_->cache_seqlens, args_->batch_size, + args_->max_block_num); + } + static std::pair + cpuinfer_interface(KVCache &kv_cache, intptr_t block_table, + intptr_t cache_seqlens, int batch_size, + int max_block_num) { + Args *args = new Args{nullptr, + &kv_cache, + (int *)block_table, + (int *)cache_seqlens, + batch_size, + max_block_num}; + return std::make_pair((intptr_t)&inner, (intptr_t)args); + } + }; + + class LoadKVCacheBindings { + public: + struct Args { + CPUInfer *cpuinfer; + KVCache *kv_cache; + std::string tensor_file_path; + }; + static void inner(void *args) { + Args *args_ = (Args *)args; + args_->cpuinfer->enqueue(&KVCache::load_kvcache, args_->kv_cache, + args_->tensor_file_path); + } + static std::pair + cpuinfer_interface(KVCache &kv_cache, std::string tensor_file_path) { + Args *args = + new Args{nullptr, &kv_cache, (std::string)tensor_file_path}; + return std::make_pair((intptr_t)&inner, (intptr_t)args); + } + }; + class DumpKVCacheBindings { + public: + struct Args { + CPUInfer *cpuinfer; + KVCache *kv_cache; + int *block_table; + int cache_total_len; + std::string tensor_file_path; + }; + static void inner(void *args) { + Args *args_ = (Args *)args; + args_->cpuinfer->enqueue(&KVCache::dump_kvcache, args_->kv_cache, + args_->block_table, args_->cache_total_len, + args_->tensor_file_path); + } + static std::pair + cpuinfer_interface(KVCache &kv_cache, intptr_t block_table, + int cache_total_len, std::string tensor_file_path) { + Args *args = + new Args{nullptr, &kv_cache, (int *)block_table, + cache_total_len, (std::string)tensor_file_path}; + return std::make_pair((intptr_t)&inner, (intptr_t)args); + } + }; +}; + +class LinearBindings { + public: + class WarmUpBindinds { + public: + struct Args { + CPUInfer *cpuinfer; + Linear *linear; + }; + static void inner(void *args) { + Args *args_ = (Args *)args; + args_->cpuinfer->enqueue(&Linear::warm_up, args_->linear); + } + static std::pair + cpuinfer_interface(Linear &linear) { + Args *args = new Args{nullptr, &linear}; + return std::make_pair((intptr_t)&inner, (intptr_t)args); + } + }; + class ForwardBindings { + public: + struct Args { + CPUInfer *cpuinfer; + Linear *linear; + int qlen; + const void *input; + void *output; + }; + static void inner(void *args) { + Args *args_ = (Args *)args; + args_->cpuinfer->enqueue(&Linear::forward, args_->linear, + args_->qlen, args_->input, args_->output); + } + static std::pair + cpuinfer_interface(Linear &linear, int qlen, intptr_t input, + intptr_t output) { + Args *args = new Args{nullptr, &linear, qlen, (const void *)input, + (void *)output}; + return std::make_pair((intptr_t)&inner, (intptr_t)args); + } + }; +}; + +class MLPBindings { + public: + class WarmUpBindinds { + public: + struct Args { + CPUInfer *cpuinfer; + MLP *mlp; + }; + static void inner(void *args) { + Args *args_ = (Args *)args; + args_->cpuinfer->enqueue(&MLP::warm_up, args_->mlp); + } + static std::pair cpuinfer_interface(MLP &mlp) { + Args *args = new Args{nullptr, &mlp}; + return std::make_pair((intptr_t)&inner, (intptr_t)args); + } + }; + class ForwardBindings { + public: + struct Args { + CPUInfer *cpuinfer; + MLP *mlp; + int qlen; + const void *input; + void *output; + }; + static void inner(void *args) { + Args *args_ = (Args *)args; + args_->cpuinfer->enqueue(&MLP::forward, args_->mlp, args_->qlen, + args_->input, args_->output); + } + static std::pair + cpuinfer_interface(MLP &mlp, int qlen, intptr_t input, + intptr_t output) { + Args *args = new Args{nullptr, &mlp, qlen, (const void *)input, + (void *)output}; + return std::make_pair((intptr_t)&inner, (intptr_t)args); + } + }; +}; + +class MOEBindings { + public: + class WarmUpBindinds { + public: + struct Args { + CPUInfer *cpuinfer; + MOE *moe; + }; + static void inner(void *args) { + Args *args_ = (Args *)args; + args_->cpuinfer->enqueue(&MOE::warm_up, args_->moe); + } + static std::pair cpuinfer_interface(MOE &moe) { + Args *args = new Args{nullptr, &moe}; + return std::make_pair((intptr_t)&inner, (intptr_t)args); + } + }; + class ForwardBindings { + public: + struct Args { + CPUInfer *cpuinfer; + MOE *moe; + int qlen; + int k; + const uint64_t *expert_ids; + const float *weights; + const void *input; + void *output; + int *batch_size_tensor; + }; + static void inner(void *args) { + Args *args_ = (Args *)args; + args_->cpuinfer->enqueue( + &MOE::forward, args_->moe, args_->qlen, args_->k, + args_->expert_ids, args_->weights, args_->input, args_->output, args_->batch_size_tensor); + } + static std::pair + cpuinfer_interface(MOE &moe, int qlen, int k, intptr_t expert_ids, + intptr_t weights, intptr_t input, intptr_t output, intptr_t batch_size_tensor) { + Args *args = new Args{nullptr, + &moe, + qlen, + k, + (const uint64_t *)expert_ids, + (const float *)weights, + (const void *)input, + (void *)output, + (int *)batch_size_tensor}; + return std::make_pair((intptr_t)&inner, (intptr_t)args); + } + }; +}; + + +#if defined(__x86_64__) && defined(__HAS_AVX512F__) && defined(__HAS_AMX__) +template +class AMX_MOEBindings { + public: + class WarmUpBindings { + public: + struct Args { + CPUInfer *cpuinfer; + AMX_MOE *moe; + }; + static void inner(void *args) { + Args *args_ = (Args *)args; + args_->cpuinfer->enqueue(&AMX_MOE::warm_up, args_->moe); + } + static std::pair cpuinfer_interface(AMX_MOE &moe) { + Args *args = new Args{nullptr, &moe}; + return std::make_pair((intptr_t)&inner, (intptr_t)args); + } + }; + class LoadWeightsBindings { + public: + struct Args { + CPUInfer *cpuinfer; + AMX_MOE *moe; + }; + static void inner(void *args) { + Args *args_ = (Args *)args; + args_->cpuinfer->enqueue(&AMX_MOE::load_weights, args_->moe); + } + static std::pair cpuinfer_interface(AMX_MOE &moe) { + Args *args = new Args{nullptr, &moe}; + return std::make_pair((intptr_t)&inner, (intptr_t)args); + } + }; + class ForwardBindings { + public: + struct Args { + CPUInfer *cpuinfer; + AMX_MOE *moe; + int qlen; + int k; + const uint64_t *expert_ids; + const float *weights; + const void *input; + void *output; + int *batch_size_tensor; + }; + static void inner(void *args) { + Args *args_ = (Args *)args; + args_->cpuinfer->enqueue( + &AMX_MOE::forward, args_->moe, args_->qlen, args_->k, + args_->expert_ids, args_->weights, args_->input, args_->output, args_->batch_size_tensor); + } + static std::pair + cpuinfer_interface(AMX_MOE &moe, int qlen, int k, intptr_t expert_ids, + intptr_t weights, intptr_t input, intptr_t output, intptr_t batch_size_tensor) { + Args *args = new Args{nullptr, + &moe, + qlen, + k, + (const uint64_t *)expert_ids, + (const float *)weights, + (const void *)input, + (void *)output, + (int *)batch_size_tensor}; + return std::make_pair((intptr_t)&inner, (intptr_t)args); + } + }; +}; +#endif + +PYBIND11_MODULE(cpuinfer_ext, m) { + py::class_(m, "CPUInfer") + .def(py::init()) + .def("submit", &CPUInfer::submit) + .def("submit_with_cuda_stream", &CPUInfer::submit_with_cuda_stream) + .def("sync", &CPUInfer::sync) + .def("sync_with_cuda_stream", &CPUInfer::sync_with_cuda_stream); + + auto linear_module = m.def_submodule("linear"); + py::class_(linear_module, "LinearConfig") + .def(py::init([](int hidden_size, int intermediate_size, int stride, + int group_max_len, intptr_t proj, int proj_type, + int hidden_type) { + return LinearConfig(hidden_size, intermediate_size, stride, + group_max_len, (void *)proj, + (ggml_type)proj_type, (ggml_type)hidden_type); + })); + py::class_(linear_module, "Linear") + .def(py::init()) + .def("warm_up", &LinearBindings::WarmUpBindinds::cpuinfer_interface) + .def("forward", &LinearBindings::ForwardBindings::cpuinfer_interface); + + auto mlp_module = m.def_submodule("mlp"); + py::class_(mlp_module, "MLPConfig") + .def(py::init([](int hidden_size, int intermediate_size, int stride, + int group_max_len, intptr_t gate_proj, + intptr_t up_proj, intptr_t down_proj, int gate_type, + int up_type, int down_type, int hidden_type) { + return MLPConfig(hidden_size, intermediate_size, stride, + group_max_len, (void *)gate_proj, (void *)up_proj, + (void *)down_proj, (ggml_type)gate_type, + (ggml_type)up_type, (ggml_type)down_type, + (ggml_type)hidden_type); + })); + py::class_(mlp_module, "MLP") + .def(py::init()) + .def("warm_up", &MLPBindings::WarmUpBindinds::cpuinfer_interface) + .def("forward", &MLPBindings::ForwardBindings::cpuinfer_interface); + + auto moe_module = m.def_submodule("moe"); + py::class_(moe_module, "MOEConfig") + .def(py::init([](int expert_num, int routed_expert_num, int hidden_size, + int intermediate_size, int stride, int group_min_len, + int group_max_len, bool use_silu, intptr_t gate_proj, + intptr_t up_proj, intptr_t down_proj, int gate_type, + int up_type, int down_type, int hidden_type) { + return MOEConfig(expert_num, routed_expert_num, hidden_size, + intermediate_size, stride, group_min_len, + group_max_len, use_silu, (void *)gate_proj, (void *)up_proj, + (void *)down_proj, (ggml_type)gate_type, + (ggml_type)up_type, (ggml_type)down_type, + (ggml_type)hidden_type); + })); + py::class_(moe_module, "MOE") + .def(py::init()) + .def("warm_up", &MOEBindings::WarmUpBindinds::cpuinfer_interface) + .def("forward", &MOEBindings::ForwardBindings::cpuinfer_interface); + + + #if defined(__x86_64__) && defined(__HAS_AVX512F__) && defined(__HAS_AMX__) + py::class_(moe_module, "AMX_MOEConfig") + .def(py::init([](int expert_num, int routed_expert_num, int hidden_size, + int intermediate_size, + int max_len, bool use_silu, intptr_t gate_proj, + intptr_t up_proj, intptr_t down_proj) { + return AMX_MOEConfig(expert_num, routed_expert_num, hidden_size, + intermediate_size, + max_len, use_silu, (void *)gate_proj, + (void *)up_proj, (void *)down_proj); + })); + + py::class_>(moe_module, "AMXBF16_MOE") + .def(py::init()) + .def("warm_up", &AMX_MOEBindings::WarmUpBindings::cpuinfer_interface) + .def("load_weights", &AMX_MOEBindings::LoadWeightsBindings::cpuinfer_interface) + .def("forward", &AMX_MOEBindings::ForwardBindings::cpuinfer_interface); + py::class_>(moe_module, "AMXInt8_MOE") + .def(py::init()) + .def("warm_up", &AMX_MOEBindings::WarmUpBindings::cpuinfer_interface) + .def("load_weights", &AMX_MOEBindings::LoadWeightsBindings::cpuinfer_interface) + .def("forward", &AMX_MOEBindings::ForwardBindings::cpuinfer_interface); + + #endif + + auto kvcache_module = m.def_submodule("kvcache"); + + py::enum_(kvcache_module, "AnchorType") + .value("FIXED", AnchorType::FIXED_ANCHOR) + .value("DYNAMIC", AnchorType::DYNAMIC) + .value("QUEST", AnchorType::QUEST) + .value("BLOCK_MAX", AnchorType::BLOCK_MAX) + .value("BLOCK_MEAN", AnchorType::BLOCK_MEAN); + py::enum_(kvcache_module, "ggml_type") + .value("FP16", ggml_type::GGML_TYPE_F16) + .value("FP32", ggml_type::GGML_TYPE_F32) + .value("Q4_0", ggml_type::GGML_TYPE_Q4_0) + .value("Q8_0", ggml_type::GGML_TYPE_Q8_0); + py::enum_(kvcache_module, "RetrievalType") + .value("LAYER", RetrievalType::LAYER) + .value("KVHEAD", RetrievalType::KVHEAD) + .value("QHEAD", RetrievalType::QHEAD); + + py::class_(kvcache_module, "KVCacheConfig") + .def(py::init()) + .def_readwrite("layer_num", &KVCacheConfig::layer_num) + .def_readwrite("kv_head_num", &KVCacheConfig::kv_head_num) + .def_readwrite("q_head_num", &KVCacheConfig::q_head_num) + .def_readwrite("head_dim", &KVCacheConfig::head_dim) + .def_readwrite("block_len", &KVCacheConfig::block_len) + .def_readwrite("anchor_num", &KVCacheConfig::anchor_num) + .def_readwrite("anchor_type", &KVCacheConfig::anchor_type) + .def_readwrite("kv_type", &KVCacheConfig::kv_type) + .def_readwrite("retrieval_type", &KVCacheConfig::retrieval_type) + .def_readwrite("layer_step", &KVCacheConfig::layer_step) + .def_readwrite("token_step", &KVCacheConfig::token_step) + .def_readwrite("layer_offset", &KVCacheConfig::layer_offset) + .def_readwrite("max_block_num", &KVCacheConfig::max_block_num) + .def_readwrite("max_batch_size", &KVCacheConfig::max_batch_size) + .def_readwrite("max_thread_num", &KVCacheConfig::max_thread_num); + py::class_(kvcache_module, "KVCache") + .def(py::init()) + .def("get_cache_total_len", &KVCache::get_cache_total_len) + .def("update_cache_total_len", + [](KVCache &kvcache, int cache_total_len) { + kvcache.update_cache_total_len(cache_total_len); + }) + .def("attn", &KVCacheBindings::AttnBindings::cpuinfer_interface) + .def( + "get_all_kvcache_one_layer", + &KVCacheBindings::GetAllKVCacheOneLayerBindings::cpuinfer_interface) + .def("get_and_update_kvcache_fp16", + &KVCacheBindings::GetAndUpdateKVCacheFp16Bindings:: + cpuinfer_interface) + .def("get_kvcache_fp16", + &KVCacheBindings::GetKVCacheFp16Bindings::cpuinfer_interface) + .def("update_kvcache_fp16", + &KVCacheBindings::UpdateKVCacheFp16Bindings::cpuinfer_interface) + .def("update_importance", + &KVCacheBindings::UpdateImportanceBindings::cpuinfer_interface) + .def("attn_with_kvcache", + &KVCacheBindings::AttnWithKVCacheBindings::cpuinfer_interface) + .def("clear_importance_all_layers", + &KVCacheBindings::ClearImportanceAllLayersBindings:: + cpuinfer_interface) + .def("calc_anchor_all_layers", + &KVCacheBindings::CalcAnchorAllLayersBindinds::cpuinfer_interface); +} \ No newline at end of file diff --git a/archive/csrc/ktransformers_ext/operators/amx/la/amx.hpp b/archive/csrc/ktransformers_ext/operators/amx/la/amx.hpp new file mode 100644 index 0000000..866300d --- /dev/null +++ b/archive/csrc/ktransformers_ext/operators/amx/la/amx.hpp @@ -0,0 +1,974 @@ +/** + * @Description : + * @Author : chenht2022 + * @Date : 2025-04-25 18:28:12 + * @Version : 1.0.0 + * @LastEditors : chenht2022 + * @LastEditTime : 2025-04-25 18:28:12 + * @Copyright (c) 2024 by KVCache.AI, All Rights Reserved. + **/ + +#pragma once +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "utils.hpp" +#include + +#if (defined(_WIN32) || defined(_WIN64)) +#define RESTRICT __restrict +#else +#define RESTRICT __restrict__ +#endif + +#if (defined(_WIN32) || defined(_WIN64)) +#define ALWAYS_INLINE __forceinline +#elif __has_attribute(always_inline) || defined(__GNUC__) +#define ALWAYS_INLINE __attribute__((__always_inline__)) inline +#else +#define ALWAYS_INLINE inline +#endif + +namespace amx { + +#define ARCH_GET_XCOMP_PERM 0x1022 +#define ARCH_REQ_XCOMP_PERM 0x1023 +#define XFEATURE_XTILECFG 17 +#define XFEATURE_XTILEDATA 18 + +const int TMMCount = 8; +const int MaxTileHeight = 16; +const int MaxTileWidth = 64; + +const int AMX_BLK_SIZE = 32; + +#define TMM0 0 +#define TMM1 1 +#define TMM2 2 +#define TMM3 3 +#define TMM4 4 +#define TMM5 5 +#define TMM6 6 +#define TMM7 7 + +inline bool enable_amx() { + static thread_local bool initialized = false; + if (initialized) { + return true; + } + initialized = true; + + if (syscall(SYS_arch_prctl, ARCH_REQ_XCOMP_PERM, XFEATURE_XTILEDATA)) { + printf("\n Fail to do XFEATURE_XTILEDATA \n\n"); + return false; + } else { + // printf("\n TILE DATA USE SET - OK \n\n"); + return true; + } + return true; +} + +struct alignas(64) TileConfig { + uint8_t palette; + uint8_t start_row; + std::array __0 = {}; + std::array colsb; + std::array __1 = {}; + std::array rows; + std::array __2 = {}; + + TileConfig() { + palette = 1; + start_row = 0; + for (int i = 0; i < 8; i++) { + set_row_col(i, 0, 0); + } + } + + void set_row_col(int i, uint8_t row, uint16_t col) { + colsb[i] = col; + rows[i] = row; + } + + void set_config() { _tile_loadconfig(this); } + + static void load_data(int to, void *from, size_t stride) { + switch (to) { + case 0: + _tile_loadd(0, from, stride); + break; + case 1: + _tile_loadd(1, from, stride); + break; + case 2: + _tile_loadd(2, from, stride); + break; + case 3: + _tile_loadd(3, from, stride); + break; + case 4: + _tile_loadd(4, from, stride); + break; + case 5: + _tile_loadd(5, from, stride); + break; + case 6: + _tile_loadd(6, from, stride); + break; + case 7: + _tile_loadd(7, from, stride); + break; + default: + throw std::runtime_error("no such tile"); + } + } + + static void store_data(int from, void *to, size_t stride) { + switch (from) { + case 0: + _tile_stored(0, to, stride); + break; + case 1: + _tile_stored(1, to, stride); + break; + case 2: + _tile_stored(2, to, stride); + break; + case 3: + _tile_stored(3, to, stride); + break; + case 4: + _tile_stored(4, to, stride); + break; + case 5: + _tile_stored(5, to, stride); + break; + case 6: + _tile_stored(6, to, stride); + break; + case 7: + _tile_stored(7, to, stride); + break; + default: + throw std::runtime_error("no such tile"); + } + } +}; + +static_assert(sizeof(TileConfig) == 64); + +inline void debug_tile(int t) { + printf("Tile %d\n", t); + uint8_t data[16][64] = {}; + TileConfig::store_data(t, data, 64); + for (int i = 0; i < 16; i++) { + for (int j = 0; j < 64; j++) { + printf("%3d ", data[i][j]); + } + printf("\n"); + } + printf("\n"); +} + +inline void debug_tiles(int to = 8) { + for (int i = 0; i < to; i++) { + debug_tile(i); + } +} + +inline void debug_m512(__m512 x) { + float data[16]; + _mm512_storeu_ps(data, x); + for (int i = 0; i < 16; i++) { + printf("%f ", data[i]); + } + printf("\n"); +} + +// transpose utils +inline void transpose_16x16_32bit(__m512i *v) { + __m512i v1[16]; + v1[0] = _mm512_unpacklo_epi32(v[0], v[1]); + v1[1] = _mm512_unpackhi_epi32(v[0], v[1]); + v1[2] = _mm512_unpacklo_epi32(v[2], v[3]); + v1[3] = _mm512_unpackhi_epi32(v[2], v[3]); + v1[4] = _mm512_unpacklo_epi32(v[4], v[5]); + v1[5] = _mm512_unpackhi_epi32(v[4], v[5]); + v1[6] = _mm512_unpacklo_epi32(v[6], v[7]); + v1[7] = _mm512_unpackhi_epi32(v[6], v[7]); + v1[8] = _mm512_unpacklo_epi32(v[8], v[9]); + v1[9] = _mm512_unpackhi_epi32(v[8], v[9]); + v1[10] = _mm512_unpacklo_epi32(v[10], v[11]); + v1[11] = _mm512_unpackhi_epi32(v[10], v[11]); + v1[12] = _mm512_unpacklo_epi32(v[12], v[13]); + v1[13] = _mm512_unpackhi_epi32(v[12], v[13]); + v1[14] = _mm512_unpacklo_epi32(v[14], v[15]); + v1[15] = _mm512_unpackhi_epi32(v[14], v[15]); + + v[0] = _mm512_unpacklo_epi64(v1[0], v1[2]); + v[1] = _mm512_unpackhi_epi64(v1[0], v1[2]); + v[2] = _mm512_unpacklo_epi64(v1[1], v1[3]); + v[3] = _mm512_unpackhi_epi64(v1[1], v1[3]); + v[4] = _mm512_unpacklo_epi64(v1[4], v1[6]); + v[5] = _mm512_unpackhi_epi64(v1[4], v1[6]); + v[6] = _mm512_unpacklo_epi64(v1[5], v1[7]); + v[7] = _mm512_unpackhi_epi64(v1[5], v1[7]); + v[8] = _mm512_unpacklo_epi64(v1[8], v1[10]); + v[9] = _mm512_unpackhi_epi64(v1[8], v1[10]); + v[10] = _mm512_unpacklo_epi64(v1[9], v1[11]); + v[11] = _mm512_unpackhi_epi64(v1[9], v1[11]); + v[12] = _mm512_unpacklo_epi64(v1[12], v1[14]); + v[13] = _mm512_unpackhi_epi64(v1[12], v1[14]); + v[14] = _mm512_unpacklo_epi64(v1[13], v1[15]); + v[15] = _mm512_unpackhi_epi64(v1[13], v1[15]); + + v1[0] = _mm512_shuffle_i32x4(v[0], v[4], 0x88); + v1[1] = _mm512_shuffle_i32x4(v[1], v[5], 0x88); + v1[2] = _mm512_shuffle_i32x4(v[2], v[6], 0x88); + v1[3] = _mm512_shuffle_i32x4(v[3], v[7], 0x88); + v1[4] = _mm512_shuffle_i32x4(v[0], v[4], 0xdd); + v1[5] = _mm512_shuffle_i32x4(v[1], v[5], 0xdd); + v1[6] = _mm512_shuffle_i32x4(v[2], v[6], 0xdd); + v1[7] = _mm512_shuffle_i32x4(v[3], v[7], 0xdd); + v1[8] = _mm512_shuffle_i32x4(v[8], v[12], 0x88); + v1[9] = _mm512_shuffle_i32x4(v[9], v[13], 0x88); + v1[10] = _mm512_shuffle_i32x4(v[10], v[14], 0x88); + v1[11] = _mm512_shuffle_i32x4(v[11], v[15], 0x88); + v1[12] = _mm512_shuffle_i32x4(v[8], v[12], 0xdd); + v1[13] = _mm512_shuffle_i32x4(v[9], v[13], 0xdd); + v1[14] = _mm512_shuffle_i32x4(v[10], v[14], 0xdd); + v1[15] = _mm512_shuffle_i32x4(v[11], v[15], 0xdd); + + v[0] = _mm512_shuffle_i32x4(v1[0], v1[8], 0x88); + v[1] = _mm512_shuffle_i32x4(v1[1], v1[9], 0x88); + v[2] = _mm512_shuffle_i32x4(v1[2], v1[10], 0x88); + v[3] = _mm512_shuffle_i32x4(v1[3], v1[11], 0x88); + v[4] = _mm512_shuffle_i32x4(v1[4], v1[12], 0x88); + v[5] = _mm512_shuffle_i32x4(v1[5], v1[13], 0x88); + v[6] = _mm512_shuffle_i32x4(v1[6], v1[14], 0x88); + v[7] = _mm512_shuffle_i32x4(v1[7], v1[15], 0x88); + v[8] = _mm512_shuffle_i32x4(v1[0], v1[8], 0xdd); + v[9] = _mm512_shuffle_i32x4(v1[1], v1[9], 0xdd); + v[10] = _mm512_shuffle_i32x4(v1[2], v1[10], 0xdd); + v[11] = _mm512_shuffle_i32x4(v1[3], v1[11], 0xdd); + v[12] = _mm512_shuffle_i32x4(v1[4], v1[12], 0xdd); + v[13] = _mm512_shuffle_i32x4(v1[5], v1[13], 0xdd); + v[14] = _mm512_shuffle_i32x4(v1[6], v1[14], 0xdd); + v[15] = _mm512_shuffle_i32x4(v1[7], v1[15], 0xdd); +} + +/* + Transpose 16x16 32-bit elements + Note that v must be 64 byte aligned +*/ +inline void transpose_16x16_32bit(__m512i *v, size_t stride) { + assert(reinterpret_cast(v) % 64 == 0 && "v must be 64 aligned"); + + auto stride_v = [=](int i) { return offset_pointer(v, i * stride); }; + __m512i v1[16]; + + v1[0] = _mm512_unpacklo_epi32(*stride_v(0), *stride_v(1)); + v1[1] = _mm512_unpackhi_epi32(*stride_v(0), *stride_v(1)); + v1[2] = _mm512_unpacklo_epi32(*stride_v(2), *stride_v(3)); + v1[3] = _mm512_unpackhi_epi32(*stride_v(2), *stride_v(3)); + v1[4] = _mm512_unpacklo_epi32(*stride_v(4), *stride_v(5)); + v1[5] = _mm512_unpackhi_epi32(*stride_v(4), *stride_v(5)); + v1[6] = _mm512_unpacklo_epi32(*stride_v(6), *stride_v(7)); + v1[7] = _mm512_unpackhi_epi32(*stride_v(6), *stride_v(7)); + v1[8] = _mm512_unpacklo_epi32(*stride_v(8), *stride_v(9)); + v1[9] = _mm512_unpackhi_epi32(*stride_v(8), *stride_v(9)); + v1[10] = _mm512_unpacklo_epi32(*stride_v(10), *stride_v(11)); + v1[11] = _mm512_unpackhi_epi32(*stride_v(10), *stride_v(11)); + v1[12] = _mm512_unpacklo_epi32(*stride_v(12), *stride_v(13)); + v1[13] = _mm512_unpackhi_epi32(*stride_v(12), *stride_v(13)); + v1[14] = _mm512_unpacklo_epi32(*stride_v(14), *stride_v(15)); + v1[15] = _mm512_unpackhi_epi32(*stride_v(14), *stride_v(15)); + + *stride_v(0) = _mm512_unpacklo_epi64(v1[0], v1[2]); + *stride_v(1) = _mm512_unpackhi_epi64(v1[0], v1[2]); + *stride_v(2) = _mm512_unpacklo_epi64(v1[1], v1[3]); + *stride_v(3) = _mm512_unpackhi_epi64(v1[1], v1[3]); + *stride_v(4) = _mm512_unpacklo_epi64(v1[4], v1[6]); + *stride_v(5) = _mm512_unpackhi_epi64(v1[4], v1[6]); + *stride_v(6) = _mm512_unpacklo_epi64(v1[5], v1[7]); + *stride_v(7) = _mm512_unpackhi_epi64(v1[5], v1[7]); + *stride_v(8) = _mm512_unpacklo_epi64(v1[8], v1[10]); + *stride_v(9) = _mm512_unpackhi_epi64(v1[8], v1[10]); + *stride_v(10) = _mm512_unpacklo_epi64(v1[9], v1[11]); + *stride_v(11) = _mm512_unpackhi_epi64(v1[9], v1[11]); + *stride_v(12) = _mm512_unpacklo_epi64(v1[12], v1[14]); + *stride_v(13) = _mm512_unpackhi_epi64(v1[12], v1[14]); + *stride_v(14) = _mm512_unpacklo_epi64(v1[13], v1[15]); + *stride_v(15) = _mm512_unpackhi_epi64(v1[13], v1[15]); + + v1[0] = _mm512_shuffle_i32x4(*stride_v(0), *stride_v(4), 0x88); + v1[1] = _mm512_shuffle_i32x4(*stride_v(1), *stride_v(5), 0x88); + v1[2] = _mm512_shuffle_i32x4(*stride_v(2), *stride_v(6), 0x88); + v1[3] = _mm512_shuffle_i32x4(*stride_v(3), *stride_v(7), 0x88); + v1[4] = _mm512_shuffle_i32x4(*stride_v(0), *stride_v(4), 0xdd); + v1[5] = _mm512_shuffle_i32x4(*stride_v(1), *stride_v(5), 0xdd); + v1[6] = _mm512_shuffle_i32x4(*stride_v(2), *stride_v(6), 0xdd); + v1[7] = _mm512_shuffle_i32x4(*stride_v(3), *stride_v(7), 0xdd); + v1[8] = _mm512_shuffle_i32x4(*stride_v(8), *stride_v(12), 0x88); + v1[9] = _mm512_shuffle_i32x4(*stride_v(9), *stride_v(13), 0x88); + v1[10] = _mm512_shuffle_i32x4(*stride_v(10), *stride_v(14), 0x88); + v1[11] = _mm512_shuffle_i32x4(*stride_v(11), *stride_v(15), 0x88); + v1[12] = _mm512_shuffle_i32x4(*stride_v(8), *stride_v(12), 0xdd); + v1[13] = _mm512_shuffle_i32x4(*stride_v(9), *stride_v(13), 0xdd); + v1[14] = _mm512_shuffle_i32x4(*stride_v(10), *stride_v(14), 0xdd); + v1[15] = _mm512_shuffle_i32x4(*stride_v(11), *stride_v(15), 0xdd); + + *stride_v(0) = _mm512_shuffle_i32x4(v1[0], v1[8], 0x88); + *stride_v(1) = _mm512_shuffle_i32x4(v1[1], v1[9], 0x88); + *stride_v(2) = _mm512_shuffle_i32x4(v1[2], v1[10], 0x88); + *stride_v(3) = _mm512_shuffle_i32x4(v1[3], v1[11], 0x88); + *stride_v(4) = _mm512_shuffle_i32x4(v1[4], v1[12], 0x88); + *stride_v(5) = _mm512_shuffle_i32x4(v1[5], v1[13], 0x88); + *stride_v(6) = _mm512_shuffle_i32x4(v1[6], v1[14], 0x88); + *stride_v(7) = _mm512_shuffle_i32x4(v1[7], v1[15], 0x88); + *stride_v(8) = _mm512_shuffle_i32x4(v1[0], v1[8], 0xdd); + *stride_v(9) = _mm512_shuffle_i32x4(v1[1], v1[9], 0xdd); + *stride_v(10) = _mm512_shuffle_i32x4(v1[2], v1[10], 0xdd); + *stride_v(11) = _mm512_shuffle_i32x4(v1[3], v1[11], 0xdd); + *stride_v(12) = _mm512_shuffle_i32x4(v1[4], v1[12], 0xdd); + *stride_v(13) = _mm512_shuffle_i32x4(v1[5], v1[13], 0xdd); + *stride_v(14) = _mm512_shuffle_i32x4(v1[6], v1[14], 0xdd); + *stride_v(15) = _mm512_shuffle_i32x4(v1[7], v1[15], 0xdd); +} + +struct GemmKernel224BF { + using dt = ggml_bf16_t; + using output_t = float; + static const int TILE_M = 16; + static const int TILE_K = 32; + static const int TILE_N = 16; + static const int VNNI_BLK = 2; + + static const int M_STEP = TILE_M * 2; + static const int N_STEP = TILE_N * 2; + static const int K_STEP = TILE_K; + + static inline const int N_BLOCK = 256; + static inline const int K_BLOCK = 1792; + + static int recommended_nth(int n) { return (n + N_BLOCK - 1) / N_BLOCK; } + + static std::pair split_range_n(int n, int ith, int nth) { + int n_start = N_BLOCK * ith; + int n_end = std::min(n, N_BLOCK * (ith + 1)); + return {n_start, n_end}; + } + + static void config() { + enable_amx(); + TileConfig tile_config; + + // size is 16 x 32 + for (int i = 0; i < 2; i++) + tile_config.set_row_col(i, TILE_M, TILE_K * sizeof(dt)); + + // size is 16 x 32 + for (int i = 2; i < 4; i++) + tile_config.set_row_col(i, TILE_K / VNNI_BLK, TILE_N * VNNI_BLK * sizeof(dt)); + + // size is 16 x 16 + for (int i = 4; i < 8; i++) + tile_config.set_row_col(i, TILE_M, TILE_N * sizeof(output_t)); + + tile_config.set_config(); + } + + static void load_a(dt *a, size_t lda) { + _tile_loadd(0, a, lda); + _tile_loadd(1, offset_pointer(a, lda * TILE_M), lda); + } + + static void load_b(dt *b, size_t ldb) { + _tile_loadd(2, b, ldb); + _tile_loadd(3, offset_pointer(b, ldb * TILE_N), ldb); + } + + static void clean_c() { + _tile_zero(4); + _tile_zero(5); + _tile_zero(6); + _tile_zero(7); + } + + static void load_c(output_t *c, size_t ldc) { + _tile_loadd(4, c, ldc); + _tile_loadd(5, offset_pointer(c, TILE_N * sizeof(output_t)), ldc); + _tile_loadd(6, offset_pointer(c, ldc * TILE_M), ldc); + _tile_loadd(7, offset_pointer(c, ldc * TILE_M + TILE_N * sizeof(output_t)), ldc); + } + + static void store_c(output_t *c, size_t ldc) { + _tile_stored(4, c, ldc); + _tile_stored(5, offset_pointer(c, TILE_N * sizeof(output_t)), ldc); + _tile_stored(6, offset_pointer(c, ldc * TILE_M), ldc); + _tile_stored(7, offset_pointer(c, ldc * TILE_M + TILE_N * sizeof(output_t)), ldc); + } + + static void run_tile() { + _tile_dpbf16ps(4, 0, 2); + _tile_dpbf16ps(5, 0, 3); + _tile_dpbf16ps(6, 1, 2); + _tile_dpbf16ps(7, 1, 3); + } + + struct BufferA { + ggml_bf16_t *a; + int max_m, k; + + static size_t required_size(int max_m, int k) { return max_m * k * sizeof(ggml_bf16_t); } + + BufferA(int max_m, int k, void *ptr) : max_m(max_m), k(k) { + assert(reinterpret_cast(ptr) % 64 == 0); + assert(max_m % M_STEP == 0); + assert(k % K_STEP == 0); + a = reinterpret_cast(ptr); + } + + void from_mat(int m, ggml_bf16_t *src, int ith, int nth) { + assert(m <= max_m); + assert(ith == 0 && nth == 1); + int m_block_size = (m + M_STEP - 1) / M_STEP * M_STEP; + for (int m_begin = 0; m_begin < m; m_begin += M_STEP) { + for (int k_block_begin = 0; k_block_begin < k; k_block_begin += K_BLOCK) { + int k_block_size = std::min(K_BLOCK, k - k_block_begin); + for (int k_begin = 0; k_begin < k_block_size; k_begin += K_STEP) { + for (int i = 0; i < M_STEP && m_begin + i < m; i++) { + __m512i *s = (__m512i *)(src + (m_begin + i) * k + k_block_begin + k_begin); + __m512i *d = (__m512i *)(a + k_block_begin * m_block_size + m_begin * k_block_size + k_begin * M_STEP + + i * K_STEP); + avx512_copy_32xbf16(s, d); + } + } + } + } + } + + ggml_bf16_t *get_submat(int m, int k, int m_begin, int k_begin) { + int m_block_size = (m + M_STEP - 1) / M_STEP * M_STEP; + int k_block_begin = k_begin / K_BLOCK * K_BLOCK; + k_begin -= k_block_begin; + int k_block_size = std::min(K_BLOCK, k - k_block_begin); + return a + k_block_begin * m_block_size + m_begin * k_block_size + k_begin * M_STEP; + } + }; + + struct BufferB { + ggml_bf16_t *b; + int n, k; + + static size_t required_size(int n, int k) { return n * k * sizeof(ggml_bf16_t); } + + BufferB(int n, int k, void *ptr) : n(n), k(k) { + assert(reinterpret_cast(ptr) % 64 == 0); + assert(n % N_STEP == 0); + assert(k % K_STEP == 0); + b = reinterpret_cast(ptr); + } + + void from_mat(ggml_bf16_t *src, int ith, int nth) { + auto [n_start, n_end] = split_range_n(n, ith, nth); + int n_block_begin = n_start; + int n_block_size = n_end - n_block_begin; + for (int n_begin = 0; n_begin < n_block_size; n_begin += N_STEP) { + for (int k_block_begin = 0; k_block_begin < k; k_block_begin += K_BLOCK) { + int k_block_size = std::min(K_BLOCK, k - k_block_begin); + for (int k_begin = 0; k_begin < k_block_size; k_begin += K_STEP) { + for (int i = 0; i < N_STEP; i++) { + __m512i *s = (__m512i *)(src + (n_block_begin + n_begin + i) * k + k_block_begin + k_begin); + __m512i *d = (__m512i *)(b + n_block_begin * k + k_block_begin * n_block_size + n_begin * k_block_size + + k_begin * N_STEP + i * K_STEP); + avx512_copy_32xbf16(s, d); + } + transpose_16x16_32bit((__m512i *)(b + n_block_begin * k + k_block_begin * n_block_size + + n_begin * k_block_size + k_begin * N_STEP)); + transpose_16x16_32bit((__m512i *)(b + n_block_begin * k + k_block_begin * n_block_size + + n_begin * k_block_size + k_begin * N_STEP + TILE_N * K_STEP)); + } + } + } + } + + ggml_bf16_t *get_submat(int n, int k, int n_begin, int k_begin) { + int n_block_begin = n_begin / N_BLOCK * N_BLOCK; + n_begin -= n_block_begin; + int n_block_size = std::min(N_BLOCK, n - n_block_begin); + int k_block_begin = k_begin / K_BLOCK * K_BLOCK; + k_begin -= k_block_begin; + int k_block_size = std::min(K_BLOCK, k - k_block_begin); + return b + n_block_begin * k + k_block_begin * n_block_size + n_begin * k_block_size + k_begin * N_STEP; + } + }; + + struct BufferC { + float *c; + int max_m, n; + + static size_t required_size(int max_m, int n) { return max_m * n * sizeof(float); } + + BufferC(int max_m, int n, void *ptr) : max_m(max_m), n(n) { + assert(reinterpret_cast(ptr) % 64 == 0); + assert(max_m % M_STEP == 0); + assert(n % N_STEP == 0); + c = reinterpret_cast(ptr); + } + + void to_mat(int m, ggml_bf16_t *dst, int ith, int nth) { + assert(m <= max_m); + auto [n_start, n_end] = split_range_n(n, ith, nth); + int m_block_size = (m + M_STEP - 1) / M_STEP * M_STEP; + int n_block_begin = n_start; + int n_block_size = n_end - n_block_begin; + for (int m_begin = 0; m_begin < m; m_begin += M_STEP) { + for (int n_begin = 0; n_begin < n_block_size; n_begin += N_STEP) { + for (int i = 0; i < M_STEP && m_begin + i < m; i++) { + __m512 *x0 = + (__m512 *)(c + m_block_size * n_block_begin + m_begin * n_block_size + n_begin * M_STEP + i * N_STEP); + __m512 *x1 = (__m512 *)(c + m_block_size * n_block_begin + m_begin * n_block_size + n_begin * M_STEP + + i * N_STEP + 16); + avx512_32xfp32_to_32xbf16(x0, x1, (__m512i *)(dst + (m_begin + i) * n + n_block_begin + n_begin)); + } + } + } + } + + float *get_submat(int m, int n, int m_begin, int n_begin) { + int m_block_size = (m + M_STEP - 1) / M_STEP * M_STEP; + int n_block_begin = n_begin / N_BLOCK * N_BLOCK; + int n_block_size = std::min(N_BLOCK, n - n_block_begin); + n_begin -= n_block_begin; + return c + m_block_size * n_block_begin + m_begin * n_block_size + n_begin * M_STEP; + } + }; +}; + +struct GemmKernel224Int8 { + using dt = int8_t; + using output_t = int32_t; + static const int TILE_M = 16; + static const int TILE_K = 64; + static const int TILE_N = 16; + static const int VNNI_BLK = 4; + + static const int M_STEP = TILE_M * 2; + static const int N_STEP = TILE_N * 2; + static const int K_STEP = TILE_K; + + static inline const int N_BLOCK = 256; + static inline const int K_BLOCK = 3584; + + static int recommended_nth(int n) { return (n + N_BLOCK - 1) / N_BLOCK; } + + static std::pair split_range_n(int n, int ith, int nth) { + int n_start = N_BLOCK * ith; + int n_end = std::min(n, N_BLOCK * (ith + 1)); + return {n_start, n_end}; + } + + static void config() { + enable_amx(); + TileConfig tile_config; + + // size is 16 x 64 + for (int i = 0; i < 2; i++) + tile_config.set_row_col(i, TILE_M, TILE_K * sizeof(dt)); + + // size is 16 x 64 + for (int i = 2; i < 4; i++) + tile_config.set_row_col(i, TILE_K / VNNI_BLK, TILE_N * VNNI_BLK * sizeof(dt)); + + // size is 16 x 16 + for (int i = 4; i < 8; i++) + tile_config.set_row_col(i, TILE_M, TILE_N * sizeof(output_t)); + + tile_config.set_config(); + } + + static void load_a(dt *a, size_t lda) { + _tile_loadd(0, a, lda); + _tile_loadd(1, offset_pointer(a, lda * TILE_M), lda); + } + + static void load_b(dt *b, size_t ldb) { + _tile_loadd(2, b, ldb); + _tile_loadd(3, offset_pointer(b, ldb * TILE_N), ldb); + } + + static void clean_c() { + _tile_zero(4); + _tile_zero(5); + _tile_zero(6); + _tile_zero(7); + } + + static void load_c(output_t *c, size_t ldc) { + _tile_loadd(4, c, ldc); + _tile_loadd(5, offset_pointer(c, TILE_N * sizeof(output_t)), ldc); + _tile_loadd(6, offset_pointer(c, ldc * TILE_M), ldc); + _tile_loadd(7, offset_pointer(c, ldc * TILE_M + TILE_N * sizeof(output_t)), ldc); + } + + static void store_c(output_t *c, size_t ldc) { + _tile_stored(4, c, ldc); + _tile_stored(5, offset_pointer(c, TILE_N * sizeof(output_t)), ldc); + _tile_stored(6, offset_pointer(c, ldc * TILE_M), ldc); + _tile_stored(7, offset_pointer(c, ldc * TILE_M + TILE_N * sizeof(output_t)), ldc); + } + + static void run_tile() { + _tile_dpbssd(4, 0, 2); + _tile_dpbssd(5, 0, 3); + _tile_dpbssd(6, 1, 2); + _tile_dpbssd(7, 1, 3); + } + + struct BufferA { + int8_t *a; + float *d; + int max_m, k; + + static size_t required_size(int max_m, int k) { return max_m * k * sizeof(int8_t) + max_m * sizeof(float); } + + BufferA(int max_m, int k, void *ptr) : max_m(max_m), k(k) { + assert(reinterpret_cast(ptr) % 64 == 0); + assert(max_m % M_STEP == 0); + assert(k % K_STEP == 0); + a = reinterpret_cast(ptr); + d = reinterpret_cast(a + max_m * k); + } + + void from_mat(int m, ggml_bf16_t *src, int ith, int nth) { + assert(m <= max_m); + assert(ith == 0 && nth == 1); + for (int m_begin = 0; m_begin < m; m_begin += M_STEP) { + for (int i = 0; i < M_STEP && m_begin + i < m; i++) { + float amax = 0.0f; + for (int j = 0; j < k; j += 32) { + __m512 f0, f1; + avx512_32xbf16_to_32xfp32((__m512i *)(src + (m_begin + i) * k + j), &f0, &f1); + amax = MAX(amax, _mm512_reduce_max_ps(_mm512_abs_ps(f0))); + amax = MAX(amax, _mm512_reduce_max_ps(_mm512_abs_ps(f1))); + } + d[m_begin + i] = amax / ((1 << 7) - 1); + } + } + int m_block_size = (m + M_STEP - 1) / M_STEP * M_STEP; + for (int m_begin = 0; m_begin < m; m_begin += M_STEP) { + for (int k_block_begin = 0; k_block_begin < k; k_block_begin += K_BLOCK) { + int k_block_size = std::min(K_BLOCK, k - k_block_begin); + for (int k_begin = 0; k_begin < k_block_size; k_begin += K_STEP) { + for (int i = 0; i < M_STEP && m_begin + i < m; i++) { + __m512 id = _mm512_set1_ps(d[m_begin + i] ? 1.0f / d[m_begin + i] : 0.0f); + int8_t *dst = a + k_block_begin * m_block_size + m_begin * k_block_size + k_begin * M_STEP + i * K_STEP; + __m512 f0, f1, f2, f3; + avx512_32xbf16_to_32xfp32((__m512i *)(src + (m_begin + i) * k + k_block_begin + k_begin), &f0, &f1); + avx512_32xbf16_to_32xfp32((__m512i *)(src + (m_begin + i) * k + k_block_begin + k_begin) + 1, &f2, &f3); + __m512i i0 = _mm512_cvtps_epi32(_mm512_mul_ps(f0, id)); + __m512i i1 = _mm512_cvtps_epi32(_mm512_mul_ps(f1, id)); + __m512i i2 = _mm512_cvtps_epi32(_mm512_mul_ps(f2, id)); + __m512i i3 = _mm512_cvtps_epi32(_mm512_mul_ps(f3, id)); + __m128i s0 = _mm512_cvtsepi32_epi8(i0); + __m128i s1 = _mm512_cvtsepi32_epi8(i1); + __m128i s2 = _mm512_cvtsepi32_epi8(i2); + __m128i s3 = _mm512_cvtsepi32_epi8(i3); + _mm_storeu_si128((__m128i *)dst, s0); + _mm_storeu_si128((__m128i *)(dst + 16), s1); + _mm_storeu_si128((__m128i *)(dst + 32), s2); + _mm_storeu_si128((__m128i *)(dst + 48), s3); + } + } + } + } + } + + int8_t *get_submat(int m, int k, int m_begin, int k_begin) { + int m_block_size = (m + M_STEP - 1) / M_STEP * M_STEP; + int k_block_begin = k_begin / K_BLOCK * K_BLOCK; + k_begin -= k_block_begin; + int k_block_size = std::min(K_BLOCK, k - k_block_begin); + return a + k_block_begin * m_block_size + m_begin * k_block_size + k_begin * M_STEP; + } + + float *get_scale(int m, int m_begin) { return d + m_begin; } + }; + + struct BufferB { + int8_t *b; + float *d; + int n, k; + + static size_t required_size(int n, int k) { return n * k * sizeof(int8_t) + n * sizeof(float); } + + BufferB(int n, int k, void *ptr) : n(n), k(k) { + assert(reinterpret_cast(ptr) % 64 == 0); + assert(n % N_STEP == 0); + assert(k % K_STEP == 0); + b = reinterpret_cast(ptr); + d = reinterpret_cast(b + n * k); + } + + void from_mat(ggml_bf16_t *src, int ith, int nth) { + auto [n_start, n_end] = split_range_n(n, ith, nth); + int n_block_begin = n_start; + int n_block_size = n_end - n_block_begin; + for (int n_begin = 0; n_begin < n_block_size; n_begin += N_STEP) { + for (int i = 0; i < N_STEP; i++) { + float amax = 0.0f; + for (int j = 0; j < k; j += 32) { + __m512 f0, f1; + avx512_32xbf16_to_32xfp32((__m512i *)(src + (n_block_begin + n_begin + i) * k + j), &f0, &f1); + amax = MAX(amax, _mm512_reduce_max_ps(_mm512_abs_ps(f0))); + amax = MAX(amax, _mm512_reduce_max_ps(_mm512_abs_ps(f1))); + } + d[n_block_begin + n_begin + i] = amax / ((1 << 7) - 1); + } + } + for (int n_begin = 0; n_begin < n_block_size; n_begin += N_STEP) { + for (int k_block_begin = 0; k_block_begin < k; k_block_begin += K_BLOCK) { + int k_block_size = std::min(K_BLOCK, k - k_block_begin); + for (int k_begin = 0; k_begin < k_block_size; k_begin += K_STEP) { + for (int i = 0; i < N_STEP; i++) { + __m512 id = _mm512_set1_ps(d[n_block_begin + n_begin + i] ? 1.0f / d[n_block_begin + n_begin + i] : 0.0f); + int8_t *dst = b + n_block_begin * k + k_block_begin * n_block_size + n_begin * k_block_size + + k_begin * N_STEP + i * K_STEP; + __m512 f0, f1, f2, f3; + avx512_32xbf16_to_32xfp32((__m512i *)(src + (n_block_begin + n_begin + i) * k + k_block_begin + k_begin), + &f0, &f1); + avx512_32xbf16_to_32xfp32( + (__m512i *)(src + (n_block_begin + n_begin + i) * k + k_block_begin + k_begin) + 1, &f2, &f3); + __m512i i0 = _mm512_cvtps_epi32(_mm512_mul_ps(f0, id)); + __m512i i1 = _mm512_cvtps_epi32(_mm512_mul_ps(f1, id)); + __m512i i2 = _mm512_cvtps_epi32(_mm512_mul_ps(f2, id)); + __m512i i3 = _mm512_cvtps_epi32(_mm512_mul_ps(f3, id)); + __m128i s0 = _mm512_cvtsepi32_epi8(i0); + __m128i s1 = _mm512_cvtsepi32_epi8(i1); + __m128i s2 = _mm512_cvtsepi32_epi8(i2); + __m128i s3 = _mm512_cvtsepi32_epi8(i3); + _mm_storeu_si128((__m128i *)dst, s0); + _mm_storeu_si128((__m128i *)(dst + 16), s1); + _mm_storeu_si128((__m128i *)(dst + 32), s2); + _mm_storeu_si128((__m128i *)(dst + 48), s3); + } + transpose_16x16_32bit((__m512i *)(b + n_block_begin * k + k_block_begin * n_block_size + + n_begin * k_block_size + k_begin * N_STEP)); + transpose_16x16_32bit((__m512i *)(b + n_block_begin * k + k_block_begin * n_block_size + + n_begin * k_block_size + k_begin * N_STEP + TILE_N * K_STEP)); + } + } + } + } + + int8_t *get_submat(int n, int k, int n_begin, int k_begin) { + int n_block_begin = n_begin / N_BLOCK * N_BLOCK; + n_begin -= n_block_begin; + int n_block_size = std::min(N_BLOCK, n - n_block_begin); + int k_block_begin = k_begin / K_BLOCK * K_BLOCK; + k_begin -= k_block_begin; + int k_block_size = std::min(K_BLOCK, k - k_block_begin); + return b + n_block_begin * k + k_block_begin * n_block_size + n_begin * k_block_size + k_begin * N_STEP; + } + + float *get_scale(int n, int n_begin) { return d + n_begin; } + }; + + struct BufferC { + float *c; + int max_m, n; + + static size_t required_size(int max_m, int n) { return max_m * n * sizeof(float); } + + BufferC(int max_m, int n, void *ptr) : max_m(max_m), n(n) { + assert(reinterpret_cast(ptr) % 64 == 0); + assert(max_m % M_STEP == 0); + assert(n % N_STEP == 0); + c = reinterpret_cast(ptr); + } + + void to_mat(int m, ggml_bf16_t *dst, int ith, int nth) { + assert(m <= max_m); + auto [n_start, n_end] = split_range_n(n, ith, nth); + int m_block_size = (m + M_STEP - 1) / M_STEP * M_STEP; + int n_block_begin = n_start; + int n_block_size = n_end - n_block_begin; + for (int m_begin = 0; m_begin < m; m_begin += M_STEP) { + for (int n_begin = 0; n_begin < n_block_size; n_begin += N_STEP) { + for (int i = 0; i < M_STEP && m_begin + i < m; i++) { + __m512 *x0 = + (__m512 *)(c + m_block_size * n_block_begin + m_begin * n_block_size + n_begin * M_STEP + i * N_STEP); + __m512 *x1 = (__m512 *)(c + m_block_size * n_block_begin + m_begin * n_block_size + n_begin * M_STEP + + i * N_STEP + 16); + avx512_32xfp32_to_32xbf16(x0, x1, (__m512i *)(dst + (m_begin + i) * n + n_block_begin + n_begin)); + } + } + } + } + + float *get_submat(int m, int n, int m_begin, int n_begin) { + int m_block_size = (m + M_STEP - 1) / M_STEP * M_STEP; + int n_block_begin = n_begin / N_BLOCK * N_BLOCK; + int n_block_size = std::min(N_BLOCK, n - n_block_begin); + n_begin -= n_block_begin; + return c + m_block_size * n_block_begin + m_begin * n_block_size + n_begin * M_STEP; + } + }; +}; + +inline void mat_mul(int m, int n, int k, std::shared_ptr ba, + std::shared_ptr bb, std::shared_ptr bc, int ith, + int nth, bool use_amx) { + using K = GemmKernel224BF; + assert(n % K::N_STEP == 0); + assert(k % K::K_STEP == 0); + + auto [n_start, n_end] = K::split_range_n(n, ith, nth); + + for (int k_block_begin = 0; k_block_begin < k; k_block_begin += K::K_BLOCK) { + for (int m_begin = 0; m_begin < m; m_begin += K::M_STEP) { + for (int n_begin = n_start; n_begin < n_end; n_begin += K::N_STEP) { + + float *c = bc->get_submat(m, n, m_begin, n_begin); + if (!use_amx) { + __m512 *c512 = (__m512 *)c; + if (k_block_begin == 0) { + for (int m_i = 0; m_i < m && m_i < K::M_STEP; m_i++) { + c512[m_i * 2] = _mm512_setzero_ps(); + c512[m_i * 2 + 1] = _mm512_setzero_ps(); + } + } + + for (int k_begin = 0; k_begin < K::K_BLOCK && k_block_begin + k_begin < k; k_begin += K::K_STEP) { + int32_t *a32 = (int32_t *)ba->get_submat(m, k, m_begin, k_block_begin + k_begin); + __m512bh *b512 = (__m512bh *)bb->get_submat(n, k, n_begin, k_block_begin + k_begin); + for (int m_i = 0; m_i < m && m_i < K::M_STEP; m_i++) { + for (int k_i = 0; k_i < 16; k_i++) { + __m512bh ma = (__m512bh)_mm512_set1_epi32(a32[m_i * 16 + k_i]); + for (int n_i = 0; n_i < 2; n_i++) { + c512[m_i * 2 + n_i] = _mm512_dpbf16_ps(c512[m_i * 2 + n_i], ma, b512[n_i * 16 + k_i]); + } + } + } + } + + } else { + if (k_block_begin == 0) { + K::clean_c(); + } else { + K::load_c(c, K::N_STEP * sizeof(float)); + } + for (int k_begin = 0; k_begin < K::K_BLOCK && k_block_begin + k_begin < k; k_begin += K::K_STEP) { + K::load_a(ba->get_submat(m, k, m_begin, k_block_begin + k_begin), K::K_STEP * sizeof(ggml_bf16_t)); + K::load_b(bb->get_submat(n, k, n_begin, k_block_begin + k_begin), K::K_STEP * sizeof(ggml_bf16_t)); + K::run_tile(); + } + K::store_c(c, K::N_STEP * sizeof(float)); + } + } + } + } +} + +inline __m512i _mm512_dpbssd_epi32(__m512i src, __m512i a, __m512i b) { + __m256i a_lo = _mm512_extracti64x4_epi64(a, 0); + __m256i a_hi = _mm512_extracti64x4_epi64(a, 1); + __m256i b_lo = _mm512_extracti64x4_epi64(b, 0); + __m256i b_hi = _mm512_extracti64x4_epi64(b, 1); + + b_lo = _mm256_sign_epi8(b_lo, a_lo); + b_hi = _mm256_sign_epi8(b_hi, a_hi); + + b = _mm512_inserti64x4(b, b_lo, 0); + b = _mm512_inserti64x4(b, b_hi, 1); + + a = _mm512_abs_epi8(a); + + return _mm512_dpbusd_epi32(src, a, b); +} + +inline void mat_mul(int m, int n, int k, std::shared_ptr ba, + std::shared_ptr bb, std::shared_ptr bc, + int ith, int nth, bool use_amx) { + using K = GemmKernel224Int8; + assert(n % K::N_STEP == 0); + assert(k % K::K_STEP == 0); + + auto [n_start, n_end] = K::split_range_n(n, ith, nth); + + for (int k_block_begin = 0; k_block_begin < k; k_block_begin += K::K_BLOCK) { + for (int m_begin = 0; m_begin < m; m_begin += K::M_STEP) { + for (int n_begin = n_start; n_begin < n_end; n_begin += K::N_STEP) { + float *c = bc->get_submat(m, n, m_begin, n_begin); + + if (!use_amx) { + __m512i *c512 = (__m512i *)c; + if (k_block_begin == 0) { + for (int m_i = 0; m_i < m && m_i < K::M_STEP; m_i++) { + c512[m_i * 2] = _mm512_setzero_si512(); + c512[m_i * 2 + 1] = _mm512_setzero_si512(); + } + } + + for (int k_begin = 0; k_begin < K::K_BLOCK && k_block_begin + k_begin < k; k_begin += K::K_STEP) { + static_assert(K::K_STEP * sizeof(int8_t) == sizeof(__m512i)); + static_assert(K::N_STEP / K::TILE_N == 2, "Must be lke this"); + + int32_t *a32 = (int32_t *)ba->get_submat(m, k, m_begin, k_block_begin + k_begin); + __m512i *b512 = (__m512i *)bb->get_submat(n, k, n_begin, k_block_begin + k_begin); + for (int m_i = 0; m_i < m && m_i < K::M_STEP; m_i++) { + for (int k_i = 0; k_i < 16; k_i++) { + __m512i ma = _mm512_set1_epi32(a32[m_i * 16 + k_i]); + for (int n_i = 0; n_i < 2; n_i++) { + c512[m_i * 2 + n_i] = _mm512_dpbssd_epi32(c512[m_i * 2 + n_i], ma, b512[n_i * 16 + k_i]); + } + } + } + } + } else { + if (k_block_begin == 0) { + K::clean_c(); + } else { + K::load_c((int32_t *)c, K::N_STEP * sizeof(int32_t)); + } + for (int k_begin = 0; k_begin < K::K_BLOCK && k_block_begin + k_begin < k; k_begin += K::K_STEP) { + K::load_a(ba->get_submat(m, k, m_begin, k_block_begin + k_begin), K::K_STEP * sizeof(int8_t)); + K::load_b(bb->get_submat(n, k, n_begin, k_block_begin + k_begin), K::K_STEP * sizeof(int8_t)); + K::run_tile(); + } + K::store_c((int32_t *)c, K::N_STEP * sizeof(int32_t)); + } + + if (k_block_begin + K::K_BLOCK >= k) { + int to = m - m_begin; + if (m - m_begin > K::M_STEP) { + to = K::M_STEP; + } + for (int i = 0; i < to; i++) { + __m512 as = _mm512_set1_ps(*ba->get_scale(m, m_begin + i)); + __m512 bs = _mm512_load_ps(bb->get_scale(n, n_begin)); + __m512i now = _mm512_load_si512((__m512i *)(c + i * K::N_STEP)); + __m512 result = _mm512_mul_ps(_mm512_mul_ps(as, bs), _mm512_cvtepi32_ps(now)); + _mm512_store_ps((__m512 *)(c + i * K::N_STEP), result); + bs = _mm512_load_ps(bb->get_scale(n, n_begin) + K::TILE_N); + now = _mm512_load_si512((__m512i *)(c + i * K::N_STEP + K::TILE_N)); + result = _mm512_mul_ps(_mm512_mul_ps(as, bs), _mm512_cvtepi32_ps(now)); + _mm512_store_ps((__m512 *)(c + i * K::N_STEP + K::TILE_N), result); + } + } + } + } + } +} + +} // namespace amx \ No newline at end of file diff --git a/archive/csrc/ktransformers_ext/operators/amx/la/utils.hpp b/archive/csrc/ktransformers_ext/operators/amx/la/utils.hpp new file mode 100644 index 0000000..a95a2e7 --- /dev/null +++ b/archive/csrc/ktransformers_ext/operators/amx/la/utils.hpp @@ -0,0 +1,46 @@ +/** + * @Description : + * @Author : chenht2022 + * @Date : 2025-04-25 18:28:12 + * @Version : 1.0.0 + * @LastEditors : chenht2022 + * @LastEditTime : 2025-04-25 18:28:12 + * @Copyright (c) 2024 by KVCache.AI, All Rights Reserved. + **/ + +#pragma once +#include + + +template +T* offset_pointer(T* ptr, std::size_t byte_offset) { + return reinterpret_cast(reinterpret_cast(ptr) + byte_offset); +} + +template +const T* offset_pointer(const T* ptr, std::size_t byte_offset) { + return reinterpret_cast(reinterpret_cast(ptr) + byte_offset); +} + +template +T* offset_pointer_row_major(T* t, int row, int col, std::size_t ld) { + return offset_pointer(t, row * ld) + col; +} + +template +T* offset_pointer_col_major(T* t, int row, int col, std::size_t ld) { + return offset_pointer(t, col * ld) + row; +} + +static inline void avx512_copy_32xbf16(__m512i* src, __m512i* dst) { + _mm512_storeu_si512(dst, _mm512_loadu_si512(src)); +} + +static inline void avx512_32xfp32_to_32xbf16(__m512* src0, __m512* src1, __m512i* dst) { + _mm512_storeu_si512(dst, __m512i(_mm512_cvtne2ps_pbh(*src1, *src0))); +} + +static inline void avx512_32xbf16_to_32xfp32(__m512i* src, __m512* dst0, __m512* dst1) { + _mm512_storeu_ps(dst0, _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu16_epi32(_mm256_loadu_si256((const __m256i *)(src))), 16))); + _mm512_storeu_ps(dst1, _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu16_epi32(_mm256_loadu_si256((const __m256i *)(src) + 1)), 16))); +} \ No newline at end of file diff --git a/archive/csrc/ktransformers_ext/operators/amx/moe.hpp b/archive/csrc/ktransformers_ext/operators/amx/moe.hpp new file mode 100644 index 0000000..662189f --- /dev/null +++ b/archive/csrc/ktransformers_ext/operators/amx/moe.hpp @@ -0,0 +1,422 @@ +/** + * @Description : + * @Author : chenht2022 + * @Date : 2025-04-25 18:28:12 + * @Version : 1.0.0 + * @LastEditors : chenht2022 + * @LastEditTime : 2025-04-25 18:28:12 + * @Copyright (c) 2024 by KVCache.AI, All Rights Reserved. + **/ +#ifndef CPUINFER_OPERATOR_AMX_MOE_H +#define CPUINFER_OPERATOR_AMX_MOE_H + +#include +#include +#include +#include +#include + +#include "../../cpu_backend/backend.h" +#include "../../cpu_backend/shared_mem_buffer.h" +#include "llama.cpp/ggml-impl.h" +#include "llama.cpp/ggml-quants.h" +#include "llama.cpp/ggml.h" +#include "llamafile/sgemm.h" + +#include "la/amx.hpp" + +#ifdef USE_NUMA +#include +#include +void *numa_alloc_aligned(size_t size, int node, size_t alignment) { + void *ptr = numa_alloc_onnode(size, node); + assert(reinterpret_cast(ptr) % 64 == 0); + return ptr; +} +#endif + +static inline __m512 exp_avx512(__m512 x) { + const __m512 log2e = _mm512_set1_ps(1.44269504089f); + const __m512 c1 = _mm512_set1_ps(0.69314718056f); + + __m512 y = _mm512_mul_ps(x, log2e); + __m512i int_part = _mm512_cvtps_epi32(y); + __m512 frac_part = _mm512_sub_ps(y, _mm512_cvtepi32_ps(int_part)); + + const __m512 poly_1 = _mm512_set1_ps(0.9999999995f); + const __m512 poly_2 = _mm512_set1_ps(0.6931471805f); + const __m512 poly_3 = _mm512_set1_ps(0.2402265069f); + const __m512 poly_4 = _mm512_set1_ps(0.0555041087f); + const __m512 poly_5 = _mm512_set1_ps(0.0096181291f); + const __m512 poly_6 = _mm512_set1_ps(0.0013333558f); + + __m512 frac_exp = _mm512_fmadd_ps( + frac_part, poly_6, + _mm512_fmadd_ps(frac_part, poly_5, + _mm512_fmadd_ps(frac_part, poly_4, + _mm512_fmadd_ps(frac_part, poly_3, _mm512_fmadd_ps(frac_part, poly_2, poly_1))))); + + __m512 two_pow_i = _mm512_scalef_ps(_mm512_set1_ps(1.0f), _mm512_cvtepi32_ps(int_part)); + return _mm512_mul_ps(two_pow_i, frac_exp); +} + +static inline __m512 act_fn(__m512 gate_val, __m512 up_val) { + __m512 neg_gate_val = _mm512_sub_ps(_mm512_setzero_ps(), gate_val); + __m512 exp_neg_gate = exp_avx512(neg_gate_val); + __m512 denom = _mm512_add_ps(_mm512_set1_ps(1.0f), exp_neg_gate); + __m512 act_val = _mm512_div_ps(gate_val, denom); + + return _mm512_mul_ps(act_val, up_val); +} + +static inline __m512 relu_act_fn(__m512 gate_val, __m512 up_val) { + __m512 zero_vec = _mm512_setzero_ps(); + __m512 act_val = _mm512_max_ps(zero_vec, gate_val); + return _mm512_mul_ps(act_val, up_val); +} + +struct AMX_MOEConfig { + int expert_num; + int routed_expert_num; + int hidden_size; + int intermediate_size; + int max_len; + bool use_silu; + void *gate_proj; + void *up_proj; + void *down_proj; + + AMX_MOEConfig() {} + + AMX_MOEConfig(int expert_num, int routed_expert_num, int hidden_size, int intermediate_size, int max_len, bool use_silu, + void *gate_proj, void *up_proj, void *down_proj) + : expert_num(expert_num), routed_expert_num(routed_expert_num), hidden_size(hidden_size), + intermediate_size(intermediate_size), max_len(max_len), use_silu(use_silu), gate_proj(gate_proj), up_proj(up_proj), + down_proj(down_proj) {} +}; + +template class AMX_MOE { +private: + AMX_MOEConfig config_; + void *gate_proj_; // [expert_num * intermediate_size * hidden_size ( /32 if quantized)] + void *up_proj_; // [expert_num * intermediate_size * hidden_size ( /32 if quantized)] + void *down_proj_; // [expert_num * hidden_size * intermediate_size ( /32 if quantized)] + + ggml_bf16_t *m_local_input_; // [routed_expert_num * max_len * hidden_size] + ggml_bf16_t *m_local_gate_output_; // [routed_expert_num * max_len * intermediate_size] + ggml_bf16_t *m_local_up_output_; // [routed_expert_num * max_len * intermediate_size] + ggml_bf16_t *m_local_down_output_; // [routed_expert_num * max_len * hidden_size] + + std::vector> m_local_pos_; // [max_len, routed_expert_num] + std::vector m_local_num_; // [expert_num] + std::vector m_expert_id_map_; // [expert_num] + std::vector m_local_input_ptr_; // [expert_num] + std::vector m_local_gate_output_ptr_; // [expert_num] + std::vector m_local_up_output_ptr_; // [expert_num] + std::vector m_local_down_output_ptr_; // [expert_num] + + std::vector> gate_up_ba_; + std::vector> gate_bc_; + std::vector> up_bc_; + std::vector> down_ba_; + std::vector> down_bc_; + +#ifdef USE_NUMA + std::vector>> gate_bb_numa_; + std::vector>> up_bb_numa_; + std::vector>> down_bb_numa_; +#else + std::vector> gate_bb_; + std::vector> up_bb_; + std::vector> down_bb_; +#endif + +public: + AMX_MOE(AMX_MOEConfig config) { + config_ = config; + gate_proj_ = config_.gate_proj; + up_proj_ = config_.up_proj; + down_proj_ = config_.down_proj; + + std::vector> m_mem_requests; + m_mem_requests.push_back({(void **)&m_local_input_, + sizeof(ggml_bf16_t) * config_.routed_expert_num * config_.max_len * config_.hidden_size}); + m_mem_requests.push_back({(void **)&m_local_gate_output_, sizeof(ggml_bf16_t) * config_.routed_expert_num * + config_.max_len * config_.intermediate_size}); + m_mem_requests.push_back({(void **)&m_local_up_output_, sizeof(ggml_bf16_t) * config_.routed_expert_num * + config_.max_len * config_.intermediate_size}); + m_mem_requests.push_back({(void **)&m_local_down_output_, + sizeof(ggml_bf16_t) * config_.routed_expert_num * config_.max_len * config_.hidden_size}); + std::vector gate_up_ba_ptr(config_.expert_num); + std::vector gate_bc_ptr(config_.expert_num); + std::vector up_bc_ptr(config_.expert_num); + std::vector down_ba_ptr(config_.expert_num); + std::vector down_bc_ptr(config_.expert_num); + for (int i = 0; i < config_.expert_num; i++) { + m_mem_requests.push_back( + {(void **)&gate_up_ba_ptr[i], T::BufferA::required_size(config_.max_len, config_.hidden_size)}); + m_mem_requests.push_back( + {(void **)&gate_bc_ptr[i], T::BufferC::required_size(config_.max_len, config_.intermediate_size)}); + m_mem_requests.push_back( + {(void **)&up_bc_ptr[i], T::BufferC::required_size(config_.max_len, config_.intermediate_size)}); + m_mem_requests.push_back( + {(void **)&down_ba_ptr[i], T::BufferA::required_size(config_.max_len, config_.intermediate_size)}); + m_mem_requests.push_back( + {(void **)&down_bc_ptr[i], T::BufferC::required_size(config_.max_len, config_.hidden_size)}); + } + shared_mem_buffer.alloc(this, m_mem_requests); + + m_local_pos_.resize(config_.max_len); + for (int i = 0; i < config_.max_len; i++) { + m_local_pos_[i].resize(config_.routed_expert_num); + } + m_expert_id_map_.resize(config_.expert_num); + m_local_num_.resize(config_.expert_num); + m_local_input_ptr_.resize(config_.expert_num); + m_local_gate_output_ptr_.resize(config_.expert_num); + m_local_up_output_ptr_.resize(config_.expert_num); + m_local_down_output_ptr_.resize(config_.expert_num); + + for (uint64_t i = 0; i < config_.expert_num; i++) { + gate_up_ba_.push_back( + std::make_shared(config_.max_len, config_.hidden_size, gate_up_ba_ptr[i])); + gate_bc_.push_back( + std::make_shared(config_.max_len, config_.intermediate_size, gate_bc_ptr[i])); + up_bc_.push_back(std::make_shared(config_.max_len, config_.intermediate_size, up_bc_ptr[i])); + down_ba_.push_back( + std::make_shared(config_.max_len, config_.intermediate_size, down_ba_ptr[i])); + down_bc_.push_back(std::make_shared(config_.max_len, config_.hidden_size, down_bc_ptr[i])); + +#ifdef USE_NUMA + int numa_nodes = numa_num_configured_nodes(); + gate_bb_numa_.resize(numa_nodes); + up_bb_numa_.resize(numa_nodes); + down_bb_numa_.resize(numa_nodes); + for (int j = 0; j < numa_nodes; j++) { + void *gate_bb_ptr = + numa_alloc_aligned(T::BufferB::required_size(config_.intermediate_size, config_.hidden_size), j, 64); + gate_bb_numa_[j].push_back( + std::make_shared(config_.intermediate_size, config_.hidden_size, gate_bb_ptr)); + void *up_bb_ptr = + numa_alloc_aligned(T::BufferB::required_size(config_.intermediate_size, config_.hidden_size), j, 64); + up_bb_numa_[j].push_back( + std::make_shared(config_.intermediate_size, config_.hidden_size, up_bb_ptr)); + void *down_bb_ptr = + numa_alloc_aligned(T::BufferB::required_size(config_.hidden_size, config_.intermediate_size), j, 64); + down_bb_numa_[j].push_back( + std::make_shared(config_.hidden_size, config_.intermediate_size, down_bb_ptr)); + } +#else + void *gate_bb_ptr = + std::aligned_alloc(64, T::BufferB::required_size(config_.intermediate_size, config_.hidden_size)); + gate_bb_.push_back( + std::make_shared(config_.intermediate_size, config_.hidden_size, gate_bb_ptr)); + + void *up_bb_ptr = + std::aligned_alloc(64, T::BufferB::required_size(config_.intermediate_size, config_.hidden_size)); + up_bb_.push_back( + std::make_shared(config_.intermediate_size, config_.hidden_size, up_bb_ptr)); + + void *down_bb_ptr = + std::aligned_alloc(64, T::BufferB::required_size(config_.hidden_size, config_.intermediate_size)); + down_bb_.push_back( + std::make_shared(config_.hidden_size, config_.intermediate_size, down_bb_ptr)); +#endif + } + } + + ~AMX_MOE() { shared_mem_buffer.dealloc(this); } + + void load_weights(Backend *backend) { + int nth = T::recommended_nth(config_.intermediate_size); + backend->do_work_stealing_job( + nth * config_.expert_num, nullptr, + [&](int task_id) { + uint64_t expert_idx = task_id / nth; + int ith = task_id % nth; +#ifdef USE_NUMA + int numa_nodes = numa_num_configured_nodes(); + for (int j = 0; j < numa_nodes; j++) { + gate_bb_numa_[j][expert_idx]->from_mat((ggml_bf16_t *)config_.gate_proj + + expert_idx * config_.intermediate_size * config_.hidden_size, + ith, nth); + up_bb_numa_[j][expert_idx]->from_mat((ggml_bf16_t *)config_.up_proj + + expert_idx * config_.intermediate_size * config_.hidden_size, + ith, nth); + } +#else + gate_bb_[expert_idx]->from_mat((ggml_bf16_t *)config_.gate_proj + + expert_idx * config_.intermediate_size * config_.hidden_size, + ith, nth); + up_bb_[expert_idx]->from_mat( + (ggml_bf16_t *)config_.up_proj + expert_idx * config_.intermediate_size * config_.hidden_size, ith, nth); +#endif + }, + nullptr); + nth = T::recommended_nth(config_.hidden_size); + backend->do_work_stealing_job( + nth * config_.expert_num, nullptr, + [&](int task_id) { + uint64_t expert_idx = task_id / nth; + int ith = task_id % nth; +#ifdef USE_NUMA + int numa_nodes = numa_num_configured_nodes(); + for (int j = 0; j < numa_nodes; j++) { + down_bb_numa_[j][expert_idx]->from_mat((ggml_bf16_t *)config_.down_proj + + expert_idx * config_.hidden_size * config_.intermediate_size, + ith, nth); + } +#else + down_bb_[expert_idx]->from_mat((ggml_bf16_t *)config_.down_proj + + expert_idx * config_.hidden_size * config_.intermediate_size, + ith, nth); +#endif + }, + nullptr); + } + + void warm_up(Backend *backend) {} + + void forward(int qlen, int k, const uint64_t *expert_ids, const float *weights, const void *input, void *output, + int *batch_size_tensor, Backend *backend) { + qlen = batch_size_tensor[0]; + bool use_amx = (qlen > 4 * config_.expert_num / config_.routed_expert_num); + int activated_expert = 0; + for (int i = 0; i < config_.expert_num; i++) { + m_local_num_[i] = 0; + } + for (int i = 0; i < qlen; i++) { + for (int j = 0; j < k; j++) { + m_local_pos_[i][j] = m_local_num_[expert_ids[i * k + j]]++; + } + } + for (int i = 0; i < config_.expert_num; i++) { + if (m_local_num_[i] > 0) { + m_expert_id_map_[activated_expert] = i; + activated_expert++; + } + } + uint64_t offset = 0; + for (int i = 0; i < config_.expert_num; i++) { + m_local_input_ptr_[i] = m_local_input_ + offset * config_.hidden_size; + m_local_gate_output_ptr_[i] = m_local_gate_output_ + offset * config_.intermediate_size; + m_local_up_output_ptr_[i] = m_local_up_output_ + offset * config_.intermediate_size; + m_local_down_output_ptr_[i] = m_local_down_output_ + offset * config_.hidden_size; + offset += m_local_num_[i]; + } + backend->do_work_stealing_job( + qlen, nullptr, + [&](int i) { + for (int j = 0; j < k; j++) { + memcpy(m_local_input_ptr_[expert_ids[i * k + j]] + m_local_pos_[i][j] * config_.hidden_size, + (ggml_bf16_t *)input + i * config_.hidden_size, sizeof(ggml_bf16_t) * config_.hidden_size); + } + }, + nullptr); + backend->do_work_stealing_job( + activated_expert, nullptr, + [&](int task_id) { + int expert_idx = m_expert_id_map_[task_id]; + + gate_up_ba_[expert_idx]->from_mat(m_local_num_[expert_idx], m_local_input_ptr_[expert_idx], 0, 1); + }, + nullptr); + int nth = T::recommended_nth(config_.intermediate_size); + backend->do_work_stealing_job( + nth * activated_expert, [&](int _) { T::config(); }, + [&](int task_id) { + int expert_idx = m_expert_id_map_[task_id / nth]; + int ith = task_id % nth; +#ifdef USE_NUMA + amx::mat_mul(m_local_num_[expert_idx], config_.intermediate_size, config_.hidden_size, + gate_up_ba_[expert_idx], gate_bb_numa_[Backend::numa_node][expert_idx], gate_bc_[expert_idx], + ith, nth, use_amx); + amx::mat_mul(m_local_num_[expert_idx], config_.intermediate_size, config_.hidden_size, + gate_up_ba_[expert_idx], up_bb_numa_[Backend::numa_node][expert_idx], up_bc_[expert_idx], ith, + nth, use_amx); +#else + amx::mat_mul(m_local_num_[expert_idx], config_.intermediate_size, config_.hidden_size, + gate_up_ba_[expert_idx], gate_bb_[expert_idx], gate_bc_[expert_idx], ith, nth, use_amx); + amx::mat_mul(m_local_num_[expert_idx], config_.intermediate_size, config_.hidden_size, + gate_up_ba_[expert_idx], up_bb_[expert_idx], up_bc_[expert_idx], ith, nth, use_amx); +#endif + gate_bc_[expert_idx]->to_mat(m_local_num_[expert_idx], m_local_gate_output_ptr_[expert_idx], ith, nth); + up_bc_[expert_idx]->to_mat(m_local_num_[expert_idx], m_local_up_output_ptr_[expert_idx], ith, nth); + auto [n_start, n_end] = T::split_range_n(config_.intermediate_size, ith, nth); + if (config_.use_silu) { + for (int i = 0; i < m_local_num_[expert_idx]; i++) { + ggml_bf16_t *gate_output_ptr = &m_local_gate_output_ptr_[expert_idx][i * config_.intermediate_size]; + ggml_bf16_t *up_output_ptr = &m_local_up_output_ptr_[expert_idx][i * config_.intermediate_size]; + for (int j = n_start; j < n_end; j += 32) { + __m512 gate_val0, gate_val1, up_val0, up_val1; + avx512_32xbf16_to_32xfp32((__m512i *)(gate_output_ptr + j), &gate_val0, &gate_val1); + avx512_32xbf16_to_32xfp32((__m512i *)(up_output_ptr + j), &up_val0, &up_val1); + __m512 result0 = act_fn(gate_val0, up_val0); + __m512 result1 = act_fn(gate_val1, up_val1); + avx512_32xfp32_to_32xbf16(&result0, &result1, (__m512i *)(gate_output_ptr + j)); + } + } + } + else { + for (int i = 0; i < m_local_num_[expert_idx]; i++) { + ggml_bf16_t *gate_output_ptr = &m_local_gate_output_ptr_[expert_idx][i * config_.intermediate_size]; + ggml_bf16_t *up_output_ptr = &m_local_up_output_ptr_[expert_idx][i * config_.intermediate_size]; + for (int j = n_start; j < n_end; j += 32) { + __m512 gate_val0, gate_val1, up_val0, up_val1; + avx512_32xbf16_to_32xfp32((__m512i *)(gate_output_ptr + j), &gate_val0, &gate_val1); + avx512_32xbf16_to_32xfp32((__m512i *)(up_output_ptr + j), &up_val0, &up_val1); + __m512 result0 = relu_act_fn(gate_val0, up_val0); + __m512 result1 = relu_act_fn(gate_val1, up_val1); + avx512_32xfp32_to_32xbf16(&result0, &result1, (__m512i *)(gate_output_ptr + j)); + } + } + } + + }, + nullptr); + backend->do_work_stealing_job( + activated_expert, nullptr, + [&](int task_id) { + int expert_idx = m_expert_id_map_[task_id]; + down_ba_[expert_idx]->from_mat(m_local_num_[expert_idx], m_local_gate_output_ptr_[expert_idx], 0, 1); + }, + nullptr); + nth = T::recommended_nth(config_.hidden_size); + backend->do_work_stealing_job( + nth * activated_expert, [&](int _) { T::config(); }, + [&](int task_id) { + int expert_idx = m_expert_id_map_[task_id / nth]; + int ith = task_id % nth; +#ifdef USE_NUMA + amx::mat_mul(m_local_num_[expert_idx], config_.hidden_size, config_.intermediate_size, down_ba_[expert_idx], + down_bb_numa_[Backend::numa_node][expert_idx], down_bc_[expert_idx], ith, nth, use_amx); +#else + amx::mat_mul(m_local_num_[expert_idx], config_.hidden_size, config_.intermediate_size, down_ba_[expert_idx], + down_bb_[expert_idx], down_bc_[expert_idx], ith, nth, use_amx); +#endif + down_bc_[expert_idx]->to_mat(m_local_num_[expert_idx], m_local_down_output_ptr_[expert_idx], ith, nth); + }, + nullptr); + backend->do_work_stealing_job( + qlen, nullptr, + [&](int i) { + for (int e = 0; e < config_.hidden_size; e += 32) { + __m512 x0 = _mm512_setzero_ps(); + __m512 x1 = _mm512_setzero_ps(); + for (int j = 0; j < k; j++) { + __m512 weight = _mm512_set1_ps(weights[i * k + j]); + __m512 down_output0, down_output1; + avx512_32xbf16_to_32xfp32((__m512i *)(m_local_down_output_ptr_[expert_ids[i * k + j]] + + m_local_pos_[i][j] * config_.hidden_size + e), + &down_output0, &down_output1); + x0 = _mm512_fmadd_ps(down_output0, weight, x0); + x1 = _mm512_fmadd_ps(down_output1, weight, x1); + } + avx512_32xfp32_to_32xbf16(&x0, &x1, (__m512i *)((ggml_bf16_t *)output + i * config_.hidden_size + e)); + } + }, + nullptr); + } +}; + +#endif diff --git a/archive/csrc/ktransformers_ext/operators/kvcache/kvcache.h b/archive/csrc/ktransformers_ext/operators/kvcache/kvcache.h new file mode 100644 index 0000000..ac91778 --- /dev/null +++ b/archive/csrc/ktransformers_ext/operators/kvcache/kvcache.h @@ -0,0 +1,727 @@ +/** + * @Description : + * @Author : Jianwei Dong + * @Date : 2024-08-26 22:47:06 + * @Version : 1.0.0 + * @LastEditors : Jianwei Dong + * @LastEditTime : 2024-08-26 22:47:06 + * @Copyright (c) 2024 by KVCache.AI, All Rights Reserved. + **/ + +#ifndef CPUINFER_OPERATOR_KVCACHE_H +#define CPUINFER_OPERATOR_KVCACHE_H + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "../../cpu_backend/backend.h" +#include "llama.cpp/ggml-common.h" +#include "llama.cpp/ggml-impl.h" +#include "llama.cpp/ggml-quants.h" +#include "llama.cpp/ggml.h" +#include "llamafile/sgemm.h" + +#define CHUNK_SIZE 32 + +/** + * @brief Converts a ggml_type enum value to its corresponding string + * representation. + * + * This function provides a human-readable string representation for a given + * ggml_type enum value. The string can be used for logging, debugging, or + * displaying information in a user interface. + * + * @param type The ggml_type enum value to convert. + * @return A string representation of the enum value. + */ +std::string ggml_type_to_string(ggml_type type); + +/** + * @enum AnchorType + * @brief Defines the types of anchors used in attention mechanisms. + * + * This enum specifies different types of anchors that can be used in attention + * mechanisms, such as fixed anchors, dynamic anchors, or special anchors like + * QUEST, BLOCK_MEAN, or BLOCK_MAX. + */ +enum AnchorType { + FIXED_ANCHOR, /**< A fixed anchor that does not change. */ + DYNAMIC, /**< A dynamic anchor that can change over time. */ + QUEST, /**< A special anchor type used for QUEST (Query and Embedding Space + Transformation). */ + BLOCK_MEAN, /**< An anchor based on the mean of a block of data. */ + BLOCK_MAX /**< An anchor based on the maximum value within a block of data. + */ +}; + +/** + * @brief Converts an AnchorType enum value to its corresponding string + * representation. + * + * This function provides a human-readable string representation for a given + * AnchorType enum value. The string can be used for logging, debugging, or + * displaying information in a user interface. + * + * @param anchor_type The AnchorType enum value to convert. + * @return A string representation of the enum value. + */ +std::string AnchorTypeToString(AnchorType anchor_type); + +/** + * @enum RetrievalType + * @brief Defines the types of retrieval strategies in attention mechanisms. + * + * This enum specifies different retrieval strategies that can be used in + * attention mechanisms, such as layer-level retrieval, key-value head-level + * retrieval, or query head-level retrieval. + */ +enum RetrievalType { + LAYER, /**< Retrieval at the layer level. */ + KVHEAD, /**< Retrieval at the key-value head level. */ + QHEAD /**< Retrieval at the query head level. */ +}; + +/** + * @brief Converts a RetrievalType enum value to its corresponding string + * representation. + * + * This function provides a human-readable string representation for a given + * RetrievalType enum value. The string can be used for logging, debugging, or + * displaying information in a user interface. + * + * @param retrieval_type The RetrievalType enum value to convert. + * @return A string representation of the enum value. + */ +std::string RetrievalTypeToString(RetrievalType retrieval_type); + +/** + * @struct KVCacheConfig + * @brief Configuration structure for Key-Value (KV) Cache. + * + * This structure holds configuration parameters for setting up and managing + * a Key-Value (KV) Cache used in various attention mechanisms. It includes + * parameters such as the number of layers, the number of heads, the dimension + * of each head, block length, anchor information, and memory-related settings. + */ +struct KVCacheConfig { + int layer_num; /**< Number of layers in the model. */ + int kv_head_num; /**< Number of heads in the KV Cache. */ + int q_head_num; /**< Number of heads in the query. */ + int head_dim; /**< Dimension of each head. */ + int block_len; /**< Length of each block in the cache. */ + int anchor_num; /**< Number of anchors used in attention. */ + + ggml_type kv_type; /**< Data type of the KV Cache (e.g., fp16, q8_0). */ + + // Controls the pre-allocated memory size + int max_block_num; /**< Maximum number of blocks that can be allocated. */ + int max_batch_size; /**< Maximum batch size that can be processed. */ + int max_thread_num; /**< Maximum number of threads that can be used. */ + + AnchorType + anchor_type; /**< Type of anchors used in the attention mechanism. */ + RetrievalType + retrieval_type; /**< Type of retrieval strategy used in the cache. */ + + int layer_step; /**< Step size between layers. */ + int token_step; /**< Step size between tokens. */ + int layer_offset; /**< Offset value for layers. */ + + /** + * @brief Default constructor for KVCacheConfig. + * + * Initializes the configuration with default values. This constructor + * does not initialize any member variables explicitly. + */ + KVCacheConfig() = default; + + /** + * @brief Parameterized constructor for KVCacheConfig. + * + * This constructor initializes the configuration with specific values + * for all member variables. + * + * @param layer_num The number of layers in the model. + * @param kv_head_num The number of heads in the KV Cache. + * @param q_head_num The number of heads in the query. + * @param head_dim The dimension of each head. + * @param block_len The length of each block in the cache. + * @param anchor_num The number of anchors used in attention. + * @param anchor_type The type of anchors used in the attention mechanism. + * @param kv_type The data type of the KV Cache (e.g., fp16, q8_0). + * @param retrieval_type The type of retrieval strategy used in the cache. + * @param layer_step The step size between layers. + * @param token_step The step size between tokens. + * @param layer_offset The offset value for layers. + * @param max_block_num The maximum number of blocks that can be allocated. + * @param max_batch_size The maximum batch size that can be processed. + * @param max_thread_num The maximum number of threads that can be used. + */ + KVCacheConfig(int layer_num, int kv_head_num, int q_head_num, int head_dim, + int block_len, int anchor_num, AnchorType anchor_type, + ggml_type kv_type, RetrievalType retrieval_type, + int layer_step, int token_step, int layer_offset, + int max_block_num, int max_batch_size, int max_thread_num); +}; + +/** + * @class KVCache + * @brief Manages the Key-Value (KV) Cache used in attention mechanisms. + * + * The KVCache class provides functionality for managing the Key-Value Cache, + * including resizing the cache, retrieving configuration parameters, and + * updating internal states. This class is typically used in transformer models + * to store and manage past key and value states for efficient attention + * computations. + */ +class KVCache { + public: + /** + * @brief Constructs a KVCache object with the given configuration. + * + * Initializes the KVCache with the specified configuration parameters, + * such as the number of layers, heads, head dimensions, and other + * relevant settings. + * + * @param config The configuration object containing initialization + * parameters. + */ + KVCache(KVCacheConfig config); + + /** + * @brief Resizes the number of threads used by the cache. + * + * This function adjusts the number of threads that the cache can utilize. + * It allows dynamic reconfiguration of the parallel processing capabilities + * based on the current workload or system resources. + * + * @param thread_num The new number of threads to use. + */ + void ThreadResize(int thread_num); + + /** + * @brief Resizes the batch size managed by the cache. + * + * This function adjusts the batch size that the cache can handle. It + * is useful when the input batch size changes dynamically, allowing + * the cache to be reconfigured accordingly. + * + * @param batch_size The new batch size. + */ + void BatchResize(int batch_size); + + /** + * @brief Resizes the number of blocks managed by the cache. + * + * This function adjusts the number of blocks that the cache can manage. + * It allows dynamic reconfiguration of the block structure based on the + * current sequence length or other factors. + * + * @param block_num The new number of blocks. + */ + void BlockResize(int block_num); + + /** + * @brief Gets the number of layers in the cache. + * + * @return The number of layers configured in the cache. + */ + int get_layer_num() { return config_.layer_num; } + + /** + * @brief Gets the number of KV heads in the cache. + * + * @return The number of KV heads configured in the cache. + */ + int get_kv_head_num() { return config_.kv_head_num; } + + /** + * @brief Gets the number of query heads in the cache. + * + * @return The number of query heads configured in the cache. + */ + int get_q_head_num() { return config_.q_head_num; } + + /** + * @brief Gets the dimension of each head in the cache. + * + * @return The dimension of each head. + */ + int get_head_dim() { return config_.head_dim; } + + /** + * @brief Gets the length of each block in the cache. + * + * @return The length of each block. + */ + int get_block_len() { return config_.block_len; } + + /** + * @brief Gets the number of blocks for a specific layer. + * + * @param layer_id The ID of the layer for which to retrieve the block + * number. + * @return The number of blocks in the specified layer. + */ + int get_block_num(int layer_id) { return past_block_num_[layer_id]; } + + /** + * @brief Gets the number of anchors in the cache. + * + * @return The number of anchors configured in the cache. + */ + int get_anchor_num() { return config_.anchor_num; } + + /** + * @brief Gets the total length of the cache. + * + * @return The total length of the cache. + */ + int get_cache_total_len() { return cache_total_len_; } + + /** + * @brief Gets the total number of blocks in the cache. + * + * This function computes and returns the total number of blocks in the + * cache based on the total cache length and the block length configuration. + * + * @return The total number of blocks in the cache. + */ + int get_cache_total_block_num() { + return (cache_total_len_ + config_.block_len - 1) / config_.block_len; + } + + /** + * @brief Updates the total length of the cache. + * + * This function sets a new total length for the cache, allowing dynamic + * adjustment of the cache size during runtime. + * + * @param cache_total_len The new total length of the cache. + */ + void update_cache_total_len(int cache_total_len) { + cache_total_len_ = cache_total_len; + } + void attn(const ggml_fp16_t *q_in, ggml_fp16_t *output, float *attn_lse, + int layer_idx, int generate_token_idx, int q_len, int batch_size, + int max_block_num, int *block_table, int *cache_seqlens, + int pick_block_num, int init_block_num, int local_block_num, + Backend *backend); + + void update_kvcache_one_block_fp16(const ggml_fp16_t *k_in, + const ggml_fp16_t *v_in, int layer_id, + int block_idx, Backend *backend); + + void get_kvcache_one_block_fp16(ggml_fp16_t *k_in, ggml_fp16_t *v_in, + int layer_id, int block_idx, + Backend *backend); + + void update_importance_one_block(const ggml_fp16_t *importance, + int layer_id, int block_idx, + Backend *backend); + void get_importance_one_block(ggml_fp16_t *importance, int layer_id, + int block_idx, Backend *backend); + + void get_anchor_one_block(ggml_fp16_t *anchor, int layer_id, int block_idx, + Backend *backend); + + void update_anchor_one_block(const ggml_fp16_t *anchor, int layer_id, + int block_idx, Backend *backend); + + void calc_anchor_all_layers(int *block_table, int *cache_seqlens, + int batch_size, int max_block_num, + Backend *backend); + + void load_kvcache(std::string tensor_file_path, Backend *backend); + void dump_kvcache(int *block_table, int cache_total_len, + std::string tensor_file_path, Backend *backend); + + void get_and_update_kvcache_fp16(ggml_fp16_t *k_in, ggml_fp16_t *v_in, + int layer_id, int *block_table, + int batch_size, int max_block_num, + int *cache_seqlens, int q_len, + Backend *backend); + + void get_kvcache_fp16(ggml_fp16_t *k_in, ggml_fp16_t *v_in, int layer_id, + int *block_table, int batch_size, int max_block_num, + int *cache_seqlens, Backend *backend); + + void update_kvcache_fp16(const ggml_fp16_t *k_in, const ggml_fp16_t *v_in, + int layer_id, int *block_table, int batch_size, + int max_block_num, int *cache_seqlens, int q_len, + Backend *backend); + + void update_importance(const ggml_fp16_t *importance, int layer_id, + int *block_table, int batch_size, int max_block_num, + int *offset, int width, Backend *backend); + + void attn_with_kvcache(const ggml_fp16_t *q_in, const ggml_fp16_t *k_in, + const ggml_fp16_t *v_in, ggml_fp16_t *output, + float *attn_lse, int layer_idx, + int generate_token_idx, int q_len, int batch_size, + int max_block_num, int *block_table, + int *cache_seqlens, int topk, int local, + Backend *backend); + + void clear_importance_all_layers(int *block_table, int *cache_seqlens, + int batch_size, int max_block_num, + Backend *backend); + + void clear_kvcache_all_layers(int *block_table, int *cache_seqlens, + int batch_size, int max_block_num, + Backend *backend); + + void get_sincos(ggml_fp16_t *sin, ggml_fp16_t *cos, int seqlen); + + void get_attn_sparsity(const ggml_fp16_t *q_in, float *attn_sparsity, + int layer_idx, int generate_token_idx, int q_len, + int batch_size, int max_block_num, int *block_table, + int *cache_seqlens, int *block_table_origin, + int *cache_seqlens_origin, int max_block_num_origin, + int topk, int local, Backend *backend); + + void get_all_kvcache_one_layer(int layer_id, ggml_fp16_t *k_in, + ggml_fp16_t *v_in, Backend *backend); + + private: + // Persistent data + KVCacheConfig config_; + int n_gqa_; // q_head_num / kv_head_num + int cache_total_len_; // Number of tokens in cache + std::vector past_block_num_; // [layer_num] + std::vector>>> + k_cache_q4; // [layer_num, kv_head_num, past_block_num, block_len * + // (head_dim / QK_4)] + std::vector>>> + v_cache_q4; // [layer_num, kv_head_num, past_block_num, head_dim * + // (block_len / QK_4)] + std::vector>>> + k_cache_q8; // [layer_num, kv_head_num, past_block_num, block_len * + // (head_dim / QK_8)] + std::vector>>> + v_cache_q8; // [layer_num, kv_head_num, past_block_num, head_dim * + // (block_len / QK_8)] + + std::vector>>> + k_cache_fp16_; // [layer_num, kv_head_num, past_block_num, block_len * + // head_dim] + std::vector>>> + v_cache_fp16_; // [layer_num, kv_head_num, past_block_num, head_dim * + // block_len] + + std::vector>>> + importance_; // [layer_num, past_block_num, block_len, + // attention_head_num] + + std::vector + anchor_; // [layer_num * past_block_num * anchor_num * + // attention_head_num * head_dim] + + // Runtime data + int64_t layer_id_; + int64_t block_idx_; + int *block_table_; + uint64_t block_num_; + int max_block_num_after_retrieval_; + + // Rotary positional embeddings + std::vector> sin_; // [seq_len, head_dim] + std::vector> cos_; // [seq_len, head_dim] + + // update/get + int seq_len_; + uint16_t *k_scales_; // q4_0 + uint8_t *k_in_; // q4_0 + uint16_t *v_scales_; // q4_0 + uint8_t *v_in_; // q4_0 + uint16_t *k_data_; // fp16 + uint16_t *v_data_; // fp16 + uint16_t *importance_data_; // fp16 + uint16_t *anchor_data_; // fp16 + + // sparsity = (sigma(block lse / lse)) + std::vector>> + block_lse_; // [batch_size, max_block_num, q_head_num] + std::vector> attn_sparsity_; // [batch_size, q_head_num] + + // attn + std::vector> + avg_q; // [batch_size, q_head_num * head_dim] + + std::vector> + avg_q_fp16; // [batch_size, q_head_num * head_dim] + std::vector< + std::priority_queue, + std::vector>, std::greater<>>> + top_similar_block_; + + std::vector> block_similar_; + std::vector>> block_similar_kv_head_; + std::vector>> block_similar_q_head_; + + std::vector cache_seqlens_; // [batch_size] + std::vector selected_blocks_num_history_; // [layer_num // layer_step] + + std::vector>> selected_blocks_history_; + // [layer_num // layer_step, batch_size, max_block_num] + + std::vector>>> + selected_blocks_history_kvhead_; // [layer_num // layer_step, + // batch_size, max_block_num, + // kv_head_num] + + std::vector> + block_table_before_retrieval_; // [batch_size, max_block_num] + std::vector> + block_table_after_retrieval_; // [batch_size, pick_block_num] + + std::vector>> + block_table_before_retrieval_qhead_; // [batch_size, max_block_num, + // q_head_num] + std::vector>> + block_table_after_retrieval_qhead_; // [batch_size, pick_block_num, + // q_head_num] + + std::vector>> + block_table_before_retrieval_kvhead_; // [batch_size, max_block_num, + // kv_head_num] + std::vector>> + block_table_after_retrieval_kvhead_; // [batch_size, pick_block_num, + // kv_head_num] + + std::vector>> + mutex_; // [batch_size, kv_head_num] + std::vector>> + q_q8_0_; // [batch_size, kv_head_num, n_gqa * head_dim / QK8_0] + std::vector>> + q_fp32_; // [batch_size, kv_head_num, n_gqa * head_dim] + + std::vector>> + output_fp32_; // [batch_size, kv_head_num, n_gqa * head_dim] + std::vector>> + attn_lse_; // [batch_size, kv_head_num, n_gqa] + + std::vector> thread_cur_head_idx_; // [thread_num] + + std::vector> + thread_local_output_q8_0_; // [thread_num, n_gqa * head_dim / QK8_0] + std::vector> + thread_local_attn_score_; // [thread_num, n_gqa * block_len] + std::vector> + thread_local_output_fp32_; // [thread_num, n_gqa * head_dim] + std::vector> + thread_local_attn_lse_; // [thread_num, n_gqa] + std::vector> + thread_local_cur_output_fp32_; // [thread_num, n_gqa * head_dim] + std::vector> + thread_local_cur_attn_lse_; // [thread_num, n_gqa] + std::vector> + thread_local_attn_mask_; // [thread_num, block_len // 8] + std::vector> + thread_local_draft_; // [thread_num, 2 * n_gqa * block_len + 6 * n_gqa * + // head_dim + 2 * block_len * head_dim] + + // tmp space + std::vector q_fp32; // [n_gqa * head_dim] + + void quantize_q_(const uint16_t *q_in_data, int batch_size); + void attn_initialize_layer_(int batch_size, int layer_idx, int *block_table, + int &max_block_num, int *cache_seqlens); + void attn_initialize_kvhead_(int batch_size, int layer_idx, + int *block_table, int &max_block_num, + int *cache_seqlens); + void retrieval_kvcache_layer_(const uint16_t *q_in_data, int init_block_num, + int local_block_num, int pick_block_num, + int q_len, int generate_token_idx, + int batch_size, int layer_idx, + int *cache_seqlens, int &max_block_num, + Backend *backend); + void retrieval_kvcache_kvhead_(const uint16_t *q_in_data, + int init_block_num, int local_block_num, + int pick_block_num, int q_len, + int generate_token_idx, int batch_size, + int layer_idx, int *cache_seqlens, + int &max_block_num, Backend *backend); + + void calculate_block_similarity_layer_( + const uint16_t *q_in_data, int batch_size, int layer_idx, int q_len, + int max_block_num, int *cache_seqlens, int init_block_num, + int local_block_num, int pick_block_num, Backend *backend); + void calculate_block_similarity_kvhead_( + const uint16_t *q_in_data, int batch_size, int layer_idx, int q_len, + int max_block_num, int *cache_seqlens, int init_block_num, + int local_block_num, int pick_block_num, Backend *backend); + + void select_block_layer_(int batch_size, int layer_idx, int max_block_num, + int init_block_num, int local_block_num, + int pick_block_num); + void select_block_kvhead_(int batch_size, int layer_idx, int max_block_num, + int init_block_num, int local_block_num, + int pick_block_num); + + void calculate_sparsity_layer_(const uint16_t *q_in_data, + float *attn_sparsity, int batch_size, + int max_block_num, int *block_table, + int *cache_seqlens, Backend *backend); + void calculate_sparsity_kvhead_(const uint16_t *q_in_data, + float *attn_sparsity, int batch_size, + int max_block_num, int *block_table, + int *cache_seqlens, Backend *backend); + + void attention_kvhead_(const uint16_t *q_in_data, ggml_fp16_t *output, + float *attn_lse, int batch_size, Backend *backend); + void attention_layer_(const uint16_t *q_in_data, ggml_fp16_t *output, + float *attn_lse, int batch_size, Backend *backend); + + /** + * @brief Computes attention with KV cache for one block. + * + * This function performs attention computation for one block using KV + * cache. The function supports different data types for Q, K, and V caches, + * and provides options for quantization. The function does not perform any + * dynamic memory allocation internally, so all necessary buffers must be + * pre-allocated externally. + * + * @param head_dim The dimension of the head. + * @param bsz The batch size. + * @param q_type The data type of Q (GGML data type). Only supports fp16 and + * q8_0. + * @param q Pointer to the Q tensor [bsz, head_dim]. The quantization is + * always applied along the head_dim dimension. The size must be + * bsz * head_dim/32 * qtype_size. If head_dim % 32 != 0, an error + * will be raised. + * @param past_kv_len The length of the past KV cache. + * @param past_kv_offset The offset in the past KV cache. + * @param is_full_attn Boolean flag indicating whether to use full attention + * (true for full 1 mask). + * @param attn_mask Pointer to the attention mask [bsz, past_kv_len]. If + * is_full_attn = false, a bit matrix is passed to + * represent the mask. + * @param k_type The data type of K cache (GGML data type). Only supports + * fp16, q4_0, and q8_0. + * @param k_quant_type Quantization type for K cache. 0 for per_token, 1 for + * per_channel. Other values will raise an error. + * @param k_cache Pointer to the K cache tensor [seq_len, head_dim]. If + * quant_type == 0, head_dim % 32 must be 0. If quant_type == + * 1, seq_len % 32 must be 0. + * @param num_k_anchor The number of K anchors. If num_k_anchor == 0, it + * means no anchor is present. + * @param k_cache_anchors Pointer to the K cache anchors [num_k_anchor, + * head_dim]. The k_anchor_type must be fp16. + * @param k_cache_anchor_pos Pointer to the K cache anchor positions. Each + * token is associated with the nearest previous anchor position. + * @param v_type The data type of V cache (GGML data type). + * @param v_quant_type Quantization type for V cache. + * @param v_cache Pointer to the V cache tensor [head_dim, seq_len]. + * @param num_v_anchor The number of V anchors. + * @param v_cache_anchors Pointer to the V cache anchors. + * @param v_cache_anchor_pos Pointer to the V cache anchor positions. + * @param attn_score Pre-allocated buffer for attention scores [bsz, + * past_kv_len]. + * @param output Output tensor [bsz, head_dim] with the same type as q_type. + * @param lse Pre-allocated buffer [bsz] for the log-sum-exp of the + * attention scores. + * @param draft Pre-allocated temporary buffer. The buffer size should be + * enough to hold (2 * bsz * past_kv_len + 6 * bsz * head_dim + 2 * + * past_kv_len * head_dim + past_kv_len * head_dim / 32) bytes. + * @param rotary_angle Pointer to the rotary angle tensor. + * @param rotary_cos Pointer to the cosine values for rotary embedding. + * @param rotary_sin Pointer to the sine values for rotary embedding. + */ + void attn_with_kvcache_one_block_( + int head_dim, int bsz, + ggml_type q_type, // GGML data type of `Q`, only supports fp16 and q8_0 + // [bsz, head_dim] + // Quantization is always on the head_dim dimension (per_token). If + // head_dim % 32 != 0, an error will be raised. The size must be bsz * + // head_dim/32 * qtype_size. + const void *q, + + int past_kv_len, int past_kv_offset, + bool is_full_attn, // true indicates a full 1 mask + // If is_full_attn = false, a bit matrix representing the mask is + // passed. [bsz, past_kv_len] + const uint8_t *attn_mask, + + ggml_type k_type, // GGML data type of `K Cache`, only supports fp16, + // q4_0, q8_0 + int k_quant_type, // 0 for per_token, 1 for per_channel, others raise an + // error + // [seq_len, head_dim] + // If quant_type == 0, head_dim % 32 must be 0. + // If quant_type == 1, seq_len % 32 must be 0. + const void *k_cache, + + // k_anchor_type must be fp16 + int num_k_anchor, // num_k_anchor == 0 indicates no anchor + // [num_k_anchor, head_dim] + const void *k_cache_anchors, + // Each token is associated with the nearest previous position's anchor, + // with the same distance. + const int *k_cache_anchor_pos, + + // v_cache similar to k_cache + ggml_type v_type, int v_quant_type, + // [head_dim, seq_len] + const void *v_cache, int num_v_anchor, const void *v_cache_anchors, + const int *v_cache_anchor_pos, + + // Pre-allocated buffer for intermediate calculations [bsz, + // past_kv_len]. No malloc is performed inside this function. + float *attn_score, + + // Output: [bsz, head_dim], with the same type as q_type + void *output, + // [bsz] + float *lse, + + // Pre-allocated temporary buffer with sufficient size: + // (2 * bsz * past_kv_len + 6 * bsz * head_dim + 2 * past_kv_len * + // head_dim + past_kv_len * head_dim / 32) bytes. + void *draft, + + // Apply rotary embedding online + const int *rotary_angle, const void *rotary_cos, const void *rotary_sin + // rotary_cos=None, + // rotary_sin=None, + // cache_seqlens: Optional[Union[(int, torch.Tensor)]] = None, + // cache_batch_idx: Optional[torch.Tensor] = None, + // rotary_interleaved=True, + + // // Not supported for now + // window_size=(-1, -1), # -1 means infinite context window + // alibi_slopes=None, + ); +}; + +/** + * @brief Scales a float32 vector by a given scalar value. + * + * This function multiplies each element of the input vector `y` by a scalar + * `v`. It uses platform-specific optimizations if available, such as Apple's + * Accelerate framework or SIMD instructions. If no specific optimization is + * available, the function falls back to a simple scalar multiplication loop. + * + * @param n The number of elements in the vector `y`. + * @param y The input vector to be scaled. The result will be stored in the same + * vector. + * @param v The scalar value by which to scale the vector. + */ +void ggml_vec_scale_f32(const int n, float *y, const float v); +#endif \ No newline at end of file diff --git a/archive/csrc/ktransformers_ext/operators/kvcache/kvcache_attn.cpp b/archive/csrc/ktransformers_ext/operators/kvcache/kvcache_attn.cpp new file mode 100644 index 0000000..4190c03 --- /dev/null +++ b/archive/csrc/ktransformers_ext/operators/kvcache/kvcache_attn.cpp @@ -0,0 +1,2535 @@ +/** + * @Description : + * @Author : Jianwei Dong + * @Date : 2024-08-26 22:47:06 + * @Version : 1.0.0 + * @LastEditors : Jianwei Dong + * @LastEditTime : 2024-08-26 22:47:06 + * @Copyright (c) 2024 by KVCache.AI, All Rights Reserved. + **/ + +#include "kvcache.h" + +#include + +void KVCache::attention_kvhead_(const uint16_t *q_in_data, ggml_fp16_t *output, + float *attn_lse, int batch_size, + Backend *backend) { + // Timer start + auto start = std::chrono::high_resolution_clock::now(); + seq_len_ = config_.block_len; + + backend->do_work_stealing_job( + batch_size * config_.kv_head_num * max_block_num_after_retrieval_, + [&](int thread_id) { + thread_cur_head_idx_[thread_id].first = -1; + thread_cur_head_idx_[thread_id].second = -1; + }, + [&](int task_id) { + int batch_id = task_id / (config_.kv_head_num * + max_block_num_after_retrieval_); + int head_id = (task_id % (config_.kv_head_num * + max_block_num_after_retrieval_)) / + max_block_num_after_retrieval_; + int block_id = task_id % max_block_num_after_retrieval_; + int thread_id = Backend::thread_local_id; + + // If the block is out of the sequence length, skip it. + if (cache_seqlens_[batch_id] / config_.block_len < block_id) { + return; + } + int block_idx = + block_table_after_retrieval_kvhead_[batch_id][block_id] + [head_id]; + if (cache_seqlens_[batch_id] / config_.block_len == block_id) { + int seq_len = cache_seqlens_[batch_id] % config_.block_len; + if (seq_len == 0) + return; + + // Prepare the attention mask for the last block. + int full_blocks = seq_len / 8; + int remaining_bits = seq_len % 8; + // Fill full blocks with 1s + for (int i = 0; i < full_blocks; ++i) { + thread_local_attn_mask_[thread_id][i] = 0xFF; + } + // Fill the remaining bits in the next block + if (remaining_bits > 0 && full_blocks < seq_len_ / 8) { + thread_local_attn_mask_[thread_id][full_blocks] = + (1 << remaining_bits) - 1; + } else { + thread_local_attn_mask_[thread_id][full_blocks] = 0; + } + + for (int i = full_blocks + 1; i < seq_len_ / 8; ++i) { + thread_local_attn_mask_[thread_id][i] = 0; + } + if (config_.kv_type == ggml_type::GGML_TYPE_F16) { + attn_with_kvcache_one_block_( + config_.head_dim, + config_.q_head_num / config_.kv_head_num, GGML_TYPE_F16, + (void *)&q_in_data[batch_id * config_.kv_head_num * + n_gqa_ * config_.head_dim + + head_id * n_gqa_ * config_.head_dim], + seq_len_, 0, false, + thread_local_attn_mask_[thread_id].data(), + GGML_TYPE_F16, 0, + k_cache_fp16_[layer_id_][head_id][block_idx].data(), 0, + nullptr, nullptr, GGML_TYPE_F16, 1, + v_cache_fp16_[layer_id_][head_id][block_idx].data(), 0, + nullptr, nullptr, + thread_local_attn_score_[thread_id].data(), + thread_local_output_fp32_[thread_id].data(), + thread_local_attn_lse_[thread_id].data(), + thread_local_draft_[thread_id].data(), nullptr, + cos_.data(), sin_.data()); + } else if (config_.kv_type == ggml_type::GGML_TYPE_Q4_0) { + attn_with_kvcache_one_block_( + config_.head_dim, + config_.q_head_num / config_.kv_head_num, + GGML_TYPE_Q8_0, q_q8_0_[batch_id][head_id].data(), + seq_len_, 0, false, + thread_local_attn_mask_[thread_id].data(), + GGML_TYPE_Q4_0, 0, + k_cache_q4[layer_id_][head_id][block_idx].data(), 0, + nullptr, nullptr, GGML_TYPE_Q4_0, 1, + v_cache_q4[layer_id_][head_id][block_idx].data(), 0, + nullptr, nullptr, + thread_local_attn_score_[thread_id].data(), + thread_local_output_q8_0_[thread_id].data(), + thread_local_attn_lse_[thread_id].data(), + thread_local_draft_[thread_id].data(), nullptr, + cos_.data(), sin_.data()); + dequantize_row_q8_0( + thread_local_output_q8_0_[thread_id].data(), + thread_local_output_fp32_[thread_id].data(), + n_gqa_ * config_.head_dim); + } else if (config_.kv_type == ggml_type::GGML_TYPE_Q8_0) { + attn_with_kvcache_one_block_( + config_.head_dim, + config_.q_head_num / config_.kv_head_num, + GGML_TYPE_Q8_0, q_q8_0_[batch_id][head_id].data(), + seq_len_, 0, false, + thread_local_attn_mask_[thread_id].data(), + GGML_TYPE_Q8_0, 0, + k_cache_q8[layer_id_][head_id][block_idx].data(), 0, + nullptr, nullptr, GGML_TYPE_Q8_0, 1, + v_cache_q8[layer_id_][head_id][block_idx].data(), 0, + nullptr, nullptr, + thread_local_attn_score_[thread_id].data(), + thread_local_output_q8_0_[thread_id].data(), + thread_local_attn_lse_[thread_id].data(), + thread_local_draft_[thread_id].data(), nullptr, + cos_.data(), sin_.data()); + dequantize_row_q8_0( + thread_local_output_q8_0_[thread_id].data(), + thread_local_output_fp32_[thread_id].data(), + n_gqa_ * config_.head_dim); + } + } else { + if (config_.kv_type == ggml_type::GGML_TYPE_F16) { + attn_with_kvcache_one_block_( + config_.head_dim, + config_.q_head_num / config_.kv_head_num, GGML_TYPE_F16, + (void *)&q_in_data[batch_id * config_.kv_head_num * + n_gqa_ * config_.head_dim + + head_id * n_gqa_ * config_.head_dim], + seq_len_, 0, true, nullptr, GGML_TYPE_F16, 0, + k_cache_fp16_[layer_id_][head_id][block_idx].data(), 0, + nullptr, nullptr, GGML_TYPE_F16, 1, + v_cache_fp16_[layer_id_][head_id][block_idx].data(), 0, + nullptr, nullptr, + thread_local_attn_score_[thread_id].data(), + thread_local_output_fp32_[thread_id].data(), + thread_local_attn_lse_[thread_id].data(), + thread_local_draft_[thread_id].data(), nullptr, + cos_.data(), sin_.data()); + + } else if (config_.kv_type == ggml_type::GGML_TYPE_Q4_0) { + attn_with_kvcache_one_block_( + config_.head_dim, + config_.q_head_num / config_.kv_head_num, + GGML_TYPE_Q8_0, q_q8_0_[batch_id][head_id].data(), + seq_len_, 0, true, nullptr, GGML_TYPE_Q4_0, 0, + k_cache_q4[layer_id_][head_id][block_idx].data(), 0, + nullptr, nullptr, GGML_TYPE_Q4_0, 1, + v_cache_q4[layer_id_][head_id][block_idx].data(), 0, + nullptr, nullptr, + thread_local_attn_score_[thread_id].data(), + thread_local_output_q8_0_[thread_id].data(), + thread_local_attn_lse_[thread_id].data(), + thread_local_draft_[thread_id].data(), nullptr, + cos_.data(), sin_.data()); + dequantize_row_q8_0( + thread_local_output_q8_0_[thread_id].data(), + thread_local_output_fp32_[thread_id].data(), + n_gqa_ * config_.head_dim); + } else if (config_.kv_type == ggml_type::GGML_TYPE_Q8_0) { + attn_with_kvcache_one_block_( + config_.head_dim, + config_.q_head_num / config_.kv_head_num, + GGML_TYPE_Q8_0, q_q8_0_[batch_id][head_id].data(), + seq_len_, 0, true, nullptr, GGML_TYPE_Q8_0, 0, + k_cache_q8[layer_id_][head_id][block_idx].data(), 0, + nullptr, nullptr, GGML_TYPE_Q8_0, 1, + v_cache_q8[layer_id_][head_id][block_idx].data(), 0, + nullptr, nullptr, + thread_local_attn_score_[thread_id].data(), + thread_local_output_q8_0_[thread_id].data(), + thread_local_attn_lse_[thread_id].data(), + thread_local_draft_[thread_id].data(), nullptr, + cos_.data(), sin_.data()); + dequantize_row_q8_0( + thread_local_output_q8_0_[thread_id].data(), + thread_local_output_fp32_[thread_id].data(), + n_gqa_ * config_.head_dim); + } + } + int cur_batch_idx = thread_cur_head_idx_[thread_id].first; + int cur_head_id = thread_cur_head_idx_[thread_id].second; + if (batch_id == cur_batch_idx && head_id == cur_head_id) { + for (int i = 0; i < n_gqa_; i++) { + float new_attn_lse = + thread_local_cur_attn_lse_[thread_id][i] + + std::log( + 1.0 + + std::exp(thread_local_attn_lse_[thread_id][i] - + thread_local_cur_attn_lse_[thread_id][i])); + ggml_vec_scale_f32( + config_.head_dim, + thread_local_cur_output_fp32_[thread_id].data() + + i * config_.head_dim, + std::exp(thread_local_cur_attn_lse_[thread_id][i] - + new_attn_lse)); + ggml_vec_scale_f32( + config_.head_dim, + thread_local_output_fp32_[thread_id].data() + + i * config_.head_dim, + std::exp(thread_local_attn_lse_[thread_id][i] - + new_attn_lse)); + for (int j = 0; j < config_.head_dim; j++) { + thread_local_cur_output_fp32_[thread_id] + [i * config_.head_dim + + j] += + thread_local_output_fp32_[thread_id] + [i * config_.head_dim + j]; + } + thread_local_cur_attn_lse_[thread_id][i] = new_attn_lse; + } + } else { + if (cur_batch_idx != -1) { + mutex_[cur_batch_idx][cur_head_id]->lock(); + for (int i = 0; i < n_gqa_; i++) { + if (std::abs(attn_lse_[cur_batch_idx][cur_head_id][i]) < + 1e-6) { + attn_lse_[cur_batch_idx][cur_head_id][i] = + thread_local_cur_attn_lse_[thread_id][i]; + for (int j = 0; j < config_.head_dim; j++) { + output_fp32_[cur_batch_idx][cur_head_id] + [i * config_.head_dim + j] = + thread_local_cur_output_fp32_ + [thread_id] + [i * config_.head_dim + j]; + } + continue; + } + float new_attn_lse = + attn_lse_[cur_batch_idx][cur_head_id][i] + + std::log( + 1.0 + + std::exp( + thread_local_cur_attn_lse_[thread_id][i] - + attn_lse_[cur_batch_idx][cur_head_id][i])); + ggml_vec_scale_f32( + config_.head_dim, + output_fp32_[cur_batch_idx][cur_head_id].data() + + i * config_.head_dim, + std::exp(attn_lse_[cur_batch_idx][cur_head_id][i] - + new_attn_lse)); + ggml_vec_scale_f32( + config_.head_dim, + thread_local_cur_output_fp32_[thread_id].data() + + i * config_.head_dim, + std::exp(thread_local_cur_attn_lse_[thread_id][i] - + new_attn_lse)); + for (int j = 0; j < config_.head_dim; j++) { + output_fp32_[cur_batch_idx][cur_head_id] + [i * config_.head_dim + j] += + thread_local_cur_output_fp32_ + [thread_id][i * config_.head_dim + j]; + } + attn_lse_[cur_batch_idx][cur_head_id][i] = new_attn_lse; + } + mutex_[cur_batch_idx][cur_head_id]->unlock(); + } + thread_cur_head_idx_[thread_id].first = batch_id; + thread_cur_head_idx_[thread_id].second = head_id; + for (int i = 0; i < n_gqa_; i++) { + thread_local_cur_attn_lse_[thread_id][i] = + thread_local_attn_lse_[thread_id][i]; + for (int j = 0; j < config_.head_dim; j++) { + thread_local_cur_output_fp32_ + [thread_id][i * config_.head_dim + j] = + thread_local_output_fp32_[thread_id] + [i * config_.head_dim + + j]; + } + } + } + }, + // Merge the results of the remaining blocks. + [&](int thread_id) { + int cur_batch_idx = thread_cur_head_idx_[thread_id].first; + int cur_head_id = thread_cur_head_idx_[thread_id].second; + if (cur_head_id != -1) { + mutex_[cur_batch_idx][cur_head_id]->lock(); + for (int i = 0; i < n_gqa_; i++) { + float new_attn_lse; + if (std::abs(attn_lse_[cur_batch_idx][cur_head_id][i]) < + 1e-6) { + attn_lse_[cur_batch_idx][cur_head_id][i] = + thread_local_cur_attn_lse_[thread_id][i]; + for (int j = 0; j < config_.head_dim; j++) { + output_fp32_[cur_batch_idx][cur_head_id] + [i * config_.head_dim + j] = + thread_local_cur_output_fp32_ + [thread_id] + [i * config_.head_dim + j]; + } + continue; + } + new_attn_lse = + attn_lse_[cur_batch_idx][cur_head_id][i] + + std::log( + 1.0 + + std::exp(thread_local_cur_attn_lse_[thread_id][i] - + attn_lse_[cur_batch_idx][cur_head_id][i])); + ggml_vec_scale_f32( + config_.head_dim, + output_fp32_[cur_batch_idx][cur_head_id].data() + + i * config_.head_dim, + std::exp(attn_lse_[cur_batch_idx][cur_head_id][i] - + new_attn_lse)); + ggml_vec_scale_f32( + config_.head_dim, + thread_local_cur_output_fp32_[thread_id].data() + + i * config_.head_dim, + std::exp(thread_local_cur_attn_lse_[thread_id][i] - + new_attn_lse)); + for (int j = 0; j < config_.head_dim; j++) { + output_fp32_[cur_batch_idx][cur_head_id] + [i * config_.head_dim + j] += + thread_local_cur_output_fp32_[thread_id] + [i * config_.head_dim + + j]; + } + attn_lse_[cur_batch_idx][cur_head_id][i] = new_attn_lse; + } + mutex_[cur_batch_idx][cur_head_id]->unlock(); + } + }); + // move the results to output and attn_lse + uint16_t *output_data = reinterpret_cast(output); + float *attn_lse_data = attn_lse; + for (int batch_idx = 0; batch_idx < batch_size; batch_idx++) { + for (int i = 0; i < config_.kv_head_num; i++) { + for (int j = 0; j < n_gqa_ * config_.head_dim; j++) { + output_data[batch_idx * config_.kv_head_num * n_gqa_ * + config_.head_dim + + i * n_gqa_ * config_.head_dim + j] = + GGML_FP32_TO_FP16(output_fp32_[batch_idx][i][j]); + } + for (int j = 0; j < n_gqa_; j++) { + attn_lse_data[batch_idx * config_.kv_head_num * n_gqa_ + + i * n_gqa_ + j] = attn_lse_[batch_idx][i][j]; + } + } + } + + // Timer end + auto end = std::chrono::high_resolution_clock::now(); + std::chrono::duration diff = end - start; + // printf("layer %d time of computing attention: %f s\n", layer_idx, + // diff.count()); +} + +void KVCache::attention_layer_(const uint16_t *q_in_data, ggml_fp16_t *output, + float *attn_lse, int batch_size, + Backend *backend) { + // Timer start + auto start = std::chrono::high_resolution_clock::now(); + seq_len_ = config_.block_len; + backend->do_work_stealing_job( + batch_size * config_.kv_head_num * max_block_num_after_retrieval_, + [&](int thread_id) { + thread_cur_head_idx_[thread_id].first = -1; + thread_cur_head_idx_[thread_id].second = -1; + }, + [&](int task_id) { + int batch_id = task_id / (config_.kv_head_num * + max_block_num_after_retrieval_); + int head_id = (task_id % (config_.kv_head_num * + max_block_num_after_retrieval_)) / + max_block_num_after_retrieval_; + int block_id = task_id % max_block_num_after_retrieval_; + int thread_id = Backend::thread_local_id; + // If the block is out of the sequence length, skip it. + if (cache_seqlens_[batch_id] / config_.block_len < block_id) { + return; + } + int block_idx = block_table_after_retrieval_[batch_id][block_id]; + if (cache_seqlens_[batch_id] / config_.block_len == block_id) { + int seq_len = cache_seqlens_[batch_id] % config_.block_len; + if (seq_len == 0) + return; + + // Prepare the attention mask for the last block. + int full_blocks = seq_len / 8; + int remaining_bits = seq_len % 8; + + // Fill full blocks with 1s + for (int i = 0; i < full_blocks; ++i) { + thread_local_attn_mask_[thread_id][i] = 0xFF; + } + // Fill the remaining bits in the next block + if (remaining_bits > 0 && full_blocks < seq_len_ / 8) { + thread_local_attn_mask_[thread_id][full_blocks] = + (1 << remaining_bits) - 1; + } else { + thread_local_attn_mask_[thread_id][full_blocks] = 0; + } + + for (int i = full_blocks + 1; i < seq_len_ / 8; ++i) { + thread_local_attn_mask_[thread_id][i] = 0; + } + if (config_.kv_type == ggml_type::GGML_TYPE_F16) { + attn_with_kvcache_one_block_( + config_.head_dim, + config_.q_head_num / config_.kv_head_num, GGML_TYPE_F16, + (void *)&q_in_data[batch_id * config_.kv_head_num * + n_gqa_ * config_.head_dim + + head_id * n_gqa_ * config_.head_dim], + seq_len_, 0, false, + thread_local_attn_mask_[thread_id].data(), + GGML_TYPE_F16, 0, + k_cache_fp16_[layer_id_][head_id][block_idx].data(), 0, + nullptr, nullptr, GGML_TYPE_F16, 1, + v_cache_fp16_[layer_id_][head_id][block_idx].data(), 0, + nullptr, nullptr, + thread_local_attn_score_[thread_id].data(), + thread_local_output_fp32_[thread_id].data(), + thread_local_attn_lse_[thread_id].data(), + thread_local_draft_[thread_id].data(), nullptr, + cos_.data(), sin_.data()); + } else if (config_.kv_type == ggml_type::GGML_TYPE_Q4_0) { + attn_with_kvcache_one_block_( + config_.head_dim, + config_.q_head_num / config_.kv_head_num, + GGML_TYPE_Q8_0, q_q8_0_[batch_id][head_id].data(), + seq_len_, 0, false, + thread_local_attn_mask_[thread_id].data(), + GGML_TYPE_Q4_0, 0, + k_cache_q4[layer_id_][head_id][block_idx].data(), 0, + nullptr, nullptr, GGML_TYPE_Q4_0, 1, + v_cache_q4[layer_id_][head_id][block_idx].data(), 0, + nullptr, nullptr, + thread_local_attn_score_[thread_id].data(), + thread_local_output_q8_0_[thread_id].data(), + thread_local_attn_lse_[thread_id].data(), + thread_local_draft_[thread_id].data(), nullptr, + cos_.data(), sin_.data()); + dequantize_row_q8_0( + thread_local_output_q8_0_[thread_id].data(), + thread_local_output_fp32_[thread_id].data(), + n_gqa_ * config_.head_dim); + } else if (config_.kv_type == ggml_type::GGML_TYPE_Q8_0) { + attn_with_kvcache_one_block_( + config_.head_dim, + config_.q_head_num / config_.kv_head_num, + GGML_TYPE_Q8_0, q_q8_0_[batch_id][head_id].data(), + seq_len_, 0, false, + thread_local_attn_mask_[thread_id].data(), + GGML_TYPE_Q8_0, 0, + k_cache_q8[layer_id_][head_id][block_idx].data(), 0, + nullptr, nullptr, GGML_TYPE_Q8_0, 1, + v_cache_q8[layer_id_][head_id][block_idx].data(), 0, + nullptr, nullptr, + thread_local_attn_score_[thread_id].data(), + thread_local_output_q8_0_[thread_id].data(), + thread_local_attn_lse_[thread_id].data(), + thread_local_draft_[thread_id].data(), nullptr, + cos_.data(), sin_.data()); + dequantize_row_q8_0( + thread_local_output_q8_0_[thread_id].data(), + thread_local_output_fp32_[thread_id].data(), + n_gqa_ * config_.head_dim); + } + } else { + if (config_.kv_type == ggml_type::GGML_TYPE_F16) { + attn_with_kvcache_one_block_( + config_.head_dim, + config_.q_head_num / config_.kv_head_num, GGML_TYPE_F16, + (void *)&q_in_data[batch_id * config_.kv_head_num * + n_gqa_ * config_.head_dim + + head_id * n_gqa_ * config_.head_dim], + seq_len_, 0, true, nullptr, GGML_TYPE_F16, 0, + k_cache_fp16_[layer_id_][head_id][block_idx].data(), 0, + nullptr, nullptr, GGML_TYPE_F16, 1, + v_cache_fp16_[layer_id_][head_id][block_idx].data(), 0, + nullptr, nullptr, + thread_local_attn_score_[thread_id].data(), + thread_local_output_fp32_[thread_id].data(), + thread_local_attn_lse_[thread_id].data(), + thread_local_draft_[thread_id].data(), nullptr, + cos_.data(), sin_.data()); + + } else if (config_.kv_type == ggml_type::GGML_TYPE_Q4_0) { + attn_with_kvcache_one_block_( + config_.head_dim, + config_.q_head_num / config_.kv_head_num, + GGML_TYPE_Q8_0, q_q8_0_[batch_id][head_id].data(), + seq_len_, 0, true, nullptr, GGML_TYPE_Q4_0, 0, + k_cache_q4[layer_id_][head_id][block_idx].data(), 0, + nullptr, nullptr, GGML_TYPE_Q4_0, 1, + v_cache_q4[layer_id_][head_id][block_idx].data(), 0, + nullptr, nullptr, + thread_local_attn_score_[thread_id].data(), + thread_local_output_q8_0_[thread_id].data(), + thread_local_attn_lse_[thread_id].data(), + thread_local_draft_[thread_id].data(), nullptr, + cos_.data(), sin_.data()); + dequantize_row_q8_0( + thread_local_output_q8_0_[thread_id].data(), + thread_local_output_fp32_[thread_id].data(), + n_gqa_ * config_.head_dim); + } else if (config_.kv_type == ggml_type::GGML_TYPE_Q8_0) { + attn_with_kvcache_one_block_( + config_.head_dim, + config_.q_head_num / config_.kv_head_num, + GGML_TYPE_Q8_0, q_q8_0_[batch_id][head_id].data(), + seq_len_, 0, true, nullptr, GGML_TYPE_Q8_0, 0, + k_cache_q8[layer_id_][head_id][block_idx].data(), 0, + nullptr, nullptr, GGML_TYPE_Q8_0, 1, + v_cache_q8[layer_id_][head_id][block_idx].data(), 0, + nullptr, nullptr, + thread_local_attn_score_[thread_id].data(), + thread_local_output_q8_0_[thread_id].data(), + thread_local_attn_lse_[thread_id].data(), + thread_local_draft_[thread_id].data(), nullptr, + cos_.data(), sin_.data()); + dequantize_row_q8_0( + thread_local_output_q8_0_[thread_id].data(), + thread_local_output_fp32_[thread_id].data(), + n_gqa_ * config_.head_dim); + } + } + int cur_batch_idx = thread_cur_head_idx_[thread_id].first; + int cur_head_id = thread_cur_head_idx_[thread_id].second; + if (batch_id == cur_batch_idx && head_id == cur_head_id) { + for (int i = 0; i < n_gqa_; i++) { + float new_attn_lse = + thread_local_cur_attn_lse_[thread_id][i] + + std::log( + 1.0 + + std::exp(thread_local_attn_lse_[thread_id][i] - + thread_local_cur_attn_lse_[thread_id][i])); + ggml_vec_scale_f32( + config_.head_dim, + thread_local_cur_output_fp32_[thread_id].data() + + i * config_.head_dim, + std::exp(thread_local_cur_attn_lse_[thread_id][i] - + new_attn_lse)); + ggml_vec_scale_f32( + config_.head_dim, + thread_local_output_fp32_[thread_id].data() + + i * config_.head_dim, + std::exp(thread_local_attn_lse_[thread_id][i] - + new_attn_lse)); + for (int j = 0; j < config_.head_dim; j++) { + thread_local_cur_output_fp32_[thread_id] + [i * config_.head_dim + + j] += + thread_local_output_fp32_[thread_id] + [i * config_.head_dim + j]; + } + thread_local_cur_attn_lse_[thread_id][i] = new_attn_lse; + } + } else { + if (cur_batch_idx != -1) { + mutex_[cur_batch_idx][cur_head_id]->lock(); + for (int i = 0; i < n_gqa_; i++) { + if (std::abs(attn_lse_[cur_batch_idx][cur_head_id][i]) < + 1e-6) { + attn_lse_[cur_batch_idx][cur_head_id][i] = + thread_local_cur_attn_lse_[thread_id][i]; + for (int j = 0; j < config_.head_dim; j++) { + output_fp32_[cur_batch_idx][cur_head_id] + [i * config_.head_dim + j] = + thread_local_cur_output_fp32_ + [thread_id] + [i * config_.head_dim + j]; + } + continue; + } + float new_attn_lse = + attn_lse_[cur_batch_idx][cur_head_id][i] + + std::log( + 1.0 + + std::exp( + thread_local_cur_attn_lse_[thread_id][i] - + attn_lse_[cur_batch_idx][cur_head_id][i])); + ggml_vec_scale_f32( + config_.head_dim, + output_fp32_[cur_batch_idx][cur_head_id].data() + + i * config_.head_dim, + std::exp(attn_lse_[cur_batch_idx][cur_head_id][i] - + new_attn_lse)); + ggml_vec_scale_f32( + config_.head_dim, + thread_local_cur_output_fp32_[thread_id].data() + + i * config_.head_dim, + std::exp(thread_local_cur_attn_lse_[thread_id][i] - + new_attn_lse)); + for (int j = 0; j < config_.head_dim; j++) { + output_fp32_[cur_batch_idx][cur_head_id] + [i * config_.head_dim + j] += + thread_local_cur_output_fp32_ + [thread_id][i * config_.head_dim + j]; + } + attn_lse_[cur_batch_idx][cur_head_id][i] = new_attn_lse; + } + mutex_[cur_batch_idx][cur_head_id]->unlock(); + } + thread_cur_head_idx_[thread_id].first = batch_id; + thread_cur_head_idx_[thread_id].second = head_id; + for (int i = 0; i < n_gqa_; i++) { + thread_local_cur_attn_lse_[thread_id][i] = + thread_local_attn_lse_[thread_id][i]; + for (int j = 0; j < config_.head_dim; j++) { + thread_local_cur_output_fp32_ + [thread_id][i * config_.head_dim + j] = + thread_local_output_fp32_[thread_id] + [i * config_.head_dim + + j]; + } + } + } + }, + // Merge the results of the remaining blocks. + [&](int thread_id) { + int cur_batch_idx = thread_cur_head_idx_[thread_id].first; + int cur_head_id = thread_cur_head_idx_[thread_id].second; + if (cur_head_id != -1) { + mutex_[cur_batch_idx][cur_head_id]->lock(); + for (int i = 0; i < n_gqa_; i++) { + float new_attn_lse; + if (std::abs(attn_lse_[cur_batch_idx][cur_head_id][i]) < + 1e-6) { + attn_lse_[cur_batch_idx][cur_head_id][i] = + thread_local_cur_attn_lse_[thread_id][i]; + for (int j = 0; j < config_.head_dim; j++) { + output_fp32_[cur_batch_idx][cur_head_id] + [i * config_.head_dim + j] = + thread_local_cur_output_fp32_ + [thread_id] + [i * config_.head_dim + j]; + } + continue; + } + new_attn_lse = + attn_lse_[cur_batch_idx][cur_head_id][i] + + std::log( + 1.0 + + std::exp(thread_local_cur_attn_lse_[thread_id][i] - + attn_lse_[cur_batch_idx][cur_head_id][i])); + ggml_vec_scale_f32( + config_.head_dim, + output_fp32_[cur_batch_idx][cur_head_id].data() + + i * config_.head_dim, + std::exp(attn_lse_[cur_batch_idx][cur_head_id][i] - + new_attn_lse)); + ggml_vec_scale_f32( + config_.head_dim, + thread_local_cur_output_fp32_[thread_id].data() + + i * config_.head_dim, + std::exp(thread_local_cur_attn_lse_[thread_id][i] - + new_attn_lse)); + for (int j = 0; j < config_.head_dim; j++) { + output_fp32_[cur_batch_idx][cur_head_id] + [i * config_.head_dim + j] += + thread_local_cur_output_fp32_[thread_id] + [i * config_.head_dim + + j]; + } + attn_lse_[cur_batch_idx][cur_head_id][i] = new_attn_lse; + } + mutex_[cur_batch_idx][cur_head_id]->unlock(); + } + }); + + // move the results to output and attn_lse + uint16_t *output_data = reinterpret_cast(output); + float *attn_lse_data = attn_lse; + for (int batch_idx = 0; batch_idx < batch_size; batch_idx++) { + for (int i = 0; i < config_.kv_head_num; i++) { + for (int j = 0; j < n_gqa_ * config_.head_dim; j++) { + output_data[batch_idx * config_.kv_head_num * n_gqa_ * + config_.head_dim + + i * n_gqa_ * config_.head_dim + j] = + GGML_FP32_TO_FP16(output_fp32_[batch_idx][i][j]); + } + for (int j = 0; j < n_gqa_; j++) { + attn_lse_data[batch_idx * config_.kv_head_num * n_gqa_ + + i * n_gqa_ + j] = attn_lse_[batch_idx][i][j]; + } + } + } + // Timer end + auto end = std::chrono::high_resolution_clock::now(); + std::chrono::duration diff = end - start; + // printf("layer %d time of computing attention: %f s\n", layer_id_, + // diff.count()); +} + +void KVCache::attn(const ggml_fp16_t *q_in, ggml_fp16_t *output, + float *attn_lse, int layer_idx, int generate_token_idx, + int q_len, int batch_size, int max_block_num, + int *block_table, int *cache_seqlens, int pick_block_num, + int init_block_num, int local_block_num, Backend *backend) { + + // Timer start + auto start = std::chrono::high_resolution_clock::now(); + layer_id_ = layer_idx; + batch_size = batch_size * q_len; + + const uint16_t *q_in_data = const_cast(q_in); + + quantize_q_(q_in_data, batch_size); + if (config_.retrieval_type == RetrievalType::LAYER) { + attn_initialize_layer_(batch_size, layer_idx, block_table, + max_block_num, cache_seqlens); + retrieval_kvcache_layer_(q_in_data, init_block_num, local_block_num, + pick_block_num, q_len, generate_token_idx, + batch_size, layer_idx, cache_seqlens, + max_block_num, backend); + attention_layer_(q_in_data, output, attn_lse, batch_size, backend); + } else if (config_.retrieval_type == RetrievalType::KVHEAD) { + attn_initialize_kvhead_(batch_size, layer_idx, block_table, + max_block_num, cache_seqlens); + retrieval_kvcache_kvhead_(q_in_data, init_block_num, local_block_num, + pick_block_num, q_len, generate_token_idx, + batch_size, layer_idx, cache_seqlens, + max_block_num, backend); + attention_kvhead_(q_in_data, output, attn_lse, batch_size, backend); + } + + // Timer end + auto end = std::chrono::high_resolution_clock::now(); + std::chrono::duration diff = end - start; + // printf("layer %d time of computing attention: %f s\n", layer_idx, + // diff.count()); +} + +void KVCache::attn_with_kvcache( + const ggml_fp16_t *q_in, const ggml_fp16_t *k_in, const ggml_fp16_t *v_in, + ggml_fp16_t *output, float *attn_lse, int layer_idx, int generate_token_idx, + int q_len, int batch_size, int max_block_num, int *block_table, + int *cache_seqlens, int topk, int local, Backend *backend) { + // printf("attn_with_kvcache start\n"); + assert(q_len == 1); + // Timer start + auto start = std::chrono::high_resolution_clock::now(); + + layer_id_ = layer_idx; + + update_kvcache_fp16(k_in, v_in, layer_idx, block_table, batch_size, + max_block_num, cache_seqlens, q_len, backend); + // printf("update finished.\n"); + + // cache_seqlens memory is modified. + for (int i = 0; i < batch_size; i++) { + cache_seqlens[i] += q_len; + } + int init_block_num = 1; + if (config_.block_len <= 32) { + init_block_num = 64 / config_.block_len; + } + + attn(q_in, output, attn_lse, layer_idx, generate_token_idx, q_len, + batch_size, max_block_num, block_table, cache_seqlens, topk, + init_block_num, local, backend); + + // Timer end + auto end = std::chrono::high_resolution_clock::now(); + std::chrono::duration diff = end - start; + // printf("layer %d time of computing attention with kvcache: %f s\n", + // layer_idx, diff.count()); +} + +void KVCache::quantize_q_(const uint16_t *q_in_data, int batch_size) { + // Timer start + auto start = std::chrono::high_resolution_clock::now(); + for (int batch_idx = 0; batch_idx < batch_size; batch_idx++) { + if (config_.kv_type == ggml_type::GGML_TYPE_F16) { + // quantize q + for (int i = 0; i < config_.kv_head_num; i++) { + for (int j = 0; j < n_gqa_ * config_.head_dim; j++) { + q_fp32_[batch_idx][i][j] = GGML_FP16_TO_FP32( + q_in_data[batch_idx * config_.kv_head_num * n_gqa_ * + config_.head_dim + + i * n_gqa_ * config_.head_dim + j]); + } + } + } else { + // quantize q + for (int i = 0; i < config_.kv_head_num; i++) { + for (int j = 0; j < n_gqa_ * config_.head_dim; j++) { + q_fp32[j] = GGML_FP16_TO_FP32( + q_in_data[batch_idx * config_.kv_head_num * n_gqa_ * + config_.head_dim + + i * n_gqa_ * config_.head_dim + j]); + } + quantize_row_q8_0(q_fp32.data(), q_q8_0_[batch_idx][i].data(), + n_gqa_ * config_.head_dim); + } + } + } + // Timer end + auto end = std::chrono::high_resolution_clock::now(); + // printf("time of quantizing q: %f s\n", + // std::chrono::duration(end - start).count()); +} +void KVCache::attn_initialize_layer_(int batch_size, int layer_idx, + int *block_table, int &max_block_num, + int *cache_seqlens) { + // Timer start + auto start = std::chrono::high_resolution_clock::now(); + for (int batch_idx = 0; batch_idx < batch_size; batch_idx++) { + // initialize output_fp32_ and attn_lse_ + for (int i = 0; i < config_.kv_head_num; i++) { + for (int j = 0; j < n_gqa_ * config_.head_dim; j++) { + output_fp32_[batch_idx][i][j] = 0; + } + for (int j = 0; j < n_gqa_; j++) { + attn_lse_[batch_idx][i][j] = 0; + } + } + // clear top_similar_block_ + + while (!top_similar_block_[batch_idx].empty()) + top_similar_block_[batch_idx].pop(); + } + + // get block_table_before_retrieval_ and cache_seqlens_ + if (block_table == nullptr) { + max_block_num = past_block_num_[layer_idx]; + for (int batch_idx = 0; batch_idx < batch_size; batch_idx++) { + if (cache_total_len_ != 0) + cache_seqlens_[batch_idx] = cache_total_len_; + else + cache_seqlens_[batch_idx] = max_block_num * config_.block_len; + for (int i = 0; i < max_block_num; i++) { + block_table_before_retrieval_[batch_idx][i] = i; + block_similar_[batch_idx][i] = 0; + } + } + } else { + for (int batch_idx = 0; batch_idx < batch_size; batch_idx++) { + cache_seqlens_[batch_idx] = cache_seqlens[batch_idx]; + for (int i = 0; i < max_block_num; i++) { + block_table_before_retrieval_[batch_idx][i] = + block_table[batch_idx * max_block_num + i]; + block_similar_[batch_idx][i] = 0; + } + } + } + // Timer end + auto end = std::chrono::high_resolution_clock::now(); + // printf("layer %d time of initializing attention: %f s\n", layer_idx, + // std::chrono::duration(end - start).count()); +} + +void KVCache::calculate_block_similarity_layer_( + const uint16_t *q_in_data, int batch_size, int layer_idx, int q_len, + int max_block_num, int *cache_seqlens, int init_block_num, + int local_block_num, int pick_block_num, Backend *backend) { + // Timer start + auto start = std::chrono::high_resolution_clock::now(); + + if (batch_size == 1 && + config_.anchor_num == 1) { // TODO: improve batch_size > 1 + for (int batch_id = 0; batch_id < batch_size; batch_id++) { + if (q_len == 1) { + for (int j = 0; j < config_.head_dim * config_.q_head_num; + j++) { + avg_q[batch_id][j] = GGML_FP16_TO_FP32( + q_in_data[batch_id * q_len * config_.q_head_num * + config_.head_dim + + j]); + avg_q_fp16[batch_id][j] = + q_in_data[batch_id * q_len * config_.q_head_num * + config_.head_dim + + j]; + } + } else { + for (int j = 0; j < config_.head_dim * config_.q_head_num; + j++) { + avg_q[batch_id][j] = 0; + } + for (int i = 0; i < q_len; i++) { + for (int j = 0; j < config_.head_dim; j++) { + avg_q[batch_id][j] += GGML_FP16_TO_FP32( + q_in_data[batch_id * q_len * config_.q_head_num * + config_.head_dim + + i * config_.q_head_num * + config_.head_dim + + j]); + } + } + for (int j = 0; j < config_.head_dim * config_.q_head_num; + j++) { + avg_q[batch_id][j] /= q_len; + avg_q_fp16[batch_id][j] = + GGML_FP32_TO_FP16(avg_q[batch_id][j]); + } + } + int seq_len = cache_seqlens_[batch_id]; + int block_num = (seq_len / config_.block_len) - local_block_num - + init_block_num; + if (block_num <= 0) { + continue; + } + bool is_seq = true; + for (int i = init_block_num + 1; + i < (seq_len / config_.block_len) - local_block_num; i++) { + if (block_table_before_retrieval_[batch_id][i] != + block_table_before_retrieval_[batch_id][i - 1] + 1) { + is_seq = false; + break; + } + } + if (is_seq) { + int nth = backend->get_thread_num(); + backend->do_work_stealing_job( + nth, nullptr, + [&](int task_id) { + int ith = task_id; + bool ok = llamafile_sgemm( + block_num, 1, config_.q_head_num * config_.head_dim, + anchor_.data() + + (layer_idx * config_.max_block_num + + block_table_before_retrieval_ + [batch_id][init_block_num]) * + config_.anchor_num * config_.q_head_num * + config_.head_dim, + config_.q_head_num * config_.head_dim, + avg_q_fp16[batch_id].data(), + config_.q_head_num * config_.head_dim, + block_similar_[batch_id].data() + init_block_num, + block_num, ith, nth, GGML_TASK_TYPE_COMPUTE, + GGML_TYPE_F16, GGML_TYPE_F16, GGML_TYPE_F32, + GGML_PREC_DEFAULT); + if (!ok) { + printf("llamafile_sgemm failed\n"); + } + }, + nullptr); + } else { + backend->do_work_stealing_job( + block_num, nullptr, + [&](int task_id) { + int block_id = task_id + init_block_num; + int block_idx = + block_table_before_retrieval_[batch_id][block_id]; + bool ok = llamafile_sgemm( + 1, 1, config_.q_head_num * config_.head_dim, + anchor_.data() + + (layer_idx * config_.max_block_num + + block_table_before_retrieval_[batch_id] + [block_idx]) * + config_.anchor_num * config_.q_head_num * + config_.head_dim, + config_.q_head_num * config_.head_dim, + avg_q_fp16[batch_id].data(), + config_.q_head_num * config_.head_dim, + block_similar_[batch_id].data() + block_id, 1, 0, 1, + GGML_TASK_TYPE_COMPUTE, GGML_TYPE_F16, + GGML_TYPE_F16, GGML_TYPE_F32, GGML_PREC_DEFAULT); + if (!ok) { + printf("llamafile_sgemm failed\n"); + } + }, + nullptr); + } + } + } else { + backend->do_work_stealing_job( + batch_size * max_block_num, nullptr, + [&](int task_id) { + int batch_id = task_id / max_block_num; + int block_id = task_id % max_block_num; + int seq_len = cache_seqlens_[batch_id]; + + if (block_id < init_block_num || + block_id >= + (seq_len / config_.block_len) - local_block_num) { + return; + } + + int block_idx = + block_table_before_retrieval_[batch_id][block_id]; + float sim = 0; + + for (int head_id = 0; head_id < config_.q_head_num; head_id++) { + for (int i = 0; i < config_.head_dim; i++) { + float q_i = 0, + qa_i = std::numeric_limits::lowest(); + for (int q_id = 0; q_id < q_len; q_id++) { + q_i += GGML_FP16_TO_FP32( + q_in_data[batch_id * q_len * + config_.q_head_num * + config_.head_dim + + q_id * config_.q_head_num * + config_.head_dim + + head_id * config_.head_dim + i]); + } + q_i /= q_len; + for (int anchor_id = 0; anchor_id < config_.anchor_num; + anchor_id++) { + qa_i = std::max( + qa_i, + GGML_FP16_TO_FP32( + anchor_[(long long)layer_idx * + config_.max_block_num * + config_.anchor_num * + config_.q_head_num * + config_.head_dim + + block_idx * config_.anchor_num * + config_.q_head_num * + config_.head_dim + + anchor_id * config_.q_head_num * + config_.head_dim + + head_id * config_.head_dim + i]) * + q_i); + } + sim += qa_i; + } + } + block_similar_[batch_id][block_id] = sim; + }, + nullptr); + } + // Timer end + auto end = std::chrono::high_resolution_clock::now(); + std::chrono::duration diff = end - start; + // printf("layer %d time of calculating similarity: %f s\n", layer_idx, + // diff.count()); +} + +void KVCache::select_block_layer_(int batch_size, int layer_idx, + int max_block_num, int init_block_num, + int local_block_num, int pick_block_num) { + // Timer start + auto start = std::chrono::high_resolution_clock::now(); + + for (int batch_idx = 0; batch_idx < batch_size; batch_idx++) { + + if (cache_seqlens_[batch_idx] / config_.block_len <= + init_block_num + pick_block_num + local_block_num) { + block_table_after_retrieval_[batch_idx].swap( + block_table_before_retrieval_[batch_idx]); + selected_blocks_num_history_[(layer_idx - config_.layer_offset) / + config_.layer_step] = 0; + continue; + } + + for (int block_id = init_block_num; + block_id < + (cache_seqlens_[batch_idx] / config_.block_len) - local_block_num; + block_id++) { + top_similar_block_[batch_idx].push(std::make_pair( + block_similar_[batch_idx][block_id], + block_table_before_retrieval_[batch_idx][block_id])); + if (top_similar_block_[batch_idx].size() > pick_block_num) { + top_similar_block_[batch_idx].pop(); + } + } + + int i = 0; + for (; i < init_block_num; i++) { + block_table_after_retrieval_[batch_idx][i] = + block_table_before_retrieval_[batch_idx][i]; + } + while (!top_similar_block_[batch_idx].empty()) { + block_table_after_retrieval_[batch_idx][i] = + top_similar_block_[batch_idx].top().second; + top_similar_block_[batch_idx].pop(); + i++; + } + for (; i < init_block_num + pick_block_num + local_block_num; i++) { + block_table_after_retrieval_[batch_idx][i] = + block_table_before_retrieval_[batch_idx] + [(cache_seqlens_[batch_idx] / + config_.block_len) - + local_block_num + i - + init_block_num - pick_block_num]; + } + if (cache_seqlens_[batch_idx] % config_.block_len != 0) { + block_table_after_retrieval_[batch_idx][i] = + block_table_before_retrieval_[batch_idx][( + cache_seqlens_[batch_idx] / config_.block_len)]; + cache_seqlens_[batch_idx] = + (cache_seqlens_[batch_idx] % config_.block_len) + + i * config_.block_len; + i++; + } else { + cache_seqlens_[batch_idx] = + (cache_seqlens_[batch_idx] % config_.block_len) + + i * config_.block_len; + } + for (int j = 0; j < i; j++) { + selected_blocks_history_[(layer_idx - config_.layer_offset) / + config_.layer_step][batch_idx][j] = + block_table_after_retrieval_[batch_idx][j]; + } + selected_blocks_num_history_[(layer_idx - config_.layer_offset) / + config_.layer_step] = i; + } + + // Timer end + auto end = std::chrono::high_resolution_clock::now(); + std::chrono::duration diff = end - start; + // printf("layer %d time of selecting blocks: %f s\n", layer_idx, + // diff.count()); +} + +// retrieval kvcache, get the init_block_num block at beginning, top +// pick_block_num similar and last local_block_num blocks. Each task +// calculates the simlarity of a certain block with the query, then push +// the block into the priority queue. Finally, the required blocks are +// pushed into the block_table_after_retrieval_. +void KVCache::retrieval_kvcache_layer_(const uint16_t *q_in_data, + int init_block_num, int local_block_num, + int pick_block_num, int q_len, + int generate_token_idx, int batch_size, + int layer_idx, int *cache_seqlens, + int &max_block_num, Backend *backend) { + // Timer start + auto start = std::chrono::high_resolution_clock::now(); + max_block_num_after_retrieval_ = 0; + if (pick_block_num != -1 && + (generate_token_idx % config_.token_step != 0 || + (layer_idx % config_.layer_step != config_.layer_offset))) { + + if (selected_blocks_num_history_[(layer_idx - config_.layer_offset) / + config_.layer_step] == 0) { + max_block_num_after_retrieval_ = max_block_num; + block_table_after_retrieval_.swap(block_table_before_retrieval_); + } else { + max_block_num_after_retrieval_ = selected_blocks_num_history_ + [(layer_idx - config_.layer_offset) / config_.layer_step]; + for (int batch_idx = 0; batch_idx < batch_size; batch_idx++) { + for (int i = 0; i < max_block_num_after_retrieval_; i++) { + block_table_after_retrieval_[batch_idx][i] = + selected_blocks_history_[(layer_idx - + config_.layer_offset) / + config_.layer_step][batch_idx] + [i]; + } + + if (cache_seqlens[batch_idx] % config_.block_len == 1) { + selected_blocks_num_history_[(layer_idx - + config_.layer_offset) / + config_.layer_step] += 1; + int x = + selected_blocks_num_history_[(layer_idx - + config_.layer_offset) / + config_.layer_step]; + int last_block_idx = + block_table_before_retrieval_[batch_idx] + [cache_seqlens[batch_idx] / + config_.block_len]; + selected_blocks_history_[(layer_idx - + config_.layer_offset) / + config_.layer_step][batch_idx] + [x - 1] = last_block_idx; + block_table_after_retrieval_[batch_idx][x - 1] = + last_block_idx; + } + cache_seqlens_[batch_idx] = + (cache_seqlens_[batch_idx] % config_.block_len) + + selected_blocks_num_history_[(layer_idx - + config_.layer_offset) / + config_.layer_step] * + config_.block_len - + config_.block_len; + } + } + } else if (pick_block_num != -1) { + max_block_num_after_retrieval_ = + std::min(max_block_num, + init_block_num + pick_block_num + local_block_num + 1); + calculate_block_similarity_layer_(q_in_data, batch_size, layer_idx, + q_len, max_block_num, cache_seqlens, + init_block_num, local_block_num, + pick_block_num, backend); + select_block_layer_(batch_size, layer_idx, max_block_num, + init_block_num, local_block_num, pick_block_num); + } else { + selected_blocks_num_history_[(layer_idx - config_.layer_offset) / + config_.layer_step] = 0; + max_block_num_after_retrieval_ = max_block_num; + block_table_after_retrieval_.swap(block_table_before_retrieval_); + } + + // Timer end + auto end = std::chrono::high_resolution_clock::now(); + // printf("layer %d time of retrieval kvcache: %f s\n", layer_idx, + // std::chrono::duration(end - start).count()); +} +void KVCache::calculate_sparsity_layer_(const uint16_t *q_in_data, + float *attn_sparsity, int batch_size, + int max_block_num, int *block_table, + int *cache_seqlens, Backend *backend + +) { + // Timer start + auto start = std::chrono::high_resolution_clock::now(); + seq_len_ = config_.block_len; + backend->do_work_stealing_job( + batch_size * config_.kv_head_num * max_block_num, + [&](int thread_id) { + thread_cur_head_idx_[thread_id].first = -1; + thread_cur_head_idx_[thread_id].second = -1; + }, + [&](int task_id) { + int batch_id = task_id / (config_.kv_head_num * max_block_num); + int head_id = (task_id % (config_.kv_head_num * max_block_num)) / + max_block_num; + int block_id = task_id % max_block_num; + int thread_id = Backend::thread_local_id; + // If the block is out of the sequence length, skip it. + if (cache_seqlens[batch_id] / config_.block_len < block_id) { + return; + } + int block_idx = block_table[batch_id * max_block_num + block_id]; + if (cache_seqlens_[batch_id] / config_.block_len == block_id) { + int seq_len = cache_seqlens_[batch_id] % config_.block_len; + if (seq_len == 0) + return; + + // Prepare the attention mask for the last block. + int full_blocks = seq_len / 8; + int remaining_bits = seq_len % 8; + // Fill full blocks with 1s + for (int i = 0; i < full_blocks; ++i) { + thread_local_attn_mask_[thread_id][i] = 0xFF; + } + // Fill the remaining bits in the next block + if (remaining_bits > 0 && full_blocks < seq_len_ / 8) { + thread_local_attn_mask_[thread_id][full_blocks] = + (1 << remaining_bits) - 1; + } else { + thread_local_attn_mask_[thread_id][full_blocks] = 0; + } + + for (int i = full_blocks + 1; i < seq_len_ / 8; ++i) { + thread_local_attn_mask_[thread_id][i] = 0; + } + if (config_.kv_type == ggml_type::GGML_TYPE_F16) { + attn_with_kvcache_one_block_( + config_.head_dim, + config_.q_head_num / config_.kv_head_num, GGML_TYPE_F16, + (void *)&q_in_data[batch_id * config_.kv_head_num * + n_gqa_ * config_.head_dim + + head_id * n_gqa_ * config_.head_dim], + seq_len_, 0, false, + thread_local_attn_mask_[thread_id].data(), + GGML_TYPE_F16, 0, + k_cache_fp16_[layer_id_][head_id][block_idx].data(), 0, + nullptr, nullptr, GGML_TYPE_F16, 1, + v_cache_fp16_[layer_id_][head_id][block_idx].data(), 0, + nullptr, nullptr, + thread_local_attn_score_[thread_id].data(), + thread_local_output_fp32_[thread_id].data(), + thread_local_attn_lse_[thread_id].data(), + thread_local_draft_[thread_id].data(), nullptr, + cos_.data(), sin_.data()); + } else if (config_.kv_type == ggml_type::GGML_TYPE_Q4_0) { + attn_with_kvcache_one_block_( + config_.head_dim, + config_.q_head_num / config_.kv_head_num, + GGML_TYPE_Q8_0, q_q8_0_[batch_id][head_id].data(), + seq_len_, 0, false, + thread_local_attn_mask_[thread_id].data(), + GGML_TYPE_Q4_0, 0, + k_cache_q4[layer_id_][head_id][block_idx].data(), 0, + nullptr, nullptr, GGML_TYPE_Q4_0, 1, + v_cache_q4[layer_id_][head_id][block_idx].data(), 0, + nullptr, nullptr, + thread_local_attn_score_[thread_id].data(), + thread_local_output_q8_0_[thread_id].data(), + thread_local_attn_lse_[thread_id].data(), + thread_local_draft_[thread_id].data(), nullptr, + cos_.data(), sin_.data()); + dequantize_row_q8_0( + thread_local_output_q8_0_[thread_id].data(), + thread_local_output_fp32_[thread_id].data(), + n_gqa_ * config_.head_dim); + } else if (config_.kv_type == ggml_type::GGML_TYPE_Q8_0) { + attn_with_kvcache_one_block_( + config_.head_dim, + config_.q_head_num / config_.kv_head_num, + GGML_TYPE_Q8_0, q_q8_0_[batch_id][head_id].data(), + seq_len_, 0, false, + thread_local_attn_mask_[thread_id].data(), + GGML_TYPE_Q8_0, 0, + k_cache_q8[layer_id_][head_id][block_idx].data(), 0, + nullptr, nullptr, GGML_TYPE_Q8_0, 1, + v_cache_q8[layer_id_][head_id][block_idx].data(), 0, + nullptr, nullptr, + thread_local_attn_score_[thread_id].data(), + thread_local_output_q8_0_[thread_id].data(), + thread_local_attn_lse_[thread_id].data(), + thread_local_draft_[thread_id].data(), nullptr, + cos_.data(), sin_.data()); + dequantize_row_q8_0( + thread_local_output_q8_0_[thread_id].data(), + thread_local_output_fp32_[thread_id].data(), + n_gqa_ * config_.head_dim); + } + } else { + if (config_.kv_type == ggml_type::GGML_TYPE_F16) { + attn_with_kvcache_one_block_( + config_.head_dim, + config_.q_head_num / config_.kv_head_num, GGML_TYPE_F16, + (void *)&q_in_data[batch_id * config_.kv_head_num * + n_gqa_ * config_.head_dim + + head_id * n_gqa_ * config_.head_dim], + seq_len_, 0, true, nullptr, GGML_TYPE_F16, 0, + k_cache_fp16_[layer_id_][head_id][block_idx].data(), 0, + nullptr, nullptr, GGML_TYPE_F16, 1, + v_cache_fp16_[layer_id_][head_id][block_idx].data(), 0, + nullptr, nullptr, + thread_local_attn_score_[thread_id].data(), + thread_local_output_fp32_[thread_id].data(), + thread_local_attn_lse_[thread_id].data(), + thread_local_draft_[thread_id].data(), nullptr, + cos_.data(), sin_.data()); + + } else if (config_.kv_type == ggml_type::GGML_TYPE_Q4_0) { + attn_with_kvcache_one_block_( + config_.head_dim, + config_.q_head_num / config_.kv_head_num, + GGML_TYPE_Q8_0, q_q8_0_[batch_id][head_id].data(), + seq_len_, 0, true, nullptr, GGML_TYPE_Q4_0, 0, + k_cache_q4[layer_id_][head_id][block_idx].data(), 0, + nullptr, nullptr, GGML_TYPE_Q4_0, 1, + v_cache_q4[layer_id_][head_id][block_idx].data(), 0, + nullptr, nullptr, + thread_local_attn_score_[thread_id].data(), + thread_local_output_q8_0_[thread_id].data(), + thread_local_attn_lse_[thread_id].data(), + thread_local_draft_[thread_id].data(), nullptr, + cos_.data(), sin_.data()); + dequantize_row_q8_0( + thread_local_output_q8_0_[thread_id].data(), + thread_local_output_fp32_[thread_id].data(), + n_gqa_ * config_.head_dim); + } else if (config_.kv_type == ggml_type::GGML_TYPE_Q8_0) { + attn_with_kvcache_one_block_( + config_.head_dim, + config_.q_head_num / config_.kv_head_num, + GGML_TYPE_Q8_0, q_q8_0_[batch_id][head_id].data(), + seq_len_, 0, true, nullptr, GGML_TYPE_Q8_0, 0, + k_cache_q8[layer_id_][head_id][block_idx].data(), 0, + nullptr, nullptr, GGML_TYPE_Q8_0, 1, + v_cache_q8[layer_id_][head_id][block_idx].data(), 0, + nullptr, nullptr, + thread_local_attn_score_[thread_id].data(), + thread_local_output_q8_0_[thread_id].data(), + thread_local_attn_lse_[thread_id].data(), + thread_local_draft_[thread_id].data(), nullptr, + cos_.data(), sin_.data()); + dequantize_row_q8_0( + thread_local_output_q8_0_[thread_id].data(), + thread_local_output_fp32_[thread_id].data(), + n_gqa_ * config_.head_dim); + } + } + for (int i = 0; i < n_gqa_; i++) { + block_lse_[batch_id][block_idx][head_id * n_gqa_ + i] = + thread_local_attn_lse_[thread_id][i]; + } + int cur_batch_idx = thread_cur_head_idx_[thread_id].first; + int cur_head_id = thread_cur_head_idx_[thread_id].second; + if (batch_id == cur_batch_idx && head_id == cur_head_id) { + for (int i = 0; i < n_gqa_; i++) { + float new_attn_lse = + thread_local_cur_attn_lse_[thread_id][i] + + std::log( + 1.0 + + std::exp(thread_local_attn_lse_[thread_id][i] - + thread_local_cur_attn_lse_[thread_id][i])); + ggml_vec_scale_f32( + config_.head_dim, + thread_local_cur_output_fp32_[thread_id].data() + + i * config_.head_dim, + std::exp(thread_local_cur_attn_lse_[thread_id][i] - + new_attn_lse)); + ggml_vec_scale_f32( + config_.head_dim, + thread_local_output_fp32_[thread_id].data() + + i * config_.head_dim, + std::exp(thread_local_attn_lse_[thread_id][i] - + new_attn_lse)); + for (int j = 0; j < config_.head_dim; j++) { + thread_local_cur_output_fp32_[thread_id] + [i * config_.head_dim + + j] += + thread_local_output_fp32_[thread_id] + [i * config_.head_dim + j]; + } + thread_local_cur_attn_lse_[thread_id][i] = new_attn_lse; + } + } else { + if (cur_batch_idx != -1) { + mutex_[cur_batch_idx][cur_head_id]->lock(); + for (int i = 0; i < n_gqa_; i++) { + if (std::abs(attn_lse_[cur_batch_idx][cur_head_id][i]) < + 1e-6) { + attn_lse_[cur_batch_idx][cur_head_id][i] = + thread_local_cur_attn_lse_[thread_id][i]; + for (int j = 0; j < config_.head_dim; j++) { + output_fp32_[cur_batch_idx][cur_head_id] + [i * config_.head_dim + j] = + thread_local_cur_output_fp32_ + [thread_id] + [i * config_.head_dim + j]; + } + continue; + } + float new_attn_lse = + attn_lse_[cur_batch_idx][cur_head_id][i] + + std::log( + 1.0 + + std::exp( + thread_local_cur_attn_lse_[thread_id][i] - + attn_lse_[cur_batch_idx][cur_head_id][i])); + ggml_vec_scale_f32( + config_.head_dim, + output_fp32_[cur_batch_idx][cur_head_id].data() + + i * config_.head_dim, + std::exp(attn_lse_[cur_batch_idx][cur_head_id][i] - + new_attn_lse)); + ggml_vec_scale_f32( + config_.head_dim, + thread_local_cur_output_fp32_[thread_id].data() + + i * config_.head_dim, + std::exp(thread_local_cur_attn_lse_[thread_id][i] - + new_attn_lse)); + for (int j = 0; j < config_.head_dim; j++) { + output_fp32_[cur_batch_idx][cur_head_id] + [i * config_.head_dim + j] += + thread_local_cur_output_fp32_ + [thread_id][i * config_.head_dim + j]; + } + attn_lse_[cur_batch_idx][cur_head_id][i] = new_attn_lse; + } + mutex_[cur_batch_idx][cur_head_id]->unlock(); + } + thread_cur_head_idx_[thread_id].first = batch_id; + thread_cur_head_idx_[thread_id].second = head_id; + for (int i = 0; i < n_gqa_; i++) { + thread_local_cur_attn_lse_[thread_id][i] = + thread_local_attn_lse_[thread_id][i]; + for (int j = 0; j < config_.head_dim; j++) { + thread_local_cur_output_fp32_ + [thread_id][i * config_.head_dim + j] = + thread_local_output_fp32_[thread_id] + [i * config_.head_dim + + j]; + } + } + } + }, + // Merge the results of the remaining blocks. + [&](int thread_id) { + int cur_batch_idx = thread_cur_head_idx_[thread_id].first; + int cur_head_id = thread_cur_head_idx_[thread_id].second; + if (cur_head_id != -1) { + mutex_[cur_batch_idx][cur_head_id]->lock(); + for (int i = 0; i < n_gqa_; i++) { + float new_attn_lse; + if (std::abs(attn_lse_[cur_batch_idx][cur_head_id][i]) < + 1e-6) { + attn_lse_[cur_batch_idx][cur_head_id][i] = + thread_local_cur_attn_lse_[thread_id][i]; + for (int j = 0; j < config_.head_dim; j++) { + output_fp32_[cur_batch_idx][cur_head_id] + [i * config_.head_dim + j] = + thread_local_cur_output_fp32_ + [thread_id] + [i * config_.head_dim + j]; + } + continue; + } + new_attn_lse = + attn_lse_[cur_batch_idx][cur_head_id][i] + + std::log( + 1.0 + + std::exp(thread_local_cur_attn_lse_[thread_id][i] - + attn_lse_[cur_batch_idx][cur_head_id][i])); + ggml_vec_scale_f32( + config_.head_dim, + output_fp32_[cur_batch_idx][cur_head_id].data() + + i * config_.head_dim, + std::exp(attn_lse_[cur_batch_idx][cur_head_id][i] - + new_attn_lse)); + ggml_vec_scale_f32( + config_.head_dim, + thread_local_cur_output_fp32_[thread_id].data() + + i * config_.head_dim, + std::exp(thread_local_cur_attn_lse_[thread_id][i] - + new_attn_lse)); + for (int j = 0; j < config_.head_dim; j++) { + output_fp32_[cur_batch_idx][cur_head_id] + [i * config_.head_dim + j] += + thread_local_cur_output_fp32_[thread_id] + [i * config_.head_dim + + j]; + } + attn_lse_[cur_batch_idx][cur_head_id][i] = new_attn_lse; + } + mutex_[cur_batch_idx][cur_head_id]->unlock(); + } + }); + + for (int i = 0; i < batch_size; i++) { + for (int j = 0; j < max_block_num_after_retrieval_; j++) { + int block_idx = block_table_after_retrieval_[i][j]; + for (int k = 0; k < config_.q_head_num; k++) { + attn_sparsity[i * config_.q_head_num + k] += + std::exp(block_lse_[i][block_idx][k] - + attn_lse_[i][k / n_gqa_][k % n_gqa_]); + } + } + } + + // Timer end + auto end = std::chrono::high_resolution_clock::now(); + std::chrono::duration diff = end - start; + // printf("layer %d time of calculating sparsity: %f s\n", layer_id_, + // diff.count()); +} + +void KVCache::attn_initialize_kvhead_(int batch_size, int layer_idx, + int *block_table, int &max_block_num, + int *cache_seqlens) { + // Timer start + auto start = std::chrono::high_resolution_clock::now(); + for (int batch_idx = 0; batch_idx < batch_size; batch_idx++) { + // initialize output_fp32_ and attn_lse_ + for (int i = 0; i < config_.kv_head_num; i++) { + for (int j = 0; j < n_gqa_ * config_.head_dim; j++) { + output_fp32_[batch_idx][i][j] = 0; + } + for (int j = 0; j < n_gqa_; j++) { + attn_lse_[batch_idx][i][j] = 0; + } + } + + // clear top_similar_block_ + while (!top_similar_block_[batch_idx].empty()) + top_similar_block_[batch_idx].pop(); + } + + for (int batch_idx = 0; batch_idx < batch_size; batch_idx++) { + cache_seqlens_[batch_idx] = cache_seqlens[batch_idx]; + for (int i = 0; i < max_block_num; i++) { + for (int j = 0; j < config_.kv_head_num; j++) { + block_table_before_retrieval_kvhead_[batch_idx][i][j] = + block_table[batch_idx * max_block_num + i]; + block_similar_kv_head_[batch_idx][i][j] = 0; + } + } + } + + // Timer end + auto end = std::chrono::high_resolution_clock::now(); + // printf("layer %d time of initializing attn: %f s\n", layer_idx, + // std::chrono::duration(end - start).count()); +} +void KVCache::retrieval_kvcache_kvhead_(const uint16_t *q_in_data, + int init_block_num, int local_block_num, + int pick_block_num, int q_len, + int generate_token_idx, int batch_size, + int layer_idx, int *cache_seqlens, + int &max_block_num, Backend *backend) { + // Timer start + auto start = std::chrono::high_resolution_clock::now(); + max_block_num_after_retrieval_ = 0; + if (pick_block_num != -1 && + (generate_token_idx % config_.token_step != 0 || + (layer_idx % config_.layer_step != config_.layer_offset))) { + + if (selected_blocks_num_history_[(layer_idx - config_.layer_offset) / + config_.layer_step] == 0) { + max_block_num_after_retrieval_ = max_block_num; + for (int batch_idx = 0; batch_idx < batch_size; batch_idx++) { + for (int i = 0; i < max_block_num; i++) { + for (int j = 0; j < config_.kv_head_num; j++) { + block_table_after_retrieval_kvhead_[batch_idx][i][j] = + block_table_before_retrieval_kvhead_[batch_idx][i] + [j]; + } + } + } + } else { + + max_block_num_after_retrieval_ = selected_blocks_num_history_ + [(layer_idx - config_.layer_offset) / config_.layer_step]; + + for (int batch_idx = 0; batch_idx < batch_size; batch_idx++) { + for (int i = 0; i < max_block_num_after_retrieval_; i++) { + for (int j = 0; j < config_.kv_head_num; j++) { + block_table_after_retrieval_kvhead_[batch_idx][i][j] = + selected_blocks_history_kvhead_ + [(layer_idx - config_.layer_offset) / + config_.layer_step][batch_idx][i][j]; + } + } + + if (cache_seqlens[batch_idx] % config_.block_len == 1) { + selected_blocks_num_history_[(layer_idx - + config_.layer_offset) / + config_.layer_step] += 1; + int x = + selected_blocks_num_history_[(layer_idx - + config_.layer_offset) / + config_.layer_step]; + for (int i = 0; i < config_.kv_head_num; i++) { + int last_block_idx = + block_table_before_retrieval_kvhead_ + [batch_idx][cache_seqlens[batch_idx] / + config_.block_len][i]; + selected_blocks_history_kvhead_[(layer_idx - + config_.layer_offset) / + config_.layer_step] + [batch_idx][x - 1][i] = + last_block_idx; + block_table_after_retrieval_kvhead_[batch_idx][x - 1] + [i] = last_block_idx; + } + } + cache_seqlens_[batch_idx] = std::min( + cache_seqlens_[batch_idx], + (cache_seqlens_[batch_idx] % config_.block_len) + + (init_block_num + pick_block_num + local_block_num) * + config_.block_len); + } + } + } else if (pick_block_num != -1) { + max_block_num_after_retrieval_ = + std::min(max_block_num, + init_block_num + pick_block_num + local_block_num + 1); + calculate_block_similarity_kvhead_(q_in_data, batch_size, layer_idx, + q_len, max_block_num, cache_seqlens, + init_block_num, local_block_num, + pick_block_num, backend); + select_block_kvhead_(batch_size, layer_idx, max_block_num, + init_block_num, local_block_num, pick_block_num); + } else { + selected_blocks_num_history_[(layer_idx - config_.layer_offset) / + config_.layer_step] = 0; + max_block_num_after_retrieval_ = max_block_num; + for (int batch_idx = 0; batch_idx < batch_size; batch_idx++) { + for (int i = 0; i < max_block_num; i++) { + for (int j = 0; j < config_.kv_head_num; j++) { + block_table_after_retrieval_kvhead_[batch_idx][i][j] = + block_table_before_retrieval_kvhead_[batch_idx][i][j]; + } + } + } + } + + // Timer end + auto end = std::chrono::high_resolution_clock::now(); + // printf("layer %d time of retrieval kvcache: %f s\n", layer_idx, + // std::chrono::duration(end - start).count()); +} +void KVCache::calculate_sparsity_kvhead_(const uint16_t *q_in_data, + float *attn_sparsity, int batch_size, + int max_block_num, int *block_table, + int *cache_seqlens, Backend *backend) { + // Timer start + auto start = std::chrono::high_resolution_clock::now(); + seq_len_ = config_.block_len; + backend->do_work_stealing_job( + batch_size * config_.kv_head_num * max_block_num, + [&](int thread_id) { + thread_cur_head_idx_[thread_id].first = -1; + thread_cur_head_idx_[thread_id].second = -1; + }, + [&](int task_id) { + int batch_id = task_id / (config_.kv_head_num * max_block_num); + int head_id = (task_id % (config_.kv_head_num * max_block_num)) / + max_block_num; + int block_id = task_id % max_block_num; + int thread_id = Backend::thread_local_id; + // If the block is out of the sequence length, skip it. + if (cache_seqlens[batch_id] / config_.block_len < block_id) { + return; + } + int block_idx = block_table[batch_id * max_block_num + block_id]; + if (cache_seqlens_[batch_id] / config_.block_len == block_id) { + int seq_len = cache_seqlens_[batch_id] % config_.block_len; + if (seq_len == 0) + return; + + // Prepare the attention mask for the last block. + int full_blocks = seq_len / 8; + int remaining_bits = seq_len % 8; + + // Fill full blocks with 1s + for (int i = 0; i < full_blocks; ++i) { + thread_local_attn_mask_[thread_id][i] = 0xFF; + } + // Fill the remaining bits in the next block + if (remaining_bits > 0 && full_blocks < seq_len_ / 8) { + thread_local_attn_mask_[thread_id][full_blocks] = + (1 << remaining_bits) - 1; + } else { + thread_local_attn_mask_[thread_id][full_blocks] = 0; + } + + for (int i = full_blocks + 1; i < seq_len_ / 8; ++i) { + thread_local_attn_mask_[thread_id][i] = 0; + } + if (config_.kv_type == ggml_type::GGML_TYPE_F16) { + attn_with_kvcache_one_block_( + config_.head_dim, + config_.q_head_num / config_.kv_head_num, GGML_TYPE_F16, + (void *)&q_in_data[batch_id * config_.kv_head_num * + n_gqa_ * config_.head_dim + + head_id * n_gqa_ * config_.head_dim], + seq_len_, 0, false, + thread_local_attn_mask_[thread_id].data(), + GGML_TYPE_F16, 0, + k_cache_fp16_[layer_id_][head_id][block_idx].data(), 0, + nullptr, nullptr, GGML_TYPE_F16, 1, + v_cache_fp16_[layer_id_][head_id][block_idx].data(), 0, + nullptr, nullptr, + thread_local_attn_score_[thread_id].data(), + thread_local_output_fp32_[thread_id].data(), + thread_local_attn_lse_[thread_id].data(), + thread_local_draft_[thread_id].data(), nullptr, + cos_.data(), sin_.data()); + } else if (config_.kv_type == ggml_type::GGML_TYPE_Q4_0) { + attn_with_kvcache_one_block_( + config_.head_dim, + config_.q_head_num / config_.kv_head_num, + GGML_TYPE_Q8_0, q_q8_0_[batch_id][head_id].data(), + seq_len_, 0, false, + thread_local_attn_mask_[thread_id].data(), + GGML_TYPE_Q4_0, 0, + k_cache_q4[layer_id_][head_id][block_idx].data(), 0, + nullptr, nullptr, GGML_TYPE_Q4_0, 1, + v_cache_q4[layer_id_][head_id][block_idx].data(), 0, + nullptr, nullptr, + thread_local_attn_score_[thread_id].data(), + thread_local_output_q8_0_[thread_id].data(), + thread_local_attn_lse_[thread_id].data(), + thread_local_draft_[thread_id].data(), nullptr, + cos_.data(), sin_.data()); + dequantize_row_q8_0( + thread_local_output_q8_0_[thread_id].data(), + thread_local_output_fp32_[thread_id].data(), + n_gqa_ * config_.head_dim); + } else if (config_.kv_type == ggml_type::GGML_TYPE_Q8_0) { + attn_with_kvcache_one_block_( + config_.head_dim, + config_.q_head_num / config_.kv_head_num, + GGML_TYPE_Q8_0, q_q8_0_[batch_id][head_id].data(), + seq_len_, 0, false, + thread_local_attn_mask_[thread_id].data(), + GGML_TYPE_Q8_0, 0, + k_cache_q8[layer_id_][head_id][block_idx].data(), 0, + nullptr, nullptr, GGML_TYPE_Q8_0, 1, + v_cache_q8[layer_id_][head_id][block_idx].data(), 0, + nullptr, nullptr, + thread_local_attn_score_[thread_id].data(), + thread_local_output_q8_0_[thread_id].data(), + thread_local_attn_lse_[thread_id].data(), + thread_local_draft_[thread_id].data(), nullptr, + cos_.data(), sin_.data()); + dequantize_row_q8_0( + thread_local_output_q8_0_[thread_id].data(), + thread_local_output_fp32_[thread_id].data(), + n_gqa_ * config_.head_dim); + } + } else { + if (config_.kv_type == ggml_type::GGML_TYPE_F16) { + attn_with_kvcache_one_block_( + config_.head_dim, + config_.q_head_num / config_.kv_head_num, GGML_TYPE_F16, + (void *)&q_in_data[batch_id * config_.kv_head_num * + n_gqa_ * config_.head_dim + + head_id * n_gqa_ * config_.head_dim], + seq_len_, 0, true, nullptr, GGML_TYPE_F16, 0, + k_cache_fp16_[layer_id_][head_id][block_idx].data(), 0, + nullptr, nullptr, GGML_TYPE_F16, 1, + v_cache_fp16_[layer_id_][head_id][block_idx].data(), 0, + nullptr, nullptr, + thread_local_attn_score_[thread_id].data(), + thread_local_output_fp32_[thread_id].data(), + thread_local_attn_lse_[thread_id].data(), + thread_local_draft_[thread_id].data(), nullptr, + cos_.data(), sin_.data()); + + } else if (config_.kv_type == ggml_type::GGML_TYPE_Q4_0) { + attn_with_kvcache_one_block_( + config_.head_dim, + config_.q_head_num / config_.kv_head_num, + GGML_TYPE_Q8_0, q_q8_0_[batch_id][head_id].data(), + seq_len_, 0, true, nullptr, GGML_TYPE_Q4_0, 0, + k_cache_q4[layer_id_][head_id][block_idx].data(), 0, + nullptr, nullptr, GGML_TYPE_Q4_0, 1, + v_cache_q4[layer_id_][head_id][block_idx].data(), 0, + nullptr, nullptr, + thread_local_attn_score_[thread_id].data(), + thread_local_output_q8_0_[thread_id].data(), + thread_local_attn_lse_[thread_id].data(), + thread_local_draft_[thread_id].data(), nullptr, + cos_.data(), sin_.data()); + dequantize_row_q8_0( + thread_local_output_q8_0_[thread_id].data(), + thread_local_output_fp32_[thread_id].data(), + n_gqa_ * config_.head_dim); + } else if (config_.kv_type == ggml_type::GGML_TYPE_Q8_0) { + attn_with_kvcache_one_block_( + config_.head_dim, + config_.q_head_num / config_.kv_head_num, + GGML_TYPE_Q8_0, q_q8_0_[batch_id][head_id].data(), + seq_len_, 0, true, nullptr, GGML_TYPE_Q8_0, 0, + k_cache_q8[layer_id_][head_id][block_idx].data(), 0, + nullptr, nullptr, GGML_TYPE_Q8_0, 1, + v_cache_q8[layer_id_][head_id][block_idx].data(), 0, + nullptr, nullptr, + thread_local_attn_score_[thread_id].data(), + thread_local_output_q8_0_[thread_id].data(), + thread_local_attn_lse_[thread_id].data(), + thread_local_draft_[thread_id].data(), nullptr, + cos_.data(), sin_.data()); + dequantize_row_q8_0( + thread_local_output_q8_0_[thread_id].data(), + thread_local_output_fp32_[thread_id].data(), + n_gqa_ * config_.head_dim); + } + } + for (int i = 0; i < n_gqa_; i++) { + block_lse_[batch_id][block_idx][head_id * n_gqa_ + i] = + thread_local_attn_lse_[thread_id][i]; + } + int cur_batch_idx = thread_cur_head_idx_[thread_id].first; + int cur_head_id = thread_cur_head_idx_[thread_id].second; + if (batch_id == cur_batch_idx && head_id == cur_head_id) { + for (int i = 0; i < n_gqa_; i++) { + float new_attn_lse = + thread_local_cur_attn_lse_[thread_id][i] + + std::log( + 1.0 + + std::exp(thread_local_attn_lse_[thread_id][i] - + thread_local_cur_attn_lse_[thread_id][i])); + ggml_vec_scale_f32( + config_.head_dim, + thread_local_cur_output_fp32_[thread_id].data() + + i * config_.head_dim, + std::exp(thread_local_cur_attn_lse_[thread_id][i] - + new_attn_lse)); + ggml_vec_scale_f32( + config_.head_dim, + thread_local_output_fp32_[thread_id].data() + + i * config_.head_dim, + std::exp(thread_local_attn_lse_[thread_id][i] - + new_attn_lse)); + for (int j = 0; j < config_.head_dim; j++) { + thread_local_cur_output_fp32_[thread_id] + [i * config_.head_dim + + j] += + thread_local_output_fp32_[thread_id] + [i * config_.head_dim + j]; + } + thread_local_cur_attn_lse_[thread_id][i] = new_attn_lse; + } + } else { + if (cur_batch_idx != -1) { + mutex_[cur_batch_idx][cur_head_id]->lock(); + for (int i = 0; i < n_gqa_; i++) { + if (std::abs(attn_lse_[cur_batch_idx][cur_head_id][i]) < + 1e-6) { + attn_lse_[cur_batch_idx][cur_head_id][i] = + thread_local_cur_attn_lse_[thread_id][i]; + for (int j = 0; j < config_.head_dim; j++) { + output_fp32_[cur_batch_idx][cur_head_id] + [i * config_.head_dim + j] = + thread_local_cur_output_fp32_ + [thread_id] + [i * config_.head_dim + j]; + } + continue; + } + float new_attn_lse = + attn_lse_[cur_batch_idx][cur_head_id][i] + + std::log( + 1.0 + + std::exp( + thread_local_cur_attn_lse_[thread_id][i] - + attn_lse_[cur_batch_idx][cur_head_id][i])); + ggml_vec_scale_f32( + config_.head_dim, + output_fp32_[cur_batch_idx][cur_head_id].data() + + i * config_.head_dim, + std::exp(attn_lse_[cur_batch_idx][cur_head_id][i] - + new_attn_lse)); + ggml_vec_scale_f32( + config_.head_dim, + thread_local_cur_output_fp32_[thread_id].data() + + i * config_.head_dim, + std::exp(thread_local_cur_attn_lse_[thread_id][i] - + new_attn_lse)); + for (int j = 0; j < config_.head_dim; j++) { + output_fp32_[cur_batch_idx][cur_head_id] + [i * config_.head_dim + j] += + thread_local_cur_output_fp32_ + [thread_id][i * config_.head_dim + j]; + } + attn_lse_[cur_batch_idx][cur_head_id][i] = new_attn_lse; + } + mutex_[cur_batch_idx][cur_head_id]->unlock(); + } + thread_cur_head_idx_[thread_id].first = batch_id; + thread_cur_head_idx_[thread_id].second = head_id; + for (int i = 0; i < n_gqa_; i++) { + thread_local_cur_attn_lse_[thread_id][i] = + thread_local_attn_lse_[thread_id][i]; + for (int j = 0; j < config_.head_dim; j++) { + thread_local_cur_output_fp32_ + [thread_id][i * config_.head_dim + j] = + thread_local_output_fp32_[thread_id] + [i * config_.head_dim + + j]; + } + } + } + }, + // Merge the results of the remaining blocks. + [&](int thread_id) { + int cur_batch_idx = thread_cur_head_idx_[thread_id].first; + int cur_head_id = thread_cur_head_idx_[thread_id].second; + if (cur_head_id != -1) { + mutex_[cur_batch_idx][cur_head_id]->lock(); + for (int i = 0; i < n_gqa_; i++) { + float new_attn_lse; + if (std::abs(attn_lse_[cur_batch_idx][cur_head_id][i]) < + 1e-6) { + attn_lse_[cur_batch_idx][cur_head_id][i] = + thread_local_cur_attn_lse_[thread_id][i]; + for (int j = 0; j < config_.head_dim; j++) { + output_fp32_[cur_batch_idx][cur_head_id] + [i * config_.head_dim + j] = + thread_local_cur_output_fp32_ + [thread_id] + [i * config_.head_dim + j]; + } + continue; + } + new_attn_lse = + attn_lse_[cur_batch_idx][cur_head_id][i] + + std::log( + 1.0 + + std::exp(thread_local_cur_attn_lse_[thread_id][i] - + attn_lse_[cur_batch_idx][cur_head_id][i])); + ggml_vec_scale_f32( + config_.head_dim, + output_fp32_[cur_batch_idx][cur_head_id].data() + + i * config_.head_dim, + std::exp(attn_lse_[cur_batch_idx][cur_head_id][i] - + new_attn_lse)); + ggml_vec_scale_f32( + config_.head_dim, + thread_local_cur_output_fp32_[thread_id].data() + + i * config_.head_dim, + std::exp(thread_local_cur_attn_lse_[thread_id][i] - + new_attn_lse)); + for (int j = 0; j < config_.head_dim; j++) { + output_fp32_[cur_batch_idx][cur_head_id] + [i * config_.head_dim + j] += + thread_local_cur_output_fp32_[thread_id] + [i * config_.head_dim + + j]; + } + attn_lse_[cur_batch_idx][cur_head_id][i] = new_attn_lse; + } + mutex_[cur_batch_idx][cur_head_id]->unlock(); + } + }); + + for (int i = 0; i < batch_size; i++) { + for (int j = 0; j < max_block_num_after_retrieval_; j++) { + for (int k = 0; k < config_.q_head_num; k++) { + int block_idx = + block_table_after_retrieval_kvhead_[i][j][k / n_gqa_]; + attn_sparsity[i * config_.q_head_num + k] += + std::exp(block_lse_[i][block_idx][k] - + attn_lse_[i][k / n_gqa_][k % n_gqa_]); + } + } + } + + // Timer end + auto end = std::chrono::high_resolution_clock::now(); + std::chrono::duration diff = end - start; + // printf("layer %d time of calculating sparsity: %f s\n", layer_id_, + // diff.count()); +} +void KVCache::calculate_block_similarity_kvhead_( + const uint16_t *q_in_data, int batch_size, int layer_idx, int q_len, + int max_block_num, int *cache_seqlens, int init_block_num, + int local_block_num, int pick_block_num, Backend *backend) { + // Timer start + auto start = std::chrono::high_resolution_clock::now(); + backend->do_work_stealing_job( + batch_size * max_block_num, nullptr, + [&](int task_id) { + int batch_id = task_id / max_block_num; + int block_id = task_id % max_block_num; + int seq_len = cache_seqlens_[batch_id]; + + if (block_id < init_block_num || + block_id >= (seq_len / config_.block_len) - local_block_num) { + return; + } + int block_idx = + block_table_before_retrieval_kvhead_[batch_id][block_id][0]; + + for (int head_id = 0; head_id < config_.q_head_num; head_id++) { + for (int i = 0; i < config_.head_dim; i++) { + float q_i = 0, qa_i = std::numeric_limits::lowest(); + for (int q_id = 0; q_id < q_len; q_id++) { + q_i += GGML_FP16_TO_FP32( + q_in_data[batch_id * q_len * config_.q_head_num * + config_.head_dim + + q_id * config_.q_head_num * + config_.head_dim + + head_id * config_.head_dim + i]); + } + q_i /= q_len; + for (int anchor_id = 0; anchor_id < config_.anchor_num; + anchor_id++) { + qa_i = std::max( + qa_i, + GGML_FP16_TO_FP32( + anchor_[layer_idx * config_.max_block_num * + config_.anchor_num * + config_.q_head_num * + config_.head_dim + + block_idx * config_.anchor_num * + config_.q_head_num * + config_.head_dim + + anchor_id * config_.q_head_num * + config_.head_dim + + head_id * config_.head_dim + i]) * + q_i); + } + block_similar_kv_head_[batch_id][block_id] + [head_id / n_gqa_] += qa_i; + } + } + }, + nullptr); + + // Timer end + auto end = std::chrono::high_resolution_clock::now(); + std::chrono::duration diff = end - start; + // printf("layer %d time of calculating similarity: %f s\n", layer_idx, + // diff.count()); +} +void KVCache::select_block_kvhead_(int batch_size, int layer_idx, + int max_block_num, int init_block_num, + int local_block_num, int pick_block_num) { + // Timer start + auto start = std::chrono::high_resolution_clock::now(); + + for (int batch_idx = 0; batch_idx < batch_size; batch_idx++) { + int cache_len_after_retrieval = 0; + if (cache_seqlens_[batch_idx] / config_.block_len <= + init_block_num + pick_block_num + local_block_num) { + selected_blocks_num_history_[(layer_idx - config_.layer_offset) / + config_.layer_step] = 0; + for (int i = 0; i < max_block_num; i++) { + for (int j = 0; j < config_.kv_head_num; j++) { + block_table_after_retrieval_kvhead_[batch_idx][i][j] = + block_table_before_retrieval_kvhead_[batch_idx][i][j]; + } + } + continue; + } + for (int head_id = 0; head_id < config_.kv_head_num; head_id++) { + + for (int block_id = init_block_num; + block_id < (cache_seqlens_[batch_idx] / config_.block_len) - + local_block_num; + block_id++) { + + top_similar_block_[batch_idx].push(std::make_pair( + block_similar_kv_head_[batch_idx][block_id][head_id], + block_table_before_retrieval_kvhead_[batch_idx][block_id] + [head_id])); + if (top_similar_block_[batch_idx].size() > pick_block_num) { + top_similar_block_[batch_idx].pop(); + } + } + + int i = 0; + for (; i < init_block_num; i++) { + block_table_after_retrieval_kvhead_[batch_idx][i][head_id] = + block_table_before_retrieval_kvhead_[batch_idx][i][head_id]; + } + while (!top_similar_block_[batch_idx].empty()) { + block_table_after_retrieval_kvhead_[batch_idx][i][head_id] = + top_similar_block_[batch_idx].top().second; + top_similar_block_[batch_idx].pop(); + i++; + } + for (; i < init_block_num + pick_block_num + local_block_num; i++) { + block_table_after_retrieval_kvhead_[batch_idx][i][head_id] = + block_table_before_retrieval_kvhead_ + [batch_idx] + [(cache_seqlens_[batch_idx] / config_.block_len) - + local_block_num + i - init_block_num - pick_block_num] + [head_id]; + } + if (cache_seqlens_[batch_idx] % config_.block_len != 0) { + block_table_after_retrieval_kvhead_[batch_idx][i][head_id] = + block_table_before_retrieval_kvhead_[batch_idx][( + cache_seqlens_[batch_idx] / config_.block_len)] + [head_id]; + cache_len_after_retrieval = + (cache_seqlens_[batch_idx] % config_.block_len) + + i * config_.block_len; + i++; + } else { + cache_len_after_retrieval = + (cache_seqlens_[batch_idx] % config_.block_len) + + i * config_.block_len; + } + for (int j = 0; j < i; j++) { + selected_blocks_history_kvhead_ + [(layer_idx - config_.layer_offset) / config_.layer_step] + [batch_idx][j][head_id] = + block_table_after_retrieval_kvhead_[batch_idx][j] + [head_id]; + } + } + cache_seqlens_[batch_idx] = cache_len_after_retrieval; + selected_blocks_num_history_[(layer_idx - config_.layer_offset) / + config_.layer_step] = + (cache_len_after_retrieval + config_.block_len - 1) / + config_.block_len; + } + + // Timer end + auto end = std::chrono::high_resolution_clock::now(); + std::chrono::duration diff = end - start; + // printf("layer %d time of selecting block: %f s\n", layer_idx, + // diff.count()) +} + +void KVCache::get_attn_sparsity(const ggml_fp16_t *q_in, float *attn_sparsity, + int layer_idx, int generate_token_idx, + int q_len, int batch_size, int max_block_num, + int *block_table, int *cache_seqlens, + int *block_table_origin, + int *cache_seqlens_origin, + int max_block_num_origin, int topk, int local, + Backend *backend) { + // Timer start + auto start = std::chrono::high_resolution_clock::now(); + layer_id_ = layer_idx; + int thread_num = backend->get_thread_num(); + batch_size = 1; + + const uint16_t *q_in_data = const_cast(q_in); + + quantize_q_(q_in_data, batch_size); + if (config_.retrieval_type == RetrievalType::LAYER) { + attn_initialize_layer_(batch_size, layer_idx, block_table, + max_block_num, cache_seqlens); + retrieval_kvcache_layer_(q_in_data, 1, local, topk, q_len, + generate_token_idx, batch_size, layer_idx, + cache_seqlens, max_block_num, backend); + calculate_sparsity_layer_(q_in_data, attn_sparsity, batch_size, + max_block_num_origin, block_table_origin, + cache_seqlens_origin, backend); + } else if (config_.retrieval_type == RetrievalType::KVHEAD) { + attn_initialize_kvhead_(batch_size, layer_idx, block_table, + max_block_num, cache_seqlens); + retrieval_kvcache_kvhead_(q_in_data, 1, local, topk, q_len, + generate_token_idx, batch_size, layer_idx, + cache_seqlens, max_block_num, backend); + calculate_sparsity_kvhead_(q_in_data, attn_sparsity, batch_size, + max_block_num_origin, block_table_origin, + cache_seqlens_origin, backend); + } +} + +void KVCache::attn_with_kvcache_one_block_( + int head_dim, int bsz, + ggml_type q_type, // GGML data type of `Q`, only supports fp16 and q8_0 + // [bsz, head_dim] + // Quantization is always on the head_dim dimension (per_token). If + // head_dim % 32 != 0, an error will be raised. The size must be bsz * + // head_dim/32 * qtype_size. + const void *q, + + int past_kv_len, int past_kv_offset, + bool is_full_attn, // true indicates a full 1 mask + // If is_full_attn = false, a bit matrix representing the mask is + // passed. [bsz, past_kv_len] + const uint8_t *attn_mask, + + ggml_type k_type, // GGML data type of `K Cache`, only supports fp16, + // q4_0, q8_0 + int k_quant_type, // 0 for per_token, 1 for per_channel, others raise an + // error + // [seq_len, head_dim] + // If quant_type == 0, head_dim % 32 must be 0. + // If quant_type == 1, seq_len % 32 must be 0. + const void *k_cache, + + // k_anchor_type must be fp16 + int num_k_anchor, // num_k_anchor == 0 indicates no anchor + // [num_k_anchor, head_dim] + const void *k_cache_anchors, + // Each token is associated with the nearest previous position's anchor, + // with the same distance. + const int *k_cache_anchor_pos, + + // v_cache similar to k_cache + ggml_type v_type, int v_quant_type, + // [head_dim, seq_len] + const void *v_cache, int num_v_anchor, const void *v_cache_anchors, + const int *v_cache_anchor_pos, + + // Pre-allocated buffer for intermediate calculations [bsz, + // past_kv_len]. No malloc is performed inside this function. + float *attn_score, + + // Output: [bsz, head_dim], with the same type as q_type + void *output, + // [bsz] + float *lse, + + // Pre-allocated temporary buffer with sufficient size: + // (2 * bsz * past_kv_len + 6 * bsz * head_dim + 2 * past_kv_len * + // head_dim + past_kv_len * head_dim / 32) bytes. + void *draft, + + // Apply rotary embedding online + const int *rotary_angle, const void *rotary_cos, const void *rotary_sin + // rotary_cos=None, + // rotary_sin=None, + // cache_seqlens: Optional[Union[(int, torch.Tensor)]] = None, + // cache_batch_idx: Optional[torch.Tensor] = None, + // rotary_interleaved=True, + + // // Not supported for now + // window_size=(-1, -1), # -1 means infinite context window + // alibi_slopes=None, +) { + assert(head_dim % 32 == 0); + assert(k_quant_type == 0); + assert(v_quant_type == 1); + assert(q_type == GGML_TYPE_F16 || q_type == GGML_TYPE_Q8_0); + if (q_type == GGML_TYPE_F16) { + assert(k_type == GGML_TYPE_F16); + assert(v_type == GGML_TYPE_F16); + + // attn = q * k + q * k_anchor + // TODO: anchor + assert(num_k_anchor == 0); + + if (rotary_angle != nullptr) { + ggml_fp16_t *k_cache_with_rope_fp16 = + (reinterpret_cast(draft) + + sizeof(block_q8_0) * bsz * past_kv_len / QK8_0 + + sizeof(float) * bsz * head_dim); + // dequant k_cache and apply rope + // k_rope(i) = k(i) * cos(i) - k(i+l) * sin(i) + // k_rope(i+l) = k(i+l) * cos(i+l) + k(i) * sin(i) + + // k(i)cos(i) -> k_rope(i) + // k(i)sin(i+l) -> k_rope(i+l) + + // k(i)cos(i) -> k_rope(i) + // -k(i)sin(i-l) -> k_rope(i-l) + + std::vector block_fp32(32); + for (int k = 0; k < past_kv_len; k++) { + int angle = rotary_angle[k]; + for (int l = 0; l < head_dim / 32; l++) { + for (int m = 0; m < 32; m++) { + float x = GGML_FP16_TO_FP32(( + (ggml_fp16_t *)k_cache)[k * head_dim + l * 32 + m]); + float sin_val = GGML_FP16_TO_FP32( + ((ggml_fp16_t *) + rotary_sin)[angle * head_dim + l * 32 + m]); + float cos_val = GGML_FP16_TO_FP32( + ((ggml_fp16_t *) + rotary_cos)[angle * head_dim + l * 32 + m]); + + if (l * 32 + m < head_dim / 2) { + k_cache_with_rope_fp16[k * head_dim + l * 32 + m] = + GGML_FP32_TO_FP16(x * cos_val); + k_cache_with_rope_fp16[k * head_dim + l * 32 + m + + head_dim / 2] = + GGML_FP32_TO_FP16(-x * sin_val); + } else { + k_cache_with_rope_fp16[k * head_dim + l * 32 + m] = + GGML_FP32_TO_FP16( + GGML_FP16_TO_FP32( + k_cache_with_rope_fp16[k * head_dim + + l * 32 + m]) + + x * sin_val); + k_cache_with_rope_fp16[k * head_dim + l * 32 + m - + head_dim / 2] = + GGML_FP32_TO_FP16( + GGML_FP16_TO_FP32( + k_cache_with_rope_fp16[k * head_dim + + l * 32 + m - + head_dim / 2]) - + x * cos_val); + } + } + } + } + + llamafile_sgemm(past_kv_len, bsz, head_dim, + (ggml_fp16_t *)k_cache_with_rope_fp16, head_dim, + (ggml_fp16_t *)q, head_dim, attn_score, past_kv_len, + 0, 1, GGML_TASK_TYPE_COMPUTE, k_type, GGML_TYPE_F16, + GGML_TYPE_F32, GGML_PREC_DEFAULT); + } else { + bool ok = llamafile_sgemm( + past_kv_len, bsz, head_dim, (ggml_fp16_t *)k_cache, head_dim, + (ggml_fp16_t *)q, head_dim, attn_score, past_kv_len, 0, 1, + GGML_TASK_TYPE_COMPUTE, k_type, GGML_TYPE_F16, GGML_TYPE_F32, + GGML_PREC_DEFAULT); + + if (!ok) { + printf("llamafile_sgemm failed\n"); + } + } + // attn = attn * scale + float scale_factor = 1.0 / std::sqrt(float(head_dim)); + ggml_vec_scale_f32(bsz * past_kv_len, attn_score, scale_factor); + + // attn = attn & mask + if (!is_full_attn) { + for (int i = 0; i < bsz; i++) { + for (int j = 0; j < past_kv_len; j++) { + int index = i * past_kv_len + j; + if (!(attn_mask[j / 8] & (1 << (j % 8)))) { + attn_score[index] = + std::numeric_limits::lowest(); + } + } + } + } + + // attn = softmax(attn) + for (int i = 0; i < bsz; i++) { + float sum_exp = 0; + for (int j = 0; j < past_kv_len; j++) { + attn_score[i * past_kv_len + j] = + std::exp(attn_score[i * past_kv_len + j]); + sum_exp += attn_score[i * past_kv_len + j]; + } + for (int j = 0; j < past_kv_len; j++) { + attn_score[i * past_kv_len + j] /= sum_exp; + } + if (lse != nullptr) { + lse[i] = std::log(sum_exp); + } + } + + // output = attn * v + attn * v_anchor + // std::vector sum(bsz * head_dim); + float *sum = reinterpret_cast(reinterpret_cast(draft) + + sizeof(block_q8_0) * bsz * + past_kv_len / QK8_0); + + // float* attn_score_fp16(bsz, past_kv_len) + ggml_fp16_t *attn_score_fp16 = (reinterpret_cast( + reinterpret_cast(draft) + + sizeof(block_q8_0) * bsz * past_kv_len / QK8_0 + + sizeof(float) * bsz * head_dim)); + + for (int i = 0; i < bsz * past_kv_len; i++) { + attn_score_fp16[i] = GGML_FP32_TO_FP16(attn_score[i]); + } + + // TODO: anchor + assert(num_v_anchor == 0); + bool ok = llamafile_sgemm( + head_dim, bsz, past_kv_len, (ggml_fp16_t *)v_cache, past_kv_len, + (ggml_fp16_t *)attn_score_fp16, past_kv_len, sum, head_dim, 0, 1, + GGML_TASK_TYPE_COMPUTE, v_type, GGML_TYPE_F16, GGML_TYPE_F32, + GGML_PREC_DEFAULT); + if (!ok) { + printf("llamafile_sgemm failed\n"); + } + + // copy to output + for (int i = 0; i < bsz; i++) { + for (int j = 0; j < head_dim; j++) { + ((float *)output)[i * head_dim + j] = sum[i * head_dim + j]; + } + } + } else { + assert(k_type == GGML_TYPE_Q4_0 || k_type == GGML_TYPE_Q8_0); + assert(v_type == GGML_TYPE_Q4_0 || v_type == GGML_TYPE_Q8_0); + + // attn = q * k + q * k_anchor + // TODO: anchor + assert(num_k_anchor == 0); + + if (rotary_angle != nullptr) { + ggml_fp16_t *k_cache_with_rope_fp16 = + (reinterpret_cast(draft) + + sizeof(block_q8_0) * bsz * past_kv_len / QK8_0 + + sizeof(float) * bsz * head_dim); + block_q4_0 *k_cache_with_rope_q4 = + (reinterpret_cast(draft) + + sizeof(block_q8_0) * bsz * past_kv_len / QK8_0 + + sizeof(float) * bsz * head_dim) + + sizeof(ggml_fp16_t) * bsz * head_dim; + // dequant k_cache and apply rope + // k_rope(i) = k(i) * cos(i) - k(i+l) * sin(i) + // k_rope(i+l) = k(i+l) * cos(i+l) + k(i) * sin(i) + + // k(i)cos(i) -> k_rope(i) + // k(i)sin(i+l) -> k_rope(i+l) + + // k(i)cos(i) -> k_rope(i) + // -k(i)sin(i-l) -> k_rope(i-l) + + std::vector block_fp32(32); + for (int k = 0; k < past_kv_len; k++) { + int angle = rotary_angle[k]; + for (int l = 0; l < head_dim / 32; l++) { + block_q4_0 block = + ((block_q4_0 *)k_cache)[k * head_dim / 32 + l]; + dequantize_row_q4_0(&block, block_fp32.data(), 32); + for (int m = 0; m < 32; m++) { + float sin_val = GGML_FP16_TO_FP32( + ((ggml_fp16_t *) + rotary_sin)[angle * head_dim + l * 32 + m]); + float cos_val = GGML_FP16_TO_FP32( + ((ggml_fp16_t *) + rotary_cos)[angle * head_dim + l * 32 + m]); + + if (l * 32 + m < head_dim / 2) { + k_cache_with_rope_fp16[k * head_dim + l * 32 + m] = + GGML_FP32_TO_FP16(block_fp32[m] * cos_val); + k_cache_with_rope_fp16[k * head_dim + l * 32 + m + + head_dim / 2] = + GGML_FP32_TO_FP16(-block_fp32[m] * sin_val); + } else { + k_cache_with_rope_fp16[k * head_dim + l * 32 + m] += + GGML_FP32_TO_FP16(block_fp32[m] * sin_val); + k_cache_with_rope_fp16[k * head_dim + l * 32 + m - + head_dim / 2] -= + GGML_FP32_TO_FP16(block_fp32[m] * cos_val); + } + } + } + } + // quantize k_cache_with_rope_fp16 + for (int k = 0; k < past_kv_len; k++) { + for (int l = 0; l < head_dim / 32; l++) { + for (int m = 0; m < 32; m++) { + block_fp32[m] = GGML_FP16_TO_FP32( + k_cache_with_rope_fp16[k * head_dim + l * 32 + m]); + } + quantize_row_q4_0( + block_fp32.data(), + &k_cache_with_rope_q4[k * head_dim / 32 + l], 32); + } + } + + llamafile_sgemm(past_kv_len, bsz, head_dim / 32, + (block_q4_0 *)k_cache_with_rope_q4, head_dim / 32, + (block_q8_0 *)q, head_dim / 32, attn_score, + past_kv_len, 0, 1, GGML_TASK_TYPE_COMPUTE, k_type, + GGML_TYPE_Q8_0, GGML_TYPE_F32, GGML_PREC_DEFAULT); + } else { + llamafile_sgemm(past_kv_len, bsz, head_dim / 32, + (block_q4_0 *)k_cache, head_dim / 32, + (block_q8_0 *)q, head_dim / 32, attn_score, + past_kv_len, 0, 1, GGML_TASK_TYPE_COMPUTE, k_type, + GGML_TYPE_Q8_0, GGML_TYPE_F32, GGML_PREC_DEFAULT); + } + + // attn = attn * scale + float scale_factor = 1.0 / std::sqrt(float(head_dim)); + ggml_vec_scale_f32(bsz * past_kv_len, attn_score, scale_factor); + + // attn = attn & mask + if (!is_full_attn) { + for (int i = 0; i < bsz; i++) { + for (int j = 0; j < past_kv_len; j++) { + int index = i * past_kv_len + j; + if (!(attn_mask[j / 8] & (1 << (j % 8)))) { + attn_score[index] = + std::numeric_limits::lowest(); + } + } + } + } + + // attn = softmax(attn) + for (int i = 0; i < bsz; i++) { + float sum_exp = 0; + for (int j = 0; j < past_kv_len; j++) { + attn_score[i * past_kv_len + j] = + std::exp(attn_score[i * past_kv_len + j]); + sum_exp += attn_score[i * past_kv_len + j]; + } + for (int j = 0; j < past_kv_len; j++) { + attn_score[i * past_kv_len + j] /= sum_exp; + } + if (lse != nullptr) { + lse[i] = std::log(sum_exp); + } + } + + // output = attn * v + attn * v_anchor + // std::vector attn_q8_0(bsz * past_kv_len / QK8_0); + block_q8_0 *attn_q8_0 = reinterpret_cast(draft); + quantize_row_q8_0(attn_score, attn_q8_0, bsz * past_kv_len); + // std::vector sum(bsz * head_dim); + float *sum = reinterpret_cast(reinterpret_cast(draft) + + sizeof(block_q8_0) * bsz * + past_kv_len / QK8_0); + // TODO: anchor + assert(num_v_anchor == 0); + llamafile_sgemm(head_dim, bsz, past_kv_len / 32, (block_q4_0 *)v_cache, + past_kv_len / 32, attn_q8_0, past_kv_len / 32, sum, + head_dim, 0, 1, GGML_TASK_TYPE_COMPUTE, v_type, + GGML_TYPE_Q8_0, GGML_TYPE_F32, GGML_PREC_DEFAULT); + + quantize_row_q8_0(sum, (block_q8_0 *)output, bsz * head_dim); + } +} diff --git a/archive/csrc/ktransformers_ext/operators/kvcache/kvcache_load_dump.cpp b/archive/csrc/ktransformers_ext/operators/kvcache/kvcache_load_dump.cpp new file mode 100644 index 0000000..4de217f --- /dev/null +++ b/archive/csrc/ktransformers_ext/operators/kvcache/kvcache_load_dump.cpp @@ -0,0 +1,126 @@ +/** + * @Description : + * @Author : Jianwei Dong + * @Date : 2024-08-26 22:47:06 + * @Version : 1.0.0 + * @LastEditors : Jianwei Dong + * @LastEditTime : 2024-08-26 22:47:06 + * @Copyright (c) 2024 by KVCache.AI, All Rights Reserved. + **/ + +#include "kvcache.h" + +#include + +void KVCache::load_kvcache(std::string tensor_file_path, Backend *backend) { + // Timer start + auto start = std::chrono::high_resolution_clock::now(); + std::ifstream ifs_tensor(tensor_file_path, std::ios::binary); + if (!ifs_tensor) { + throw std::runtime_error("Failed to open tensor file"); + } + ifs_tensor.read(reinterpret_cast(&cache_total_len_), + sizeof(cache_total_len_)); + int past_block_num = + (cache_total_len_ + config_.block_len - 1) / config_.block_len; + printf("cache_total_len: %d, past_block_num: %d\n", cache_total_len_, + past_block_num); + for (int i = 0; i < config_.layer_num; ++i) { + past_block_num_[i] = past_block_num; + } + ifs_tensor.read(reinterpret_cast(anchor_.data()), + anchor_.size() * sizeof(ggml_fp16_t)); + for (int i = 0; i < config_.layer_num; ++i) { + for (int j = 0; j < config_.kv_head_num; ++j) { + for (int k = 0; k < past_block_num_[i]; ++k) { + if (config_.kv_type == GGML_TYPE_F16) { + ifs_tensor.read( + reinterpret_cast(k_cache_fp16_[i][j][k].data()), + k_cache_fp16_[i][j][k].size() * sizeof(ggml_fp16_t)); + ifs_tensor.read( + reinterpret_cast(v_cache_fp16_[i][j][k].data()), + v_cache_fp16_[i][j][k].size() * sizeof(ggml_fp16_t)); + } else if (config_.kv_type == GGML_TYPE_Q4_0) { + ifs_tensor.read( + reinterpret_cast(k_cache_q4[i][j][k].data()), + k_cache_q4[i][j][k].size() * sizeof(block_q4_0)); + ifs_tensor.read( + reinterpret_cast(v_cache_q4[i][j][k].data()), + v_cache_q4[i][j][k].size() * sizeof(block_q4_0)); + } + } + } + for (int k = 0; k < past_block_num_[i]; ++k) { + for (int l = 0; l < config_.block_len; l++) { + ifs_tensor.read( + reinterpret_cast(importance_[i][k][l].data()), + importance_[i][k][l].size() * sizeof(ggml_fp16_t)); + } + } + } + ifs_tensor.close(); + // Timer end + auto end = std::chrono::high_resolution_clock::now(); + std::chrono::duration diff = end - start; + printf("time of load: %f s\n", diff.count()); +} +void KVCache::dump_kvcache(int *block_table, int cache_total_len, + std::string tensor_file_path, Backend *backend) { + // Timer start + auto start = std::chrono::high_resolution_clock::now(); + std::ofstream ofs(tensor_file_path, std::ios::binary); + printf("dump_kvcache: %s\n", tensor_file_path.c_str()); + if (!ofs.is_open()) { + std::cerr << "Cannot open file " << tensor_file_path << std::endl; + return; + } + ofs.write(reinterpret_cast(&cache_total_len), + sizeof(cache_total_len)); + int past_block_num = + (cache_total_len + config_.block_len - 1) / config_.block_len; + printf("cache_total_len: %d, past_block_num: %d\n", cache_total_len, + past_block_num); + ofs.write(reinterpret_cast(anchor_.data()), + anchor_.size() * sizeof(ggml_fp16_t)); + for (int i = 0; i < config_.layer_num; ++i) { + for (int j = 0; j < config_.kv_head_num; ++j) { + for (int k = 0; k < past_block_num; ++k) { + int block_idx = block_table[k]; + if (config_.kv_type == GGML_TYPE_F16) { + ofs.write(reinterpret_cast( + k_cache_fp16_[i][j][block_idx].data()), + k_cache_fp16_[i][j][block_idx].size() * + sizeof(ggml_fp16_t)); + ofs.write(reinterpret_cast( + v_cache_fp16_[i][j][block_idx].data()), + v_cache_fp16_[i][j][block_idx].size() * + sizeof(ggml_fp16_t)); + + } else if (config_.kv_type == GGML_TYPE_Q4_0) { + ofs.write(reinterpret_cast( + k_cache_q4[i][j][block_idx].data()), + k_cache_q4[i][j][block_idx].size() * + sizeof(block_q4_0)); + ofs.write(reinterpret_cast( + v_cache_q4[i][j][block_idx].data()), + v_cache_q4[i][j][block_idx].size() * + sizeof(block_q4_0)); + } + } + } + for (int k = 0; k < past_block_num; ++k) { + int block_idx = block_table[k]; + for (int l = 0; l < config_.block_len; l++) { + ofs.write(reinterpret_cast( + importance_[i][block_idx][l].data()), + importance_[i][block_idx][l].size() * + sizeof(ggml_fp16_t)); + } + } + } + ofs.close(); + // Timer end + auto end = std::chrono::high_resolution_clock::now(); + std::chrono::duration diff = end - start; + printf("time of dump: %f s\n", diff.count()); +} \ No newline at end of file diff --git a/archive/csrc/ktransformers_ext/operators/kvcache/kvcache_read_write.cpp b/archive/csrc/ktransformers_ext/operators/kvcache/kvcache_read_write.cpp new file mode 100644 index 0000000..0104905 --- /dev/null +++ b/archive/csrc/ktransformers_ext/operators/kvcache/kvcache_read_write.cpp @@ -0,0 +1,1021 @@ +/** + * @Description : + * @Author : Jianwei Dong + * @Date : 2024-08-26 22:47:06 + * @Version : 1.0.0 + * @LastEditors : Jianwei Dong + * @LastEditTime : 2024-08-26 22:47:06 + * @Copyright (c) 2024 by KVCache.AI, All Rights Reserved. + **/ + +#include "kvcache.h" + +#include + +void KVCache::get_anchor_one_block(ggml_fp16_t *anchor, int layer_id, + int block_idx, Backend *backend) { + // Timer start + auto start = std::chrono::high_resolution_clock::now(); + + layer_id_ = layer_id; + block_idx = block_idx; + seq_len_ = config_.block_len; + anchor_data_ = const_cast(anchor); + + // Timer end + auto end = std::chrono::high_resolution_clock::now(); + std::chrono::duration duration = end - start; + printf("layer %d block %d time of reading anchor: %f s\n", layer_id, + block_idx, duration.count()); +} + +void KVCache::update_anchor_one_block(const ggml_fp16_t *anchor, int layer_id, + int block_idx, Backend *backend) { + // Timer start + auto start = std::chrono::high_resolution_clock::now(); + + layer_id_ = layer_id; + block_idx = block_idx; + seq_len_ = config_.block_len; + anchor_data_ = const_cast(anchor); + + // Each task updates the anchor of a certain position + // backend->do_work_stealing_job(config_.anchor_num, [&](int task_id) { + // int k = task_id % config_.anchor_num; + // int head_id = task_id / config_.anchor_num; + // memcpy(anchor_[layer_id_][head_id][block_idx].data() + + // k * config_.head_dim, + // anchor_data_ + k * config_.head_dim, + // sizeof(uint16_t) * config_.head_dim); + // }); + + // Timer end + auto end = std::chrono::high_resolution_clock::now(); + std::chrono::duration duration = end - start; + printf("layer %d block %d time of writting anchor: %f s\n", layer_id, + block_idx, duration.count()); +} + +void KVCache::update_importance_one_block(const ggml_fp16_t *importance, + int layer_id, int block_idx, + Backend *backend) { + // Timer start + auto start = std::chrono::high_resolution_clock::now(); + + layer_id_ = layer_id; + block_idx = block_idx; + seq_len_ = config_.block_len; + importance_data_ = const_cast(importance); + + // Each task updates the importance of a certain position + backend->do_work_stealing_job( + config_.block_len, nullptr, + [&](int task_id) { + int k = task_id; + memcpy(importance_[layer_id_][block_idx].data() + k, + importance_data_ + k, sizeof(uint16_t)); + }, + nullptr); + + // Timer end + auto end = std::chrono::high_resolution_clock::now(); + std::chrono::duration duration = end - start; + printf("layer %d block %d time of writting importance: %f s\n", layer_id, + block_idx, duration.count()); +} + +void KVCache::get_importance_one_block(ggml_fp16_t *importance, int layer_id, + int block_idx, Backend *backend) { + // Timer start + auto start = std::chrono::high_resolution_clock::now(); + + layer_id_ = layer_id; + block_idx = block_idx; + seq_len_ = config_.block_len; + importance_data_ = const_cast(importance); + + // Each task updates the importance of a certain position + backend->do_work_stealing_job( + config_.block_len, nullptr, + [&](int task_id) { + int k = task_id; + memcpy(importance_data_ + k, + importance_[layer_id_][block_idx].data() + k, + sizeof(uint16_t)); + }, + nullptr); + + // Timer end + auto end = std::chrono::high_resolution_clock::now(); + std::chrono::duration duration = end - start; + printf("layer %d block %d time of reading importance: %f s\n", layer_id, + block_idx, duration.count()); +} + +void KVCache::update_kvcache_one_block_fp16(const ggml_fp16_t *k_in, + const ggml_fp16_t *v_in, + int layer_id, int block_idx, + Backend *backend) { + // Timer start + auto start = std::chrono::high_resolution_clock::now(); + + layer_id_ = layer_id; + block_idx = block_idx; + seq_len_ = config_.block_len; + k_data_ = const_cast(k_in); + v_data_ = const_cast(v_in); + + int new_block_num = std::max((int)past_block_num_[layer_id], block_idx + 1); + + importance_[layer_id_].resize(new_block_num); + + for (int i = 0; i < config_.kv_head_num; i++) { + k_cache_q4[layer_id][i].resize(new_block_num); + v_cache_q4[layer_id][i].resize(new_block_num); + // anchor_[layer_id][i].resize(new_block_num); + } + + for (int i = 0; i < new_block_num; i++) { + importance_[layer_id][i].resize(config_.block_len); + } + + // Each task updates the k cache or v cache of a certain header + backend->do_work_stealing_job( + config_.kv_head_num * 2, nullptr, + [&](int task_id) { + std::vector block_fp32(32); + int head_id = task_id / 2; + if (task_id & 1) { + // fill k_cache_ + k_cache_q4[layer_id_][head_id][block_idx].resize( + config_.block_len * config_.head_dim / 32); + for (int k = 0; k < config_.block_len; k++) { + for (int l = 0; l < config_.head_dim / 32; l++) { + block_q4_0 block; + for (int m = 0; m < 32; m++) { + + block_fp32[m] = GGML_FP16_TO_FP32( + k_data_[((0 * config_.kv_head_num + head_id) * + seq_len_ + + 0 * config_.block_len + k) * + config_.head_dim + + l * 32 + m]); + } + quantize_row_q4_0(block_fp32.data(), &block, 32); + k_cache_q4[layer_id_][head_id][block_idx] + [k * config_.head_dim / 32 + l] = block; + } + } + } else { + // fill v_cache_ + v_cache_q4[layer_id_][head_id][block_idx].resize( + config_.head_dim * config_.block_len / 32); + for (int k = 0; k < config_.block_len / 32; k++) { + for (int l = 0; l < config_.head_dim; l++) { + block_q4_0 block; + for (int m = 0; m < 32; m++) { + + block_fp32[m] = GGML_FP16_TO_FP32( + v_data_[((0 * config_.kv_head_num + head_id) * + seq_len_ + + 0 * config_.block_len + k * 32 + m) * + config_.head_dim + + l]); + } + quantize_row_q4_0(block_fp32.data(), &block, 32); + v_cache_q4[layer_id_][head_id][block_idx] + [l * config_.block_len / 32 + k] = block; + } + } + } + }, + nullptr); + past_block_num_[layer_id] = new_block_num; + + // Timer end + auto end = std::chrono::high_resolution_clock::now(); + std::chrono::duration duration = end - start; + printf("layer %d block %d time of writting KV Cache: %f s\n", layer_id, + block_idx, duration.count()); + // printf("get_one_block_fp16 duration: %ld\n", duration); +} + +void KVCache::get_kvcache_one_block_fp16(ggml_fp16_t *k_in, ggml_fp16_t *v_in, + int layer_id, int block_idx, + Backend *backend) { + // Timer start + auto start = std::chrono::high_resolution_clock::now(); + + layer_id_ = layer_id; + seq_len_ = config_.block_len; + k_data_ = reinterpret_cast(k_in); + v_data_ = reinterpret_cast(v_in); + + // printf("layer_id: %d, block_idx: %d\n", layer_id, block_idx); + // Each task gets the k cache or v cache of a certain header + backend->do_work_stealing_job( + config_.kv_head_num * 2, nullptr, + [&](int task_id) { + std::vector block_fp32(32); + int head_id = task_id / 2; + if (task_id & 1) { + // get k_cache_ + for (int k = 0; k < config_.block_len; k++) { + for (int l = 0; l < config_.head_dim / 32; l++) { + block_q4_0 block = + k_cache_q4[layer_id_][head_id][block_idx] + [k * config_.head_dim / 32 + l]; + dequantize_row_q4_0(&block, block_fp32.data(), 32); + for (int m = 0; m < 32; m++) { + + k_data_[((0 * config_.kv_head_num + head_id) * + seq_len_ + + 0 * config_.block_len + k) * + config_.head_dim + + l * 32 + m] = + GGML_FP32_TO_FP16(block_fp32[m]); + } + } + } + } else { + // get v_cache_ + for (int k = 0; k < config_.block_len / 32; k++) { + for (int l = 0; l < config_.head_dim; l++) { + block_q4_0 block = + v_cache_q4[layer_id_][head_id][block_idx] + [l * config_.block_len / 32 + k]; + dequantize_row_q4_0(&block, block_fp32.data(), 32); + for (int m = 0; m < 32; m++) { + + v_data_[((0 * config_.kv_head_num + head_id) * + seq_len_ + + 0 * config_.block_len + k * 32 + m) * + config_.head_dim + + l] = GGML_FP32_TO_FP16(block_fp32[m]); + } + } + } + } + }, + nullptr); + + // Timer end + auto end = std::chrono::high_resolution_clock::now(); + std::chrono::duration duration = end - start; + printf("layer %d block %d time of reading KV Cache: %f s\n", layer_id, + block_idx, duration.count()); + // printf("get_one_block_fp16 duration: %ld\n", duration); +} + +// k_in: (batch_size, seq_len, head_num, head_dim) +// v_in: (batch_size, seq_len, head_num, head_dim) +void KVCache::get_and_update_kvcache_fp16(ggml_fp16_t *k_in, ggml_fp16_t *v_in, + int layer_id, int *block_table, + int batch_size, int max_block_num, + int *cache_seqlens, int q_len, + Backend *backend) { + // Timer start + auto start = std::chrono::high_resolution_clock::now(); + + layer_id_ = layer_id; + k_data_ = const_cast(k_in); + v_data_ = const_cast(v_in); + + // Each task updates the k cache and v cache of a certain header + backend->do_work_stealing_job( + config_.kv_head_num * max_block_num * batch_size, nullptr, + [&](int task_id) { + // printf("block_idx: %d, task_id: %d\n", block_idx, task_id); + std::vector block_fp32(32); + int batch_id = task_id / (config_.kv_head_num * max_block_num); + int block_id = (task_id / config_.kv_head_num) % max_block_num; + int head_id = task_id % config_.kv_head_num; + int block_idx = block_table[batch_id * max_block_num + block_id]; + int seq_len = cache_seqlens[batch_id]; + int block_l = block_id * config_.block_len; + int block_r = block_id * config_.block_len + config_.block_len; + + if (block_l < seq_len) { + if (config_.kv_type == ggml_type::GGML_TYPE_F16) { + for (int k = 0; k < config_.block_len; k++) { + if (block_id * config_.block_len + k >= seq_len) + break; + for (int l = 0; l < config_.head_dim; l++) { + k_data_ + [batch_id * + (max_block_num * config_.block_len * + config_.kv_head_num * config_.head_dim) + + block_id * + (config_.block_len * config_.kv_head_num * + config_.head_dim) + + k * (config_.kv_head_num * config_.head_dim) + + head_id * config_.head_dim + l] = + k_cache_fp16_[layer_id_][head_id][block_idx] + [k * config_.head_dim + l]; + v_data_ + [batch_id * + (max_block_num * config_.block_len * + config_.kv_head_num * config_.head_dim) + + block_id * + (config_.block_len * config_.kv_head_num * + config_.head_dim) + + k * (config_.kv_head_num * config_.head_dim) + + head_id * config_.head_dim + l] = + v_cache_fp16_[layer_id_][head_id][block_idx] + [l * config_.block_len + k]; + } + } + } else if (config_.kv_type == ggml_type::GGML_TYPE_Q4_0) { + // get k_cache_ + for (int k = 0; k < config_.block_len; k++) { + if (block_id * config_.block_len + k >= seq_len) + break; + for (int l = 0; l < config_.head_dim / 32; l++) { + block_q4_0 block = + k_cache_q4[layer_id_][head_id][block_idx] + [k * config_.head_dim / 32 + l]; + dequantize_row_q4_0(&block, block_fp32.data(), 32); + for (int m = 0; m < 32; m++) { + + k_data_[batch_id * + (max_block_num * config_.block_len * + config_.kv_head_num * + config_.head_dim) + + block_id * (config_.block_len * + config_.kv_head_num * + config_.head_dim) + + k * (config_.kv_head_num * + config_.head_dim) + + head_id * config_.head_dim + l * 32 + + m] = GGML_FP32_TO_FP16(block_fp32[m]); + } + } + } + // get v_cache_ + for (int k = 0; k < config_.block_len / 32; k++) { + for (int l = 0; l < config_.head_dim; l++) { + block_q4_0 block = + v_cache_q4[layer_id_][head_id][block_idx] + [l * config_.block_len / 32 + k]; + dequantize_row_q4_0(&block, block_fp32.data(), 32); + for (int m = 0; m < 32; m++) { + + if (block_id * config_.block_len + k * 32 + m >= + seq_len) + break; + v_data_[batch_id * + (max_block_num * config_.block_len * + config_.kv_head_num * + config_.head_dim) + + block_id * (config_.block_len * + config_.kv_head_num * + config_.head_dim) + + (k * 32 + m) * config_.kv_head_num * + config_.head_dim + + head_id * config_.head_dim + l] = + GGML_FP32_TO_FP16(block_fp32[m]); + } + } + } + } else if (config_.kv_type == ggml_type::GGML_TYPE_Q8_0) { + // get k_cache_ + for (int k = 0; k < config_.block_len; k++) { + if (block_id * config_.block_len + k >= seq_len) + break; + for (int l = 0; l < config_.head_dim / 32; l++) { + block_q8_0 block = + k_cache_q8[layer_id_][head_id][block_idx] + [k * config_.head_dim / 32 + l]; + dequantize_row_q8_0(&block, block_fp32.data(), 32); + for (int m = 0; m < 32; m++) { + + k_data_[batch_id * + (max_block_num * config_.block_len * + config_.kv_head_num * + config_.head_dim) + + block_id * (config_.block_len * + config_.kv_head_num * + config_.head_dim) + + k * (config_.kv_head_num * + config_.head_dim) + + head_id * config_.head_dim + l * 32 + + m] = GGML_FP32_TO_FP16(block_fp32[m]); + } + } + } + // get v_cache_ + for (int k = 0; k < config_.block_len / 32; k++) { + for (int l = 0; l < config_.head_dim; l++) { + block_q8_0 block = + v_cache_q8[layer_id_][head_id][block_idx] + [l * config_.block_len / 32 + k]; + dequantize_row_q8_0(&block, block_fp32.data(), 32); + for (int m = 0; m < 32; m++) { + + if (block_id * config_.block_len + k * 32 + m >= + seq_len) + break; + v_data_[batch_id * + (max_block_num * config_.block_len * + config_.kv_head_num * + config_.head_dim) + + block_id * (config_.block_len * + config_.kv_head_num * + config_.head_dim) + + (k * 32 + m) * config_.kv_head_num * + config_.head_dim + + head_id * config_.head_dim + l] = + GGML_FP32_TO_FP16(block_fp32[m]); + } + } + } + } + } + if (block_r > seq_len && block_l < seq_len + q_len) { + if (config_.kv_type == ggml_type::GGML_TYPE_F16) { + for (int k = 0; k < config_.block_len; k++) { + if (block_id * config_.block_len + k >= + seq_len + q_len || + block_id * config_.block_len + k < seq_len) + continue; + for (int l = 0; l < config_.head_dim; l++) { + k_cache_fp16_[layer_id_][head_id][block_idx] + [k * config_.head_dim + l] = k_data_ + [batch_id * (max_block_num * + config_.block_len * + config_.kv_head_num * + config_.head_dim) + + block_id * (config_.block_len * + config_.kv_head_num * + config_.head_dim) + + k * (config_.kv_head_num * + config_.head_dim) + + head_id * config_.head_dim + l]; + v_cache_fp16_[layer_id_][head_id][block_idx] + [l * config_.block_len + k] = v_data_ + [batch_id * (max_block_num * + config_.block_len * + config_.kv_head_num * + config_.head_dim) + + block_id * (config_.block_len * + config_.kv_head_num * + config_.head_dim) + + k * (config_.kv_head_num * + config_.head_dim) + + head_id * config_.head_dim + l]; + } + } + } else if (config_.kv_type == ggml_type::GGML_TYPE_Q4_0) { + // fill k_cache_ + for (int k = 0; k < config_.block_len; k++) { + if (block_id * config_.block_len + k >= + seq_len + q_len || + block_id * config_.block_len + k < seq_len) + continue; + for (int l = 0; l < config_.head_dim / 32; l++) { + block_q4_0 block; + for (int m = 0; m < 32; m++) { + + block_fp32[m] = GGML_FP16_TO_FP32( + k_data_[batch_id * (max_block_num * + config_.block_len * + config_.kv_head_num * + config_.head_dim) + + block_id * (config_.block_len * + config_.kv_head_num * + config_.head_dim) + + k * (config_.kv_head_num * + config_.head_dim) + + head_id * config_.head_dim + + l * 32 + m]); + } + quantize_row_q4_0(block_fp32.data(), &block, 32); + k_cache_q4[layer_id_][head_id][block_idx] + [k * config_.head_dim / 32 + l] = block; + } + } + + // fill v_cache_ + for (int k = 0; k < config_.block_len / 32; k++) { + for (int l = 0; l < config_.head_dim; l++) { + block_q4_0 block; + for (int m = 0; m < 32; m++) { + + if (block_id * config_.block_len + k * 32 + m >= + seq_len + q_len) { + block_fp32[m] = 0; + continue; + } + block_fp32[m] = GGML_FP16_TO_FP32( + v_data_[batch_id * (max_block_num * + config_.block_len * + config_.kv_head_num * + config_.head_dim) + + block_id * (config_.block_len * + config_.kv_head_num * + config_.head_dim) + + (k * 32 + m) * config_.kv_head_num * + config_.head_dim + + head_id * config_.head_dim + l]); + } + quantize_row_q4_0(block_fp32.data(), &block, 32); + v_cache_q4[layer_id_][head_id][block_idx] + [l * config_.block_len / 32 + k] = block; + } + } + } else if (config_.kv_type == ggml_type::GGML_TYPE_Q8_0) { + // fill k_cache_ + for (int k = 0; k < config_.block_len; k++) { + if (block_id * config_.block_len + k >= + seq_len + q_len || + block_id * config_.block_len + k < seq_len) + continue; + for (int l = 0; l < config_.head_dim / 32; l++) { + block_q8_0 block; + for (int m = 0; m < 32; m++) { + + block_fp32[m] = GGML_FP16_TO_FP32( + k_data_[batch_id * (max_block_num * + config_.block_len * + config_.kv_head_num * + config_.head_dim) + + block_id * (config_.block_len * + config_.kv_head_num * + config_.head_dim) + + k * (config_.kv_head_num * + config_.head_dim) + + head_id * config_.head_dim + + l * 32 + m]); + } + quantize_row_q8_0(block_fp32.data(), &block, 32); + k_cache_q8[layer_id_][head_id][block_idx] + [k * config_.head_dim / 32 + l] = block; + } + } + + // fill v_cache_ + for (int k = 0; k < config_.block_len / 32; k++) { + for (int l = 0; l < config_.head_dim; l++) { + block_q8_0 block; + for (int m = 0; m < 32; m++) { + + if (block_id * config_.block_len + k * 32 + m >= + seq_len + q_len) { + block_fp32[m] = 0; + continue; + } + block_fp32[m] = GGML_FP16_TO_FP32( + v_data_[batch_id * (max_block_num * + config_.block_len * + config_.kv_head_num * + config_.head_dim) + + block_id * (config_.block_len * + config_.kv_head_num * + config_.head_dim) + + (k * 32 + m) * config_.kv_head_num * + config_.head_dim + + head_id * config_.head_dim + l]); + } + quantize_row_q8_0(block_fp32.data(), &block, 32); + v_cache_q8[layer_id_][head_id][block_idx] + [l * config_.block_len / 32 + k] = block; + } + } + } + } + }, + nullptr); + + // Timer end + auto end = std::chrono::high_resolution_clock::now(); + std::chrono::duration duration = end - start; + + // printf("layer %d time of reading and updating KV Cache: %f s\n", + // layer_id, + // duration.count()); +} + +void KVCache::update_importance(const ggml_fp16_t *importance, int layer_id, + int *block_table, int batch_size, + int max_block_num, int *offset, int width, + Backend *backend) { + // Timer start + auto start = std::chrono::high_resolution_clock::now(); + + layer_id_ = layer_id; + importance_data_ = const_cast(importance); + + // Each task updates the importance of a certain position + backend->do_work_stealing_job( + max_block_num * batch_size, nullptr, + [&](int task_id) { + int block_id = task_id % max_block_num; + int batch_id = task_id / max_block_num; + int block_idx = block_table[batch_id * max_block_num + block_id]; + if (block_id > (offset[batch_id] + width) / config_.block_len) { + return; + } + for (int k = 0; k < config_.block_len; k++) { + for (int head_id = 0; head_id < config_.q_head_num; head_id++) { + importance_[layer_id_][block_idx][k][head_id] = + GGML_FP32_TO_FP16( + GGML_FP16_TO_FP32( + importance_data_[batch_id * max_block_num * + config_.block_len * + config_.q_head_num + + (block_id * config_.block_len + + k) * + config_.q_head_num + + head_id]) + + GGML_FP16_TO_FP32( + importance_[layer_id_][block_idx][k][head_id])); + } + } + }, + nullptr); + + // Timer end + auto end = std::chrono::high_resolution_clock::now(); + std::chrono::duration duration = end - start; + + // printf("layer %d time of updating importance: %f s\n", layer_id, + // duration.count()); +} + +void KVCache::get_kvcache_fp16(ggml_fp16_t *k_in, ggml_fp16_t *v_in, + int layer_id, int *block_table, int batch_size, + int max_block_num, int *cache_seqlens, + Backend *backend) { + // Timer start + auto start = std::chrono::high_resolution_clock::now(); + + layer_id_ = layer_id; + k_data_ = const_cast(k_in); + v_data_ = const_cast(v_in); + + // Each task updates the k cache and v cache of a certain header + backend->do_work_stealing_job( + config_.kv_head_num * max_block_num * batch_size, nullptr, + [&](int task_id) { + // printf("block_idx: %d, task_id: %d\n", block_idx, task_id); + std::vector block_fp32(32); + int batch_id = task_id / (config_.kv_head_num * max_block_num); + int block_id = (task_id / config_.kv_head_num) % max_block_num; + int head_id = task_id % config_.kv_head_num; + int block_idx = block_table[batch_id * max_block_num + block_id]; + int seq_len = cache_seqlens[batch_id]; + int block_l = block_id * config_.block_len; + int block_r = block_id * config_.block_len + config_.block_len; + + if (block_l < seq_len) { + if (config_.kv_type == ggml_type::GGML_TYPE_F16) { + for (int k = 0; k < config_.block_len; k++) { + if (block_id * config_.block_len + k >= seq_len) + break; + for (int l = 0; l < config_.head_dim; l++) { + k_data_ + [batch_id * + (max_block_num * config_.block_len * + config_.kv_head_num * config_.head_dim) + + block_id * + (config_.block_len * config_.kv_head_num * + config_.head_dim) + + k * (config_.kv_head_num * config_.head_dim) + + head_id * config_.head_dim + l] = + k_cache_fp16_[layer_id_][head_id][block_idx] + [k * config_.head_dim + l]; + v_data_ + [batch_id * + (max_block_num * config_.block_len * + config_.kv_head_num * config_.head_dim) + + block_id * + (config_.block_len * config_.kv_head_num * + config_.head_dim) + + k * (config_.kv_head_num * config_.head_dim) + + head_id * config_.head_dim + l] = + v_cache_fp16_[layer_id_][head_id][block_idx] + [l * config_.block_len + k]; + } + } + } else if (config_.kv_type == ggml_type::GGML_TYPE_Q4_0) { + // get k_cache_ + for (int k = 0; k < config_.block_len; k++) { + if (block_id * config_.block_len + k >= seq_len) + break; + for (int l = 0; l < config_.head_dim / 32; l++) { + block_q4_0 block = + k_cache_q4[layer_id_][head_id][block_idx] + [k * config_.head_dim / 32 + l]; + dequantize_row_q4_0(&block, block_fp32.data(), 32); + for (int m = 0; m < 32; m++) { + + k_data_[batch_id * + (max_block_num * config_.block_len * + config_.kv_head_num * + config_.head_dim) + + block_id * (config_.block_len * + config_.kv_head_num * + config_.head_dim) + + k * (config_.kv_head_num * + config_.head_dim) + + head_id * config_.head_dim + l * 32 + + m] = GGML_FP32_TO_FP16(block_fp32[m]); + } + } + } + // get v_cache_ + for (int k = 0; k < config_.block_len / 32; k++) { + for (int l = 0; l < config_.head_dim; l++) { + block_q4_0 block = + v_cache_q4[layer_id_][head_id][block_idx] + [l * config_.block_len / 32 + k]; + dequantize_row_q4_0(&block, block_fp32.data(), 32); + for (int m = 0; m < 32; m++) { + + if (block_id * config_.block_len + k * 32 + m >= + seq_len) + break; + v_data_[batch_id * + (max_block_num * config_.block_len * + config_.kv_head_num * + config_.head_dim) + + block_id * (config_.block_len * + config_.kv_head_num * + config_.head_dim) + + (k * 32 + m) * config_.kv_head_num * + config_.head_dim + + head_id * config_.head_dim + l] = + GGML_FP32_TO_FP16(block_fp32[m]); + } + } + } + } else if (config_.kv_type == ggml_type::GGML_TYPE_Q8_0) { + // get k_cache_ + for (int k = 0; k < config_.block_len; k++) { + if (block_id * config_.block_len + k >= seq_len) + break; + for (int l = 0; l < config_.head_dim / 32; l++) { + block_q8_0 block = + k_cache_q8[layer_id_][head_id][block_idx] + [k * config_.head_dim / 32 + l]; + dequantize_row_q8_0(&block, block_fp32.data(), 32); + for (int m = 0; m < 32; m++) { + + k_data_[batch_id * + (max_block_num * config_.block_len * + config_.kv_head_num * + config_.head_dim) + + block_id * (config_.block_len * + config_.kv_head_num * + config_.head_dim) + + k * (config_.kv_head_num * + config_.head_dim) + + head_id * config_.head_dim + l * 32 + + m] = GGML_FP32_TO_FP16(block_fp32[m]); + } + } + } + // get v_cache_ + for (int k = 0; k < config_.block_len / 32; k++) { + for (int l = 0; l < config_.head_dim; l++) { + block_q8_0 block = + v_cache_q8[layer_id_][head_id][block_idx] + [l * config_.block_len / 32 + k]; + dequantize_row_q8_0(&block, block_fp32.data(), 32); + for (int m = 0; m < 32; m++) { + + if (block_id * config_.block_len + k * 32 + m >= + seq_len) + break; + v_data_[batch_id * + (max_block_num * config_.block_len * + config_.kv_head_num * + config_.head_dim) + + block_id * (config_.block_len * + config_.kv_head_num * + config_.head_dim) + + (k * 32 + m) * config_.kv_head_num * + config_.head_dim + + head_id * config_.head_dim + l] = + GGML_FP32_TO_FP16(block_fp32[m]); + } + } + } + } + } + }, + nullptr); + + // Timer end + auto end = std::chrono::high_resolution_clock::now(); + std::chrono::duration duration = end - start; +} + +void KVCache::update_kvcache_fp16(const ggml_fp16_t *k_in, + const ggml_fp16_t *v_in, int layer_id, + int *block_table, int batch_size, + int max_block_num, int *cache_seqlens, + int q_len, Backend *backend) { + // Timer start + auto start = std::chrono::high_resolution_clock::now(); + + layer_id_ = layer_id; + k_data_ = const_cast(k_in); + v_data_ = const_cast(v_in); + // Each task updates the k cache and v cache of a certain header + backend->do_work_stealing_job( + batch_size * config_.kv_head_num * q_len, nullptr, + [&](int task_id) { + int batch_id = task_id / (config_.kv_head_num * q_len); + int head_id = task_id / q_len % config_.kv_head_num; + int seq_len = cache_seqlens[batch_id] + task_id % q_len; + int q_offset = task_id % q_len; + + int block_id = seq_len / config_.block_len; + int block_idx = block_table[batch_id * max_block_num + block_id]; + int pos_in_block = seq_len % config_.block_len; + + if (config_.kv_type == ggml_type::GGML_TYPE_F16) { + for (int l = 0; l < config_.head_dim; l++) { + k_cache_fp16_[layer_id_][head_id][block_idx] + [pos_in_block * config_.head_dim + l] = + k_data_[batch_id * + (q_len * config_.kv_head_num * + config_.head_dim) + + q_offset * config_.kv_head_num * + config_.head_dim + + head_id * config_.head_dim + l]; + v_cache_fp16_[layer_id_][head_id][block_idx] + [l * config_.block_len + pos_in_block] = + v_data_[batch_id * + (q_len * config_.kv_head_num * + config_.head_dim) + + q_offset * config_.kv_head_num * + config_.head_dim + + head_id * config_.head_dim + l]; + } + } else if (config_.kv_type == ggml_type::GGML_TYPE_Q4_0) { + std::vector block_fp32(32); + // fill k_cache_ + for (int l = 0; l < config_.head_dim / 32; l++) { + block_q4_0 block; + for (int m = 0; m < 32; m++) { + + block_fp32[m] = GGML_FP16_TO_FP32( + k_data_[batch_id * (q_len * config_.kv_head_num * + config_.head_dim) + + head_id * config_.head_dim + l * 32 + m]); + } + quantize_row_q4_0(block_fp32.data(), &block, 32); + + k_cache_q4[layer_id_][head_id][block_idx] + [pos_in_block * config_.head_dim / 32 + l] = + block; + } + + // fill v_cache_ + for (int l = 0; l < config_.head_dim; l++) { + block_q4_0 block = v_cache_q4[layer_id_][head_id][block_idx] + [l * config_.block_len / 32 + + pos_in_block / 32]; + dequantize_row_q4_0(&block, block_fp32.data(), 32); + block_fp32[pos_in_block % 32] = GGML_FP16_TO_FP32( + v_data_[batch_id * (q_len * config_.kv_head_num * + config_.head_dim) + + head_id * config_.head_dim + l]); + quantize_row_q4_0(block_fp32.data(), &block, 32); + v_cache_q4[layer_id_][head_id][block_idx] + [l * config_.block_len / 32 + pos_in_block / 32] = + block; + } + } else if (config_.kv_type == ggml_type::GGML_TYPE_Q8_0) { + std::vector block_fp32(32); + // fill k_cache_ + for (int l = 0; l < config_.head_dim / 32; l++) { + block_q8_0 block; + for (int m = 0; m < 32; m++) { + + block_fp32[m] = GGML_FP16_TO_FP32( + k_data_[batch_id * (q_len * config_.kv_head_num * + config_.head_dim) + + head_id * config_.head_dim + l * 32 + m]); + } + quantize_row_q8_0(block_fp32.data(), &block, 32); + + k_cache_q8[layer_id_][head_id][block_idx] + [pos_in_block * config_.head_dim / 32 + l] = + block; + } + + // fill v_cache_ + for (int l = 0; l < config_.head_dim; l++) { + block_q8_0 block = v_cache_q8[layer_id_][head_id][block_idx] + [l * config_.block_len / 32 + + pos_in_block / 32]; + dequantize_row_q8_0(&block, block_fp32.data(), 32); + block_fp32[pos_in_block % 32] = GGML_FP16_TO_FP32( + v_data_[batch_id * (q_len * config_.kv_head_num * + config_.head_dim) + + head_id * config_.head_dim + l]); + quantize_row_q8_0(block_fp32.data(), &block, 32); + v_cache_q8[layer_id_][head_id][block_idx] + [l * config_.block_len / 32 + pos_in_block / 32] = + block; + } + } + }, + nullptr); + + // Timer end + auto end = std::chrono::high_resolution_clock::now(); + std::chrono::duration duration = end - start; + // printf("layer %d time of reading KV Cache: %f s\n", layer_id, + // duration.count()); +} + +void KVCache::get_all_kvcache_one_layer(int layer_id, ggml_fp16_t *k_in, + ggml_fp16_t *v_in, Backend *backend) { + // Timer start + auto start = std::chrono::high_resolution_clock::now(); + + layer_id_ = layer_id; + seq_len_ = config_.block_len; + block_num_ = get_cache_total_block_num(); + k_data_ = reinterpret_cast(k_in); + v_data_ = reinterpret_cast(v_in); + + // Each task gets the k cache or v cache of a certain header + backend->do_work_stealing_job( + config_.kv_head_num * past_block_num_[layer_id] * 2, nullptr, + [&](int task_id) { + std::vector block_fp32(32); + int head_id = task_id / 2 / past_block_num_[layer_id]; + int block_idx = task_id / 2 % past_block_num_[layer_id]; + if (block_idx >= block_num_) + return; + + int max_offset = 0; + if (task_id & 1) { + // get k_cache_ + for (int k = 0; k < config_.block_len; k++) { + if (block_idx * seq_len_ + k >= cache_total_len_) + break; + for (int l = 0; l < config_.head_dim / 32; l++) { + block_q4_0 block = + k_cache_q4[layer_id_][head_id][block_idx] + [k * config_.head_dim / 32 + l]; + dequantize_row_q4_0(&block, block_fp32.data(), 32); + for (int m = 0; m < 32; m++) { + + k_data_[(head_id * cache_total_len_ + + block_idx * config_.block_len + k) * + config_.head_dim + + l * 32 + m] = + GGML_FP32_TO_FP16(block_fp32[m]); + max_offset = std::max( + max_offset, + (int)(head_id * cache_total_len_ + + block_idx * config_.block_len + k) * + config_.head_dim + + l * 32 + m); + } + } + } + } else { + // get v_cache_ + for (int k = 0; k < config_.block_len / 32; k++) { + for (int l = 0; l < config_.head_dim; l++) { + block_q4_0 block = + v_cache_q4[layer_id_][head_id][block_idx] + [l * config_.block_len / 32 + k]; + dequantize_row_q4_0(&block, block_fp32.data(), 32); + for (int m = 0; m < 32; m++) { + + if (block_idx * seq_len_ + k * 32 + m >= + cache_total_len_) + break; + v_data_[(head_id * cache_total_len_ + + block_idx * config_.block_len + k * 32 + + m) * + config_.head_dim + + l] = GGML_FP32_TO_FP16(block_fp32[m]); + max_offset = + std::max(max_offset, + (int)((head_id * cache_total_len_ + + block_idx * config_.block_len + + k * 32 + m) * + config_.head_dim + + l)); + } + } + } + } + }, + nullptr); + + // Timer end + auto end = std::chrono::high_resolution_clock::now(); + std::chrono::duration duration = end - start; + // printf("layer %d block num %d time of reading all KV Cache: %f s\n", + // layer_id, block_num_, duration.count()); +} diff --git a/archive/csrc/ktransformers_ext/operators/kvcache/kvcache_utils.cpp b/archive/csrc/ktransformers_ext/operators/kvcache/kvcache_utils.cpp new file mode 100644 index 0000000..c57d475 --- /dev/null +++ b/archive/csrc/ktransformers_ext/operators/kvcache/kvcache_utils.cpp @@ -0,0 +1,1159 @@ +/** + * @Description : + * @Author : Jianwei Dong + * @Date : 2024-08-26 22:47:06 + * @Version : 1.0.0 + * @LastEditors : Jianwei Dong + * @LastEditTime : 2024-08-26 22:47:06 + * @Copyright (c) 2024 by KVCache.AI, All Rights Reserved. + **/ + +#include "kvcache.h" + +#include + +std::string ggml_type_to_string(ggml_type type) { + switch (type) { + case GGML_TYPE_F32: + return "GGML_TYPE_F32"; + case GGML_TYPE_F16: + return "GGML_TYPE_F16"; + case GGML_TYPE_Q4_0: + return "GGML_TYPE_Q4_0"; + case GGML_TYPE_Q8_0: + return "GGML_TYPE_Q8_0"; + } + return "UNDIFINED"; +} +std::string AnchorTypeToString(AnchorType type) { + switch (type) { + case AnchorType::DYNAMIC: + return "DYNAMIC"; + case AnchorType::BLOCK_MEAN: + return "BLOCK_MEAN"; + case AnchorType::BLOCK_MAX: + return "BLOCK_MAX"; + case AnchorType::FIXED_ANCHOR: + return "FIXED_ANCHOR"; + case AnchorType::QUEST: + return "QUEST"; + } + return "UNDIFINED"; +} +std::string RetrievalTypeToString(RetrievalType type) { + switch (type) { + case RetrievalType::LAYER: + return "SHARED"; + case RetrievalType::KVHEAD: + return "SEPARATE"; + case RetrievalType::QHEAD: + return "INDIVIDUAL"; + } + return "UNDIFINED"; +} +KVCacheConfig::KVCacheConfig(int layer_num, int kv_head_num, int q_head_num, + int head_dim, int block_len, int anchor_num, + AnchorType anchor_type, ggml_type kv_type, + RetrievalType retrieval_type, int layer_step, + int token_step, int layer_offset, + int max_block_num, int max_batch_size, + int max_thread_num) + : layer_num(layer_num), kv_head_num(kv_head_num), q_head_num(q_head_num), + head_dim(head_dim), block_len(block_len), anchor_num(anchor_num), + anchor_type(anchor_type), kv_type(kv_type), + retrieval_type(retrieval_type), layer_step(layer_step), + token_step(token_step), layer_offset(layer_offset), + max_block_num(max_block_num), max_batch_size(max_batch_size), + max_thread_num(max_thread_num) { + printf( + "layer_num: %d, kv_head_num: %d, q_head_num: %d, head_dim: %d, " + "block_len: %d, anchor_num: %d, anchor_type: %s, kv_type: %s, " + "retrieval_type: %s, layer_step: %d, token_step: %d, layer_offset: %d," + "max_block_num: %d, max_batch_size: %d, max_thread_num: %d\n", + layer_num, kv_head_num, q_head_num, head_dim, block_len, anchor_num, + AnchorTypeToString(anchor_type).c_str(), + ggml_type_to_string(kv_type).c_str(), + RetrievalTypeToString(retrieval_type).c_str(), layer_step, token_step, + layer_offset, max_block_num, max_batch_size, max_thread_num); + assert(q_head_num % kv_head_num == 0); +} +KVCache::KVCache(KVCacheConfig config) { + this->config_ = config; + + n_gqa_ = config_.q_head_num / config_.kv_head_num; + if (config_.kv_type == ggml_type::GGML_TYPE_F16) { + // TODO: Elegant implement + k_cache_fp16_.resize(config_.layer_num); + v_cache_fp16_.resize(config_.layer_num); + selected_blocks_num_history_.resize(config_.layer_num / + config_.layer_step); + if (config_.retrieval_type == RetrievalType::LAYER) { + selected_blocks_history_.resize(config_.layer_num / + config_.layer_step); + } else if (config_.retrieval_type == RetrievalType::KVHEAD) { + selected_blocks_history_kvhead_.resize(config_.layer_num / + config_.layer_step); + } else if (config_.retrieval_type == RetrievalType::QHEAD) { + } + } else if (config_.kv_type == ggml_type::GGML_TYPE_Q4_0) { + k_cache_q4.resize(config.layer_num); + v_cache_q4.resize(config.layer_num); + } else if (config_.kv_type == ggml_type::GGML_TYPE_Q8_0) { + k_cache_q8.resize(config.layer_num); + v_cache_q8.resize(config.layer_num); + } else { + assert(false); + } + anchor_.resize(config.layer_num * config.max_block_num * config.anchor_num * + config.q_head_num * config.head_dim); + importance_.resize(config.layer_num); + past_block_num_.resize(config.layer_num); + for (int i = 0; i < config.layer_num; i++) { + past_block_num_[i] = 0; + } + + ThreadResize(config.max_thread_num); + BatchResize(config.max_batch_size); + BlockResize(config.max_block_num); + q_fp32.resize(n_gqa_ * config.head_dim); +} + +void KVCache::ThreadResize(int thread_num) { + thread_local_output_q8_0_.resize(thread_num); + thread_local_attn_score_.resize(thread_num); + thread_local_output_fp32_.resize(thread_num); + thread_local_attn_lse_.resize(thread_num); + thread_local_cur_output_fp32_.resize(thread_num); + thread_local_cur_attn_lse_.resize(thread_num); + thread_local_draft_.resize(thread_num); + thread_cur_head_idx_.resize(thread_num); + thread_local_attn_mask_.resize(thread_num); + for (int i = 0; i < thread_num; i++) { + thread_local_output_q8_0_[i].resize(n_gqa_ * config_.head_dim / QK8_0); + thread_local_attn_score_[i].resize(n_gqa_ * config_.block_len); + thread_local_output_fp32_[i].resize(n_gqa_ * config_.head_dim); + thread_local_attn_lse_[i].resize(n_gqa_); + thread_local_cur_output_fp32_[i].resize(n_gqa_ * config_.head_dim); + thread_local_cur_attn_lse_[i].resize(n_gqa_); + thread_local_draft_[i].resize( + 2 * n_gqa_ * config_.block_len + 6 * n_gqa_ * config_.head_dim + + 2 * config_.block_len * config_.head_dim + + config_.block_len * config_.head_dim / QK4_0); + thread_local_attn_mask_[i].resize(config_.block_len / 8); + } +} +void KVCache::BatchResize(int batch_size) { + mutex_.resize(batch_size); + q_q8_0_.resize(batch_size); + q_fp32_.resize(batch_size); + output_fp32_.resize(batch_size); + attn_lse_.resize(batch_size); + block_lse_.resize(batch_size); + attn_sparsity_.resize(batch_size); + + if (config_.retrieval_type == RetrievalType::LAYER) { + block_table_before_retrieval_.resize(batch_size); + block_table_after_retrieval_.resize(batch_size); + + for (int i = 0; i < config_.layer_num / config_.layer_step; i++) { + selected_blocks_history_[i].resize(batch_size); + } + + } else if (config_.retrieval_type == RetrievalType::KVHEAD) { + block_table_before_retrieval_kvhead_.resize(batch_size); + block_table_after_retrieval_kvhead_.resize(batch_size); + for (int i = 0; i < config_.layer_num / config_.layer_step; i++) { + selected_blocks_history_kvhead_[i].resize(batch_size); + } + } else if (config_.retrieval_type == RetrievalType::QHEAD) { + block_table_before_retrieval_qhead_.resize(batch_size); + block_table_after_retrieval_qhead_.resize(batch_size); + } + cache_seqlens_.resize(batch_size); + if (config_.retrieval_type == RetrievalType::LAYER) { + block_similar_.resize(batch_size); + } else if (config_.retrieval_type == RetrievalType::KVHEAD) { + block_similar_kv_head_.resize(batch_size); + } else if (config_.retrieval_type == RetrievalType::QHEAD) { + block_similar_q_head_.resize(batch_size); + } + for (int i = 0; i < batch_size; i++) { + top_similar_block_.resize(batch_size); + + mutex_[i].resize(config_.kv_head_num); + q_q8_0_[i].resize(config_.kv_head_num); + q_fp32_[i].resize(config_.kv_head_num); + output_fp32_[i].resize(config_.kv_head_num); + attn_lse_[i].resize(config_.kv_head_num); + + for (int j = 0; j < config_.kv_head_num; j++) { + if (!mutex_[i][j]) { + mutex_[i][j] = std::make_unique(); + } + q_q8_0_[i][j].resize(n_gqa_ * config_.head_dim / QK8_0); + q_fp32_[i][j].resize(n_gqa_ * config_.head_dim); + output_fp32_[i][j].resize(n_gqa_ * config_.head_dim); + attn_lse_[i][j].resize(n_gqa_); + } + } + avg_q.resize(batch_size); + avg_q_fp16.resize(batch_size); + for (int i = 0; i < batch_size; i++) { + attn_sparsity_[i].resize(config_.q_head_num); + avg_q[i].resize(config_.q_head_num * config_.head_dim); + avg_q_fp16[i].resize(config_.q_head_num * config_.head_dim); + } +} + +void KVCache::BlockResize(int max_block_num) { + sin_.resize(max_block_num * config_.block_len); + cos_.resize(max_block_num * config_.block_len); + for (int i = 0; i < max_block_num * config_.block_len; i++) { + sin_[i].resize(config_.head_dim); + cos_[i].resize(config_.head_dim); + } + + for (int i = 0; i < config_.layer_num / config_.layer_step; i++) { + for (int j = 0; j < config_.max_batch_size; j++) { + if (config_.retrieval_type == RetrievalType::LAYER) { + selected_blocks_history_[i][j].resize(max_block_num); + } else if (config_.retrieval_type == RetrievalType::KVHEAD) { + selected_blocks_history_kvhead_[i][j].resize(max_block_num); + for (int k = 0; k < config_.max_block_num; k++) { + selected_blocks_history_kvhead_[i][j][k].resize( + config_.kv_head_num); + } + } else if (config_.retrieval_type == RetrievalType::QHEAD) { + } + } + } + + for (int layer_id = 0; layer_id < config_.layer_num; layer_id++) { + importance_[layer_id].resize(max_block_num); + + if (config_.kv_type == ggml_type::GGML_TYPE_F16) { + // TODO: Elegant implement + k_cache_fp16_[layer_id].resize(config_.kv_head_num); + v_cache_fp16_[layer_id].resize(config_.kv_head_num); + + for (int i = 0; i < config_.kv_head_num; i++) { + k_cache_fp16_[layer_id][i].resize(max_block_num); + v_cache_fp16_[layer_id][i].resize(max_block_num); + + for (int j = 0; j < max_block_num; j++) { + k_cache_fp16_[layer_id][i][j].resize(config_.block_len * + config_.head_dim); + v_cache_fp16_[layer_id][i][j].resize(config_.block_len * + config_.head_dim); + } + } + + } else if (config_.kv_type == ggml_type::GGML_TYPE_Q4_0) { + k_cache_q4[layer_id].resize(config_.kv_head_num); + v_cache_q4[layer_id].resize(config_.kv_head_num); + for (int i = 0; i < config_.kv_head_num; i++) { + k_cache_q4[layer_id][i].resize(max_block_num); + v_cache_q4[layer_id][i].resize(max_block_num); + + for (int j = 0; j < max_block_num; j++) { + k_cache_q4[layer_id][i][j].resize(config_.block_len * + config_.head_dim / 32); + v_cache_q4[layer_id][i][j].resize(config_.block_len * + config_.head_dim / 32); + } + } + } else if (config_.kv_type == ggml_type::GGML_TYPE_Q8_0) { + k_cache_q8[layer_id].resize(config_.kv_head_num); + v_cache_q8[layer_id].resize(config_.kv_head_num); + for (int i = 0; i < config_.kv_head_num; i++) { + k_cache_q8[layer_id][i].resize(max_block_num); + v_cache_q8[layer_id][i].resize(max_block_num); + + for (int j = 0; j < max_block_num; j++) { + k_cache_q8[layer_id][i][j].resize(config_.block_len * + config_.head_dim / 32); + v_cache_q8[layer_id][i][j].resize(config_.block_len * + config_.head_dim / 32); + } + } + } else { + assert(false); + } + for (int i = 0; i < config_.max_batch_size; i++) { + if (config_.retrieval_type == RetrievalType::LAYER) { + block_similar_[i].resize(max_block_num); + block_table_before_retrieval_[i].resize(max_block_num); + block_table_after_retrieval_[i].resize(max_block_num); + } else if (config_.retrieval_type == RetrievalType::KVHEAD) { + block_similar_kv_head_[i].resize(max_block_num); + block_table_before_retrieval_kvhead_[i].resize(max_block_num); + block_table_after_retrieval_kvhead_[i].resize(max_block_num); + for (int j = 0; j < max_block_num; j++) { + block_similar_kv_head_[i][j].resize(config_.kv_head_num); + block_table_before_retrieval_kvhead_[i][j].resize( + config_.kv_head_num); + block_table_after_retrieval_kvhead_[i][j].resize( + config_.kv_head_num); + } + } else if (config_.retrieval_type == RetrievalType::QHEAD) { + block_similar_q_head_[i].resize(max_block_num); + block_table_before_retrieval_qhead_[i].resize(max_block_num); + block_table_after_retrieval_qhead_[i].resize(max_block_num); + for (int j = 0; j < max_block_num; j++) { + block_similar_q_head_[i][j].resize(config_.q_head_num); + block_table_before_retrieval_qhead_[i][j].resize( + config_.q_head_num); + block_table_after_retrieval_qhead_[i][j].resize( + config_.q_head_num); + } + } + block_lse_[i].resize(max_block_num); + for (int j = 0; j < max_block_num; j++) { + block_lse_[i][j].resize(config_.q_head_num); + } + } + + for (int i = 0; i < max_block_num; i++) { + importance_[layer_id][i].resize(config_.block_len); + for (int j = 0; j < config_.block_len; j++) { + importance_[layer_id][i][j].resize(config_.q_head_num); + } + } + } +} + +void KVCache::calc_anchor_all_layers(int *block_table, int *cache_seqlens, + int batch_size, int max_block_num, + Backend *backend) { + // Timer start + auto start = std::chrono::high_resolution_clock::now(); + + // Each task updates the importance of a certain block + seq_len_ = config_.block_len; + backend->do_work_stealing_job( + config_.layer_num * batch_size * max_block_num, nullptr, + [&](int task_id) { + int layer_id = task_id / (batch_size * max_block_num); + int batch_id = (task_id / max_block_num) % batch_size; + int block_id = task_id % max_block_num; + // If the block is out of the sequence length, skip it. In + // particular, the last block of the sequence that is shorter than + // the block length should be skipped. + + if (cache_seqlens[batch_id] / config_.block_len < block_id) { + return; + } + int block_idx = block_table[batch_id * max_block_num + block_id]; + + std::vector block_fp32(32); + if (config_.anchor_type == AnchorType::DYNAMIC) { + + // clear anchor_ + for (int anchor_id = 0; anchor_id < 1; anchor_id++) { + for (int head_id = 0; head_id < config_.q_head_num; + head_id++) { + for (int l = 0; l < config_.head_dim; l++) { + anchor_[layer_id * config_.max_block_num * + config_.anchor_num * + config_.q_head_num * config_.head_dim + + block_idx * config_.anchor_num * + config_.q_head_num * config_.head_dim + + anchor_id * config_.q_head_num * + config_.head_dim + + head_id * config_.head_dim + l] = 0; + } + } + } + + // find top anchor_num importances and their corresponding + // positions in the importance_ tensor + // TODO: Move top_importances to the class member to avoid + // repeated memory allocation + std::priority_queue< + std::pair>, + std::vector>>, + std::greater<>> + top_importances; + for (int head_id = 0; head_id < config_.q_head_num; head_id++) { + for (int k = 0; k < seq_len_; k++) { + top_importances.push(std::make_pair( + GGML_FP16_TO_FP32( + importance_[layer_id][block_idx][k][head_id]), + std::make_pair(block_idx, k))); + // TODO: change to config_ item + if (top_importances.size() > config_.anchor_num) { + top_importances.pop(); + } + } + + // fill anchor_ + + for (int l = 0; l < config_.head_dim; l++) { + anchor_[layer_id * config_.max_block_num * + config_.anchor_num * config_.q_head_num * + config_.head_dim + + block_idx * config_.anchor_num * + config_.q_head_num * config_.head_dim + + 0 * config_.q_head_num * config_.head_dim + + head_id * config_.head_dim + l] = 0; + } + for (int k = 0; k < config_.anchor_num; k++) { + int top_indice = top_importances.top().second.second; + int top_block_idx = top_importances.top().second.first; + + if (config_.kv_type == ggml_type::GGML_TYPE_F16) { + + for (int l = 0; l < config_.head_dim; l++) { + anchor_[layer_id * config_.max_block_num * + config_.anchor_num * + config_.q_head_num * + config_.head_dim + + top_block_idx * config_.anchor_num * + config_.q_head_num * + config_.head_dim + + 0 * config_.q_head_num * + config_.head_dim + + head_id * config_.head_dim + l] = + GGML_FP32_TO_FP16( + GGML_FP16_TO_FP32( + anchor_[layer_id * + config_.max_block_num * + config_.anchor_num * + config_.q_head_num * + config_.head_dim + + top_block_idx * + config_.anchor_num * + config_.q_head_num * + config_.head_dim + + 0 * config_.q_head_num * + config_.head_dim + + head_id * config_.head_dim + + l]) + + GGML_FP16_TO_FP32( + k_cache_fp16_[layer_id] + [head_id / n_gqa_] + [top_block_idx] + [top_indice * + config_.head_dim + + l])); + } + + } else if (config_.kv_type == + ggml_type::GGML_TYPE_Q4_0) { + for (int l = 0; l < config_.head_dim / 32; l++) { + block_q4_0 block = k_cache_q4 + [layer_id][head_id / n_gqa_][top_block_idx] + [top_indice * config_.head_dim / 32 + l]; + dequantize_row_q4_0(&block, block_fp32.data(), + 32); + for (int m = 0; m < 32; m++) { + anchor_[layer_id * config_.max_block_num * + config_.anchor_num * + config_.q_head_num * + config_.head_dim + + top_block_idx * config_.anchor_num * + config_.q_head_num * + config_.head_dim + + 0 * config_.q_head_num * + config_.head_dim + + head_id * config_.head_dim + + l * 32 + m] = + GGML_FP32_TO_FP16( + block_fp32[m] / 4 + + GGML_FP16_TO_FP32( + anchor_[layer_id * + config_ + .max_block_num * + config_.anchor_num * + config_.q_head_num * + config_.head_dim + + top_block_idx * + config_.anchor_num * + config_.q_head_num * + config_.head_dim + + 0 * config_.q_head_num * + config_.head_dim + + head_id * + config_.head_dim + + l * 32 + m])); + } + } + } else if (config_.kv_type == + ggml_type::GGML_TYPE_Q8_0) { + for (int l = 0; l < config_.head_dim / 32; l++) { + block_q8_0 block = k_cache_q8 + [layer_id][head_id / n_gqa_][top_block_idx] + [top_indice * config_.head_dim / 32 + l]; + dequantize_row_q8_0(&block, block_fp32.data(), + 32); + for (int m = 0; m < 32; m++) { + anchor_[layer_id * config_.max_block_num * + config_.anchor_num * + config_.q_head_num * + config_.head_dim + + top_block_idx * config_.anchor_num * + config_.q_head_num * + config_.head_dim + + 0 * config_.q_head_num * + config_.head_dim + + head_id * config_.head_dim + + l * 32 + m] = + GGML_FP32_TO_FP16( + block_fp32[m] / 4 + + GGML_FP16_TO_FP32( + anchor_[layer_id * + config_ + .max_block_num * + config_.anchor_num * + config_.q_head_num * + config_.head_dim + + top_block_idx * + config_.anchor_num * + config_.q_head_num * + config_.head_dim + + 0 * config_.q_head_num * + config_.head_dim + + head_id * + config_.head_dim + + l * 32 + m])); + } + } + } + top_importances.pop(); + } + } + } else if (config_.anchor_type == AnchorType::BLOCK_MEAN) { + // clear anchor_ + for (int anchor_id = 0; anchor_id < config_.anchor_num; + anchor_id++) { + for (int head_id = 0; head_id < config_.q_head_num; + head_id++) { + for (int l = 0; l < config_.head_dim; l++) { + anchor_[layer_id * config_.max_block_num * + config_.anchor_num * + config_.q_head_num * config_.head_dim + + block_idx * config_.anchor_num * + config_.q_head_num * config_.head_dim + + anchor_id * config_.q_head_num * + config_.head_dim + + head_id * config_.head_dim + l] = 0; + } + } + } + + // fill anchor_ + if (config_.kv_type == ggml_type::GGML_TYPE_F16) { + + for (int head_id = 0; head_id < config_.q_head_num; + head_id++) { + for (int k = 0; k < config_.block_len; k++) { + for (int l = 0; l < config_.head_dim; l++) { + anchor_[layer_id * config_.max_block_num * + config_.anchor_num * + config_.q_head_num * + config_.head_dim + + block_idx * config_.anchor_num * + config_.q_head_num * + config_.head_dim + + 0 * config_.q_head_num * + config_.head_dim + + head_id * config_.head_dim + l] = + GGML_FP32_TO_FP16( + GGML_FP16_TO_FP32( + anchor_[layer_id * + config_.max_block_num * + config_.anchor_num * + config_.q_head_num * + config_.head_dim + + block_idx * + config_.anchor_num * + config_.q_head_num * + config_.head_dim + + 0 * config_.q_head_num * + config_.head_dim + + head_id * config_.head_dim + + l]) + + GGML_FP16_TO_FP32( + k_cache_fp16_[layer_id] + [head_id / n_gqa_] + [block_idx] + [k * config_.head_dim + + l]) / + config_.block_len); + } + } + } + } + } else if (config_.anchor_type == AnchorType::BLOCK_MAX) { + // clear anchor_ + for (int anchor_id = 0; anchor_id < config_.anchor_num; + anchor_id++) { + for (int head_id = 0; head_id < config_.q_head_num; + head_id++) { + for (int l = 0; l < config_.head_dim; l++) { + anchor_[layer_id * config_.max_block_num * + config_.anchor_num * + config_.q_head_num * config_.head_dim + + block_idx * config_.anchor_num * + config_.q_head_num * config_.head_dim + + anchor_id * config_.q_head_num * + config_.head_dim + + head_id * config_.head_dim + l] = 0; + } + } + } + + // fill anchor_ + if (config_.kv_type == ggml_type::GGML_TYPE_F16) { + + for (int head_id = 0; head_id < config_.q_head_num; + head_id++) { + for (int k = 0; k < config_.block_len; k++) { + for (int l = 0; l < config_.head_dim; l++) { + anchor_[layer_id * config_.max_block_num * + config_.anchor_num * + config_.q_head_num * + config_.head_dim + + block_idx * config_.anchor_num * + config_.q_head_num * + config_.head_dim + + 0 * config_.q_head_num * + config_.head_dim + + head_id * config_.head_dim + l] = + GGML_FP32_TO_FP16(std::max( + GGML_FP16_TO_FP32( + anchor_[layer_id * + config_.max_block_num * + config_.anchor_num * + config_.q_head_num * + config_.head_dim + + block_idx * + config_.anchor_num * + config_.q_head_num * + config_.head_dim + + 0 * config_.q_head_num * + config_.head_dim + + head_id * config_.head_dim + + l]), + GGML_FP16_TO_FP32( + k_cache_fp16_ + [layer_id][head_id / n_gqa_] + [block_idx] + [k * config_.head_dim + l]))); + } + } + } + } + } else if (config_.anchor_type == AnchorType::FIXED_ANCHOR) { + // clear anchor_ + for (int anchor_id = 0; anchor_id < 1; anchor_id++) { + for (int head_id = 0; head_id < config_.q_head_num; + head_id++) { + for (int l = 0; l < config_.head_dim; l++) { + anchor_[layer_id * config_.max_block_num * + config_.anchor_num * + config_.q_head_num * config_.head_dim + + block_idx * config_.anchor_num * + config_.q_head_num * config_.head_dim + + anchor_id * config_.q_head_num * + config_.head_dim + + head_id * config_.head_dim + l] = 0; + } + } + } + + // fill anchor_ + if (config_.kv_type == ggml_type::GGML_TYPE_F16) { + + int stride = config_.block_len / config_.anchor_num; + for (int head_id = 0; head_id < config_.q_head_num; + head_id++) { + for (int k = 0, tot = 0; + k < config_.block_len, tot < config_.anchor_num; + k += stride, tot++) { + for (int l = 0; l < config_.head_dim; l++) { + anchor_[layer_id * config_.max_block_num * + config_.anchor_num * + config_.q_head_num * + config_.head_dim + + block_idx * config_.anchor_num * + config_.q_head_num * + config_.head_dim + + 0 * config_.q_head_num * + config_.head_dim + + head_id * config_.head_dim + l] = + GGML_FP32_TO_FP16( + GGML_FP16_TO_FP32( + anchor_[layer_id * + config_.max_block_num * + config_.anchor_num * + config_.q_head_num * + config_.head_dim + + block_idx * + config_.anchor_num * + config_.q_head_num * + config_.head_dim + + 0 * config_.q_head_num * + config_.head_dim + + head_id * config_.head_dim + + l]) + + GGML_FP16_TO_FP32( + k_cache_fp16_[layer_id] + [head_id / n_gqa_] + [block_idx] + [k * config_.head_dim + + l]) / + config_.anchor_num); + } + } + } + } + + } else if (config_.anchor_type == AnchorType::QUEST) { + // clear anchor_ + for (int head_id = 0; head_id < config_.q_head_num; head_id++) { + for (int l = 0; l < config_.head_dim; l++) { + anchor_[layer_id * config_.max_block_num * + config_.anchor_num * config_.q_head_num * + config_.head_dim + + block_idx * config_.anchor_num * + config_.q_head_num * config_.head_dim + + 1 * config_.q_head_num * config_.head_dim + + head_id * config_.head_dim + l] = + GGML_FP32_TO_FP16( + std::numeric_limits::max()); + + anchor_[layer_id * config_.max_block_num * + config_.anchor_num * config_.q_head_num * + config_.head_dim + + block_idx * config_.anchor_num * + config_.q_head_num * config_.head_dim + + 0 * config_.q_head_num * config_.head_dim + + head_id * config_.head_dim + l] = + GGML_FP32_TO_FP16( + std::numeric_limits::min()); + } + } + + // fill anchor_ + + if (config_.kv_type == ggml_type::GGML_TYPE_F16) { + for (int indice = 0; indice < seq_len_; indice++) { + for (int head_id = 0; head_id < config_.kv_head_num; + head_id++) { + for (int l = 0; l < config_.head_dim; l++) { + anchor_[layer_id * config_.max_block_num * + config_.anchor_num * + config_.q_head_num * + config_.head_dim + + block_idx * config_.anchor_num * + config_.q_head_num * + config_.head_dim + + 0 * config_.q_head_num * + config_.head_dim + + head_id * config_.head_dim + l] = + GGML_FP32_TO_FP16(std::max( + GGML_FP16_TO_FP32( + k_cache_fp16_ + [layer_id][head_id][block_idx] + [indice * config_.head_dim + + l]), + GGML_FP16_TO_FP32( + anchor_[layer_id * + config_.max_block_num * + config_.anchor_num * + config_.q_head_num * + config_.head_dim + + block_idx * + config_.anchor_num * + config_.q_head_num * + config_.head_dim + + 0 * config_.q_head_num * + config_.head_dim + + head_id * config_.head_dim + + l]))); + + anchor_[layer_id * config_.max_block_num * + config_.anchor_num * + config_.q_head_num * + config_.head_dim + + block_idx * config_.anchor_num * + config_.q_head_num * + config_.head_dim + + 1 * config_.q_head_num * + config_.head_dim + + head_id * config_.head_dim + l] = + GGML_FP32_TO_FP16(std::min( + GGML_FP16_TO_FP32( + k_cache_fp16_ + [layer_id][head_id][block_idx] + [indice * config_.head_dim + + l]), + GGML_FP16_TO_FP32( + anchor_[layer_id * + config_.max_block_num * + config_.anchor_num * + config_.q_head_num * + config_.head_dim + + block_idx * + config_.anchor_num * + config_.q_head_num * + config_.head_dim + + 1 * config_.q_head_num * + config_.head_dim + + head_id * config_.head_dim + + l]))); + } + } + } + + } else if (config_.kv_type == ggml_type::GGML_TYPE_Q4_0) { + for (int indice = 0; indice < seq_len_; indice++) { + for (int head_id = 0; head_id < config_.kv_head_num; + head_id++) { + for (int l = 0; l < config_.head_dim / 32; l++) { + block_q4_0 block = + k_cache_q4[layer_id][head_id][block_idx] + [indice * config_.head_dim / 32 + + l]; + dequantize_row_q4_0(&block, block_fp32.data(), + 32); + + for (int m = 0; m < 32; m++) { + for (int gqa_idx = 0; gqa_idx < n_gqa_; + gqa_idx++) { + + anchor_[layer_id * + config_.max_block_num * + config_.anchor_num * + config_.q_head_num * + config_.head_dim + + block_idx * config_.anchor_num * + config_.q_head_num * + config_.head_dim + + 0 * config_.q_head_num * + config_.head_dim + + head_id * config_.head_dim + + l * 32 + m] = + GGML_FP32_TO_FP16(std::max( + block_fp32[m], + GGML_FP16_TO_FP32( + anchor_ + [layer_id * + config_ + .max_block_num * + config_ + .anchor_num * + config_ + .q_head_num * + config_.head_dim + + block_idx * + config_ + .anchor_num * + config_ + .q_head_num * + config_.head_dim + + 0 * + config_ + .q_head_num * + config_.head_dim + + head_id * + config_.head_dim + + l * 32 + m]))); + + anchor_[layer_id * + config_.max_block_num * + config_.anchor_num * + config_.q_head_num * + config_.head_dim + + block_idx * config_.anchor_num * + config_.q_head_num * + config_.head_dim + + 1 * config_.q_head_num * + config_.head_dim + + head_id * config_.head_dim + + l * 32 + m] = + GGML_FP32_TO_FP16(std::min( + block_fp32[m], + GGML_FP16_TO_FP32( + anchor_ + [layer_id * + config_ + .max_block_num * + config_ + .anchor_num * + config_ + .q_head_num * + config_.head_dim + + block_idx * + config_ + .anchor_num * + config_ + .q_head_num * + config_.head_dim + + 1 * + config_ + .q_head_num * + config_.head_dim + + head_id * + config_.head_dim + + l * 32 + m]))); + } + } + } + } + } + } else if (config_.kv_type == ggml_type::GGML_TYPE_Q8_0) { + for (int indice = 0; indice < seq_len_; indice++) { + for (int head_id = 0; head_id < config_.kv_head_num; + head_id++) { + for (int l = 0; l < config_.head_dim / 32; l++) { + block_q8_0 block = + k_cache_q8[layer_id][head_id][block_idx] + [indice * config_.head_dim / 32 + + l]; + dequantize_row_q8_0(&block, block_fp32.data(), + 32); + + for (int m = 0; m < 32; m++) { + for (int gqa_idx = 0; gqa_idx < n_gqa_; + gqa_idx++) { + + anchor_[layer_id * + config_.max_block_num * + config_.anchor_num * + config_.q_head_num * + config_.head_dim + + block_idx * config_.anchor_num * + config_.q_head_num * + config_.head_dim + + 0 * config_.q_head_num * + config_.head_dim + + head_id * config_.head_dim + + l * 32 + m] = + GGML_FP32_TO_FP16(std::max( + block_fp32[m], + GGML_FP16_TO_FP32( + anchor_ + [layer_id * + config_ + .max_block_num * + config_ + .anchor_num * + config_ + .q_head_num * + config_.head_dim + + block_idx * + config_ + .anchor_num * + config_ + .q_head_num * + config_.head_dim + + 0 * + config_ + .q_head_num * + config_.head_dim + + head_id * + config_.head_dim + + l * 32 + m]))); + + anchor_[layer_id * + config_.max_block_num * + config_.anchor_num * + config_.q_head_num * + config_.head_dim + + block_idx * config_.anchor_num * + config_.q_head_num * + config_.head_dim + + 1 * config_.q_head_num * + config_.head_dim + + head_id * config_.head_dim + + l * 32 + m] = + GGML_FP32_TO_FP16(std::min( + block_fp32[m], + GGML_FP16_TO_FP32( + anchor_ + [layer_id * + config_ + .max_block_num * + config_ + .anchor_num * + config_ + .q_head_num * + config_.head_dim + + block_idx * + config_ + .anchor_num * + config_ + .q_head_num * + config_.head_dim + + 1 * + config_ + .q_head_num * + config_.head_dim + + head_id * + config_.head_dim + + l * 32 + m]))); + } + } + } + } + } + } + } else { + assert(false); + } + }, + nullptr); + + // Timer end + auto end = std::chrono::high_resolution_clock::now(); + std::chrono::duration duration = end - start; + // printf("time of calc_anchor_all_layers: %f s\n", duration.count()); +} + +void KVCache::clear_importance_all_layers(int *block_table, int *cache_seqlens, + int batch_size, int max_block_num, + Backend *backend) { + // Timer start + auto start = std::chrono::high_resolution_clock::now(); + + // Each task updates the importance of a certain block + seq_len_ = config_.block_len; + backend->do_work_stealing_job( + config_.layer_num * batch_size * max_block_num, nullptr, + [&](int task_id) { + int layer_id = task_id / (batch_size * max_block_num); + int batch_id = (task_id / max_block_num) % batch_size; + int block_id = task_id % max_block_num; + // If the block is out of the sequence length, skip it. In + // particular, the last block of the sequence that is shorter than + // the block length should be skipped. + + if (cache_seqlens[batch_id] / config_.block_len < block_id) { + return; + } + int block_idx = block_table[batch_id * max_block_num + block_id]; + + if (config_.anchor_type == AnchorType::DYNAMIC) { + + // clear anchor_ + for (int head_id = 0; head_id < config_.q_head_num; head_id++) { + for (int l = 0; l < config_.block_len; l++) { + importance_[layer_id][block_idx][l][head_id] = 0; + } + } + } + }, + nullptr); + + // Timer end + auto end = std::chrono::high_resolution_clock::now(); + std::chrono::duration duration = end - start; + // printf("time of clear_importance_all_layerssssss: %f s\n", + // duration.count()); +} + +void KVCache::clear_kvcache_all_layers(int *block_table, int *cache_seqlens, + int batch_size, int max_block_num, + Backend *backend) { + // Timer start + auto start = std::chrono::high_resolution_clock::now(); + + // Each task updates the importance of a certain block + seq_len_ = config_.block_len; + backend->do_work_stealing_job( + config_.layer_num * batch_size * max_block_num * config_.kv_head_num, + nullptr, + [&](int task_id) { + int layer_id = + task_id / (batch_size * max_block_num * config_.kv_head_num); + int batch_id = + (task_id / (max_block_num * config_.kv_head_num)) % batch_size; + int block_id = task_id / config_.kv_head_num % max_block_num; + int head_id = task_id % config_.kv_head_num; + // If the block is out of the sequence length, skip it. In + // particular, the last block of the sequence that is shorter than + // the block length should be skipped. + if (cache_seqlens[batch_id] / config_.block_len < block_id) { + return; + } + int block_idx = block_table[batch_id * max_block_num + block_id]; + + if (config_.kv_type == ggml_type::GGML_TYPE_F16) { + for (int l = 0; l < config_.block_len * config_.head_dim; l++) { + k_cache_fp16_[layer_id][head_id][block_idx][l] = 0; + v_cache_fp16_[layer_id][head_id][block_idx][l] = 0; + } + } else if (config_.kv_type == ggml_type::GGML_TYPE_Q4_0) { + for (int l = 0; l < config_.block_len * config_.head_dim / 32; + l++) { + k_cache_q4[layer_id][head_id][block_idx][l].d = 0; + v_cache_q4[layer_id][head_id][block_idx][l].d = 0; + } + } else if (config_.kv_type == ggml_type::GGML_TYPE_Q8_0) { + for (int l = 0; l < config_.block_len * config_.head_dim / 32; + l++) { + k_cache_q8[layer_id][head_id][block_idx][l].d = 0; + v_cache_q8[layer_id][head_id][block_idx][l].d = 0; + } + } + }, + nullptr); + + // Timer end + auto end = std::chrono::high_resolution_clock::now(); + std::chrono::duration duration = end - start; + // printf("time of clear_kvcache_all_layers: %f s\n", duration.count()); +} + +void KVCache::get_sincos(ggml_fp16_t *sin, ggml_fp16_t *cos, int seqlen) { + // Timer start + auto start = std::chrono::high_resolution_clock::now(); + + const uint16_t *sin_data = const_cast(sin); + const uint16_t *cos_data = const_cast(cos); + + for (int i = 0; i < seqlen; i++) { + for (int j = 0; j < config_.head_dim; j++) { + sin_[i][j] = sin_data[i * config_.head_dim + j]; + cos_[i][j] = cos_data[i * config_.head_dim + j]; + } + } + + // Timer end + auto end = std::chrono::high_resolution_clock::now(); + std::chrono::duration duration = end - start; + printf("time of get_sincos: %f s\n", duration.count()); +} + +void ggml_vec_scale_f32(const int n, float *y, const float v) { +#if defined(GGML_USE_ACCELERATE) + vDSP_vsmul(y, 1, &v, y, 1, n); +#elif defined(GGML_SIMD) + const int np = (n & ~(GGML_F32_STEP - 1)); + + GGML_F32_VEC vx = GGML_F32_VEC_SET1(v); + + GGML_F32_VEC ay[GGML_F32_ARR]; + + for (int i = 0; i < np; i += GGML_F32_STEP) { + for (int j = 0; j < GGML_F32_ARR; j++) { + ay[j] = GGML_F32_VEC_LOAD(y + i + j * GGML_F32_EPR); + ay[j] = GGML_F32_VEC_MUL(ay[j], vx); + + GGML_F32_VEC_STORE(y + i + j * GGML_F32_EPR, ay[j]); + } + } + + // leftovers + for (int i = np; i < n; ++i) { + y[i] *= v; + } +#else + // scalar + for (int i = 0; i < n; ++i) { + y[i] *= v; + } +#endif +} \ No newline at end of file diff --git a/archive/csrc/ktransformers_ext/operators/llamafile/conversion.h b/archive/csrc/ktransformers_ext/operators/llamafile/conversion.h new file mode 100644 index 0000000..2a194d6 --- /dev/null +++ b/archive/csrc/ktransformers_ext/operators/llamafile/conversion.h @@ -0,0 +1,32 @@ +/** + * @Description : + * @Author : chenht2022 + * @Date : 2024-07-12 10:07:58 + * @Version : 1.0.0 + * @LastEditors : chenht2022 + * @LastEditTime : 2024-07-25 10:34:55 + * @Copyright (c) 2024 by KVCache.AI, All Rights Reserved. + **/ +#ifndef CPUINFER_CONVERSION_H +#define CPUINFER_CONVERSION_H + +#include +#include "llama.cpp/ggml.h" + +inline void to_float(const void* input, float* output, int size, ggml_type type) { + if (type == ggml_type::GGML_TYPE_F32) { + memcpy(output, input, size * sizeof(float)); + } else { + ggml_internal_get_type_traits(type).to_float(input, output, size); + } +} + +inline void from_float(const float* input, void* output, int size, ggml_type type) { + if (type == ggml_type::GGML_TYPE_F32) { + memcpy(output, input, size * sizeof(float)); + } else { + ggml_internal_get_type_traits(type).from_float(input, output, size); + } +} + +#endif \ No newline at end of file diff --git a/archive/csrc/ktransformers_ext/operators/llamafile/linear.cpp b/archive/csrc/ktransformers_ext/operators/llamafile/linear.cpp new file mode 100644 index 0000000..d1e7967 --- /dev/null +++ b/archive/csrc/ktransformers_ext/operators/llamafile/linear.cpp @@ -0,0 +1,77 @@ +/** + * @Description : + * @Author : chenht2022 + * @Date : 2024-07-12 10:07:58 + * @Version : 1.0.0 + * @LastEditors : kkk1nak0 + * @LastEditTime : 2024-08-15 07:45:18 + * @Copyright (c) 2024 by KVCache.AI, All Rights Reserved. + **/ +#include "linear.h" + +Linear::Linear(LinearConfig config) { + config_ = config; + proj_ = config_.proj; + + std::vector> mem_requests; + mem_requests.push_back({(void**)&input_fp32_, sizeof(float) * config_.group_max_len * config_.input_size}); + mem_requests.push_back({(void**)&proj_input_, config_.group_max_len * config_.input_size * ggml_type_size(ggml_internal_get_type_traits(config_.proj_type).vec_dot_type) / ggml_blck_size(ggml_internal_get_type_traits(config_.proj_type).vec_dot_type)}); + mem_requests.push_back({(void**)&proj_output_, sizeof(float) * config_.group_max_len * config_.output_size}); + shared_mem_buffer.alloc(this, mem_requests); +} + +Linear::~Linear() { + shared_mem_buffer.dealloc(this); +} + +void Linear::warm_up(Backend *backend) { + std::vector input_fp32(config_.input_size); + std::vector input(config_.input_size * + ggml_type_size(config_.hidden_type) / + ggml_blck_size(config_.hidden_type)); + std::vector output(config_.output_size * + ggml_type_size(config_.hidden_type) / + ggml_blck_size(config_.hidden_type)); + for (int i = 0; i < config_.input_size; i++) { + input_fp32[i] = 0; + } + from_float(input_fp32.data(), input.data(), config_.input_size, config_.hidden_type); + forward_many(1, input.data(), output.data(), backend); +} + +void Linear::forward_many(int qlen, const void* input, void* output, Backend* backend) { + const void* proj_input_ptr; + if (config_.hidden_type == ggml_internal_get_type_traits(config_.proj_type).vec_dot_type) { + proj_input_ptr = input; + } else { + to_float(input, input_fp32_, qlen * config_.input_size, config_.hidden_type); + from_float(input_fp32_, proj_input_, qlen * config_.input_size, ggml_internal_get_type_traits(config_.proj_type).vec_dot_type); + proj_input_ptr = proj_input_; + } + int nth = config_.output_size / config_.stride; + backend->do_work_stealing_job(nth, nullptr, [&](int task_id) { + int ith = task_id; + void* proj_ptr = (uint8_t*)proj_ + ith * config_.stride * config_.input_size * ggml_type_size(config_.proj_type) / ggml_blck_size(config_.proj_type); + float* proj_output_ptr = proj_output_ + ith * config_.stride; + llamafile_sgemm(config_.stride, qlen, config_.input_size / ggml_blck_size(config_.proj_type), proj_ptr, config_.input_size / ggml_blck_size(config_.proj_type), proj_input_ptr, config_.input_size / ggml_blck_size(config_.proj_type), proj_output_ptr, config_.output_size, 0, 1, GGML_TASK_TYPE_COMPUTE, config_.proj_type, ggml_internal_get_type_traits(config_.proj_type).vec_dot_type, GGML_TYPE_F32, GGML_PREC_DEFAULT); + if (config_.stride % ggml_blck_size(config_.hidden_type) == 0) { + for (int i = 0; i < qlen; i++) { + float* output_fp32_ptr = proj_output_ + i * config_.output_size + ith * config_.stride; + void* output_ptr = (uint8_t*)output + i * config_.output_size * ggml_type_size(config_.hidden_type) / ggml_blck_size(config_.hidden_type) + ith * config_.stride * ggml_type_size(config_.hidden_type) / ggml_blck_size(config_.hidden_type); + from_float(output_fp32_ptr, output_ptr, config_.stride, config_.hidden_type); + } + } + }, nullptr); + if (config_.stride % ggml_blck_size(config_.hidden_type) != 0) { + from_float(proj_output_, output, qlen * config_.output_size, config_.hidden_type); + } +} + +void Linear::forward(int qlen, const void* input, void* output, Backend* backend) { + if (qlen <= 0) { + return; + } + int forward_len = std::min(qlen, config_.group_max_len); + forward_many(forward_len, input, output, backend); + forward(qlen - forward_len, (uint8_t*)input + forward_len * config_.input_size * ggml_type_size(config_.hidden_type) / ggml_blck_size(config_.hidden_type), (uint8_t*)output + forward_len * config_.output_size * ggml_type_size(config_.hidden_type) / ggml_blck_size(config_.hidden_type), backend); +} \ No newline at end of file diff --git a/archive/csrc/ktransformers_ext/operators/llamafile/linear.h b/archive/csrc/ktransformers_ext/operators/llamafile/linear.h new file mode 100644 index 0000000..f8ae7ae --- /dev/null +++ b/archive/csrc/ktransformers_ext/operators/llamafile/linear.h @@ -0,0 +1,59 @@ +/** + * @Description : + * @Author : chenht2022 + * @Date : 2024-07-12 10:07:58 + * @Version : 1.0.0 + * @LastEditors : chenht2022 + * @LastEditTime : 2024-07-25 10:35:00 + * @Copyright (c) 2024 by KVCache.AI, All Rights Reserved. + **/ +#ifndef CPUINFER_OPERATOR_LINEAR_H +#define CPUINFER_OPERATOR_LINEAR_H + +#include +#include +#include +#include +#include + +#include "../../cpu_backend/backend.h" +#include "../../cpu_backend/shared_mem_buffer.h" +#include "conversion.h" +#include "llama.cpp/ggml-impl.h" +#include "llama.cpp/ggml-quants.h" +#include "llama.cpp/ggml.h" +#include "llamafile/sgemm.h" + +struct LinearConfig { + int input_size; + int output_size; + int stride; + int group_max_len; + void* proj; + ggml_type proj_type; + ggml_type hidden_type; + + LinearConfig() {} + + LinearConfig(int input_size, int output_size, int stride, int group_max_len, void* proj, ggml_type proj_type, ggml_type hidden_type) + : input_size(input_size), output_size(output_size), stride(stride), group_max_len(group_max_len), proj(proj), proj_type(proj_type), hidden_type(hidden_type) {} +}; + +class Linear { + public: + Linear(LinearConfig); + ~Linear(); + void warm_up(Backend* backend); + void forward_many(int qlen, const void* input, void* output, Backend* backend); + void forward(int qlen, const void* input, void* output, Backend* backend); + + private: + LinearConfig config_; + void* proj_; // [output_size * input_size ( /32 if quantized)] + + float* input_fp32_; // [group_max_len * input_size] + uint8_t* proj_input_; // [group_max_len * input_size * ggml_type_size(ggml_internal_get_type_traits(proj_type).vec_dot_type) / ggml_blck_size(ggml_internal_get_type_traits(proj_type).vec_dot_type)] + float* proj_output_; // [group_max_len * output_size] +}; + +#endif \ No newline at end of file diff --git a/archive/csrc/ktransformers_ext/operators/llamafile/mlp.cpp b/archive/csrc/ktransformers_ext/operators/llamafile/mlp.cpp new file mode 100644 index 0000000..602fdcb --- /dev/null +++ b/archive/csrc/ktransformers_ext/operators/llamafile/mlp.cpp @@ -0,0 +1,125 @@ +/** + * @Description : + * @Author : chenht2022 + * @Date : 2024-07-16 10:43:18 + * @Version : 1.0.0 + * @LastEditors : kkk1nak0 + * @LastEditTime : 2024-08-15 07:44:38 + * @Copyright (c) 2024 by KVCache.AI, All Rights Reserved. + **/ +#include "mlp.h" + +MLP::MLP(MLPConfig config) { + config_ = config; + gate_proj_ = config_.gate_proj; + up_proj_ = config_.up_proj; + down_proj_ = config_.down_proj; + + std::vector> mem_requests; + mem_requests.push_back({(void**)&input_fp32_, sizeof(float) * config_.group_max_len * config_.hidden_size}); + mem_requests.push_back({(void**)&gate_input_, config_.group_max_len * config_.hidden_size * ggml_type_size(ggml_internal_get_type_traits(config_.gate_type).vec_dot_type) / ggml_blck_size(ggml_internal_get_type_traits(config_.gate_type).vec_dot_type)}); + mem_requests.push_back({(void**)&up_input_, config_.group_max_len * config_.hidden_size * ggml_type_size(ggml_internal_get_type_traits(config_.up_type).vec_dot_type) / ggml_blck_size(ggml_internal_get_type_traits(config_.up_type).vec_dot_type)}); + mem_requests.push_back({(void**)&gate_output_, sizeof(float) * config_.group_max_len * config_.intermediate_size}); + mem_requests.push_back({(void**)&up_output_, sizeof(float) * config_.group_max_len * config_.intermediate_size}); + mem_requests.push_back({(void**)&intermediate_fp32_, sizeof(float) * config_.group_max_len * config_.intermediate_size}); + mem_requests.push_back({(void**)&down_input_, config_.group_max_len * config_.intermediate_size * ggml_type_size(ggml_internal_get_type_traits(config_.down_type).vec_dot_type) / ggml_blck_size(ggml_internal_get_type_traits(config_.down_type).vec_dot_type)}); + mem_requests.push_back({(void**)&down_output_, sizeof(float) * config_.group_max_len * config_.hidden_size}); + shared_mem_buffer.alloc(this, mem_requests); +} + +MLP::~MLP() { + shared_mem_buffer.dealloc(this); +} + +void MLP::warm_up(Backend *backend) { + std::vector input_fp32(config_.hidden_size); + std::vector input(config_.hidden_size * + ggml_type_size(config_.hidden_type) / + ggml_blck_size(config_.hidden_type)); + std::vector output(config_.hidden_size * + ggml_type_size(config_.hidden_type) / + ggml_blck_size(config_.hidden_type)); + for (int i = 0; i < config_.hidden_size; i++) { + input_fp32[i] = 0; + } + from_float(input_fp32.data(), input.data(), config_.hidden_size, config_.hidden_type); + forward_many(1, input.data(), output.data(), backend); +} + +static float act_fn(float x) { return x / (1.0f + expf(-x)); } + +void MLP::forward_many(int qlen, const void* input, void* output, Backend* backend) { + const void* gate_input_ptr; + const void* up_input_ptr; + if (config_.hidden_type == ggml_internal_get_type_traits(config_.gate_type).vec_dot_type && config_.hidden_type == ggml_internal_get_type_traits(config_.up_type).vec_dot_type) { + gate_input_ptr = up_input_ptr = input; + } else { + to_float(input, input_fp32_, qlen * config_.hidden_size, config_.hidden_type); + if (ggml_internal_get_type_traits(config_.gate_type).vec_dot_type == ggml_internal_get_type_traits(config_.up_type).vec_dot_type) { + from_float(input_fp32_, gate_input_, qlen * config_.hidden_size, ggml_internal_get_type_traits(config_.gate_type).vec_dot_type); + gate_input_ptr = up_input_ptr = gate_input_; + } else { + if (config_.hidden_type != ggml_internal_get_type_traits(config_.gate_type).vec_dot_type) { + from_float(input_fp32_, gate_input_, qlen * config_.hidden_size, ggml_internal_get_type_traits(config_.gate_type).vec_dot_type); + gate_input_ptr = gate_input_; + } else { + gate_input_ptr = input; + } + if (config_.hidden_type != ggml_internal_get_type_traits(config_.up_type).vec_dot_type) { + from_float(input_fp32_, up_input_, qlen * config_.hidden_size, ggml_internal_get_type_traits(config_.up_type).vec_dot_type); + up_input_ptr = up_input_; + } else { + up_input_ptr = input; + } + } + } + int nth = config_.intermediate_size / config_.stride; + backend->do_work_stealing_job(nth, nullptr, [&](int task_id) { + int ith = task_id; + void* gate_proj_ptr = (uint8_t*)gate_proj_ + ith * config_.stride * config_.hidden_size * ggml_type_size(config_.gate_type) / ggml_blck_size(config_.gate_type); + float* gate_output_ptr = gate_output_ + ith * config_.stride; + llamafile_sgemm(config_.stride, qlen, config_.hidden_size / ggml_blck_size(config_.gate_type), gate_proj_ptr, config_.hidden_size / ggml_blck_size(config_.gate_type), gate_input_ptr, config_.hidden_size / ggml_blck_size(config_.gate_type), gate_output_ptr, config_.intermediate_size, 0, 1, GGML_TASK_TYPE_COMPUTE, config_.gate_type, ggml_internal_get_type_traits(config_.gate_type).vec_dot_type, GGML_TYPE_F32, GGML_PREC_DEFAULT); + void* up_proj_ptr = (uint8_t*)up_proj_ + ith * config_.stride * config_.hidden_size * ggml_type_size(config_.up_type) / ggml_blck_size(config_.up_type); + float* up_output_ptr = up_output_ + ith * config_.stride; + llamafile_sgemm(config_.stride, qlen, config_.hidden_size / ggml_blck_size(config_.up_type), up_proj_ptr, config_.hidden_size / ggml_blck_size(config_.up_type), up_input_ptr, config_.hidden_size / ggml_blck_size(config_.up_type), up_output_ptr, config_.intermediate_size, 0, 1, GGML_TASK_TYPE_COMPUTE, config_.up_type, ggml_internal_get_type_traits(config_.up_type).vec_dot_type, GGML_TYPE_F32, GGML_PREC_DEFAULT); + for (int i = 0; i < qlen; i++) { + for (int j = ith * config_.stride; j < (ith + 1) * config_.stride; j++) { + intermediate_fp32_[i * config_.intermediate_size + j] = act_fn(gate_output_[i * config_.intermediate_size + j]) * up_output_[i * config_.intermediate_size + j]; + } + if (config_.stride % ggml_blck_size(ggml_internal_get_type_traits(config_.down_type).vec_dot_type) == 0) { + float* intermediate_fp32_ptr = intermediate_fp32_ + i * config_.intermediate_size + ith * config_.stride; + void* down_input_ptr = (uint8_t*)down_input_ + i * config_.intermediate_size * ggml_type_size(ggml_internal_get_type_traits(config_.down_type).vec_dot_type) / ggml_blck_size(ggml_internal_get_type_traits(config_.down_type).vec_dot_type) + ith * config_.stride * ggml_type_size(ggml_internal_get_type_traits(config_.down_type).vec_dot_type) / ggml_blck_size(ggml_internal_get_type_traits(config_.down_type).vec_dot_type); + from_float(intermediate_fp32_ptr, down_input_ptr, config_.stride, ggml_internal_get_type_traits(config_.down_type).vec_dot_type); + } + } + }, nullptr); + if (config_.stride % ggml_blck_size(ggml_internal_get_type_traits(config_.down_type).vec_dot_type) != 0) { + from_float(intermediate_fp32_, down_input_, qlen * config_.intermediate_size, ggml_internal_get_type_traits(config_.down_type).vec_dot_type); + } + nth = config_.hidden_size / config_.stride; + backend->do_work_stealing_job(nth, nullptr, [&](int task_id) { + int ith = task_id; + void* down_proj_ptr = (uint8_t*)down_proj_ + ith * config_.stride * config_.intermediate_size * ggml_type_size(config_.down_type) / ggml_blck_size(config_.down_type); + float* down_output_ptr = down_output_ + ith * config_.stride; + llamafile_sgemm(config_.stride, qlen, config_.intermediate_size / ggml_blck_size(config_.down_type), down_proj_ptr, config_.intermediate_size / ggml_blck_size(config_.down_type), down_input_, config_.intermediate_size / ggml_blck_size(config_.down_type), down_output_ptr, config_.hidden_size, 0, 1, GGML_TASK_TYPE_COMPUTE, config_.down_type, ggml_internal_get_type_traits(config_.down_type).vec_dot_type, GGML_TYPE_F32, GGML_PREC_DEFAULT); + if (config_.stride % ggml_blck_size(config_.hidden_type) == 0) { + for (int i = 0; i < qlen; i++) { + float* output_fp32_ptr = down_output_ + i * config_.hidden_size + ith * config_.stride; + void* output_ptr = (uint8_t*)output + i * config_.hidden_size * ggml_type_size(config_.hidden_type) / ggml_blck_size(config_.hidden_type) + ith * config_.stride * ggml_type_size(config_.hidden_type) / ggml_blck_size(config_.hidden_type); + from_float(output_fp32_ptr, output_ptr, config_.stride, config_.hidden_type); + } + } + }, nullptr); + if (config_.stride % ggml_blck_size(config_.hidden_type) != 0) { + from_float(down_output_, output, qlen * config_.hidden_size, config_.hidden_type); + } +} + +void MLP::forward(int qlen, const void* input, void* output, Backend* backend) { + if (qlen <= 0) { + return; + } + int forward_len = std::min(qlen, config_.group_max_len); + forward_many(forward_len, input, output, backend); + forward(qlen - forward_len, (uint8_t*)input + forward_len * config_.hidden_size * ggml_type_size(config_.hidden_type) / ggml_blck_size(config_.hidden_type), (uint8_t*)output + forward_len * config_.hidden_size * ggml_type_size(config_.hidden_type) / ggml_blck_size(config_.hidden_type), backend); +} \ No newline at end of file diff --git a/archive/csrc/ktransformers_ext/operators/llamafile/mlp.h b/archive/csrc/ktransformers_ext/operators/llamafile/mlp.h new file mode 100644 index 0000000..7e6e5cc --- /dev/null +++ b/archive/csrc/ktransformers_ext/operators/llamafile/mlp.h @@ -0,0 +1,70 @@ +/** + * @Description : + * @Author : chenht2022 + * @Date : 2024-07-12 10:07:58 + * @Version : 1.0.0 + * @LastEditors : chenht2022 + * @LastEditTime : 2024-07-25 10:35:06 + * @Copyright (c) 2024 by KVCache.AI, All Rights Reserved. + **/ +#ifndef CPUINFER_OPERATOR_MLP_H +#define CPUINFER_OPERATOR_MLP_H + +#include +#include +#include +#include +#include + +#include "../../cpu_backend/backend.h" +#include "../../cpu_backend/shared_mem_buffer.h" +#include "conversion.h" +#include "llama.cpp/ggml-impl.h" +#include "llama.cpp/ggml-quants.h" +#include "llama.cpp/ggml.h" +#include "llamafile/sgemm.h" + +struct MLPConfig { + int hidden_size; + int intermediate_size; + int stride; + int group_max_len; + void* gate_proj; + void* up_proj; + void* down_proj; + ggml_type gate_type; + ggml_type up_type; + ggml_type down_type; + ggml_type hidden_type; + + MLPConfig() {} + + MLPConfig(int hidden_size, int intermediate_size, int stride, int group_max_len, void* gate_proj, void* up_proj, void* down_proj, ggml_type gate_type, ggml_type up_type, ggml_type down_type, ggml_type hidden_type) + : hidden_size(hidden_size), intermediate_size(intermediate_size), stride(stride), group_max_len(group_max_len), gate_proj(gate_proj), up_proj(up_proj), down_proj(down_proj), gate_type(gate_type), up_type(up_type), down_type(down_type), hidden_type(hidden_type) {} +}; + +class MLP { + public: + MLP(MLPConfig); + ~MLP(); + void warm_up(Backend* backend); + void forward_many(int qlen, const void* input, void* output, Backend* backend); + void forward(int qlen, const void* input, void* output, Backend* backend); + + private: + MLPConfig config_; + void* gate_proj_; // [intermediate_size * hidden_size ( /32 if quantized)] + void* up_proj_; // [intermediate_size * hidden_size ( /32 if quantized)] + void* down_proj_; // [hidden_size * intermediate_size ( /32 if quantized)] + + float* input_fp32_; // [group_max_len * hidden_size] + uint8_t* gate_input_; // [group_max_len * hidden_size * ggml_type_size(ggml_internal_get_type_traits(gate_type).vec_dot_type) / ggml_blck_size(ggml_internal_get_type_traits(gate_type).vec_dot_type)] + uint8_t* up_input_; // [group_max_len * hidden_size * ggml_type_size(ggml_internal_get_type_traits(up_type).vec_dot_type) / ggml_blck_size(ggml_internal_get_type_traits(up_type).vec_dot_type)] + float* gate_output_; // [group_max_len * intermediate_size] + float* up_output_; // [group_max_len * intermediate_size] + float* intermediate_fp32_; // [group_max_len * intermediate_size] + uint8_t* down_input_; // [group_max_len * intermediate_size * ggml_type_size(ggml_internal_get_type_traits(down_type).vec_dot_type) / ggml_blck_size(ggml_internal_get_type_traits(down_type).vec_dot_type)] + float* down_output_; // [group_max_len * hidden_size] +}; + +#endif \ No newline at end of file diff --git a/archive/csrc/ktransformers_ext/operators/llamafile/moe.cpp b/archive/csrc/ktransformers_ext/operators/llamafile/moe.cpp new file mode 100644 index 0000000..86e55a2 --- /dev/null +++ b/archive/csrc/ktransformers_ext/operators/llamafile/moe.cpp @@ -0,0 +1,380 @@ +/** + * @Description : + * @Author : chenht2022 + * @Date : 2024-07-22 02:03:22 + * @Version : 1.0.0 + * @LastEditors : kkk1nak0 + * @LastEditTime : 2024-08-15 07:43:41 + * @Copyright (c) 2024 by KVCache.AI, All Rights Reserved. + **/ +#include "moe.h" +#include +#include +#include + +#ifdef USE_NUMA +#include +#include +#endif + +MOE::MOE(MOEConfig config) { + config_ = config; + gate_proj_ = config_.gate_proj; + up_proj_ = config_.up_proj; + down_proj_ = config_.down_proj; + + #ifdef USE_NUMA + int numa_nodes = numa_num_configured_nodes(); + gate_proj_numa_.resize(numa_nodes); + up_proj_numa_.resize(numa_nodes); + down_proj_numa_.resize(numa_nodes); + size_t exp_inter_hidden_mul_ = (size_t)config.expert_num * config.intermediate_size * config.hidden_size; + for (int i = 0; i < numa_nodes; i++) { + gate_proj_numa_[i] = numa_alloc_onnode(exp_inter_hidden_mul_* ggml_type_size(config.gate_type) / ggml_blck_size(config.gate_type), i); + up_proj_numa_[i] = numa_alloc_onnode(exp_inter_hidden_mul_* ggml_type_size(config.up_type) / ggml_blck_size(config.up_type), i); + down_proj_numa_[i] = numa_alloc_onnode(exp_inter_hidden_mul_* ggml_type_size(config.down_type) / ggml_blck_size(config.down_type), i); + if (!gate_proj_numa_[i]) { + std::cout << "Memory allocation failed for gate_proj_numa_ on node " << i << std::endl; + } + if (!up_proj_numa_[i]) { + std::cout << "Memory allocation failed for up_proj_numa_ on node " << i << std::endl; + } + if (!down_proj_numa_[i]) { + std::cout << "Memory allocation failed for down_proj_numa_ on node " << i << std::endl; + } + memcpy(gate_proj_numa_[i], gate_proj_, exp_inter_hidden_mul_* ggml_type_size(config.gate_type) / ggml_blck_size(config.gate_type)); + memcpy(up_proj_numa_[i], up_proj_, exp_inter_hidden_mul_* ggml_type_size(config.up_type) / ggml_blck_size(config.up_type)); + memcpy(down_proj_numa_[i], down_proj_, exp_inter_hidden_mul_* ggml_type_size(config.down_type) / ggml_blck_size(config.down_type)); + } + #endif + + std::vector> s_mem_requests; + s_mem_requests.push_back({(void**)&s_input_fp32_, sizeof(float) * config_.hidden_size}); + s_mem_requests.push_back({(void**)&s_gate_input_, config_.hidden_size * ggml_type_size(ggml_internal_get_type_traits(config_.gate_type).vec_dot_type) / ggml_blck_size(ggml_internal_get_type_traits(config_.gate_type).vec_dot_type)}); + s_mem_requests.push_back({(void**)&s_up_input_, config_.hidden_size * ggml_type_size(ggml_internal_get_type_traits(config_.up_type).vec_dot_type) / ggml_blck_size(ggml_internal_get_type_traits(config_.up_type).vec_dot_type)}); + s_gate_output_.resize(config_.routed_expert_num); + s_up_output_.resize(config_.routed_expert_num); + s_intermediate_fp32_.resize(config_.routed_expert_num); + s_down_input_.resize(config_.routed_expert_num); + s_down_output_.resize(config_.routed_expert_num); + for (int i = 0; i < config_.routed_expert_num; i++) { + s_mem_requests.push_back({(void**)&s_gate_output_[i], sizeof(float) * config_.intermediate_size}); + s_mem_requests.push_back({(void**)&s_up_output_[i], sizeof(float) * config_.intermediate_size}); + s_mem_requests.push_back({(void**)&s_intermediate_fp32_[i], sizeof(float) * config_.intermediate_size}); + s_mem_requests.push_back({(void**)&s_down_input_[i], config_.intermediate_size * ggml_type_size(ggml_internal_get_type_traits(config_.down_type).vec_dot_type) / ggml_blck_size(ggml_internal_get_type_traits(config_.down_type).vec_dot_type)}); + s_mem_requests.push_back({(void**)&s_down_output_[i], sizeof(float) * config_.hidden_size}); + } + s_mem_requests.push_back({(void**)&s_output_fp32_, sizeof(float) * config_.hidden_size}); + shared_mem_buffer.alloc(this, s_mem_requests); + + std::vector> m_mem_requests; + m_input_fp32_.resize(config_.group_max_len); + m_gate_input_.resize(config_.group_max_len); + m_up_input_.resize(config_.group_max_len); + for (int i = 0; i < config_.group_max_len; i++) { + m_mem_requests.push_back({(void**)&m_input_fp32_[i], sizeof(float) * config_.hidden_size}); + m_mem_requests.push_back({(void**)&m_gate_input_[i], config_.hidden_size * ggml_type_size(ggml_internal_get_type_traits(config_.gate_type).vec_dot_type) / ggml_blck_size(ggml_internal_get_type_traits(config_.gate_type).vec_dot_type)}); + m_mem_requests.push_back({(void**)&m_up_input_[i], config_.hidden_size * ggml_type_size(ggml_internal_get_type_traits(config_.up_type).vec_dot_type) / ggml_blck_size(ggml_internal_get_type_traits(config_.up_type).vec_dot_type)}); + } + m_mem_requests.push_back({(void**)&m_local_gate_input_, config_.routed_expert_num * config_.group_max_len * config_.hidden_size * ggml_type_size(ggml_internal_get_type_traits(config_.gate_type).vec_dot_type) / ggml_blck_size(ggml_internal_get_type_traits(config_.gate_type).vec_dot_type)}); + m_mem_requests.push_back({(void**)&m_local_up_input_, config_.routed_expert_num * config_.group_max_len * config_.hidden_size * ggml_type_size(ggml_internal_get_type_traits(config_.up_type).vec_dot_type) / ggml_blck_size(ggml_internal_get_type_traits(config_.up_type).vec_dot_type)}); + m_mem_requests.push_back({(void**)&m_local_gate_output_, sizeof(float) * config_.routed_expert_num * config_.group_max_len * config_.intermediate_size}); + m_mem_requests.push_back({(void**)&m_local_up_output_, sizeof(float) * config_.routed_expert_num * config_.group_max_len * config_.intermediate_size}); + m_mem_requests.push_back({(void**)&m_local_intermediate_fp32_, sizeof(float) * config_.routed_expert_num * config_.group_max_len * config_.intermediate_size}); + m_mem_requests.push_back({(void**)&m_local_down_input_, config_.routed_expert_num * config_.group_max_len * config_.intermediate_size * ggml_type_size(ggml_internal_get_type_traits(config_.down_type).vec_dot_type) / ggml_blck_size(ggml_internal_get_type_traits(config_.down_type).vec_dot_type)}); + m_mem_requests.push_back({(void**)&m_local_down_output_, sizeof(float) * config_.routed_expert_num * config_.group_max_len * config_.hidden_size}); + m_output_fp32_.resize(config_.group_max_len); + for (int i = 0; i < config_.group_max_len; i++) { + m_mem_requests.push_back({(void**)&m_output_fp32_[i], sizeof(float) * config_.hidden_size}); + } + shared_mem_buffer.alloc(this, m_mem_requests); + + m_local_pos_.resize(config_.group_max_len); + for (int i = 0; i < config_.group_max_len; i++) { + m_local_pos_[i].resize(config_.routed_expert_num); + } + m_local_num_.resize(config_.expert_num); + m_local_gate_input_ptr_.resize(config_.expert_num); + m_local_up_input_ptr_.resize(config_.expert_num); + m_local_gate_output_ptr_.resize(config_.expert_num); + m_local_up_output_ptr_.resize(config_.expert_num); + m_local_intermediate_fp32_ptr_.resize(config_.expert_num); + m_local_down_input_ptr_.resize(config_.expert_num); + m_local_down_output_ptr_.resize(config_.expert_num); +} + +MOE::~MOE() { + shared_mem_buffer.dealloc(this); + + #ifdef USE_NUMA + int numa_nodes = numa_num_configured_nodes(); + for (int i = 0; i < numa_nodes; i++) { + numa_free(gate_proj_numa_[i], config_.expert_num * config_.intermediate_size * config_.hidden_size * ggml_type_size(config_.gate_type) / ggml_blck_size(config_.gate_type)); + numa_free(up_proj_numa_[i], config_.expert_num * config_.intermediate_size * config_.hidden_size * ggml_type_size(config_.up_type) / ggml_blck_size(config_.up_type)); + numa_free(down_proj_numa_[i], config_.expert_num * config_.hidden_size * config_.intermediate_size * ggml_type_size(config_.down_type) / ggml_blck_size(config_.down_type)); + } + #endif +} + +void MOE::warm_up(Backend* backend) { + std::vector input_fp32(config_.hidden_size); + std::vector input(config_.hidden_size * ggml_type_size(config_.hidden_type) / ggml_blck_size(config_.hidden_type)); + std::vector output(config_.hidden_size * ggml_type_size(config_.hidden_type) / ggml_blck_size(config_.hidden_type)); + for (int i = 0; i < config_.hidden_size; i++) { + input_fp32[i] = 0; + } + from_float(input_fp32.data(), input.data(), config_.hidden_size, config_.hidden_type); + for (int i = 0; i < config_.expert_num; i++) { + uint64_t expert_ids = i; + float weights = 0; + forward_one(1, &expert_ids, &weights, input.data(), output.data(), backend); + } +} + +static float act_fn(float x) { + return x / (1.0f + expf(-x)); +} + +static float act_fn_relu(float x) { + if(x > 0.0){ + return x; + } else { + return 0.0; + } +} + +void MOE::forward_one(int k, const uint64_t* expert_ids, const float* weights, const void* input, void* output, Backend* backend) { + const void* gate_input_ptr; + const void* up_input_ptr; + if (config_.hidden_type == ggml_internal_get_type_traits(config_.gate_type).vec_dot_type && config_.hidden_type == ggml_internal_get_type_traits(config_.up_type).vec_dot_type) { + gate_input_ptr = up_input_ptr = input; + } else { + to_float(input, s_input_fp32_, config_.hidden_size, config_.hidden_type); + if (ggml_internal_get_type_traits(config_.gate_type).vec_dot_type == ggml_internal_get_type_traits(config_.up_type).vec_dot_type) { + from_float(s_input_fp32_, s_gate_input_, config_.hidden_size, ggml_internal_get_type_traits(config_.gate_type).vec_dot_type); + gate_input_ptr = up_input_ptr = s_gate_input_; + } else { + if (config_.hidden_type != ggml_internal_get_type_traits(config_.gate_type).vec_dot_type) { + from_float(s_input_fp32_, s_gate_input_, config_.hidden_size, ggml_internal_get_type_traits(config_.gate_type).vec_dot_type); + gate_input_ptr = s_gate_input_; + } else { + gate_input_ptr = input; + } + if (config_.hidden_type != ggml_internal_get_type_traits(config_.up_type).vec_dot_type) { + from_float(s_input_fp32_, s_up_input_, config_.hidden_size, ggml_internal_get_type_traits(config_.up_type).vec_dot_type); + up_input_ptr = s_up_input_; + } else { + up_input_ptr = input; + } + } + } + int nth = config_.intermediate_size / config_.stride; + backend->do_work_stealing_job(nth * k, nullptr, [&](int task_id) { + int expert_idx = task_id / nth; + uint64_t expert_id = expert_ids[expert_idx]; + int ith = task_id % nth; + + #ifdef USE_NUMA + void* gate_proj_ptr = (uint8_t*)gate_proj_numa_[Backend::numa_node] + (expert_id * config_.intermediate_size + ith * config_.stride) * config_.hidden_size * ggml_type_size(config_.gate_type) / ggml_blck_size(config_.gate_type); + #else + void* gate_proj_ptr = (uint8_t*)gate_proj_ + (expert_id * config_.intermediate_size + ith * config_.stride) * config_.hidden_size * ggml_type_size(config_.gate_type) / ggml_blck_size(config_.gate_type); + #endif + + float* gate_output_ptr = s_gate_output_[expert_idx] + ith * config_.stride; + llamafile_sgemm(config_.stride, 1, config_.hidden_size / ggml_blck_size(config_.gate_type), gate_proj_ptr, config_.hidden_size / ggml_blck_size(config_.gate_type), gate_input_ptr, config_.hidden_size / ggml_blck_size(config_.gate_type), gate_output_ptr, config_.stride, 0, 1, GGML_TASK_TYPE_COMPUTE, config_.gate_type, ggml_internal_get_type_traits(config_.gate_type).vec_dot_type, GGML_TYPE_F32, GGML_PREC_DEFAULT); + + #ifdef USE_NUMA + void* up_proj_ptr = (uint8_t*)up_proj_numa_[Backend::numa_node] + (expert_id * config_.intermediate_size + ith * config_.stride) * config_.hidden_size * ggml_type_size(config_.up_type) / ggml_blck_size(config_.up_type); + #else + void* up_proj_ptr = (uint8_t*)up_proj_ + (expert_id * config_.intermediate_size + ith * config_.stride) * config_.hidden_size * ggml_type_size(config_.up_type) / ggml_blck_size(config_.up_type); + #endif + + float* up_output_ptr = s_up_output_[expert_idx] + ith * config_.stride; + llamafile_sgemm(config_.stride, 1, config_.hidden_size / ggml_blck_size(config_.up_type), up_proj_ptr, config_.hidden_size / ggml_blck_size(config_.up_type), up_input_ptr, config_.hidden_size / ggml_blck_size(config_.up_type), up_output_ptr, config_.stride, 0, 1, GGML_TASK_TYPE_COMPUTE, config_.up_type, ggml_internal_get_type_traits(config_.up_type).vec_dot_type, GGML_TYPE_F32, GGML_PREC_DEFAULT); + if(config_.use_silu){ + // use silu as act fn + for (int i = ith * config_.stride; i < (ith + 1) * config_.stride; i++) { + s_intermediate_fp32_[expert_idx][i] = act_fn(s_gate_output_[expert_idx][i]) * s_up_output_[expert_idx][i]; + } + } else { + // use relu as act fn + for (int i = ith * config_.stride; i < (ith + 1) * config_.stride; i++) { + s_intermediate_fp32_[expert_idx][i] = act_fn_relu(s_gate_output_[expert_idx][i]) * s_up_output_[expert_idx][i]; + } + } + if (config_.stride % ggml_blck_size(ggml_internal_get_type_traits(config_.down_type).vec_dot_type) == 0) { + float* intermediate_fp32_ptr = s_intermediate_fp32_[expert_idx] + ith * config_.stride; + void* down_input_ptr = s_down_input_[expert_idx] + ith * config_.stride * ggml_type_size(ggml_internal_get_type_traits(config_.down_type).vec_dot_type) / ggml_blck_size(ggml_internal_get_type_traits(config_.down_type).vec_dot_type); + from_float(intermediate_fp32_ptr, down_input_ptr, config_.stride, ggml_internal_get_type_traits(config_.down_type).vec_dot_type); + } + }, nullptr); + if (config_.stride % ggml_blck_size(ggml_internal_get_type_traits(config_.down_type).vec_dot_type) != 0) { + for (int i = 0; i < k; i++) { + from_float(s_intermediate_fp32_[i], s_down_input_[i], config_.intermediate_size, ggml_internal_get_type_traits(config_.down_type).vec_dot_type); + } + } + nth = config_.hidden_size / config_.stride; + backend->do_work_stealing_job(nth, nullptr, [&](int task_id) { + int ith = task_id; + for (int i = ith * config_.stride; i < (ith + 1) * config_.stride; i++) { + s_output_fp32_[i] = 0; + } + for (int expert_idx = 0; expert_idx < k; expert_idx++) { + uint64_t expert_id = expert_ids[expert_idx]; + + #ifdef USE_NUMA + void* down_proj_ptr = (uint8_t*)down_proj_numa_[Backend::numa_node] + (expert_id * config_.hidden_size + ith * config_.stride) * config_.intermediate_size * ggml_type_size(config_.down_type) / ggml_blck_size(config_.down_type); + #else + void* down_proj_ptr = (uint8_t*)down_proj_ + (expert_id * config_.hidden_size + ith * config_.stride) * config_.intermediate_size * ggml_type_size(config_.down_type) / ggml_blck_size(config_.down_type); + #endif + + float* down_output_ptr = s_down_output_[expert_idx] + ith * config_.stride; + llamafile_sgemm(config_.stride, 1, config_.intermediate_size / ggml_blck_size(config_.down_type), down_proj_ptr, config_.intermediate_size / ggml_blck_size(config_.down_type), s_down_input_[expert_idx], config_.intermediate_size / ggml_blck_size(config_.down_type), down_output_ptr, config_.stride, 0, 1, GGML_TASK_TYPE_COMPUTE, config_.down_type, ggml_internal_get_type_traits(config_.down_type).vec_dot_type, GGML_TYPE_F32, GGML_PREC_DEFAULT); + for (int i = ith * config_.stride; i < (ith + 1) * config_.stride; i++) { + s_output_fp32_[i] += s_down_output_[expert_idx][i] * weights[expert_idx]; + } + } + if (config_.stride % ggml_blck_size(config_.hidden_type) == 0) { + float* output_fp32_ptr = s_output_fp32_ + ith * config_.stride; + void* output_ptr = (uint8_t*)output + ith * config_.stride * ggml_type_size(config_.hidden_type) / ggml_blck_size(config_.hidden_type); + from_float(output_fp32_ptr, output_ptr, config_.stride, config_.hidden_type); + } + }, nullptr); + if (config_.stride % ggml_blck_size(config_.hidden_type) != 0) { + from_float(s_output_fp32_, output, config_.hidden_size, config_.hidden_type); + } +} + +void MOE::forward_many(int qlen, int k, const uint64_t* expert_ids, const float* weights, const void* input, void* output, Backend* backend) { + for (int i = 0; i < config_.expert_num; i++) { + m_local_num_[i] = 0; + } + for (int i = 0; i < qlen; i++) { + for (int j = 0; j < k; j++) { + m_local_pos_[i][j] = m_local_num_[expert_ids[i * k + j]]++; + } + } + uint64_t offset = 0; + for (int i = 0; i < config_.expert_num; i++) { + m_local_gate_input_ptr_[i] = m_local_gate_input_ + offset * config_.hidden_size * ggml_type_size(ggml_internal_get_type_traits(config_.gate_type).vec_dot_type) / ggml_blck_size(ggml_internal_get_type_traits(config_.gate_type).vec_dot_type); + m_local_up_input_ptr_[i] = m_local_up_input_ + offset * config_.hidden_size * ggml_type_size(ggml_internal_get_type_traits(config_.up_type).vec_dot_type) / ggml_blck_size(ggml_internal_get_type_traits(config_.up_type).vec_dot_type); + m_local_gate_output_ptr_[i] = m_local_gate_output_ + offset * config_.intermediate_size; + m_local_up_output_ptr_[i] = m_local_up_output_ + offset * config_.intermediate_size; + m_local_intermediate_fp32_ptr_[i] = m_local_intermediate_fp32_ + offset * config_.intermediate_size; + m_local_down_input_ptr_[i] = m_local_down_input_ + offset * config_.intermediate_size * ggml_type_size(ggml_internal_get_type_traits(config_.down_type).vec_dot_type) / ggml_blck_size(ggml_internal_get_type_traits(config_.down_type).vec_dot_type); + m_local_down_output_ptr_[i] = m_local_down_output_ + offset * config_.hidden_size; + offset += m_local_num_[i]; + } + backend->do_work_stealing_job(qlen, nullptr, [&](int i) { + const void* gate_input_ptr; + const void* up_input_ptr; + if (config_.hidden_type == ggml_internal_get_type_traits(config_.gate_type).vec_dot_type && config_.hidden_type == ggml_internal_get_type_traits(config_.up_type).vec_dot_type) { + gate_input_ptr = up_input_ptr = (uint8_t*)input + i * config_.hidden_size * ggml_type_size(config_.hidden_type) / ggml_blck_size(config_.hidden_type); + } else { + to_float((uint8_t*)input + i * config_.hidden_size * ggml_type_size(config_.hidden_type) / ggml_blck_size(config_.hidden_type), m_input_fp32_[i], config_.hidden_size, config_.hidden_type); + if (ggml_internal_get_type_traits(config_.gate_type).vec_dot_type == ggml_internal_get_type_traits(config_.up_type).vec_dot_type) { + from_float(m_input_fp32_[i], m_gate_input_[i], config_.hidden_size, ggml_internal_get_type_traits(config_.gate_type).vec_dot_type); + gate_input_ptr = up_input_ptr = m_gate_input_[i]; + } else { + if (config_.hidden_type != ggml_internal_get_type_traits(config_.gate_type).vec_dot_type) { + from_float(m_input_fp32_[i], m_gate_input_[i], config_.hidden_size, ggml_internal_get_type_traits(config_.gate_type).vec_dot_type); + gate_input_ptr = m_gate_input_[i]; + } else { + gate_input_ptr = (uint8_t*)input + i * config_.hidden_size * ggml_type_size(config_.hidden_type) / ggml_blck_size(config_.hidden_type); + } + if (config_.hidden_type != ggml_internal_get_type_traits(config_.up_type).vec_dot_type) { + from_float(m_input_fp32_[i], m_up_input_[i], config_.hidden_size, ggml_internal_get_type_traits(config_.up_type).vec_dot_type); + up_input_ptr = m_up_input_[i]; + } else { + up_input_ptr = (uint8_t*)input + i * config_.hidden_size * ggml_type_size(config_.hidden_type) / ggml_blck_size(config_.hidden_type); + } + } + } + for (int j = 0; j < k; j++) { + memcpy(m_local_gate_input_ptr_[expert_ids[i * k + j]] + m_local_pos_[i][j] * config_.hidden_size * ggml_type_size(ggml_internal_get_type_traits(config_.gate_type).vec_dot_type) / ggml_blck_size(ggml_internal_get_type_traits(config_.gate_type).vec_dot_type), gate_input_ptr, config_.hidden_size * ggml_type_size(ggml_internal_get_type_traits(config_.gate_type).vec_dot_type) / ggml_blck_size(ggml_internal_get_type_traits(config_.gate_type).vec_dot_type)); + memcpy(m_local_up_input_ptr_[expert_ids[i * k + j]] + m_local_pos_[i][j] * config_.hidden_size * ggml_type_size(ggml_internal_get_type_traits(config_.up_type).vec_dot_type) / ggml_blck_size(ggml_internal_get_type_traits(config_.up_type).vec_dot_type), up_input_ptr, config_.hidden_size * ggml_type_size(ggml_internal_get_type_traits(config_.up_type).vec_dot_type) / ggml_blck_size(ggml_internal_get_type_traits(config_.up_type).vec_dot_type)); + } + }, nullptr); + int stride = QK_K; + int nth = config_.intermediate_size / stride; + backend->do_work_stealing_job(nth * config_.expert_num, nullptr, [&](int task_id) { + uint64_t expert_idx = task_id / nth; + int ith = task_id % nth; + void* gate_input_ptr = m_local_gate_input_ptr_[expert_idx]; + + #ifdef USE_NUMA + void* gate_proj_ptr = (uint8_t*)gate_proj_numa_[Backend::numa_node] + (expert_idx * config_.intermediate_size + ith * stride) * config_.hidden_size * ggml_type_size(config_.gate_type) / ggml_blck_size(config_.gate_type); + #else + void* gate_proj_ptr = (uint8_t*)gate_proj_ + (expert_idx * config_.intermediate_size + ith * stride) * config_.hidden_size * ggml_type_size(config_.gate_type) / ggml_blck_size(config_.gate_type); + #endif + + float* gate_output_ptr = m_local_gate_output_ptr_[expert_idx] + ith * stride; + llamafile_sgemm(stride, m_local_num_[expert_idx], config_.hidden_size / ggml_blck_size(config_.gate_type), gate_proj_ptr, config_.hidden_size / ggml_blck_size(config_.gate_type), gate_input_ptr, config_.hidden_size / ggml_blck_size(config_.gate_type), gate_output_ptr, config_.intermediate_size, 0, 1, GGML_TASK_TYPE_COMPUTE, config_.gate_type, ggml_internal_get_type_traits(config_.gate_type).vec_dot_type, GGML_TYPE_F32, GGML_PREC_DEFAULT); + void* up_input_ptr = m_local_up_input_ptr_[expert_idx]; + + #ifdef USE_NUMA + void* up_proj_ptr = (uint8_t*)up_proj_numa_[Backend::numa_node] + (expert_idx * config_.intermediate_size + ith * stride) * config_.hidden_size * ggml_type_size(config_.up_type) / ggml_blck_size(config_.up_type); + #else + void* up_proj_ptr = (uint8_t*)up_proj_ + (expert_idx * config_.intermediate_size + ith * stride) * config_.hidden_size * ggml_type_size(config_.up_type) / ggml_blck_size(config_.up_type); + #endif + + float* up_output_ptr = m_local_up_output_ptr_[expert_idx] + ith * stride; + llamafile_sgemm(stride, m_local_num_[expert_idx], config_.hidden_size / ggml_blck_size(config_.up_type), up_proj_ptr, config_.hidden_size / ggml_blck_size(config_.up_type), up_input_ptr, config_.hidden_size / ggml_blck_size(config_.up_type), up_output_ptr, config_.intermediate_size, 0, 1, GGML_TASK_TYPE_COMPUTE, config_.up_type, ggml_internal_get_type_traits(config_.up_type).vec_dot_type, GGML_TYPE_F32, GGML_PREC_DEFAULT); + for (int i = 0; i < m_local_num_[expert_idx]; i++) { + if(config_.use_silu){ + for (int j = ith * stride; j < (ith + 1) * stride; j++) { + m_local_intermediate_fp32_ptr_[expert_idx][i * config_.intermediate_size + j] = act_fn(m_local_gate_output_ptr_[expert_idx][i * config_.intermediate_size + j]) * m_local_up_output_ptr_[expert_idx][i * config_.intermediate_size + j]; + } + } else { + for (int j = ith * stride; j < (ith + 1) * stride; j++) { + m_local_intermediate_fp32_ptr_[expert_idx][i * config_.intermediate_size + j] = act_fn_relu(m_local_gate_output_ptr_[expert_idx][i * config_.intermediate_size + j]) * m_local_up_output_ptr_[expert_idx][i * config_.intermediate_size + j]; + } + } + float* intermediate_fp32_ptr = m_local_intermediate_fp32_ptr_[expert_idx] + i * config_.intermediate_size + ith * stride; + void* down_input_ptr = m_local_down_input_ptr_[expert_idx] + i * config_.intermediate_size * ggml_type_size(ggml_internal_get_type_traits(config_.down_type).vec_dot_type) / ggml_blck_size(ggml_internal_get_type_traits(config_.down_type).vec_dot_type) + ith * stride * ggml_type_size(ggml_internal_get_type_traits(config_.down_type).vec_dot_type) / ggml_blck_size(ggml_internal_get_type_traits(config_.down_type).vec_dot_type); + from_float(intermediate_fp32_ptr, down_input_ptr, stride, ggml_internal_get_type_traits(config_.down_type).vec_dot_type); + } + }, nullptr); + stride = QK_K; + nth = config_.hidden_size / stride; + backend->do_work_stealing_job(nth * config_.expert_num, nullptr, [&](int task_id) { + uint64_t expert_idx = task_id / nth; + int ith = task_id % nth; + void* down_input_ptr = m_local_down_input_ptr_[expert_idx]; + + #ifdef USE_NUMA + void* down_proj_ptr = (uint8_t*)down_proj_numa_[Backend::numa_node] + (expert_idx * config_.hidden_size + ith * stride) * config_.intermediate_size * ggml_type_size(config_.down_type) / ggml_blck_size(config_.down_type); + #else + void* down_proj_ptr = (uint8_t*)down_proj_ + (expert_idx * config_.hidden_size + ith * stride) * config_.intermediate_size * ggml_type_size(config_.down_type) / ggml_blck_size(config_.down_type); + #endif + + float* down_output_ptr = m_local_down_output_ptr_[expert_idx] + ith * stride; + llamafile_sgemm(stride, m_local_num_[expert_idx], config_.intermediate_size / ggml_blck_size(config_.down_type), down_proj_ptr, config_.intermediate_size / ggml_blck_size(config_.down_type), down_input_ptr, config_.intermediate_size / ggml_blck_size(config_.down_type), down_output_ptr, config_.hidden_size, 0, 1, GGML_TASK_TYPE_COMPUTE, config_.down_type, ggml_internal_get_type_traits(config_.down_type).vec_dot_type, GGML_TYPE_F32, GGML_PREC_DEFAULT); + }, nullptr); + backend->do_work_stealing_job(qlen, nullptr, [&](int i) { + for (int e = 0; e < config_.hidden_size; e++) { + m_output_fp32_[i][e] = 0; + } + for (int j = 0; j < k; j++) { + for (int e = 0; e < config_.hidden_size; e++) { + m_output_fp32_[i][e] += m_local_down_output_ptr_[expert_ids[i * k + j]][m_local_pos_[i][j] * config_.hidden_size + e] * weights[i * k + j]; + } + } + from_float(m_output_fp32_[i], (uint8_t*)output + i * config_.hidden_size * ggml_type_size(config_.hidden_type) / ggml_blck_size(config_.hidden_type), config_.hidden_size, config_.hidden_type); + }, nullptr); +} + +void MOE::forward(int qlen, int k, const uint64_t* expert_ids, const float* weights, const void* input, void* output, int* batch_size_tensor, Backend* backend) { + qlen = batch_size_tensor[0]; + if (qlen < config_.group_min_len) { + for (int i = 0; i < qlen; i++) { + forward_one(k, expert_ids + i * k, weights + i * k, (uint8_t*)input + i * config_.hidden_size * ggml_type_size(config_.hidden_type) / ggml_blck_size(config_.hidden_type), (uint8_t*)output + i * config_.hidden_size * ggml_type_size(config_.hidden_type) / ggml_blck_size(config_.hidden_type), backend); + } + return; + } + int forward_len = std::min(config_.group_max_len, qlen); + forward_many(forward_len, k, expert_ids, weights, input, output, backend); + + batch_size_tensor[0] -= forward_len; + forward(qlen - forward_len, k, expert_ids + forward_len * k, weights + forward_len * k, (uint8_t*)input + forward_len * config_.hidden_size * ggml_type_size(config_.hidden_type) / ggml_blck_size(config_.hidden_type), (uint8_t*)output + forward_len * config_.hidden_size * ggml_type_size(config_.hidden_type) / ggml_blck_size(config_.hidden_type), batch_size_tensor, backend); +} \ No newline at end of file diff --git a/archive/csrc/ktransformers_ext/operators/llamafile/moe.h b/archive/csrc/ktransformers_ext/operators/llamafile/moe.h new file mode 100644 index 0000000..b568b51 --- /dev/null +++ b/archive/csrc/ktransformers_ext/operators/llamafile/moe.h @@ -0,0 +1,104 @@ +/** + * @Description : + * @Author : chenht2022 + * @Date : 2024-07-22 02:03:22 + * @Version : 1.0.0 + * @LastEditors : chenht2022 + * @LastEditTime : 2024-07-25 10:35:10 + * @Copyright (c) 2024 by KVCache.AI, All Rights Reserved. + **/ +#ifndef CPUINFER_OPERATOR_MOE_H +#define CPUINFER_OPERATOR_MOE_H + +#include +#include +#include +#include +#include + +#include "../../cpu_backend/backend.h" +#include "../../cpu_backend/shared_mem_buffer.h" +#include "conversion.h" +#include "llama.cpp/ggml-impl.h" +#include "llama.cpp/ggml-quants.h" +#include "llama.cpp/ggml.h" +#include "llamafile/sgemm.h" + +struct MOEConfig { + int expert_num; + int routed_expert_num; + int hidden_size; + int intermediate_size; + int stride; + int group_min_len; + int group_max_len; + bool use_silu; + void* gate_proj; + void* up_proj; + void* down_proj; + ggml_type gate_type; + ggml_type up_type; + ggml_type down_type; + ggml_type hidden_type; + + MOEConfig() {} + + MOEConfig(int expert_num, int routed_expert_num, int hidden_size, int intermediate_size, int stride, int group_min_len, int group_max_len, bool use_silu, void* gate_proj, void* up_proj, void* down_proj, ggml_type gate_type, ggml_type up_type, ggml_type down_type, ggml_type hidden_type) + : expert_num(expert_num), routed_expert_num(routed_expert_num), hidden_size(hidden_size), intermediate_size(intermediate_size), stride(stride), group_min_len(group_min_len), group_max_len(group_max_len), use_silu(use_silu), gate_proj(gate_proj), up_proj(up_proj), down_proj(down_proj), gate_type(gate_type), up_type(up_type), down_type(down_type), hidden_type(hidden_type) {} +}; + +class MOE { + public: + MOE(MOEConfig); + ~MOE(); + void warm_up(Backend* backend); + void forward_one(int k, const uint64_t* expert_ids, const float* weights, const void* input, void* output, Backend* backend); + void forward_many(int qlen, int k, const uint64_t* expert_ids, const float* weights, const void* input, void* output, Backend* backend); + void forward(int qlen, int k, const uint64_t* expert_ids, const float* weights, const void* input, void* output, int* batch_size_tensor, Backend* backend); + + private: + MOEConfig config_; + void* gate_proj_; // [expert_num * intermediate_size * hidden_size ( /32 if quantized)] + void* up_proj_; // [expert_num * intermediate_size * hidden_size ( /32 if quantized)] + void* down_proj_; // [expert_num * hidden_size * intermediate_size ( /32 if quantized)] + + #ifdef USE_NUMA + std::vector gate_proj_numa_; // [numa_num, expert_num * intermediate_size * hidden_size ( /32 if quantized)] + std::vector up_proj_numa_; // [numa_num, expert_num * intermediate_size * hidden_size ( /32 if quantized)] + std::vector down_proj_numa_; // [numa_num, expert_num * hidden_size * intermediate_size ( /32 if quantized)] + #endif + + float* s_input_fp32_; // [hidden_size] + uint8_t* s_gate_input_; // [hidden_size * ggml_type_size(ggml_internal_get_type_traits(gate_type).vec_dot_type) / ggml_blck_size(ggml_internal_get_type_traits(gate_type).vec_dot_type)] + uint8_t* s_up_input_; // [hidden_size * ggml_type_size(ggml_internal_get_type_traits(up_type).vec_dot_type) / ggml_blck_size(ggml_internal_get_type_traits(up_type).vec_dot_type)] + std::vector s_gate_output_; // [routed_expert_num, intermediate_size] + std::vector s_up_output_; // [routed_expert_num, intermediate_size] + std::vector s_intermediate_fp32_; // [routed_expert_num, intermediate_size] + std::vector s_down_input_; // [routed_expert_num, intermediate_size * ggml_type_size(ggml_internal_get_type_traits(down_type).vec_dot_type) / ggml_blck_size(ggml_internal_get_type_traits(down_type).vec_dot_type)] + std::vector s_down_output_; // [routed_expert_num, hidden_size] + float* s_output_fp32_; // [hidden_size] + + std::vector m_input_fp32_; // [group_max_len, hidden_size] + std::vector m_gate_input_; // [group_max_len, hidden_size * ggml_type_size(ggml_internal_get_type_traits(gate_type).vec_dot_type) / ggml_blck_size(ggml_internal_get_type_traits(gate_type).vec_dot_type)] + std::vector m_up_input_; // [group_max_len, hidden_size * ggml_type_size(ggml_internal_get_type_traits(up_type).vec_dot_type) / ggml_blck_size(ggml_internal_get_type_traits(up_type).vec_dot_type)] + uint8_t* m_local_gate_input_; // [routed_expert_num * group_max_len * hidden_size * ggml_type_size(ggml_internal_get_type_traits(gate_type).vec_dot_type) / ggml_blck_size(ggml_internal_get_type_traits(gate_type).vec_dot_type)] + uint8_t* m_local_up_input_; // [routed_expert_num * group_max_len * hidden_size * ggml_type_size(ggml_internal_get_type_traits(up_type).vec_dot_type) / ggml_blck_size(ggml_internal_get_type_traits(up_type).vec_dot_type)] + float* m_local_gate_output_; // [routed_expert_num * group_max_len * intermediate_size] + float* m_local_up_output_; // [routed_expert_num * group_max_len * intermediate_size] + float* m_local_intermediate_fp32_; // [routed_expert_num * group_max_len * intermediate_size] + uint8_t* m_local_down_input_; // [routed_expert_num * group_max_len * intermediate_size * ggml_type_size(ggml_internal_get_type_traits(down_type).vec_dot_type) / ggml_blck_size(ggml_internal_get_type_traits(down_type).vec_dot_type)] + float* m_local_down_output_; // [routed_expert_num * group_max_len * hidden_size] + std::vector m_output_fp32_; // [group_max_len, hidden_size] + + std::vector> m_local_pos_; // [group_max_len, routed_expert_num] + std::vector m_local_num_; // [expert_num] + std::vector m_local_gate_input_ptr_; // [expert_num] + std::vector m_local_up_input_ptr_; // [expert_num] + std::vector m_local_gate_output_ptr_; // [expert_num] + std::vector m_local_up_output_ptr_; // [expert_num] + std::vector m_local_intermediate_fp32_ptr_; // [expert_num] + std::vector m_local_down_input_ptr_; // [expert_num] + std::vector m_local_down_output_ptr_; // [expert_num] +}; + +#endif \ No newline at end of file diff --git a/archive/csrc/ktransformers_ext/vendors/cuda.h b/archive/csrc/ktransformers_ext/vendors/cuda.h new file mode 100644 index 0000000..1746b07 --- /dev/null +++ b/archive/csrc/ktransformers_ext/vendors/cuda.h @@ -0,0 +1,15 @@ +#pragma once + +#include +#include +#include +#include +#include + +#if CUDART_VERSION < 11020 +#define CU_DEVICE_ATTRIBUTE_VIRTUAL_MEMORY_MANAGEMENT_SUPPORTED CU_DEVICE_ATTRIBUTE_VIRTUAL_ADDRESS_MANAGEMENT_SUPPORTED +#define CUBLAS_TF32_TENSOR_OP_MATH CUBLAS_TENSOR_OP_MATH +#define CUBLAS_COMPUTE_16F CUDA_R_16F +#define CUBLAS_COMPUTE_32F CUDA_R_32F +#define cublasComputeType_t cudaDataType_t +#endif // CUDART_VERSION < 11020 diff --git a/archive/csrc/ktransformers_ext/vendors/hip.h b/archive/csrc/ktransformers_ext/vendors/hip.h new file mode 100644 index 0000000..abbc1e8 --- /dev/null +++ b/archive/csrc/ktransformers_ext/vendors/hip.h @@ -0,0 +1,172 @@ +#pragma once + +#define HIP_ENABLE_WARP_SYNC_BUILTINS 1 +#include +#include +#include +#include +#ifdef __HIP_PLATFORM_AMD__ +// for rocblas_initialize() +#include "rocblas/rocblas.h" +#endif // __HIP_PLATFORM_AMD__ + +#define CUBLAS_COMPUTE_16F HIPBLAS_R_16F +#define CUBLAS_COMPUTE_32F HIPBLAS_R_32F +#define CUBLAS_COMPUTE_32F_FAST_16F HIPBLAS_R_32F +#define CUBLAS_GEMM_DEFAULT HIPBLAS_GEMM_DEFAULT +#define CUBLAS_GEMM_DEFAULT_TENSOR_OP HIPBLAS_GEMM_DEFAULT +#define CUBLAS_OP_N HIPBLAS_OP_N +#define CUBLAS_OP_T HIPBLAS_OP_T +#define CUBLAS_STATUS_SUCCESS HIPBLAS_STATUS_SUCCESS +#define CUBLAS_TF32_TENSOR_OP_MATH 0 +#define CUDA_R_16F HIPBLAS_R_16F +#define CUDA_R_32F HIPBLAS_R_32F +#define CU_DEVICE_ATTRIBUTE_VIRTUAL_MEMORY_MANAGEMENT_SUPPORTED hipDeviceAttributeVirtualMemoryManagementSupported +#define CU_MEM_ALLOC_GRANULARITY_RECOMMENDED hipMemAllocationGranularityRecommended +#define CU_MEM_ALLOCATION_TYPE_PINNED hipMemAllocationTypePinned +#define CU_MEM_LOCATION_TYPE_DEVICE hipMemLocationTypeDevice +#define CU_MEM_ACCESS_FLAGS_PROT_READWRITE hipMemAccessFlagsProtReadWrite +#define CU_CHECK(fn) {hipError_t err = fn; if(err != hipSuccess) { GGML_ABORT("HipVMM Failure: %s\n", hipGetErrorString(err)); }} +#define __shfl_sync(mask, var, laneMask, width) __shfl(var, laneMask, width) +#define __shfl_xor_sync(mask, var, laneMask, width) __shfl_xor(var, laneMask, width) +#define cublasComputeType_t hipblasDatatype_t //deprecated, new hipblasComputeType_t not in 5.6 +#define cublasCreate hipblasCreate +#define cublasDestroy hipblasDestroy +#define cublasGemmEx hipblasGemmEx +#define cublasGemmBatchedEx hipblasGemmBatchedEx +#define cublasGemmStridedBatchedEx hipblasGemmStridedBatchedEx +#define cublasHandle_t hipblasHandle_t +#define cublasSetMathMode(handle, mode) CUBLAS_STATUS_SUCCESS +#define cublasSetStream hipblasSetStream +#define cublasSgemm hipblasSgemm +#define cublasStatus_t hipblasStatus_t +#define cublasOperation_t hipblasOperation_t +#define cudaDataType_t hipblasDatatype_t //deprecated, new hipblasDatatype not in 5.6 +#define cudaDeviceCanAccessPeer hipDeviceCanAccessPeer +#define cudaDeviceDisablePeerAccess hipDeviceDisablePeerAccess +#define cudaDeviceEnablePeerAccess hipDeviceEnablePeerAccess +#define cudaDeviceProp hipDeviceProp_t +#define cudaDeviceSynchronize hipDeviceSynchronize +#define cudaError_t hipError_t +#define cudaErrorPeerAccessAlreadyEnabled hipErrorPeerAccessAlreadyEnabled +#define cudaErrorPeerAccessNotEnabled hipErrorPeerAccessNotEnabled +#define cudaEventCreateWithFlags hipEventCreateWithFlags +#define cudaEventDisableTiming hipEventDisableTiming +#define cudaEventRecord hipEventRecord +#define cudaEventSynchronize hipEventSynchronize +#define cudaEvent_t hipEvent_t +#define cudaEventDestroy hipEventDestroy +#define cudaFree hipFree +#define cudaFreeHost hipHostFree +#define cudaGetDevice hipGetDevice +#define cudaGetDeviceCount hipGetDeviceCount +#define cudaGetDeviceProperties hipGetDeviceProperties +#define cudaGetErrorString hipGetErrorString +#define cudaGetLastError hipGetLastError +#define cudaHostRegister hipHostRegister +#define cudaHostRegisterPortable hipHostRegisterPortable +#define cudaHostRegisterReadOnly hipHostRegisterReadOnly +#define cudaHostUnregister hipHostUnregister +#define cudaLaunchHostFunc hipLaunchHostFunc +#define cudaMalloc hipMalloc +#define cudaMallocHost(ptr, size) hipHostMalloc(ptr, size, hipHostMallocDefault) +#define cudaMemcpy hipMemcpy +#define cudaMemcpyAsync hipMemcpyAsync +#define cudaMemcpyPeerAsync hipMemcpyPeerAsync +#define cudaMemcpy2DAsync hipMemcpy2DAsync +#define cudaMemcpyDeviceToDevice hipMemcpyDeviceToDevice +#define cudaMemcpyDeviceToHost hipMemcpyDeviceToHost +#define cudaMemcpyHostToDevice hipMemcpyHostToDevice +#define cudaMemcpyKind hipMemcpyKind +#define cudaMemset hipMemset +#define cudaMemsetAsync hipMemsetAsync +#define cudaMemGetInfo hipMemGetInfo +#define cudaOccupancyMaxPotentialBlockSize hipOccupancyMaxPotentialBlockSize +#define cudaSetDevice hipSetDevice +#define cuDeviceGet hipDeviceGet +#define CUdevice hipDevice_t +#define CUdeviceptr hipDeviceptr_t +#define cuMemUnmap hipMemUnmap +#define CUmemAccessDesc hipMemAccessDesc +#define cuMemAddressFree hipMemAddressFree +#define cuMemRelease hipMemRelease +#define CUmemGenericAllocationHandle hipMemGenericAllocationHandle_t +#define cuMemCreate hipMemCreate +#define cuMemAddressReserve hipMemAddressReserve +#define cuMemMap hipMemMap +#define cuMemSetAccess hipMemSetAccess +#define cuMemGetAllocationGranularity hipMemGetAllocationGranularity +#define CUmemAllocationProp hipMemAllocationProp +#define cuDeviceGetAttribute hipDeviceGetAttribute +#define cudaStreamCreateWithFlags hipStreamCreateWithFlags +#define cudaStreamDestroy hipStreamDestroy +#define cudaStreamFireAndForget hipStreamFireAndForget +#define cudaStreamNonBlocking hipStreamNonBlocking +#define cudaStreamPerThread hipStreamPerThread +#define cudaStreamSynchronize hipStreamSynchronize +#define cudaStreamWaitEvent(stream, event, flags) hipStreamWaitEvent(stream, event, flags) +#define cudaGraphExec_t hipGraphExec_t +#define cudaGraphNode_t hipGraphNode_t +#define cudaKernelNodeParams hipKernelNodeParams +#define cudaKernelNodeParams hipKernelNodeParams +#define cudaGraphExecDestroy hipGraphExecDestroy +#define cudaGraphLaunch hipGraphLaunch +#define cudaErrorGraphExecUpdateFailure hipErrorGraphExecUpdateFailure +#define cudaGraphExecUpdateResultInfo hipGraphExecUpdateResult +#define cudaGraphNodeType hipGraphNodeType +#define cudaGraphNodeTypeKernel hipGraphNodeTypeKernel +#define cudaGraphInstantiate hipGraphInstantiate +#define cudaStreamEndCapture hipStreamEndCapture +#define cudaGraphDestroy hipGraphDestroy +#define cudaGraphKernelNodeSetParams hipGraphKernelNodeSetParams +#define cudaErrorInvalidDeviceFunction hipErrorInvalidDeviceFunction +#define cudaGraphKernelNodeGetParams hipGraphKernelNodeGetParams +#define cudaGraphNodeGetType hipGraphNodeGetType +#define cudaGraphGetNodes hipGraphGetNodes +#define cudaGraphExecUpdate hipGraphExecUpdate +#define cudaStreamCaptureModeRelaxed hipStreamCaptureModeRelaxed +#define cudaStreamBeginCapture hipStreamBeginCapture +#define cudaGraph_t hipGraph_t +#define cudaStream_t hipStream_t +#define cudaSuccess hipSuccess +#define cudaHostFn_t hipHostFn_t +#define __trap() do { abort(); __builtin_unreachable(); } while(0) +#define CUBLAS_STATUS_SUCCESS HIPBLAS_STATUS_SUCCESS +#define CUBLAS_STATUS_NOT_INITIALIZED HIPBLAS_STATUS_NOT_INITIALIZED +#define CUBLAS_STATUS_ALLOC_FAILED HIPBLAS_STATUS_ALLOC_FAILED +#define CUBLAS_STATUS_INVALID_VALUE HIPBLAS_STATUS_INVALID_VALUE +#define CUBLAS_STATUS_ARCH_MISMATCH HIPBLAS_STATUS_ARCH_MISMATCH +#define CUBLAS_STATUS_MAPPING_ERROR HIPBLAS_STATUS_MAPPING_ERROR +#define CUBLAS_STATUS_EXECUTION_FAILED HIPBLAS_STATUS_EXECUTION_FAILED +#define CUBLAS_STATUS_INTERNAL_ERROR HIPBLAS_STATUS_INTERNAL_ERROR +#define CUBLAS_STATUS_NOT_SUPPORTED HIPBLAS_STATUS_NOT_SUPPORTED + +#define __CUDA_ARCH__ 1300 + +#if defined(__gfx803__) || defined(__gfx900__) || defined(__gfx906__) +#define GCN +#endif + +#if defined(__gfx908__) || defined(__gfx90a__) || defined(__gfx942__) +#define CDNA +#endif + +#if defined(__gfx1100__) || defined(__gfx1101__) || defined(__gfx1102__) || defined(__gfx1103__) || \ + defined(__gfx1150__) || defined(__gfx1151__) +#define RDNA3 +#endif + +#if defined(__gfx1030__) || defined(__gfx1031__) || defined(__gfx1032__) || defined(__gfx1033__) || \ + defined(__gfx1034__) || defined(__gfx1035__) || defined(__gfx1036__) || defined(__gfx1037__) +#define RDNA2 +#endif + +#if defined(__gfx1010__) || defined(__gfx1012__) +#define RDNA1 +#endif + +#ifndef __has_builtin + #define __has_builtin(x) 0 +#endif + +typedef hip_bfloat16 nv_bfloat16; diff --git a/archive/csrc/ktransformers_ext/vendors/musa.h b/archive/csrc/ktransformers_ext/vendors/musa.h new file mode 100644 index 0000000..6cc1b69 --- /dev/null +++ b/archive/csrc/ktransformers_ext/vendors/musa.h @@ -0,0 +1,137 @@ +#pragma once + +#include +#include +#include +#include +#include +#define CUBLAS_COMPUTE_16F CUDA_R_16F +#define CUBLAS_COMPUTE_32F CUDA_R_32F +#define CUBLAS_COMPUTE_32F_FAST_16F MUBLAS_COMPUTE_32F_FAST_16F +#define CUBLAS_GEMM_DEFAULT MUBLAS_GEMM_DEFAULT +#define CUBLAS_GEMM_DEFAULT_TENSOR_OP MUBLAS_GEMM_DEFAULT +#define CUBLAS_OP_N MUBLAS_OP_N +#define CUBLAS_OP_T MUBLAS_OP_T +#define CUBLAS_STATUS_SUCCESS MUBLAS_STATUS_SUCCESS +#define CUBLAS_TF32_TENSOR_OP_MATH MUBLAS_MATH_MODE_DEFAULT +#define CUDA_R_16F MUSA_R_16F +#define CUDA_R_32F MUSA_R_32F +#define cublasComputeType_t cudaDataType_t +#define cublasCreate mublasCreate +#define cublasDestroy mublasDestroy +#define cublasGemmEx mublasGemmEx +#define cublasGemmBatchedEx mublasGemmBatchedEx +#define cublasGemmStridedBatchedEx mublasGemmStridedBatchedEx +#define cublasHandle_t mublasHandle_t +#define cublasSetMathMode mublasSetMathMode +#define cublasSetStream mublasSetStream +#define cublasSgemm mublasSgemm +#define cublasStatus_t mublasStatus_t +#define cublasOperation_t mublasOperation_t +#define cublasGetStatusString mublasStatus_to_string +#define cudaDataType_t musaDataType_t +#define cudaDeviceCanAccessPeer musaDeviceCanAccessPeer +#define cudaDeviceDisablePeerAccess musaDeviceDisablePeerAccess +#define cudaDeviceEnablePeerAccess musaDeviceEnablePeerAccess +#define cudaDeviceProp musaDeviceProp +#define cudaDeviceSynchronize musaDeviceSynchronize +#define cudaError_t musaError_t +#define cudaErrorPeerAccessAlreadyEnabled musaErrorPeerAccessAlreadyEnabled +#define cudaErrorPeerAccessNotEnabled musaErrorPeerAccessNotEnabled +#define cudaEventCreateWithFlags musaEventCreateWithFlags +#define cudaEventDisableTiming musaEventDisableTiming +#define cudaEventRecord musaEventRecord +#define cudaEventSynchronize musaEventSynchronize +#define cudaEvent_t musaEvent_t +#define cudaEventDestroy musaEventDestroy +#define cudaFree musaFree +#define cudaFreeHost musaFreeHost +#define cudaGetDevice musaGetDevice +#define cudaGetDeviceCount musaGetDeviceCount +#define cudaGetDeviceProperties musaGetDeviceProperties +#define cudaGetErrorString musaGetErrorString +#define cudaGetLastError musaGetLastError +#define cudaHostRegister musaHostRegister +#define cudaHostRegisterPortable musaHostRegisterPortable +#define cudaHostRegisterReadOnly musaHostRegisterReadOnly +#define cudaHostUnregister musaHostUnregister +#define cudaLaunchHostFunc musaLaunchHostFunc +#define cudaMalloc musaMalloc +#define cudaMallocHost musaMallocHost +#define cudaMallocManaged musaMallocManaged +#define cudaMemcpy musaMemcpy +#define cudaMemcpyAsync musaMemcpyAsync +#define cudaMemcpyPeerAsync musaMemcpyPeerAsync +#define cudaMemcpy2DAsync musaMemcpy2DAsync +#define cudaMemcpyDeviceToDevice musaMemcpyDeviceToDevice +#define cudaMemcpyDeviceToHost musaMemcpyDeviceToHost +#define cudaMemcpyHostToDevice musaMemcpyHostToDevice +#define cudaMemcpyKind musaMemcpyKind +#define cudaMemset musaMemset +#define cudaMemsetAsync musaMemsetAsync +#define cudaMemGetInfo musaMemGetInfo +#define cudaOccupancyMaxPotentialBlockSize musaOccupancyMaxPotentialBlockSize +#define cudaSetDevice musaSetDevice +#define cudaStreamCreateWithFlags musaStreamCreateWithFlags +#define cudaStreamDestroy musaStreamDestroy +#define cudaStreamFireAndForget musaStreamFireAndForget +#define cudaStreamNonBlocking musaStreamNonBlocking +#define cudaStreamPerThread musaStreamPerThread +#define cudaStreamSynchronize musaStreamSynchronize +#define cudaStreamWaitEvent musaStreamWaitEvent +#define cudaStream_t musaStream_t +#define cudaSuccess musaSuccess + +// Additional mappings for MUSA virtual memory pool +#define CU_DEVICE_ATTRIBUTE_VIRTUAL_MEMORY_MANAGEMENT_SUPPORTED MU_DEVICE_ATTRIBUTE_VIRTUAL_ADDRESS_MANAGEMENT_SUPPORTED +#define CU_MEM_ACCESS_FLAGS_PROT_READWRITE MU_MEM_ACCESS_FLAGS_PROT_READWRITE +#define CU_MEM_ALLOC_GRANULARITY_RECOMMENDED MU_MEM_ALLOC_GRANULARITY_RECOMMENDED +#define CU_MEM_ALLOCATION_TYPE_PINNED MU_MEM_ALLOCATION_TYPE_PINNED +#define CU_MEM_LOCATION_TYPE_DEVICE MU_MEM_LOCATION_TYPE_DEVICE +#define CUdevice MUdevice +#define CUdeviceptr MUdeviceptr +#define CUmemAccessDesc MUmemAccessDesc +#define CUmemAllocationProp MUmemAllocationProp +#define CUmemGenericAllocationHandle MUmemGenericAllocationHandle +#define cuDeviceGet muDeviceGet +#define cuDeviceGetAttribute muDeviceGetAttribute +#define cuMemAddressFree muMemAddressFree +#define cuMemAddressReserve muMemAddressReserve +#define cuMemCreate muMemCreate +#define cuMemGetAllocationGranularity muMemGetAllocationGranularity +#define cuMemMap muMemMap +#define cuMemRelease muMemRelease +#define cuMemSetAccess muMemSetAccess +#define cuMemUnmap muMemUnmap +#define cudaFuncAttributeMaxDynamicSharedMemorySize musaFuncAttributeMaxDynamicSharedMemorySize +#define cudaFuncSetAttribute musaFuncSetAttribute +#define cudaMemcpy3DPeerParms musaMemcpy3DPeerParms +#define make_cudaExtent make_musaExtent +#define make_cudaPitchedPtr make_musaPitchedPtr + +// Additional mappings for MUSA graphs +#define CUDA_SUCCESS MUSA_SUCCESS +#define CUresult MUresult +#define cuGetErrorString muGetErrorString +#define cudaErrorGraphExecUpdateFailure musaErrorGraphExecUpdateFailure +#define cudaErrorInvalidDeviceFunction musaErrorInvalidDeviceFunction +#define cudaGraphDestroy musaGraphDestroy +#define cudaGraphExecDestroy musaGraphExecDestroy +#define cudaGraphExec_t musaGraphExec_t +#define cudaGraphExecUpdate musaGraphExecUpdate +#define cudaGraphExecUpdateResultInfo musaGraphExecUpdateResult +#define cudaGraphGetNodes musaGraphGetNodes +#define cudaGraphInstantiate musaGraphInstantiate +#define cudaGraphKernelNodeGetParams musaGraphKernelNodeGetParams +#define cudaGraphKernelNodeSetParams musaGraphKernelNodeSetParams +#define cudaGraphLaunch musaGraphLaunch +#define cudaGraphNodeGetType musaGraphNodeGetType +#define cudaGraphNode_t musaGraphNode_t +#define cudaGraphNodeType musaGraphNodeType +#define cudaGraphNodeTypeKernel musaGraphNodeTypeKernel +#define cudaGraph_t musaGraph_t +#define cudaKernelNodeParams musaKernelNodeParams +#define cudaStreamCaptureModeRelaxed musaStreamCaptureModeRelaxed +#define cudaStreamEndCapture musaStreamEndCapture + +typedef mt_bfloat16 nv_bfloat16; diff --git a/archive/csrc/ktransformers_ext/vendors/vendor.h b/archive/csrc/ktransformers_ext/vendors/vendor.h new file mode 100644 index 0000000..8470438 --- /dev/null +++ b/archive/csrc/ktransformers_ext/vendors/vendor.h @@ -0,0 +1,13 @@ +#ifndef CPUINFER_VENDOR_VENDOR_H +#define CPUINFER_VENDOR_VENDOR_H + +#ifdef USE_CUDA +#include "cuda.h" +#elif USE_HIP +#define __HIP_PLATFORM_AMD__ +#include "hip.h" +#elif USE_MUSA +#include "musa.h" +#endif + +#endif // CPUINFER_VENDOR_VENDOR_H \ No newline at end of file diff --git a/archive/install-with-cache.sh b/archive/install-with-cache.sh new file mode 100755 index 0000000..cef4341 --- /dev/null +++ b/archive/install-with-cache.sh @@ -0,0 +1,26 @@ +#!/bin/bash +set -e + +# clear build dirs +# rm -rf build +# rm -rf *.egg-info +# rm -rf csrc/build +# rm -rf csrc/ktransformers_ext/build +# rm -rf csrc/ktransformers_ext/cuda/build +# rm -rf csrc/ktransformers_ext/cuda/dist +# rm -rf csrc/ktransformers_ext/cuda/*.egg-info +rm -rf ~/.ktransformers +echo "Installing python dependencies from requirements.txt" +pip install -r requirements-local_chat.txt +pip install -r ktransformers/server/requirements.txt +echo "Installing ktransformers" +KTRANSFORMERS_FORCE_BUILD=TRUE USE_BALANCE_SERVE=1 pip install -v . --no-build-isolation +pip install third_party/custom_flashinfer/ -v + +# SITE_PACKAGES=$(python -c "import site; print(site.getsitepackages()[0])") +# echo "Copying thirdparty libs to $SITE_PACKAGES" +# cp -a csrc/balance_serve/build/third_party/prometheus-cpp/lib/libprometheus-cpp-*.so* $SITE_PACKAGES/ +# patchelf --set-rpath '$ORIGIN' $SITE_PACKAGES/sched_ext.cpython* + + +echo "Installation completed successfully" diff --git a/archive/install.bat b/archive/install.bat new file mode 100644 index 0000000..dc429e4 --- /dev/null +++ b/archive/install.bat @@ -0,0 +1,16 @@ +@echo off + +REM clear build dirs +rmdir /S /Q ktransformers\ktransformers_ext\build +rmdir /S /Q ktransformers\ktransformers_ext\cuda\build +rmdir /S /Q ktransformers\ktransformers_ext\cuda\dist +rmdir /S /Q ktransformers\ktransformers_ext\out +del /F /Q ktransformers\ktransformers_ext\cuda\*.egg-info + +echo Installing python dependencies from requirements.txt +pip install -r requirements-local_chat.txt + +echo Installing ktransformers +set KTRANSFORMERS_FORCE_BUILD=TRUE +pip install . --no-build-isolation +echo Installation completed successfully \ No newline at end of file diff --git a/archive/install.sh b/archive/install.sh new file mode 100644 index 0000000..183c9c8 --- /dev/null +++ b/archive/install.sh @@ -0,0 +1,43 @@ +#!/bin/bash +set -e + +# default backend +DEV="cuda" + +# parse --dev argument +while [[ "$#" -gt 0 ]]; do + case $1 in + --dev) DEV="$2"; shift ;; + *) echo "Unknown parameter passed: $1"; exit 1 ;; + esac + shift +done +export DEV_BACKEND="$DEV" +echo "Selected backend: $DEV_BACKEND" + +# clear build dirs +rm -rf build +rm -rf *.egg-info +rm -rf csrc/build +rm -rf csrc/ktransformers_ext/build +rm -rf csrc/ktransformers_ext/cuda/build +rm -rf csrc/ktransformers_ext/cuda/dist +rm -rf csrc/ktransformers_ext/cuda/*.egg-info +rm -rf ~/.ktransformers +echo "Installing python dependencies from requirements.txt" +pip install -r requirements-local_chat.txt +pip install -r ktransformers/server/requirements.txt + +echo "Installing ktransformers" +KTRANSFORMERS_FORCE_BUILD=TRUE pip install -v . --no-build-isolation + +if [[ "$DEV_BACKEND" == "cuda" ]]; then + echo "Installing custom_flashinfer for CUDA backend" + pip install third_party/custom_flashinfer/ +fi +# SITE_PACKAGES=$(python -c "import site; print(site.getsitepackages()[0])") +# echo "Copying thirdparty libs to $SITE_PACKAGES" +# cp -a csrc/balance_serve/build/third_party/prometheus-cpp/lib/libprometheus-cpp-*.so* $SITE_PACKAGES/ +# patchelf --set-rpath '$ORIGIN' $SITE_PACKAGES/sched_ext.cpython* + +echo "Installation completed successfully" diff --git a/archive/kt-sft/.flake8 b/archive/kt-sft/.flake8 new file mode 100644 index 0000000..fadb988 --- /dev/null +++ b/archive/kt-sft/.flake8 @@ -0,0 +1,4 @@ +[flake8] +max-line-length = 120 +extend-select = B950 +extend-ignore = E203,E501,E701, B001,B006,B007,B008,B009,B010,B011,B016,B028,B031,B950,E265,E266,E401,E402,E711,E712,E713,E721,E722,E731,F401,F403,F405,F541,F811,F821,F841,W391 \ No newline at end of file diff --git a/archive/kt-sft/.gitignore b/archive/kt-sft/.gitignore new file mode 100644 index 0000000..d47717c --- /dev/null +++ b/archive/kt-sft/.gitignore @@ -0,0 +1,89 @@ +__pycache__ +build +.vscode +*.so +*.cache +server.db +logs +node_modules +*.nsys-rep +.vs/ +*pycache* +*build/ +*/third_party/* +.DS_Store +compile_commands.json +*.egg-info* +*dist/ +ktransformers/server/local_store/ +ktransformers/server_test1.db +*.patch +img/ +tmp*.txt +tmp*.py +test.txt +book +ktransformers/tests/chat_txt.txt +mmlu_result* +ktransformers/ktransformers_ext/cuda_musa/ +test_prompt.txt +csrc/demo + +.vscode/ + +*__pycache__* +*.py[cod] +*$py.class +.pytest_cache/ + +GGUF-DeepSeek-V2-Lite-Chat +DeepSeek-V2-Lite-Chat +ktransformers/sft/adapter +tmp +graphviz/ +compute_graph* +graphviz* +third_party/ +test_adapter/demo_* +*.whl +*.svg +*_graph +tmp_package.txt +logs/ + +*.vscode/ + +__pycache__/ +*.py[cod] +*$py.class +.pytest_cache/ + +# MakeFiles for kt_ext +build/ +ktransformers/ktransformers_ext/bin +ktransformers/ktransformers_ext/CMakeFiles +ktransformers/ktransformers_ext/cmake_install.cmake +ktransformers/ktransformers_ext/CMakeCache.txt +ktransformers/ktransformers_ext/compile_commands.json +ktransformers/ktransformers_ext/Makefile +*.egg-info* +*.so + +*.txt +*.pt + +debug/* + +test_adapter/ESC_inst_all.json + +!CMakeLists.txt +!requirements-sft.txt + +*-test*.yaml + +duipai_pure_tf +data/dataset_info.json + +test_adapter/*.json + +.venv* \ No newline at end of file diff --git a/archive/kt-sft/.gitmodules b/archive/kt-sft/.gitmodules new file mode 100644 index 0000000..1864d1f --- /dev/null +++ b/archive/kt-sft/.gitmodules @@ -0,0 +1,19 @@ +[submodule "third_party/llama.cpp"] + path = third_party/llama.cpp + url = https://github.com/ggerganov/llama.cpp.git +[submodule "third_party/pybind11"] + path = third_party/pybind11 + url = https://github.com/pybind/pybind11.git +[submodule "third_party/spdlog"] + path = third_party/spdlog + url = https://github.com/gabime/spdlog.git +[submodule "third_party/custom_flashinfer"] + path = third_party/custom_flashinfer + url = https://github.com/kvcache-ai/custom_flashinfer.git + branch = fix-precision-mla-merge-main +[submodule "third_party/xxHash"] + path = third_party/xxHash + url = https://github.com/Cyan4973/xxHash.git +[submodule "third_party/prometheus-cpp"] + path = third_party/prometheus-cpp + url = https://github.com/jupp0r/prometheus-cpp.git diff --git a/archive/kt-sft/.pylintrc b/archive/kt-sft/.pylintrc new file mode 100644 index 0000000..1ed4ddb --- /dev/null +++ b/archive/kt-sft/.pylintrc @@ -0,0 +1,6 @@ +[MASTER] +extension-pkg-whitelist=pydantic +max-line-length=120 + +[MESSAGES CONTROL] +disable=missing-function-docstring \ No newline at end of file diff --git a/archive/kt-sft/Dockerfile b/archive/kt-sft/Dockerfile new file mode 100644 index 0000000..22f3f68 --- /dev/null +++ b/archive/kt-sft/Dockerfile @@ -0,0 +1,64 @@ +FROM pytorch/pytorch:2.5.1-cuda12.1-cudnn9-devel as compile_server + + +ARG CPU_INSTRUCT=NATIVE + +# 设置工作目录和 CUDA 路径 +WORKDIR /workspace +ENV CUDA_HOME=/usr/local/cuda + + + +# 安装依赖 +RUN apt update -y +RUN apt install -y --no-install-recommends \ + libtbb-dev \ + libssl-dev \ + libcurl4-openssl-dev \ + libaio1 \ + libaio-dev \ + libfmt-dev \ + libgflags-dev \ + zlib1g-dev \ + patchelf \ + git \ + wget \ + vim \ + gcc \ + g++ \ + cmake +# 拷贝代码 +RUN git clone https://github.com/kvcache-ai/ktransformers.git +# 清理 apt 缓存 +RUN rm -rf /var/lib/apt/lists/* + +# 进入项目目录 +WORKDIR /workspace/ktransformers +# 初始化子模块 +RUN git submodule update --init --recursive + +# 升级 pip +RUN pip install --upgrade pip + +# 安装构建依赖 +RUN pip install ninja pyproject numpy cpufeature aiohttp zmq openai + +# 安装 flash-attn（提前装可以避免后续某些编译依赖出错） +RUN pip install flash-attn + +# 安装 ktransformers 本体（含编译） +RUN CPU_INSTRUCT=${CPU_INSTRUCT} \ + USE_BALANCE_SERVE=1 \ + KTRANSFORMERS_FORCE_BUILD=TRUE \ + TORCH_CUDA_ARCH_LIST="8.0;8.6;8.7;8.9;9.0+PTX" \ + pip install . --no-build-isolation --verbose + +RUN pip install third_party/custom_flashinfer/ +# 清理 pip 缓存 +RUN pip cache purge + +# 拷贝 C++ 运行时库 +RUN cp /usr/lib/x86_64-linux-gnu/libstdc++.so.6 /opt/conda/lib/ + +# 保持容器运行（调试用） +ENTRYPOINT ["tail", "-f", "/dev/null"] \ No newline at end of file diff --git a/archive/kt-sft/Dockerfile.xpu b/archive/kt-sft/Dockerfile.xpu new file mode 100644 index 0000000..bb4d2dd --- /dev/null +++ b/archive/kt-sft/Dockerfile.xpu @@ -0,0 +1,68 @@ +# Base image +FROM intel/oneapi-basekit:2025.0.1-0-devel-ubuntu22.04 + +ARG http_proxy +ARG https_proxy + +ENV DEBIAN_FRONTEND=noninteractive +ENV CONDA_DIR=/opt/conda + +# Install dependencies +RUN apt-get update && apt-get install -y \ + wget \ + curl \ + bash \ + git \ + vim \ + ca-certificates \ + binutils \ + cmake \ + g++ \ + && rm -rf /var/lib/apt/lists/* + +# Install Miniforge +RUN wget https://github.com/conda-forge/miniforge/releases/latest/download/Miniforge3-Linux-x86_64.sh -O /tmp/miniforge.sh && \ + bash /tmp/miniforge.sh -b -p $CONDA_DIR && \ + rm /tmp/miniforge.sh && \ + $CONDA_DIR/bin/conda clean -afy + +# Add conda to PATH +ENV PATH=$CONDA_DIR/bin:$PATH + +RUN bash -c "\ + source /opt/conda/etc/profile.d/conda.sh && \ + conda create --name ktransformers python=3.11 -y && \ + conda activate ktransformers && \ + conda env list && \ + conda install -c conda-forge libstdcxx-ng -y && \ + strings \$(find /opt/conda/envs/ktransformers/lib -name 'libstdc++.so.6') | grep GLIBCXX | grep 3.4.32 \ +" + +RUN bash -c "\ + source /opt/conda/etc/profile.d/conda.sh && \ + conda activate ktransformers && \ + pip install ipex-llm[xpu_2.6]==2.3.0b20250518 --extra-index-url https://download.pytorch.org/whl/xpu && \ + pip uninstall -y torch torchvision torchaudio && \ + pip install torch==2.7+xpu torchvision torchaudio --index-url https://download.pytorch.org/whl/test/xpu && \ + pip uninstall -y intel-opencl-rt dpcpp-cpp-rt && \ + pip list \ +" + +# Clone and set up ktransformers repo +RUN bash -c "\ + source $CONDA_DIR/etc/profile.d/conda.sh && \ + conda activate ktransformers && \ + git clone https://github.com/kvcache-ai/ktransformers.git && \ + cd ktransformers && \ + git submodule update --init && \ + sed -i 's/torch\.xpu\.is_available()/True/g' setup.py && \ + bash install.sh --dev xpu \ +" + +# Init conda and prepare bashrc +RUN conda init bash && \ + echo "source $CONDA_DIR/etc/profile.d/conda.sh" >> ~/.bashrc && \ + echo "conda activate ktransformers" >> ~/.bashrc + +WORKDIR /ktransformers/ +CMD ["bash"] diff --git a/archive/kt-sft/LICENSE b/archive/kt-sft/LICENSE new file mode 100644 index 0000000..261eeb9 --- /dev/null +++ b/archive/kt-sft/LICENSE @@ -0,0 +1,201 @@ + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright [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/archive/kt-sft/MANIFEST.in b/archive/kt-sft/MANIFEST.in new file mode 100644 index 0000000..4097ce6 --- /dev/null +++ b/archive/kt-sft/MANIFEST.in @@ -0,0 +1,13 @@ +graft third_party +graft ktransformers +graft local_chat.py +graft csrc +include LICENSE README.md +prune ktransformers/website +prune ktransformers/logs +prune ktransformers.egg-info +prune third_party/llama.cpp/models +graft ktransformers/website/dist +global-exclude __pycache__ +include KTransformersOps.*.so +include cpuinfer_ext.*.so diff --git a/archive/kt-sft/Makefile b/archive/kt-sft/Makefile new file mode 100644 index 0000000..74cb3c9 --- /dev/null +++ b/archive/kt-sft/Makefile @@ -0,0 +1,32 @@ +flake_find: + cd ktransformers && flake8 | grep -Eo '[A-Z][0-9]{3}' | sort | uniq| paste -sd ',' - +format: + @cd ktransformers && black . + @black setup.py +dev_install: +# clear build dirs + rm -rf build + rm -rf *.egg-info + rm -rf ktransformers/ktransformers_ext/build + rm -rf ktransformers/ktransformers_ext/cuda/build + rm -rf ktransformers/ktransformers_ext/cuda/dist + rm -rf ktransformers/ktransformers_ext/cuda/*.egg-info + +# install ktransformers + echo "Installing python dependencies from requirements.txt" + pip install -r requirements-local_chat.txt + + echo "Installing ktransformers" + KTRANSFORMERS_FORCE_BUILD=TRUE pip install -e . -v --no-build-isolation + echo "Installation completed successfully" +clean: + rm -rf build + rm -rf *.egg-info + rm -rf ktransformers/ktransformers_ext/build + rm -rf ktransformers/ktransformers_ext/cuda/build + rm -rf ktransformers/ktransformers_ext/cuda/dist + rm -rf ktransformers/ktransformers_ext/cuda/*.egg-info +install_numa: + USE_NUMA=1 make dev_install +install_no_numa: + env -u USE_NUMA make dev_install \ No newline at end of file diff --git a/archive/kt-sft/README.md b/archive/kt-sft/README.md new file mode 100644 index 0000000..e94e84f --- /dev/null +++ b/archive/kt-sft/README.md @@ -0,0 +1,297 @@ +- [KTransformers Fine-Tuning × LLaMA-Factory Integration – User Guide](#ktransformers-fine-tuning-x-llama-factory-integration-–-user-guide) +- [Introduction](#introduction) + +- [Fine-Tuning Results (Examples)](#fine-tuning-results-examples) + - [Stylized Dialogue (CatGirl tone)](#stylized-dialogue-catgirl-tone) + - [Benchmarks](#benchmarks) + - [Translational-Style dataset](#translational-style-dataset) + - [AfriMed-QA (short answer)](#afrimed-qa-short-answer) + - [AfriMed-QA (multiple choice)](#afrimed-qa-multiple-choice) + +- [Quick to Start](#quick-to-start) + - [Environment Setup](#environment-setup) + - [Core Feature 1: Use KTransformers backend to fine-tune ultra-large MoE models](#core-feature-1-use-ktransformers-backend-to-fine-tune-ultra-large-moe-models) + - [Core Feature 2: Chat with the fine-tuned model (base + LoRA adapter)](#core-feature-2-chat-with-the-fine-tuned-model-base--lora-adapter) + - [Core Feature 3: Batch inference + metrics (base + LoRA adapter)](#core-feature-3-batch-inference--metrics-base--lora-adapter) + +- [KT Fine-Tuning Speed (User-Side View)](#kt-fine-tuning-speed-user-side-view) + - [End-to-End Performance](#end-to-end-performance) + - [GPU/CPU Memory Footprint](#gpucpu-memory-footprint) + +- [Conclusion](#conclusion) + + +# KTransformers Fine-Tuning × LLaMA-Factory Integration – User Guide + +**MadSys Lab, KVCache-AI Team, Approaching AI, LLaMA-Factory Team** + +## Introduction + +From **DeepSeek-V3/R1** to **Qwen3-MoE** and **Kimi-K2**, each wave of open-sourced large models brings leaps in performance and scale. However, many researchers and developers are constrained by expensive GPUs and models with tens or even hundreds of billions of parameters, making it **hard to fine-tune very large models under limited resources**. To bridge this gap, we propose a practical approach: combining **KTransformers** with **LLaMA-Factory**. With just **2–4 RTX 4090s** and a high-memory CPU, you can fine-tune ultra-large MoE models like DeepSeek-671B. + +Our goal is to give resource-constrained researchers a **local path to explore fine-tuning ultra-large models**, and also a fast way to customize smaller models (e.g., 14B/30B) for specific scenarios. We validate the setup using **stylized dialogue**, **Westernized translation tone**, and **medical Q&A** as representative tasks, showing that **personalized adaptation can be achieved within hours**. + +As shown below, LLaMA-Factory is the unified orchestration/configuration layer for the whole fine-tuning workflow—handling data, training scheduling, LoRA injection, and inference interfaces. **KTransformers** acts as a pluggable high-performance backend that takes over core operators like Attention/MoE under the same training configs, enabling efficient **GPU+CPU heterogeneous cooperation**. + +![image-20251011010558909](../doc/assets/image-20251011010558909.png) + +Within LLaMA-Factory, we compared LoRA fine-tuning with **HuggingFace**, **Unsloth**, and **KTransformers** backends. KTransformers is the **only workable 4090-class solution** for ultra-large MoE models (e.g., 671B) and also delivers higher throughput and lower GPU memory on smaller MoE models (e.g., DeepSeek-14B). + +| Under LoRA (BF16) + [NekoQA-10K stylized dialogue](https://github.com/mindsRiverPonder/LLM-practice) | HuggingFace Backend | Unsloth Backend | KTransformers Backend | +| ------------------------------------------------------------ | ---------------------------------------- | ------------------------------------ | --------------------- | +| [14B-DeepSeekV2-Lite] LoRA fine-tuning throughput | 303.58 token/s | 455.37 token/s | 530.38 token/s | +| [14B-DeepSeekV2-Lite] GPU memory | 32.12 GB | 9.64 GB | 6.08 GB | +| [671B-DeepSeekV3] LoRA fine-tuning throughput | Too Huge to run | NOT SUPPORT | 40.35 token/s | +| [671B-DeepSeekV3] GPU memory (sum across GPUs) | theoretical 1400 GB † | NOT SUPPORT | 70 GB † | + +† **1400 GB** is a **theoretical** FP16 full-parameter resident footprint (not runnable). **70 GB** is the **measured peak** with KT strategy (Attention on GPU + layered MoE offload). + +![按照模型划分的对比图_02](../doc/assets/image-compare_model.png) + +### Fine-Tuning Results (Examples) + +#### Stylized Dialogue (CatGirl tone) + +Dataset: [NekoQA-10K](https://zhuanlan.zhihu.com/p/1934983798233231689). Goal: improve style consistency and recognizability. + +The figure compares responses from the base vs. fine-tuned models. The fine-tuned model maintains the target tone and address terms more consistently (red boxes), validating the effectiveness of **style-transfer fine-tuning**. + +![image-20251016175046882](../doc/assets/image-20251016175046882.png) + +#### Benchmarks + +We use: + +(1) [Translational-Style-ChatLLM](https://github.com/Benson114/Translational-Style-ChatLLM), which asks for an exaggerated, Westernized translation tone—clear, stylized customization. + +(2) [AfriMed-QA](https://aclanthology.org/2025.acl-long.96/) (ACL 2025), a medical dataset for African contexts with strong domain specificity, including multiple-choice and short-answer sub-tasks—well-suited for vertical fine-tuning evaluation. + +The tables show metrics before vs. after LoRA fine-tuning. We observe **large improvements** across metrics, verifying fine-tuning effectiveness: + +| Translational-Style dataset | BLEU-1 | BLEU-2 | BLEU-3 | BLEU-4 | ROUGE-1 | ROUGE-2 | ROUGE-L | +| ------------------------------ | --------- | --------- | --------- | --------- | --------- | --------- | --------- | +| V2-Lite (no LoRA) | 20.66 | 8.33 | 4.54 | 2.89 | 22.71 | 4.52 | 19.19 | +| **KT-LoRA fine-tuned V2-Lite** | **35.41** | **22.44** | **15.42** | **11.18** | **42.03** | **18.38** | **33.10** | +| V3 base (no LoRA) | 8.49 | 3.34 | 1.62 | 0.96 | 15.91 | 2.55 | 10.07 | +| **KT-LoRA fine-tuned V3** | **37.02** | **23.70** | **16.21** | **11.49** | **43.43** | **18.96** | **34.54** | + +| AfriMed-QA (short answer) | BLEU-1 | BLEU-2 | BLEU-3 | BLEU-4 | ROUGE-1 | ROUGE-2 | ROUGE-L | +| ------------------------------ | --------- | --------- | --------- | --------- | --------- | --------- | --------- | +| V2-Lite (no LoRA) | 13.58 | 11.12 | 9.10 | 7.23 | 22.48 | 7.81 | 11.73 | +| **KT-LoRA fine-tuned V2-Lite** | **35.90** | **27.63** | **22.99** | **19.15** | **35.25** | **17.50** | **28.44** | +| V3 base (no LoRA) | 12.75 | 10.27 | 8.05 | 5.99 | 20.33 | 5.65 | 10.11 | +| **KT-LoRA fine-tuned V3** | **42.42** | **34.12** | **28.95** | **24.54** | **41.97** | **22.37** | **33.28** | + +| AfriMed-QA (multiple choice) | Accuracy | +| ------------------------------ | ---------- | +| V2-Lite (no LoRA) | 0.0645 | +| **KT-LoRA fine-tuned V2-Lite** | **0.4812** | +| V3 base (no LoRA) | 0.5833 | +| **KT-LoRA fine-tuned V3** | **0.7930** | + +Even for ultra-large MoE models, **KTransformers-backed fine-tuning** achieves strong task performance quickly. + + + +## Quick to Start + +This section shows how to install and use **LLaMA-Factory + KTransformers** for fine-tuning and inference: + +- Environment setup +- Fine-tune ultra-large MoE models with KTransformers backend +- Load the fine-tuned model (base + LoRA adapter) for chat/inference +- Batch inference and metric evaluation + +### Environment Setup + +According to the following example, install both the **KTransformers** and **LLaMA-Factory** environments simultaneously. + This time, to simplify the installation process of KTransformers, we have specially packaged a wheel file to avoid local compilation. + The detailed installation steps are as follows: + (Note: Make sure your local **Python version**, **Torch version**, **CUDA version**, and the **KTransformers wheel filename** correspond correctly.) + +```shell +# 1. Create a conda environment +conda create -n Kllama python=3.12 # choose from : [3.10, 3.11, 3.12, 3.13] +conda install -y -c conda-forge libstdcxx-ng gcc_impl_linux-64 +# ATTENTION: DO NOT skip this step, even if your cuda version is not 11.8! Otherwise, you will get this error: ImportError: libcudart.so.11.0: cannot open shared object file: No such file or directory. +conda install -y -c nvidia/label/cuda-11.8.0 cuda-runtime + +# 2. Install the LLaMA-Factory environment +git clone --depth 1 https://github.com/hiyouga/LLaMA-Factory.git +cd LLaMA-Factory +pip install -e ".[torch,metrics]" --no-build-isolation + +# 3. Install the KTransformers wheel that matches your Torch and Python versions, from https://github.com/kvcache-ai/ktransformers/releases/tag/v0.4.1 (Note: The CUDA version can differ from that in the wheel filename.) +pip install ktransformers-0.4.1+cu128torch27fancy-cp312-cp312-linux_x86_64.whl + +# 4. Install flash-attention, download the corresponding file based on your Python and Torch versions from: https://github.com/Dao-AILab/flash-attention/releases +pip install flash_attn-2.8.3+cu12torch2.7cxx11abiTRUE-cp312-cp312-linux_x86_64.whl +# abi=True/False can find from below +# import torch +# print(torch._C._GLIBCXX_USE_CXX11_ABI) + +# 5. (Optional) If you want to use flash_infer (otherwise it defaults to triton) +git clone https://github.com/kvcache-ai/custom_flashinfer.git +pip install custom_flashinfer/ +``` + +**Usage tip:** In LLaMA-Factory YAML, set `use_kt: true` and pick a `kt_optimize_rule` file to have KTransformers handle the core compute. The features below show typical configs. + +### Core Feature 1: Use KTransformers backend to fine-tune ultra-large MoE models + +Run the command: `USE_KT=1 llamafactory-cli train examples/train_lora/deepseek3_lora_sft_kt.yaml`. + +Note: You **must** provide a **BF16** model. DeepSeek-V3-671B is released in FP8 by default; convert with [DeepSeek-V3/inference/fp8_cast_bf16.py](https://github.com/deepseek-ai/DeepSeek-V3/blob/main/inference/fp8_cast_bf16.py). + +```yaml +### model +model_name_or_path: opensourcerelease/DeepSeek-V3-bf16 +trust_remote_code: true + +### method +stage: sft +do_train: true +finetuning_type: lora +lora_rank: 8 +lora_target: all + +### dataset +dataset: identity +template: deepseek +cutoff_len: 2048 +max_samples: 100000 +overwrite_cache: true +preprocessing_num_workers: 16 +dataloader_num_workers: 4 + +### output +output_dir: saves/Kllama_deepseekV3 +logging_steps: 10 +save_steps: 500 +plot_loss: true +overwrite_output_dir: true +save_only_model: false +report_to: none # choices: [none, wandb, tensorboard, swanlab, mlflow] + +### train +per_device_train_batch_size: 1 +gradient_accumulation_steps: 8 +learning_rate: 1.0e-4 +num_train_epochs: 3.0 +lr_scheduler_type: cosine +warmup_ratio: 0.1 +bf16: true +ddp_timeout: 180000000 +resume_from_checkpoint: null + +### ktransformers +use_kt: true # use KTransformers as LoRA sft backend +kt_optimize_rule: examples/kt_optimize_rules/DeepSeek-V3-Chat-sft-amx-multi-gpu.yaml +cpu_infer: 32 +chunk_size: 8192 +``` + +We also support RL DPO training using the KTransformers backend now. See [DPO Tutorial](../doc/en/SFT/DPO_tutorial.md) for details. + +`kt_optimize_rule` controls **placement strategy**. See also [ktransformers/optimize_rules](https://github.com/kvcache-ai/ktransformers/tree/main/ktransformers/optimize/optimize_rules). Naming hints (`*` = wildcard): + +| Pattern | Meaning | +| -------------------------------------------- | ----------------------------------------------------- | +| DeepSeek-V2-Lite-Chat-* / DeepSeek-V3-Chat-* | Target model variants | +| *-sft-* | Strategy for fine-tuning; others are for inference | +| *-amx-* | Use AMX on CPU; otherwise use **llamafile** | +| *-multi-gpu-X* | Model parallel on X GPUs (X omitted → default 2 GPUs) | + +Example: `DeepSeek-V3-Chat-sft-amx-multi-gpu.yaml` = V3-Chat fine-tuning with AMX and 2-GPU model parallel. + +We recommend **AMX acceleration** where available (`lscpu | grep amx`). AMX supports BF16/INT8. Example: + +```yaml +- match: + name: "^model\\.layers\\..*\\.mlp\\.experts$" + replace: + class: ktransformers.operators.experts.KTransformersExperts # custom MoE Kernel with expert parallelism + kwargs: + prefill_device: "cpu" + prefill_op: "KExpertsTorch" + generate_device: "cpu" + generate_op: "KSFTExpertsCPU" + out_device: "cuda" + backend: "AMXInt8" # or "AMXBF16" or "llamafile" (default) +``` + +Outputs go to `output_dir` in safetensors format plus adapter metadata for later loading. + +![image-20251016171537997](../doc/assets/image-20251016171537997.png) + +### Core Feature 2: Chat with the fine-tuned model (base + LoRA adapter) + +Run the command: `llamafactory-cli chat examples/inference/deepseek3_lora_sft_kt.yaml`. + +Use the safetensors adapter trained with KT for inference. + +```yaml +model_name_or_path: opensourcerelease/DeepSeek-V3-bf16 +adapter_name_or_path: saves/Kllama_deepseekV3 +template: deepseek +infer_backend: ktransformers # choices: [huggingface, vllm, sglang, ktransformers] +trust_remote_code: true + +use_kt: true # use KTransformers as LoRA sft backend to inference +kt_optimize_rule: examples/kt_optimize_rules/DeepSeek-V3-Chat-sft-amx-multi-gpu.yaml +cpu_infer: 32 +chunk_size: 8192 +``` + +We also support **GGUF** adapters: for safetensors, set the **directory**; for GGUF, set the **file path** in `adapter_name_or_path`. + +During loading, LLaMA-Factory maps layer names to KT’s naming. You’ll see logs like `Loaded adapter weight: XXX -> XXX`: + +![image-20251016171526210](../doc/assets/image-20251016171526210.png) + +### Core Feature 3: Batch inference + metrics (base + LoRA adapter) + +Run the command: `API_PORT=8000 llamafactory-cli api examples/inference/deepseek3_lora_sft_kt.yaml`. + Invoke the KT fine-tuned adapter to provide the API; the usage logic of other APIs is consistent with the native LLaMA-Factory approach. + +```yaml +model_name_or_path: opensourcerelease/DeepSeek-V3-bf16 +adapter_name_or_path: saves/Kllama_deepseekV3 +template: deepseek +infer_backend: ktransformers # choices: [huggingface, vllm, sglang, ktransformers] +trust_remote_code: true + +use_kt: true # use KTransformers as LoRA sft backend to inference +kt_optimize_rule: examples/kt_optimize_rules/DeepSeek-V3-Chat-sft-amx-multi-gpu.yaml +cpu_infer: 32 +chunk_size: 8192 +``` + + + +## KT Fine-Tuning Speed (User-Side View) + +### End-to-End Performance + +**Definitions** + +- `step_time`: wall-clock time for a full optimization step (tensor movement + Attention + MoE + other compute). +- `tokens_per_step = GAS × qlen`; `token/s = tokens_per_step / step_time`. + +**Settings:** `GAS=16`, `qlen=512` (→ `tokens_per_step = 8192`); LoRA (`r=8, alpha=32, dropout=0.1`); **AMX** enabled; GPU: RTX 4090, CPU: Intel Xeon Platinum 8488C. + +**Measured** + +- **DeepSeek-V3-671B:** `step_time = 203 s` → `token/s ≈ 8192 / 203 ≈ 40.35` +- **DeepSeek-V2-Lite-14B:** `step_time = 36 s` → `token/s ≈ 8192 / 36 ≈ 227.6` + +### GPU/CPU Memory Footprint + +- DeepSeek-V3 (671B; 61 layers with 58 MoE): ~**70 GB** total GPU VRAM (multi-GPU), ~**1.2–1.3 TB** RAM. +- DeepSeek-V2-Lite (14B; 27 layers with 26 MoE): ~**5.5 GB** GPU VRAM, ~**30 GB** RAM. + +## Conclusion + +By integrating **KTransformers LoRA fine-tuning** into **LLaMA-Factory**, we provide a practical guide for efficient training and deployment of MoE LLMs. KT brings cutting-edge optimizations (DeepSeek/Qwen/Kimi support with AMX-accelerated kernels), and LoRA enables customization under very low GPU memory. LLaMA-Factory offers a friendly, unified interface. + +This integration (akin to Unsloth-style speedups) means even models with tens to hundreds of billions of parameters can be fine-tuned and deployed with low latency on commodity hardware. You get **memory savings, speed-ups, and usability** together. We encourage you to try LLaMA-Factory + KT for your next MoE project and follow this guide. Feedback is welcome! diff --git a/archive/kt-sft/SECURITY.md b/archive/kt-sft/SECURITY.md new file mode 100644 index 0000000..034e848 --- /dev/null +++ b/archive/kt-sft/SECURITY.md @@ -0,0 +1,21 @@ +# Security Policy + +## Supported Versions + +Use this section to tell people about which versions of your project are +currently being supported with security updates. + +| Version | Supported | +| ------- | ------------------ | +| 5.1.x | :white_check_mark: | +| 5.0.x | :x: | +| 4.0.x | :white_check_mark: | +| < 4.0 | :x: | + +## Reporting a Vulnerability + +Use this section to tell people how to report a vulnerability. + +Tell them where to go, how often they can expect to get an update on a +reported vulnerability, what to expect if the vulnerability is accepted or +declined, etc. diff --git a/archive/kt-sft/autosetup.sh b/archive/kt-sft/autosetup.sh new file mode 100644 index 0000000..4f93a4f --- /dev/null +++ b/archive/kt-sft/autosetup.sh @@ -0,0 +1,110 @@ +#!/usr/bin/env bash +set -euo pipefail +shopt -s nullglob + +# 允许通过环境变量覆盖 +PY_LIST=${PY_LIST:-"3.13"} +TORCH_LIST=${TORCH_LIST:-"2.5.0 2.6.0 2.7.0 2.8.0 2.9.0"} +WHEELS_DIR=${WHEELS_DIR:-wheels} +FORCE=${FORCE:-0} # FORCE=1 时强制重建 +mkdir -p "$WHEELS_DIR" + +# 每个 Torch 版本选择一个存在的 CUDA 索引（可按需调整） +index_for_torch_version () { + case "$1" in + 2.3.*) echo "https://download.pytorch.org/whl/cu121" ;; + 2.4.*) echo "https://download.pytorch.org/whl/cu121" ;; + 2.5.*) echo "https://download.pytorch.org/whl/cu124" ;; + 2.6.*) echo "https://download.pytorch.org/whl/cu126" ;; + 2.7.*) echo "https://download.pytorch.org/whl/cu128" ;; + 2.8.*) echo "https://download.pytorch.org/whl/cu128" ;; # 可换 cu129 + 2.9.*) echo "https://download.pytorch.org/whl/cu128" ;; # 可换 cu129 + *) echo "https://download.pytorch.org/whl/cu121" ;; + esac +} + +# 检查指定“当前已激活环境”的组合是否已有产物 +# 依据 wheel 命名规则中的后缀：+torch 以及 ---linux_ +have_wheel_for_current_env () { + python - <<'PY' +import sys, platform, torch +from packaging.version import parse +py = f"cp{sys.version_info.major}{sys.version_info.minor}" +arch = platform.uname().machine +tver = parse(torch.__version__) +mm = f"{tver.major}{tver.minor}" +backend = "" +if torch.version.cuda: + backend = "cu" + torch.version.cuda.replace(".", "") +elif getattr(torch.version, "hip", None): + backend = "rocm" + torch.version.hip.replace(".", "") +else: + backend = "cpu" # 极少走到这里 +print(py, arch, backend, mm) +PY +} + +for py in $PY_LIST; do + PYBIN="$(command -v python${py} || true)" + if [[ ! -x "$PYBIN" ]]; then + echo ">> Skip python ${py}: not found" + continue + fi + for tv in $TORCH_LIST; do + echo "======== Build: Python ${py} × Torch ${tv} ========" + + # 1) 新建并激活 venv + ENV_DIR=".venv-py${py//./}-torch${tv%%.*}${tv#*.}" + "$PYBIN" -m venv "$ENV_DIR" + source "$ENV_DIR/bin/activate" + + # 2) 安装构建依赖 + 目标 torch（固定 CUDA 索引以避免装到 CPU 轮子） + python -m pip install -U pip + python -m pip install setuptools wheel build ninja cmake packaging cpufeature + IDX="$(index_for_torch_version "$tv")" + python -m pip install --index-url "$IDX" "torch==$tv" + + # 3) 读取当前环境的关键信息，拼出匹配的 wheel 通配符并检查是否已存在 + read -r CP_TAG ARCH BACKEND MM <<<"$(have_wheel_for_current_env)" + plat="linux_${ARCH}" + pattern="${WHEELS_DIR}/ktransformers-*+${BACKEND}torch${MM}*-${CP_TAG}-${CP_TAG}-${plat}.whl" + + if [[ "$FORCE" = "0" ]]; then + existing=( $pattern ) + if (( ${#existing[@]} > 0 )); then + echo ">> Found existing wheel, skip: ${existing[0]}" + deactivate + continue + fi + else + echo ">> FORCE=1, rebuild even if wheel exists" + fi + + # 打印对齐信息 + python - <<'PY' +import torch, sys +print(">>> torch:", torch.__version__, "cuda:", torch.version.cuda, + "cxx11abi:", torch.compiled_with_cxx11_abi()) +print(">>> python:", sys.version) +PY + + # ★ 清理所有构建产物（含内嵌 CMake build） + rm -rf build/ dist/ *.egg-info + find csrc -type d -name build -prune -exec rm -rf {} + + + # 构建 + KTRANSFORMERS_FORCE_BUILD=TRUE KTRANSFORMERS_DISABLE_PREBUILT=1 \ + python -m build --no-isolation --wheel + + # ★ 验证 wheel 内包含 cpuinfer_ext + whl="$(ls dist/*.whl)" + unzip -l "$whl" | grep -E 'cpuinfer_ext.*\.so' >/dev/null || { + echo "!! cpuinfer_ext missing in $whl"; exit 2; + } + + mv dist/*.whl wheels/ || true + deactivate + done +done + +echo "== Wheels saved in ./wheels ==" diff --git a/archive/kt-sft/book.toml b/archive/kt-sft/book.toml new file mode 100644 index 0000000..c88d9b7 --- /dev/null +++ b/archive/kt-sft/book.toml @@ -0,0 +1,18 @@ +[book] +authors = ["kvcache-ai"] +language = "zh-CN" +title = "Ktransformers" +src = "doc" + +[output.html] +git-repository-url = "https://github.com/kvcache-ai/ktransformers" +edit-url-template = "https://github.com/kvcache-ai/ktransformers/edit/main/{path}" + +[output.html.playground] +editable = true +copy-js = true +# line-numbers = true + +[output.html.fold] +enable = true +level = 0 \ No newline at end of file diff --git a/archive/kt-sft/csrc/custom_marlin/__init__.py b/archive/kt-sft/csrc/custom_marlin/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/archive/kt-sft/csrc/custom_marlin/binding.cpp b/archive/kt-sft/csrc/custom_marlin/binding.cpp new file mode 100644 index 0000000..184f3e2 --- /dev/null +++ b/archive/kt-sft/csrc/custom_marlin/binding.cpp @@ -0,0 +1,44 @@ +/** + * @Description : + * @Author : Azure-Tang + * @Date : 2024-07-25 13:38:30 + * @Version : 1.0.0 + * @LastEditors : kkk1nak0 + * @LastEditTime : 2024-08-12 03:05:04 + * @Copyright (c) 2024 by KVCache.AI, All Rights Reserved. + **/ + +#include "gptq_marlin/ops.h" +// Python bindings +#include +#include +#include +#include +#include +// namespace py = pybind11; + +PYBIND11_MODULE(vLLMMarlin, m) { + + /*m.def("dequantize_q8_0", &dequantize_q8_0, "Function to dequantize q8_0 + data.", py::arg("data"), py::arg("blk_size"), py::arg("device")); + m.def("dequantize_q6_k", &dequantize_q6_k, "Function to dequantize q6_k + data.", py::arg("data"), py::arg("blk_size"), py::arg("device")); + m.def("dequantize_q5_k", &dequantize_q5_k, "Function to dequantize q5_k + data.", py::arg("data"), py::arg("blk_size"), py::arg("device")); + m.def("dequantize_q4_k", &dequantize_q4_k, "Function to dequantize q4_k + data.", py::arg("data"), py::arg("blk_size"), py::arg("device")); + m.def("dequantize_q3_k", &dequantize_q3_k, "Function to dequantize q3_k + data.", py::arg("data"), py::arg("blk_size"), py::arg("device")); + m.def("dequantize_q2_k", &dequantize_q2_k, "Function to dequantize q2_k + data.", py::arg("data"), py::arg("blk_size"), py::arg("device")); + m.def("dequantize_iq4_xs", &dequantize_iq4_xs, "Function to dequantize + iq4_xs data.", py::arg("data"), py::arg("blk_size"), py::arg("device"));*/ + m.def("gptq_marlin_gemm", &gptq_marlin_gemm, + "Function to perform GEMM using Marlin quantization.", py::arg("a"), + py::arg("b_q_weight"), py::arg("b_scales"), py::arg("g_idx"), + py::arg("perm"), py::arg("workspace"), py::arg("num_bits"), py::arg("size_m_tensor"), + py::arg("size_m"), py::arg("size_n"), py::arg("size_k"), + py::arg("sms"), py::arg("is_k_full")); + m.def("gptq_marlin_repack", &gptq_marlin_repack, + "gptq_marlin repack from GPTQ"); +} \ No newline at end of file diff --git a/archive/kt-sft/csrc/custom_marlin/gptq_marlin/gptq_marlin.cu b/archive/kt-sft/csrc/custom_marlin/gptq_marlin/gptq_marlin.cu new file mode 100644 index 0000000..73ba3dd --- /dev/null +++ b/archive/kt-sft/csrc/custom_marlin/gptq_marlin/gptq_marlin.cu @@ -0,0 +1,2035 @@ +/* + * Modified by Neural Magic + * Copyright (C) Marlin.2024 Elias Frantar + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + + /* + * Adapted from https://github.com/IST-DASLab/marlin + */ + /* + * Adapted from + * https://github.com/vllm-project/vllm/tree/main/csrc/quantization/gptq_marlin + */ +#include "gptq_marlin.cuh" +#include "gptq_marlin_dtypes.cuh" +#include +#define STATIC_ASSERT_SCALAR_TYPE_VALID(scalar_t) \ + static_assert(std::is_same::value || \ + std::is_same::value, \ + "only float16 and bfloat16 is supported"); + +template inline std::string str(T x) { return std::to_string(x); } + +namespace gptq_marlin { + +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800 + + __global__ void permute_cols_kernel(int4 const* __restrict__ a_int4_ptr, + int const* __restrict__ perm_int_ptr, + int4* __restrict__ out_int4_ptr, int size_m, + int size_k, int block_rows) {} + + template shared + // fetch pipeline + const bool has_act_order, // whether act_order is enabled + const int group_blocks = -1 // number of consecutive 16x16 blocks + // with a separate quantization scale + > + __global__ void + Marlin(const int4* __restrict__ A, // fp16 input matrix of shape mxk + const int4* __restrict__ B, // 4bit quantized weight matrix of shape kxn + int4* __restrict__ C, // fp16 output buffer of shape mxn + const int4* __restrict__ scales_ptr, // fp16 quantization scales of shape + // (k/groupsize)xn + const int* __restrict__ g_idx, // int32 group indices of shape k + int num_groups, // number of scale groups per output channel + int prob_m, // batch dimension m + int prob_n, // output dimension n + int prob_k, // reduction dimension k + int* locks // extra global storage for barrier synchronization + ) {} + +} // namespace gptq_marlin + +torch::Tensor gptq_marlin_gemm(torch::Tensor& a, torch::Tensor& b_q_weight, + torch::Tensor& b_scales, torch::Tensor& g_idx, + torch::Tensor& perm, torch::Tensor& workspace, + int64_t num_bits, int64_t size_m, int64_t size_n, + int64_t size_k, bool is_k_full) { + TORCH_CHECK_NOT_IMPLEMENTED(false, + "marlin_gemm(..) requires CUDA_ARCH >= 8.0"); + return torch::empty({ 1, 1 }); +} + +#else + + // m16n8k16 tensor core mma instruction with fp16 inputs and fp32 + // output/accumulation. + template + __device__ inline void mma(const typename ScalarType::FragA& a_frag, + const typename ScalarType::FragB& frag_b, + typename ScalarType::FragC& frag_c) { + const uint32_t* a = reinterpret_cast(&a_frag); + const uint32_t* b = reinterpret_cast(&frag_b); + float* c = reinterpret_cast(&frag_c); + if constexpr (std::is_same::value) { + asm volatile( + "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 " + "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};\n" + : "=f"(c[0]), "=f"(c[1]), "=f"(c[2]), "=f"(c[3]) + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1]), + "f"(c[0]), "f"(c[1]), "f"(c[2]), "f"(c[3])); + } + else if constexpr (std::is_same::value) { + asm volatile( + "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 " + "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};\n" + : "=f"(c[0]), "=f"(c[1]), "=f"(c[2]), "=f"(c[3]) + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1]), + "f"(c[0]), "f"(c[1]), "f"(c[2]), "f"(c[3])); + } + else { + STATIC_ASSERT_SCALAR_TYPE_VALID(scalar_t); + } + } + + // Instruction for loading a full 16x16 matrix fragment of operand A from shared + // memory, directly in tensor core layout. + template + __device__ inline void ldsm4(typename ScalarType::FragA& frag_a, + const void* smem_ptr) { + uint32_t* a = reinterpret_cast(&frag_a); + uint32_t smem = static_cast(__cvta_generic_to_shared(smem_ptr)); + asm volatile( + "ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0,%1,%2,%3}, [%4];\n" + : "=r"(a[0]), "=r"(a[1]), "=r"(a[2]), "=r"(a[3]) + : "r"(smem)); + } + + // Lookup-table based 3-input logical operation; explicitly used for + // dequantization as the compiler does not seem to automatically recognize it in + // all cases. + template __device__ inline int lop3(int a, int b, int c) { + int res; + asm volatile("lop3.b32 %0, %1, %2, %3, %4;\n" + : "=r"(res) + : "r"(a), "r"(b), "r"(c), "n"(lut)); + return res; + } + + // Constructs destination register by taking bytes from 2 sources (based on + // mask) + template + __device__ inline uint32_t prmt(uint32_t a) { + uint32_t res; + asm volatile("prmt.b32 %0, %1, %2, %3;\n" + : "=r"(res) + : "r"(a), "n"(start_byte), "n"(mask)); + return res; + } + + // Efficiently dequantize an int32 value into a full B-fragment of 4 fp16 + // values. We mostly follow the strategy in the link below, with some small + // changes: + // - FP16: + // https://github.com/NVIDIA/FasterTransformer/blob/release/v5.3_tag/src/fastertransformer/cutlass_extensions/include/cutlass_extensions/interleaved_numeric_conversion.h#L215-L287 + // - BF16: + // https://github.com/NVIDIA/FasterTransformer/blob/release/v5.3_tag/src/fastertransformer/cutlass_extensions/include/cutlass_extensions/interleaved_numeric_conversion.h#L327-L385 + template + __device__ inline typename ScalarType::FragB dequant_4bit(int q) { + STATIC_ASSERT_SCALAR_TYPE_VALID(scalar_t); + } + + template <> + __device__ inline typename ScalarType::FragB dequant_4bit(int q) { + const int LO = 0x000f000f; + const int HI = 0x00f000f0; + const int EX = 0x64006400; + // Guarantee that the `(a & b) | c` operations are LOP3s. + int lo = lop3<(0xf0 & 0xcc) | 0xaa>(q, LO, EX); + int hi = lop3<(0xf0 & 0xcc) | 0xaa>(q, HI, EX); + // We want signed int4 outputs, hence we fuse the `-8` symmetric zero point + // directly into `SUB` and `ADD`. + const int SUB = 0x64086408; + const int MUL = 0x2c002c00; + const int ADD = 0xd480d480; + typename ScalarType::FragB frag_b; + frag_b[0] = __hsub2(*reinterpret_cast(&lo), + *reinterpret_cast(&SUB)); + frag_b[1] = __hfma2(*reinterpret_cast(&hi), + *reinterpret_cast(&MUL), + *reinterpret_cast(&ADD)); + return frag_b; + } + + template <> + __device__ inline typename ScalarType::FragB + dequant_4bit(int q) { + static constexpr uint32_t MASK = 0x000f000f; + static constexpr uint32_t EX = 0x43004300; + + // Guarantee that the `(a & b) | c` operations are LOP3s. + + int lo = lop3<(0xf0 & 0xcc) | 0xaa>(q, MASK, EX); + q >>= 4; + int hi = lop3<(0xf0 & 0xcc) | 0xaa>(q, MASK, EX); + + typename ScalarType::FragB frag_b; + static constexpr uint32_t MUL = 0x3F803F80; + static constexpr uint32_t ADD = 0xC308C308; + + frag_b[0] = __hfma2(*reinterpret_cast(&lo), + *reinterpret_cast(&MUL), + *reinterpret_cast(&ADD)); + frag_b[1] = __hfma2(*reinterpret_cast(&hi), + *reinterpret_cast(&MUL), + *reinterpret_cast(&ADD)); + return frag_b; + } + + // Fast Int8ToFp16/Int8ToBf16: Efficiently dequantize 8bit int values to fp16 or + // bf16 Reference: + // - FP16: + // https://github.com/NVIDIA/FasterTransformer/blob/release/v5.3_tag/src/fastertransformer/cutlass_extensions/include/cutlass_extensions/interleaved_numeric_conversion.h#L53-L85 + // - BF16: + // https://github.com/NVIDIA/FasterTransformer/blob/release/v5.3_tag/src/fastertransformer/cutlass_extensions/include/cutlass_extensions/interleaved_numeric_conversion.h#L125-L175 + template + __device__ inline typename ScalarType::FragB dequant_8bit(int q) { + STATIC_ASSERT_SCALAR_TYPE_VALID(scalar_t); + } + + template <> + __device__ inline typename ScalarType::FragB dequant_8bit(int q) { + static constexpr uint32_t mask_for_elt_01 = 0x5250; + static constexpr uint32_t mask_for_elt_23 = 0x5351; + static constexpr uint32_t start_byte_for_fp16 = 0x64646464; + + uint32_t lo = prmt(q); + uint32_t hi = prmt(q); + + static constexpr uint32_t I8s_TO_F16s_MAGIC_NUM = 0x64806480; + + typename ScalarType::FragB frag_b; + frag_b[0] = + __hsub2(*reinterpret_cast(&lo), + *reinterpret_cast(&I8s_TO_F16s_MAGIC_NUM)); + frag_b[1] = + __hsub2(*reinterpret_cast(&hi), + *reinterpret_cast(&I8s_TO_F16s_MAGIC_NUM)); + return frag_b; + } + + template <> + __device__ inline typename ScalarType::FragB + dequant_8bit(int q) { + typename ScalarType::FragB frag_b; + + float fp32_intermediates[4]; + uint32_t* fp32_intermediates_casted = + reinterpret_cast(fp32_intermediates); + + static constexpr uint32_t fp32_base = 0x4B000000; + fp32_intermediates_casted[0] = __byte_perm(q, fp32_base, 0x7650); + fp32_intermediates_casted[1] = __byte_perm(q, fp32_base, 0x7652); + fp32_intermediates_casted[2] = __byte_perm(q, fp32_base, 0x7651); + fp32_intermediates_casted[3] = __byte_perm(q, fp32_base, 0x7653); + + fp32_intermediates[0] -= 8388736.f; + fp32_intermediates[1] -= 8388736.f; + fp32_intermediates[2] -= 8388736.f; + fp32_intermediates[3] -= 8388736.f; + + uint32_t* bf16_result_ptr = reinterpret_cast(&frag_b); + bf16_result_ptr[0] = __byte_perm(fp32_intermediates_casted[0], + fp32_intermediates_casted[1], 0x7632); + bf16_result_ptr[1] = __byte_perm(fp32_intermediates_casted[2], + fp32_intermediates_casted[3], 0x7632); + + return frag_b; + } + + // Multiply dequantized values by the corresponding quantization scale; used + // only for grouped quantization. + template + __device__ inline void scale(typename ScalarType::FragB& frag_b, + typename ScalarType::FragS& frag_s, + int i) { + using scalar_t2 = typename ScalarType::scalar_t2; + scalar_t2 s = ScalarType::num2num2( + reinterpret_cast(&frag_s)[i]); + frag_b[0] = __hmul2(frag_b[0], s); + frag_b[1] = __hmul2(frag_b[1], s); + } + + // Same as above, but for act_order (each K is multiplied individually) + template + __device__ inline void scale4(typename ScalarType::FragB& frag_b, + typename ScalarType::FragS& frag_s_1, + typename ScalarType::FragS& frag_s_2, + typename ScalarType::FragS& frag_s_3, + typename ScalarType::FragS& frag_s_4, + int i) { + using scalar_t2 = typename ScalarType::scalar_t2; + scalar_t2 s_val_1_2; + s_val_1_2.x = reinterpret_cast(&frag_s_1)[i]; + s_val_1_2.y = reinterpret_cast(&frag_s_2)[i]; + + scalar_t2 s_val_3_4; + s_val_3_4.x = reinterpret_cast(&frag_s_3)[i]; + s_val_3_4.y = reinterpret_cast(&frag_s_4)[i]; + + frag_b[0] = __hmul2(frag_b[0], s_val_1_2); + frag_b[1] = __hmul2(frag_b[1], s_val_3_4); + } + + // Given 2 floats multiply by 2 scales (halves) + template + __device__ inline void scale_float(float* c, + typename ScalarType::FragS& s) { + scalar_t* s_ptr = reinterpret_cast(&s); + c[0] = __fmul_rn(c[0], ScalarType::num2float(s_ptr[0])); + c[1] = __fmul_rn(c[1], ScalarType::num2float(s_ptr[1])); + } + + // Wait until barrier reaches `count`, then lock for current threadblock. + __device__ inline void barrier_acquire(int* lock, int count) { + if (threadIdx.x == 0) { + int state = -1; + do + // Guarantee that subsequent writes by this threadblock will be + // visible globally. + asm volatile("ld.global.acquire.gpu.b32 %0, [%1];\n" + : "=r"(state) + : "l"(lock)); + while (state != count); + } + __syncthreads(); + } + + // Release barrier and increment visitation count. + __device__ inline void barrier_release(int* lock, bool reset = false) { + __syncthreads(); + if (threadIdx.x == 0) { + if (reset) { + lock[0] = 0; + return; + } + int val = 1; + // Make sure that all writes since acquiring this barrier are visible + // globally, while releasing the barrier. + asm volatile("fence.acq_rel.gpu;\n"); + asm volatile("red.relaxed.gpu.global.add.s32 [%0], %1;\n" + : + : "l"(lock), "r"(val)); + } + } + + // For a given "a" of size [M,K] performs a permutation of the K columns based + // on the given "perm" indices. + __global__ void permute_cols_kernel(int4 const* __restrict__ a_int4_ptr, + int const* __restrict__ perm_int_ptr, + int4* __restrict__ out_int4_ptr, int size_m, + int size_k, int block_rows) { + int start_row = block_rows * blockIdx.x; + int finish_row = start_row + block_rows; + if (finish_row > size_m) { + finish_row = size_m; + } + int cur_block_rows = finish_row - start_row; + + int row_stride = size_k * sizeof(half) / 16; + + auto permute_row = [&](int row) { + int iters = size_k / default_threads; + int rest = size_k % default_threads; + + int offset = row * row_stride; + + half const* a_row_half = + reinterpret_cast(a_int4_ptr + offset); + half* out_half = reinterpret_cast(out_int4_ptr + offset); + + int base_k = 0; + + for (int i = 0; i < iters; i++) { + int cur_k = base_k + threadIdx.x; + int src_pos = perm_int_ptr[cur_k]; + + out_half[cur_k] = a_row_half[src_pos]; + + base_k += default_threads; + } + + if (rest) { + if (threadIdx.x < rest) { + int cur_k = base_k + threadIdx.x; + int src_pos = perm_int_ptr[cur_k]; + + out_half[cur_k] = a_row_half[src_pos]; + } + } + }; + + for (int i = 0; i < cur_block_rows; i++) { + int cur_row = start_row + i; + if (cur_row < size_m) { + permute_row(cur_row); + } + } + } + + template shared + // fetch pipeline + const bool has_act_order, // whether act_order is enabled + const int group_blocks = -1 // number of consecutive 16x16 blocks + // with a separate quantization scale + > + __device__ void + Marlin(const int4* __restrict__ A, // fp16 input matrix of shape mxk + const int4* __restrict__ B, // 4bit quantized weight matrix of shape kxn + int4* __restrict__ C, // fp16 output buffer of shape mxn + const int4* __restrict__ scales_ptr, // fp16 quantization scales of shape + // (k/groupsize)xn + const int* __restrict__ g_idx, // int32 group indices of shape k + int num_groups, // number of scale groups per output channel + int prob_m, // batch dimension m, should be divisible by (16 * thread_m_blocks) if bigger than that + int prob_n, // output dimension n + int prob_k, // reduction dimension k + int* locks // extra global storage for barrier synchronization + ) { + // Each threadblock processes one "stripe" of the B matrix with (roughly) the + // same size, which might involve multiple column "slices" (of width 16 * + // `thread_n_blocks`). Stripes are defined as shown in the 3x3 matrix 5 SM + // example: + // 0 1 3 + // 0 2 3 + // 1 2 4 + // While this kind of partitioning makes things somewhat more complicated, it + // ensures good utilization of all SMs for many kinds of shape and GPU + // configurations, while requiring as few slow global cross-threadblock + // reductions as possible. + using Dtype = ScalarType; + using scalar_t2 = typename ScalarType::scalar_t2; + using FragA = typename ScalarType::FragA; + using FragB = typename ScalarType::FragB; + using FragC = typename ScalarType::FragC; + using FragS = typename ScalarType::FragS; + + constexpr int pack_factor = 32 / num_bits; + + // int prob_m = *prob_m_ptr; + // const int thread_m_blocks = min(div_ceil(prob_m, 16), template_thread_m_blocks); + // constexpr int thread_m_blocks = template_thread_m_blocks; + + // For larger GEMMs we run multiple batchsize 64 versions in parallel for a + // better partitioning with less reductions + int parallel = 1; + if (prob_m > 16 * thread_m_blocks) { + parallel = prob_m / (16 * thread_m_blocks); + prob_m = 16 * thread_m_blocks; + } + + int k_tiles = prob_k / 16 / thread_k_blocks; + int n_tiles = prob_n / 16 / thread_n_blocks; + int iters = div_ceil(k_tiles * n_tiles * parallel, gridDim.x); + + if constexpr (!has_act_order && group_blocks != -1) { + if (group_blocks >= thread_k_blocks) { + // Ensure that the number of tiles in each stripe is a multiple of the + // groupsize; this avoids an annoying special case where a stripe starts + // in the middle of group. + iters = (group_blocks / thread_k_blocks) * + div_ceil(iters, (group_blocks / thread_k_blocks)); + } + } + + int slice_row = (iters * blockIdx.x) % k_tiles; + int slice_col_par = (iters * blockIdx.x) / k_tiles; + int slice_col = slice_col_par; + int slice_iters; // number of threadblock tiles in the current slice + int slice_count = + 0; // total number of active threadblocks in the current slice + int slice_idx; // index of threadblock in current slice; numbered bottom to + // top + + // We can easily implement parallel problem execution by just remapping + // indices and advancing global pointers + if (slice_col_par >= n_tiles) { + A += (slice_col_par / n_tiles) * 16 * thread_m_blocks * prob_k / 8; + C += (slice_col_par / n_tiles) * 16 * thread_m_blocks * prob_n / 8; + locks += (slice_col_par / n_tiles) * n_tiles; + slice_col = slice_col_par % n_tiles; + } + + // Compute all information about the current slice which is required for + // synchronization. + auto init_slice = [&]() { + slice_iters = + iters * (blockIdx.x + 1) - (k_tiles * slice_col_par + slice_row); + if (slice_iters < 0 || slice_col_par >= n_tiles * parallel) slice_iters = 0; + if (slice_iters == 0) return; + if (slice_row + slice_iters > k_tiles) slice_iters = k_tiles - slice_row; + slice_count = 1; + slice_idx = 0; + int col_first = iters * div_ceil(k_tiles * slice_col_par, iters); + if (col_first <= k_tiles * (slice_col_par + 1)) { + int col_off = col_first - k_tiles * slice_col_par; + slice_count = div_ceil(k_tiles - col_off, iters); + if (col_off > 0) slice_count++; + int delta_first = iters * blockIdx.x - col_first; + if (delta_first < 0 || (col_off == 0 && delta_first == 0)) + slice_idx = slice_count - 1; + else { + slice_idx = slice_count - 1 - delta_first / iters; + if (col_off > 0) slice_idx--; + } + } + if (slice_col == n_tiles) { + A += 16 * thread_m_blocks * prob_k / 8; + C += 16 * thread_m_blocks * prob_n / 8; + locks += n_tiles; + slice_col = 0; + } + }; + init_slice(); + + // A sizes/strides + + // stride of the A matrix in global memory + int a_gl_stride = prob_k / 8; + // stride of an A matrix tile in shared memory + constexpr int a_sh_stride = 16 * thread_k_blocks / 8; + // delta between subsequent A tiles in global memory + constexpr int a_gl_rd_delta_o = 16 * thread_k_blocks / 8; + // between subsequent accesses within a tile + int a_gl_rd_delta_i = a_gl_stride * (threads / a_gl_rd_delta_o); + // between shared memory writes + constexpr int a_sh_wr_delta = a_sh_stride * (threads / a_gl_rd_delta_o); + // between shared memory tile reads + constexpr int a_sh_rd_delta_o = 2 * ((threads / 32) / (thread_n_blocks / 4)); + // within a shared memory tile + constexpr int a_sh_rd_delta_i = a_sh_stride * 16; + // overall size of a tile + constexpr int a_sh_stage = a_sh_stride * (16 * thread_m_blocks); + // number of shared write iterations for a tile + constexpr int a_sh_wr_iters = div_ceil(a_sh_stage, a_sh_wr_delta); + + // B sizes/strides + int b_gl_stride = 16 * prob_n / (pack_factor * 4); + constexpr int b_sh_stride = ((thread_n_blocks * 16) * 16 / pack_factor) / 4; + constexpr int b_thread_vecs = num_bits == 4 ? 1 : 2; + constexpr int b_sh_stride_threads = b_sh_stride / b_thread_vecs; + + int b_gl_rd_delta_o = b_gl_stride * thread_k_blocks; + int b_gl_rd_delta_i = b_gl_stride * (threads / b_sh_stride_threads); + constexpr int b_sh_wr_delta = threads * b_thread_vecs; + constexpr int b_sh_rd_delta = threads * b_thread_vecs; + constexpr int b_sh_stage = b_sh_stride * thread_k_blocks; + constexpr int b_sh_wr_iters = b_sh_stage / b_sh_wr_delta; + + // Scale sizes/strides without act_order + int s_gl_stride = prob_n / 8; + constexpr int s_sh_stride = 16 * thread_n_blocks / 8; + constexpr int s_tb_groups = + !has_act_order && group_blocks != -1 && group_blocks < thread_k_blocks + ? thread_k_blocks / group_blocks + : 1; + constexpr int s_sh_stage = s_tb_groups * s_sh_stride; + int s_gl_rd_delta = s_gl_stride; + + // Scale size/strides with act_order + constexpr int tb_k = 16 * thread_k_blocks; + constexpr int g_idx_stage = has_act_order ? (tb_k * sizeof(int)) / 16 : 0; + // constexpr int act_s_row_stride = 1; + // int act_s_col_stride = act_s_row_stride * num_groups; + int act_s_col_stride = 1; + int act_s_col_warp_stride = act_s_col_stride * 8; + int tb_n_warps = thread_n_blocks / 4; + int act_s_col_tb_stride = act_s_col_warp_stride * tb_n_warps; + + // Global A read index of current thread. + int a_gl_rd = a_gl_stride * (threadIdx.x / a_gl_rd_delta_o) + + (threadIdx.x % a_gl_rd_delta_o); + a_gl_rd += a_gl_rd_delta_o * slice_row; + // Shared write index of current thread. + int a_sh_wr = a_sh_stride * (threadIdx.x / a_gl_rd_delta_o) + + (threadIdx.x % a_gl_rd_delta_o); + // Shared read index. + int a_sh_rd = + a_sh_stride * ((threadIdx.x % 32) % 16) + (threadIdx.x % 32) / 16; + a_sh_rd += 2 * ((threadIdx.x / 32) / (thread_n_blocks / 4)); + + int b_gl_rd = b_gl_stride * (threadIdx.x / b_sh_stride_threads) + + (threadIdx.x % b_sh_stride_threads) * b_thread_vecs; + b_gl_rd += b_sh_stride * slice_col; + b_gl_rd += b_gl_rd_delta_o * slice_row; + int b_sh_wr = threadIdx.x * b_thread_vecs; + int b_sh_rd = threadIdx.x * b_thread_vecs; + + // For act_order + constexpr int k_iter_size = tb_k / b_sh_wr_iters; + int slice_k_start = tb_k * slice_row; + int slice_k_finish = slice_k_start + tb_k * slice_iters; + int slice_k_start_shared_fetch = slice_k_start; + int slice_n_offset = act_s_col_tb_stride * slice_col; + + // No act_order + int s_gl_rd; + if constexpr (!has_act_order) { + if constexpr (group_blocks == -1) { + s_gl_rd = s_sh_stride * slice_col + threadIdx.x; + } + else { + s_gl_rd = s_gl_stride * ((thread_k_blocks * slice_row) / group_blocks) + + s_sh_stride * slice_col + threadIdx.x; + } + } + int s_sh_wr = threadIdx.x; + bool s_sh_wr_pred = threadIdx.x < s_sh_stride; + + // We use a different scale layout for grouped and column-wise quantization as + // we scale a `half2` tile in column-major layout in the former and in + // row-major in the latter case. + int s_sh_rd; + if constexpr (group_blocks != -1) + s_sh_rd = 8 * ((threadIdx.x / 32) % (thread_n_blocks / 4)) + + (threadIdx.x % 32) / 4; + else + s_sh_rd = 8 * ((threadIdx.x / 32) % (thread_n_blocks / 4)) + + (threadIdx.x % 32) % 4; + + // Precompute which thread should not read memory in which iterations; this is + // needed if there are more threads than required for a certain tilesize or + // when the batchsize is not a multiple of 16. + bool a_sh_wr_pred[a_sh_wr_iters]; +#pragma unroll + for (int i = 0; i < a_sh_wr_iters; i++) { + a_sh_wr_pred[i] = a_sh_wr_delta * i + a_sh_wr < a_sh_stride * prob_m; + } + + // To ensure that writing and reading A tiles to/from shared memory, the + // latter in fragment format, is fully bank conflict free, we need to use a + // rather fancy XOR-based layout. The key here is that neither reads nor + // writes of the 16-byte `int4` blocks of 8 consecutive threads involve the + // same shared memory banks. Further, it seems (based on NSight-Compute) that + // each warp must also write a consecutive memory segment? + auto transform_a = [&](int i) { + int row = i / a_gl_rd_delta_o; + return a_gl_rd_delta_o * row + (i % a_gl_rd_delta_o) ^ row; + }; + // Since the computation of this remapping is non-trivial and, due to our main + // loop unrolls, all shared memory accesses are static, we simply precompute + // both transformed reads and writes. + int a_sh_wr_trans[a_sh_wr_iters]; +#pragma unroll + for (int i = 0; i < a_sh_wr_iters; i++) { + a_sh_wr_trans[i] = transform_a(a_sh_wr_delta * i + a_sh_wr); + } + int a_sh_rd_trans[b_sh_wr_iters][thread_m_blocks]; +#pragma unroll + for (int i = 0; i < b_sh_wr_iters; i++) { +#pragma unroll + for (int j = 0; j < thread_m_blocks; j++) + { + a_sh_rd_trans[i][j] = + transform_a(a_sh_rd_delta_o * i + a_sh_rd_delta_i * j + a_sh_rd); + } + } + + // Since B-accesses have non-constant stride they have to be computed at + // runtime; we break dependencies between subsequent accesses with a tile by + // maintining multiple pointers (we have enough registers), a tiny + // optimization. + const int4* B_ptr[b_sh_wr_iters]; +#pragma unroll + for (int i = 0; i < b_sh_wr_iters; i++) + B_ptr[i] = B + b_gl_rd_delta_i * i + b_gl_rd; + + extern __shared__ int4 sh[]; + // Shared memory storage for global fetch pipelines. + int4* sh_a = sh; + int4* sh_b = sh_a + (stages * a_sh_stage); + int4* sh_g_idx = sh_b + (stages * b_sh_stage); + int4* sh_s = sh_g_idx + (stages * g_idx_stage); + + // Register storage for double buffer of shared memory reads. + FragA frag_a[2][thread_m_blocks]; + I4 frag_b_quant[2][b_thread_vecs]; + FragC frag_c[thread_m_blocks][4][2]; + FragS frag_s[2][4]; // No act-order + FragS act_frag_s[2][4][4]; // For act-order + + // Zero accumulators. + auto zero_accums = [&]() { +#pragma unroll + for (int i = 0; i < thread_m_blocks * 4 * 2 * 4; i++) + { + reinterpret_cast(frag_c)[i] = 0; + } + }; + + int sh_first_group_id = -1; + int sh_num_groups = -1; + constexpr int sh_max_num_groups = 32; + + auto fetch_scales_to_shared = [&](bool is_async, int first_group_id, + int last_group_id) { + sh_first_group_id = first_group_id; + sh_num_groups = last_group_id - first_group_id + 1; + + if (sh_num_groups < sh_max_num_groups) { + sh_num_groups = sh_max_num_groups; + } + + if (sh_first_group_id + sh_num_groups > num_groups) { + sh_num_groups = num_groups - sh_first_group_id; + } + + int row_offset = first_group_id * s_gl_stride; + + if (is_async) { + for (int i = 0; i < sh_num_groups; i++) { + if (threadIdx.x < s_sh_stride) { + cp_async4_pred(&sh_s[(i * s_sh_stride) + threadIdx.x], + &scales_ptr[row_offset + (i * s_gl_stride) + + slice_n_offset + threadIdx.x]); + } + } + } + else { + for (int i = 0; i < sh_num_groups; i++) { + if (threadIdx.x < s_sh_stride) { + sh_s[(i * s_sh_stride) + threadIdx.x] = + scales_ptr[row_offset + (i * s_gl_stride) + slice_n_offset + + threadIdx.x]; + } + } + } + }; + // Asynchronously fetch the next A, B and s tile from global to the next + // shared memory pipeline location. + auto fetch_to_shared = [&](int pipe, int a_off, bool pred = true) { + if (pred) { + int4* sh_a_stage = sh_a + a_sh_stage * pipe; +#pragma unroll + for (int i = 0; i < a_sh_wr_iters; i++) { + cp_async4_pred( + &sh_a_stage[a_sh_wr_trans[i]], + &A[a_gl_rd_delta_i * i + a_gl_rd + a_gl_rd_delta_o * a_off], + a_sh_wr_pred[i]); + } + int4* sh_b_stage = sh_b + b_sh_stage * pipe; +#pragma unroll + for (int i = 0; i < b_sh_wr_iters; i++) { +#pragma unroll + for (int j = 0; j < b_thread_vecs; j++) { + cp_async4(&sh_b_stage[b_sh_wr_delta * i + b_sh_wr + j], B_ptr[i] + j); + } + + B_ptr[i] += b_gl_rd_delta_o; + } + + if constexpr (has_act_order) { + // Fetch g_idx thread-block portion + int full_pipe = a_off; + int cur_k = slice_k_start_shared_fetch + tb_k * full_pipe; + if (cur_k < prob_k && cur_k < slice_k_finish) { + int4* sh_g_idx_stage = sh_g_idx + g_idx_stage * pipe; + + int4 const* cur_g_idx_stage_ptr = + reinterpret_cast(&g_idx[cur_k]); + + if (threadIdx.x < g_idx_stage) { + cp_async4_pred(&sh_g_idx_stage[threadIdx.x], + &cur_g_idx_stage_ptr[threadIdx.x]); + } + } + } + else { + if constexpr (group_blocks != -1) { + int4* sh_s_stage = sh_s + s_sh_stage * pipe; + + if constexpr (group_blocks >= thread_k_blocks) { + // Only fetch scales if this tile starts a new group + if (pipe % (group_blocks / thread_k_blocks) == 0) { + if (s_sh_wr_pred) { + cp_async4(&sh_s_stage[s_sh_wr], &scales_ptr[s_gl_rd]); + } + s_gl_rd += s_gl_rd_delta; + } + } + else { + for (int i = 0; i < s_tb_groups; i++) { + if (s_sh_wr_pred) { + cp_async4(&sh_s_stage[i * s_sh_stride + s_sh_wr], + &scales_ptr[s_gl_rd]); + } + s_gl_rd += s_gl_rd_delta; + } + } + } + } + } + // Insert a fence even when we are winding down the pipeline to ensure that + // waiting is also correct at this point. + cp_async_fence(); + }; + + // Wait until the next thread tile has been loaded to shared memory. + auto wait_for_stage = [&]() { + // We only have `stages - 2` active fetches since we are double buffering + // and can only issue the next fetch when it is guaranteed that the previous + // shared memory load is fully complete (as it may otherwise be + // overwritten). + cp_async_wait(); + __syncthreads(); + }; + + // Load the next sub-tile from the current location in the shared memory pipe + // into the current register buffer. + auto fetch_to_registers = [&](int k, int pipe) { + int4* sh_a_stage = sh_a + a_sh_stage * pipe; +#pragma unroll + for (int i = 0; i < thread_m_blocks; i++) + { + ldsm4(frag_a[k % 2][i], + &sh_a_stage[a_sh_rd_trans[k % b_sh_wr_iters][i]]); + } + + int4* sh_b_stage = sh_b + b_sh_stage * pipe; + +#pragma unroll + for (int i = 0; i < b_thread_vecs; i++) { + frag_b_quant[k % 2][i] = *reinterpret_cast( + &sh_b_stage[b_sh_rd_delta * (k % b_sh_wr_iters) + b_sh_rd + i]); + } + }; + + bool is_same_group[stages]; + int same_group_id[stages]; + + auto init_same_group = [&](int pipe) { + if constexpr (!has_act_order) { + is_same_group[pipe] = false; + same_group_id[pipe] = 0; + return; + } + + int4* sh_g_idx_stage = sh_g_idx + g_idx_stage * pipe; + int* sh_g_idx_int_ptr = reinterpret_cast(sh_g_idx_stage); + + int group_id_1 = sh_g_idx_int_ptr[0]; + int group_id_2 = sh_g_idx_int_ptr[tb_k - 1]; + + is_same_group[pipe] = group_id_1 == group_id_2; + same_group_id[pipe] = group_id_1; + }; + + auto fetch_scales_to_registers = [&](int k, int full_pipe) { + int pipe = full_pipe % stages; + + if constexpr (!has_act_order) { + // No act-order case + if constexpr (group_blocks != -1) { + if constexpr (group_blocks >= thread_k_blocks) { + int4* sh_s_stage = + sh_s + s_sh_stage * ((group_blocks / thread_k_blocks) * + (pipe / (group_blocks / thread_k_blocks))); + reinterpret_cast(&frag_s[k % 2])[0] = sh_s_stage[s_sh_rd]; + } + else { + int warp_id = threadIdx.x / 32; + int n_warps = thread_n_blocks / 4; + + int warp_row = warp_id / n_warps; + + int cur_k = warp_row * 16; + cur_k += k_iter_size * (k % b_sh_wr_iters); + + int k_blocks = cur_k / 16; + int cur_group_id = k_blocks / group_blocks; + + int4* sh_s_stage = sh_s + s_sh_stage * pipe; + + reinterpret_cast(&frag_s[k % 2])[0] = + sh_s_stage[s_sh_rd + cur_group_id * s_sh_stride]; + } + } + + return; + } + + // Act-order case + + // Determine K of the "current" thread-block + int cur_k = slice_k_start + tb_k * full_pipe; + if (cur_k >= prob_k || cur_k >= slice_k_finish) { + return; + } + + // Reset (to current thread-block) since we read g_idx portion from the + // shared memory + cur_k = 0; + + // Progress to current iteration + cur_k += k_iter_size * (k % b_sh_wr_iters); + + // Determine "position" inside the thread-block (based on warp and + // thread-id) + int warp_id = threadIdx.x / 32; + int n_warps = + thread_n_blocks / 4; // Each warp processes 4 16-size tiles over N + + int warp_row = warp_id / n_warps; + int warp_col = warp_id % n_warps; + + cur_k += warp_row * 16; + + int th_id = threadIdx.x % 32; + cur_k += (th_id % 4) * 2; // Due to tensor-core layout for fp16 B matrix + + int s_col_shift = + /*slice_n_offset +*/ (act_s_col_warp_stride * warp_col) + + (th_id / 4) * act_s_col_stride; + + if (is_same_group[pipe]) { + if (k % 2 == 0) { + *(reinterpret_cast(&(act_frag_s[k % 2][0][0]))) = + sh_s[(same_group_id[pipe] - sh_first_group_id) * s_sh_stride + + s_col_shift]; + } + else { + *(reinterpret_cast(&(act_frag_s[k % 2][0][0]))) = + *(reinterpret_cast(&(act_frag_s[(k - 1) % 2][0][0]))); + } + + for (int i = 1; i < 4; i++) { + *(reinterpret_cast(&(act_frag_s[k % 2][i][0]))) = + *(reinterpret_cast(&(act_frag_s[k % 2][0][0]))); + } + return; + } + + int4* sh_g_idx_stage = sh_g_idx + g_idx_stage * pipe; + int* sh_g_idx_int_ptr = reinterpret_cast(sh_g_idx_stage); + + constexpr int k_frag_offsets[4] = { 0, 1, 8, + 9 }; // Tensor core offsets per thread + +#pragma unroll + for (int i = 0; i < 4; i++) { + int actual_k = cur_k + k_frag_offsets[i]; + + int group_id = sh_g_idx_int_ptr[actual_k]; + int rel_group_id = group_id - sh_first_group_id; + + *(reinterpret_cast(&(act_frag_s[k % 2][i][0]))) = + sh_s[rel_group_id * s_sh_stride + s_col_shift]; + } + }; + + // Execute the actual tensor core matmul of a sub-tile. + auto matmul = [&](int k) { + // We have the m dimension as the inner loop in order to encourage overlapping + // dequantization and matmul operations. +#pragma unroll + for (int j = 0; j < 4; j++) { + FragB frag_b0; + FragB frag_b1; + if constexpr (num_bits == 4) { + int b_quant = frag_b_quant[k % 2][0][j]; + int b_quant_shift = b_quant >> 8; + + frag_b0 = dequant_4bit(b_quant); + frag_b1 = dequant_4bit(b_quant_shift); + + } + else { + int* frag_b_quant_ptr = reinterpret_cast(frag_b_quant[k % 2]); + int b_quant_0 = frag_b_quant_ptr[j * 2 + 0]; + int b_quant_1 = frag_b_quant_ptr[j * 2 + 1]; + + frag_b0 = dequant_8bit(b_quant_0); + frag_b1 = dequant_8bit(b_quant_1); + } + + // Apply scale to frag_b0 + if constexpr (has_act_order) { + scale4(frag_b0, act_frag_s[k % 2][0][j], + act_frag_s[k % 2][1][j], act_frag_s[k % 2][2][j], + act_frag_s[k % 2][3][j], 0); + } + else { + if constexpr (group_blocks != -1) { + scale(frag_b0, frag_s[k % 2][j], 0); + } + } + + // Apply scale to frag_b1 + if constexpr (has_act_order) { + scale4(frag_b1, act_frag_s[k % 2][0][j], + act_frag_s[k % 2][1][j], act_frag_s[k % 2][2][j], + act_frag_s[k % 2][3][j], 1); + + } + else { + if constexpr (group_blocks != -1) { + scale(frag_b1, frag_s[k % 2][j], 1); + } + } + +#pragma unroll + for (int i = 0; i < thread_m_blocks; i++) { + mma(frag_a[k % 2][i], frag_b0, frag_c[i][j][0]); + mma(frag_a[k % 2][i], frag_b1, frag_c[i][j][1]); + } + } + }; + + // Since we slice across the k dimension of a tile in order to increase the + // number of warps while keeping the n dimension of a tile reasonable, we have + // multiple warps that accumulate their partial sums of the same output + // location; which we have to reduce over in the end. We do in shared memory. + auto thread_block_reduce = [&]() { + constexpr int red_off = threads / b_sh_stride_threads / 2; + if (red_off >= 1) { + int red_idx = threadIdx.x / b_sh_stride_threads; + constexpr int red_sh_stride = b_sh_stride_threads * 4 * 2; + constexpr int red_sh_delta = b_sh_stride_threads; + int red_sh_rd = red_sh_stride * (threadIdx.x / b_sh_stride_threads) + + (threadIdx.x % b_sh_stride_threads); + + // Parallel logarithmic shared memory reduction. We make sure to avoid any + // unnecessary read or write iterations, e.g., for two warps we write only + // once by warp 1 and read only once by warp 0. + +#pragma unroll + for (int m_block = 0; m_block < thread_m_blocks; m_block++) { +#pragma unroll + for (int i = red_off; i > 0; i /= 2) { + if (i <= red_idx && red_idx < 2 * i) { +#pragma unroll + for (int j = 0; j < 4 * 2; j++) { + int red_sh_wr = + red_sh_delta * j + (red_sh_rd - red_sh_stride * i); + if (i < red_off) { + float* c_rd = + reinterpret_cast(&sh[red_sh_delta * j + red_sh_rd]); + float* c_wr = reinterpret_cast(&sh[red_sh_wr]); +#pragma unroll + for (int k = 0; k < 4; k++) + reinterpret_cast(frag_c)[4 * 2 * m_block + j][k] += + c_rd[k] + c_wr[k]; + } + sh[red_sh_wr] = + reinterpret_cast(&frag_c)[4 * 2 * m_block + j]; + } + } + __syncthreads(); + } + if (red_idx == 0) { +#pragma unroll + for (int i = 0; i < 4 * 2; i++) { + float* c_rd = + reinterpret_cast(&sh[red_sh_delta * i + red_sh_rd]); +#pragma unroll + for (int j = 0; j < 4; j++) + reinterpret_cast(frag_c)[4 * 2 * m_block + i][j] += + c_rd[j]; + } + } + __syncthreads(); + } + } + }; + + // Since multiple threadblocks may process parts of the same column slice, we + // finally have to globally reduce over the results. As the striped + // partitioning minimizes the number of such reductions and our outputs are + // usually rather small, we perform this reduction serially in L2 cache. + auto global_reduce = [&](bool first = false, bool last = false) { + // We are very careful here to reduce directly in the output buffer to + // maximize L2 cache utilization in this step. To do this, we write out + // results in FP16 (but still reduce with FP32 compute). + constexpr int active_threads = 32 * thread_n_blocks / 4; + if (threadIdx.x < active_threads) { + int c_gl_stride = prob_n / 8; + int c_gl_wr_delta_o = 8 * c_gl_stride; + int c_gl_wr_delta_i = 4 * (active_threads / 32); + int c_gl_wr = c_gl_stride * ((threadIdx.x % 32) / 4) + + 4 * (threadIdx.x / 32) + threadIdx.x % 4; + c_gl_wr += (2 * thread_n_blocks) * slice_col; + constexpr int c_sh_wr_delta = active_threads; + int c_sh_wr = threadIdx.x; + + int row = (threadIdx.x % 32) / 4; + + if (!first) { + // Interestingly, doing direct global accesses here really seems to mess up + // the compiler and lead to slowdowns, hence we also use async-copies even + // though these fetches are not actually asynchronous. +#pragma unroll + for (int i = 0; i < thread_m_blocks * 4; i++) { + cp_async4_pred( + &sh[c_sh_wr + c_sh_wr_delta * i], + &C[c_gl_wr + c_gl_wr_delta_o * (i / 2) + + c_gl_wr_delta_i * (i % 2)], + i < (thread_m_blocks - 1) * 4 || 8 * (i / 2) + row < prob_m); + } + cp_async_fence(); + cp_async_wait<0>(); + } + +#pragma unroll + for (int i = 0; i < thread_m_blocks * 4; i++) { + if (i < (thread_m_blocks - 1) * 4 || 8 * (i / 2) + row < prob_m) { + if (!first) { + int4 c_red = sh[c_sh_wr + i * c_sh_wr_delta]; +#pragma unroll + for (int j = 0; j < 2 * 4; j++) { + reinterpret_cast( + &frag_c)[4 * 2 * 4 * (i / 4) + 4 * j + (i % 4)] += + Dtype::num2float(reinterpret_cast(&c_red)[j]); + } + } + if (!last) { + int4 c; +#pragma unroll + for (int j = 0; j < 2 * 4; j++) { + reinterpret_cast(&c)[j] = + Dtype::float2num(reinterpret_cast( + &frag_c)[4 * 2 * 4 * (i / 4) + 4 * j + (i % 4)]); + } + C[c_gl_wr + c_gl_wr_delta_o * (i / 2) + c_gl_wr_delta_i * (i % 2)] = + c; + } + } + } + } + }; + + // Write out the reduce final result in the correct layout. We only actually + // reshuffle matrix fragments in this step, the reduction above is performed + // in fragment layout. + auto write_result = [&]() { + int c_gl_stride = prob_n / 8; + constexpr int c_sh_stride = 2 * thread_n_blocks + 1; + int c_gl_wr_delta = c_gl_stride * (threads / (2 * thread_n_blocks)); + constexpr int c_sh_rd_delta = + c_sh_stride * (threads / (2 * thread_n_blocks)); + + int c_gl_wr = c_gl_stride * (threadIdx.x / (2 * thread_n_blocks)) + + (threadIdx.x % (2 * thread_n_blocks)); + c_gl_wr += (2 * thread_n_blocks) * slice_col; + int c_sh_wr = + (4 * c_sh_stride) * ((threadIdx.x % 32) / 4) + (threadIdx.x % 32) % 4; + c_sh_wr += 32 * (threadIdx.x / 32); + int c_sh_rd = c_sh_stride * (threadIdx.x / (2 * thread_n_blocks)) + + (threadIdx.x % (2 * thread_n_blocks)); + + int c_gl_wr_end = c_gl_stride * prob_m; + + // We first reorder in shared memory to guarantee the most efficient final + // global write patterns + auto write = [&](int idx, float c0, float c1, FragS& s) { + scalar_t2 res = + Dtype::nums2num2(Dtype::float2num(c0), Dtype::float2num(c1)); + + // For per-column quantization we finally apply the scale here (only for + // 4-bit) + if constexpr (!has_act_order && group_blocks == -1 && num_bits == 4) { + res = __hmul2(res, s[0]); + } + + ((scalar_t2*)sh)[idx] = res; + }; + + if (threadIdx.x / 32 < thread_n_blocks / 4) { +#pragma unroll + for (int i = 0; i < thread_m_blocks; i++) { +#pragma unroll + for (int j = 0; j < 4; j++) { + int wr = c_sh_wr + 8 * j; + write(wr + (4 * c_sh_stride) * 0 + 0, frag_c[i][j][0][0], + frag_c[i][j][0][1], frag_s[j / 2][2 * (j % 2) + 0]); + write(wr + (4 * c_sh_stride) * 8 + 0, frag_c[i][j][0][2], + frag_c[i][j][0][3], frag_s[j / 2][2 * (j % 2) + 0]); + write(wr + (4 * c_sh_stride) * 0 + 4, frag_c[i][j][1][0], + frag_c[i][j][1][1], frag_s[j / 2][2 * (j % 2) + 1]); + write(wr + (4 * c_sh_stride) * 8 + 4, frag_c[i][j][1][2], + frag_c[i][j][1][3], frag_s[j / 2][2 * (j % 2) + 1]); + } + c_sh_wr += 16 * (4 * c_sh_stride); + } + } + __syncthreads(); + +#pragma unroll + for (int i = 0; + i < div_ceil(16 * thread_m_blocks, threads / (2 * thread_n_blocks)); + i++) { + if (c_gl_wr < c_gl_wr_end) { + C[c_gl_wr] = sh[c_sh_rd]; + c_gl_wr += c_gl_wr_delta; + c_sh_rd += c_sh_rd_delta; + } + } + }; + + // Start global fetch and register load pipelines. + auto start_pipes = [&]() { + +#pragma unroll + for (int i = 0; i < stages - 1; i++) { + if (has_act_order && i == 0) { + int last_g_idx = slice_k_start + stages * tb_k * 2; + if (last_g_idx >= prob_k) { + last_g_idx = prob_k - 1; + } + fetch_scales_to_shared(true, g_idx[slice_k_start], g_idx[last_g_idx]); + } + fetch_to_shared(i, i, i < slice_iters); + } + + zero_accums(); + wait_for_stage(); + init_same_group(0); + fetch_to_registers(0, 0); + fetch_scales_to_registers(0, 0); + a_gl_rd += a_gl_rd_delta_o * (stages - 1); + slice_k_start_shared_fetch += tb_k * (stages - 1); + }; + if (slice_iters) { + start_pipes(); + } + + // Main loop. + while (slice_iters) { + // We unroll over both the global fetch and the register load pipeline to + // ensure all shared memory accesses are static. Note that both pipelines + // have even length meaning that the next iteration will always start at + // index 0. + +#pragma unroll + for (int pipe = 0; pipe < stages;) { +#pragma unroll + for (int k = 0; k < b_sh_wr_iters; k++) { + fetch_to_registers(k + 1, pipe % stages); + fetch_scales_to_registers(k + 1, pipe); + if (k == b_sh_wr_iters - 2) { + fetch_to_shared((pipe + stages - 1) % stages, pipe, + slice_iters >= stages); + pipe++; + wait_for_stage(); + init_same_group(pipe % stages); + } + matmul(k); + } + slice_iters--; + if (slice_iters == 0) { + break; + } + } + + a_gl_rd += a_gl_rd_delta_o * stages; + slice_k_start += tb_k * stages; + slice_k_start_shared_fetch += tb_k * stages; + + if constexpr (has_act_order) { + int first_group_id = g_idx[slice_k_start]; + int last_g_idx = slice_k_start + stages * tb_k * 2; + if (last_g_idx >= prob_k) { + last_g_idx = prob_k - 1; + } + int last_group_id = g_idx[last_g_idx]; + if (last_group_id >= sh_first_group_id + sh_num_groups) { + fetch_scales_to_shared(false, first_group_id, last_group_id); + __syncthreads(); + } + } + + // Process results and, if necessary, proceed to the next column slice. + // While this pattern may not be the most readable, other ways of writing + // the loop seemed to noticeably worse performance after compilation. + if (slice_iters == 0) { + cp_async_wait<0>(); + bool last = slice_idx == slice_count - 1; + // For per-column scales, we only fetch them here in the final step before + // write-out + if constexpr (!has_act_order && group_blocks == -1) { + if constexpr (num_bits == 8) { + if (s_sh_wr_pred) { + cp_async4(&sh_s[s_sh_wr], &scales_ptr[s_gl_rd]); + } + cp_async_fence(); + } + else { + if (last) { + if (s_sh_wr_pred) { + cp_async4(&sh_s[s_sh_wr], &scales_ptr[s_gl_rd]); + } + cp_async_fence(); + } + } + } + + thread_block_reduce(); + if constexpr (!has_act_order && group_blocks == -1) { + if constexpr (num_bits == 8) { + cp_async_wait<0>(); + __syncthreads(); + if (threadIdx.x / 32 < thread_n_blocks / 4) { + reinterpret_cast(&frag_s)[0] = sh_s[s_sh_rd + 0]; + reinterpret_cast(&frag_s)[1] = sh_s[s_sh_rd + 4]; + } + + } + else { + if (last) { + cp_async_wait<0>(); + __syncthreads(); + if (threadIdx.x / 32 < thread_n_blocks / 4) { + reinterpret_cast(&frag_s)[0] = sh_s[s_sh_rd + 0]; + reinterpret_cast(&frag_s)[1] = sh_s[s_sh_rd + 4]; + } + } + } + } + + // For 8-bit channelwise, we apply the scale before the global reduction + // that converts the fp32 results to fp16 (so that we avoid possible + // overflow in fp16) + if constexpr (!has_act_order && group_blocks == -1 && num_bits == 8) { + if (threadIdx.x / 32 < thread_n_blocks / 4) { +#pragma unroll + for (int i = 0; i < thread_m_blocks; i++) { +#pragma unroll + for (int j = 0; j < 4; j++) { + scale_float( + reinterpret_cast(&frag_c[i][j][0][0]), + frag_s[j / 2][2 * (j % 2) + 0]); + scale_float( + reinterpret_cast(&frag_c[i][j][0][2]), + frag_s[j / 2][2 * (j % 2) + 0]); + + scale_float( + reinterpret_cast(&frag_c[i][j][1][0]), + frag_s[j / 2][2 * (j % 2) + 1]); + scale_float( + reinterpret_cast(&frag_c[i][j][1][2]), + frag_s[j / 2][2 * (j % 2) + 1]); + } + } + } + } + + if (slice_count > 1) { // only globally reduce if there is more than one + // block in a slice + barrier_acquire(&locks[slice_col], slice_idx); + global_reduce(slice_idx == 0, last); + barrier_release(&locks[slice_col], last); + } + if (last) // only the last block in a slice actually writes the result + write_result(); + slice_row = 0; + slice_col_par++; + slice_col++; + init_slice(); + if (slice_iters) { + a_gl_rd = a_gl_stride * (threadIdx.x / a_gl_rd_delta_o) + + (threadIdx.x % a_gl_rd_delta_o); +#pragma unroll + for (int i = 0; i < b_sh_wr_iters; i++) + B_ptr[i] += b_sh_stride - b_gl_rd_delta_o * k_tiles; + if (slice_col == 0) { +#pragma unroll + for (int i = 0; i < b_sh_wr_iters; i++) B_ptr[i] -= b_gl_stride; + } + + // Update slice k/n for scales loading + if constexpr (has_act_order) { + slice_k_start = tb_k * slice_row; + slice_k_finish = slice_k_start + tb_k * slice_iters; + slice_k_start_shared_fetch = slice_k_start; + slice_n_offset = act_s_col_tb_stride * slice_col; + + } + else { + s_gl_rd = s_sh_stride * slice_col + threadIdx.x; + } + + start_pipes(); + } + } + } + } + + template shared + // fetch pipeline + const bool has_act_order, // whether act_order is enabled + const int group_blocks = -1 // number of consecutive 16x16 blocks + // with a separate quantization scale + > + __global__ void + Marlin_wrapper(const int4* __restrict__ A, // fp16 input matrix of shape mxk + const int4* __restrict__ B, // 4bit quantized weight matrix of shape kxn + int4* __restrict__ C, // fp16 output buffer of shape mxn + const int4* __restrict__ scales_ptr, // fp16 quantization scales of shape + // (k/groupsize)xn + const int* __restrict__ g_idx, // int32 group indices of shape k + int num_groups, // number of scale groups per output channel + const int* __restrict__ prob_m_ptr, // batch dimension m + int prob_n, // output dimension n + int prob_k, // reduction dimension k + int* locks // extra global storage for barrier synchronization + ) { + int prob_m = *prob_m_ptr; + prob_m = min(prob_m, 1024); + const int thread_m_blocks = min(div_ceil(prob_m, 16), template_thread_m_blocks); + if(prob_m > 16 * thread_m_blocks) + prob_m = (16 * thread_m_blocks) * div_ceil(prob_m, (16 * thread_m_blocks)); + /*if (blockIdx.x == 0 && threadIdx.x == 0) + printf("marlin prob_m %d\n", prob_m);*/ + if (thread_m_blocks == 1) { + Marlin( + A, B, C, scales_ptr, g_idx, num_groups, prob_m, prob_n, + prob_k, locks); + } + else if (thread_m_blocks == 2) { + Marlin( + A, B, C, scales_ptr, g_idx, num_groups, prob_m, prob_n, + prob_k, locks); + } + else if (thread_m_blocks == 3) { + Marlin( + A, B, C, scales_ptr, g_idx, num_groups, prob_m, prob_n, + prob_k, locks); + } + else if (thread_m_blocks == 4) { + Marlin( + A, B, C, scales_ptr, g_idx, num_groups, prob_m, prob_n, + prob_k, locks); + } + } + +#define __CALL_IF(NUM_BITS, THREAD_M_BLOCKS, THREAD_N_BLOCKS, THREAD_K_BLOCKS, \ + HAS_ACT_ORDER, GROUP_BLOCKS, NUM_THREADS) \ + else if (num_bits == NUM_BITS && thread_m_blocks == THREAD_M_BLOCKS && \ + thread_n_blocks == THREAD_N_BLOCKS && \ + thread_k_blocks == THREAD_K_BLOCKS && \ + has_act_order == HAS_ACT_ORDER && group_blocks == GROUP_BLOCKS && \ + num_threads == NUM_THREADS) { \ + cudaFuncSetAttribute( \ + Marlin_wrapper, \ + cudaFuncAttributeMaxDynamicSharedMemorySize, max_shared_mem); \ + Marlin_wrapper<<>>( \ + A_ptr, B_ptr, C_ptr, s_ptr, g_idx_ptr, num_groups, prob_m_ptr, prob_n, \ + prob_k, locks); \ + } + + typedef struct { + int thread_k; + int thread_n; + int num_threads; + } thread_config_t; + + typedef struct { + int max_m_blocks; + thread_config_t tb_cfg; + } exec_config_t; + + thread_config_t small_batch_thread_configs[] = { + // Ordered by priority + + // thread_k, thread_n, num_threads + {128, 128, 256}, + {64, 128, 128}, + {128, 64, 128}, + }; + + thread_config_t large_batch_thread_configs[] = { + // Ordered by priority + + // thread_k, thread_n, num_threads + {64, 256, 256}, + // {128, 128, 256}, + {64, 128, 128}, + {128, 64, 128}, + + }; + + int get_scales_cache_size(thread_config_t const& th_config, int prob_m, + int prob_n, int prob_k, int num_bits, int group_size, + bool has_act_order, bool is_k_full) { + bool cache_scales_chunk = has_act_order && !is_k_full; + + int tb_n = th_config.thread_n; + int tb_k = th_config.thread_k; + + // Get max scale groups per thread-block + int tb_groups; + if (group_size == -1) { + tb_groups = 1; + } + else if (group_size == 0) { + tb_groups = div_ceil(tb_k, 32); // Worst case is 32 group size + } + else { + tb_groups = div_ceil(tb_k, group_size); + } + + if (cache_scales_chunk) { + int load_groups = + tb_groups * pipe_stages * 2; // Chunk size is 2x pipeline over dim K + load_groups = max(load_groups, 32); // We load at least 32 scale groups + return load_groups * tb_n * 2; + + } + else { + int tb_scales = tb_groups * tb_n * 2; + + return tb_scales * pipe_stages; + } + } + + bool is_valid_cache_size(thread_config_t const& th_config, int max_m_blocks, + int prob_m, int prob_n, int prob_k, int num_bits, + int scales_cache_size, int max_shared_mem) { + int pack_factor = 32 / num_bits; + + // Get B size + int tb_k = th_config.thread_k; + int tb_n = th_config.thread_n; + + int b_size = (tb_k * tb_n / pack_factor) * 4; + + // Get A size + int m_blocks = div_ceil(prob_m, 16); + int tb_max_m = 16; + + // zbx: too ugly + // origin + /*while (true) { + if (m_blocks >= max_m_blocks) { + tb_max_m *= max_m_blocks; + break; + } + + max_m_blocks--; + if (max_m_blocks == 0) { + TORCH_CHECK(false, "Unexpected m_blocks = ", m_blocks); + } + }*/ + // refactor + tb_max_m *= std::min(m_blocks, max_m_blocks); + + int a_size = (tb_max_m * tb_k) * 2; + + float pipe_size = (a_size + b_size) * pipe_stages; + + TORCH_CHECK(max_shared_mem / 2 > scales_cache_size); // Sanity + return pipe_size < 0.95f * (max_shared_mem - scales_cache_size); + } + + bool is_valid_config(thread_config_t const& th_config, int max_m_blocks, + int prob_m, int prob_n, int prob_k, int num_bits, + int group_size, bool has_act_order, bool is_k_full, + int max_shared_mem) { + // Sanity + if (th_config.thread_k == -1 || th_config.thread_n == -1 || + th_config.num_threads == -1) { + return false; + } + + // Verify K/N are divisible by thread K/N + if (prob_k % th_config.thread_k != 0 || prob_n % th_config.thread_n != 0) { + return false; + } + + // Verify min for thread K/N + if (th_config.thread_n < min_thread_n || + th_config.thread_k < min_thread_k) { + return false; + } + + // num_threads must be at least 128 (= 4 warps) + if (th_config.num_threads < 128) { + return false; + } + + // Determine cache for scales + int scales_cache_size = + get_scales_cache_size(th_config, prob_m, prob_n, prob_k, num_bits, + group_size, has_act_order, is_k_full); + + // Check that pipeline fits into cache + if (!is_valid_cache_size(th_config, max_m_blocks, prob_m, prob_n, prob_k, + num_bits, scales_cache_size, max_shared_mem)) { + return false; + } + + return true; + } + + exec_config_t determine_thread_config(int prob_m, int prob_n, int prob_k, + int num_bits, int group_size, + bool has_act_order, bool is_k_full, + int max_shared_mem) { + int max_m_blocks = 4; + while (max_m_blocks > 0) { + if (prob_m <= 16) { + for (auto th_config : small_batch_thread_configs) { + if (is_valid_config(th_config, max_m_blocks, prob_m, prob_n, + prob_k, num_bits, group_size, has_act_order, + is_k_full, max_shared_mem)) { + return exec_config_t{ max_m_blocks, th_config }; + } + } + } + else { + for (auto th_config : large_batch_thread_configs) { + if (is_valid_config(th_config, max_m_blocks, prob_m, prob_n, + prob_k, num_bits, group_size, has_act_order, + is_k_full, max_shared_mem)) { + return exec_config_t{ max_m_blocks, th_config }; + } + } + } + + max_m_blocks--; // Process less M blocks per invocation to reduce cache + // usage + } + + return exec_config_t{ 0, {-1, -1, -1} }; + } + +#define CALL_IF(NUM_BITS, N_BLOCKS, K_BLOCKS, NUM_THREADS) \ + __CALL_IF(NUM_BITS, 1, N_BLOCKS, K_BLOCKS, false, 4, NUM_THREADS) \ + __CALL_IF(NUM_BITS, 2, N_BLOCKS, K_BLOCKS, false, 4, NUM_THREADS) \ + __CALL_IF(NUM_BITS, 3, N_BLOCKS, K_BLOCKS, false, 4, NUM_THREADS) \ + __CALL_IF(NUM_BITS, 4, N_BLOCKS, K_BLOCKS, false, 4, NUM_THREADS) \ + __CALL_IF(NUM_BITS, 1, N_BLOCKS, K_BLOCKS, false, 8, NUM_THREADS) \ + __CALL_IF(NUM_BITS, 2, N_BLOCKS, K_BLOCKS, false, 8, NUM_THREADS) \ + __CALL_IF(NUM_BITS, 3, N_BLOCKS, K_BLOCKS, false, 8, NUM_THREADS) \ + __CALL_IF(NUM_BITS, 4, N_BLOCKS, K_BLOCKS, false, 8, NUM_THREADS) + + template + void marlin_mm_f16i4(const void* A, const void* B, void* C, void* s, + void* g_idx, void* perm, void* a_tmp, int* prob_m_ptr, int prob_m, + int prob_n, int prob_k, void* workspace, int num_bits, + bool has_act_order, bool is_k_full, int num_groups, + int group_size, int dev, cudaStream_t stream, int thread_k, + int thread_n, int sms, int max_par) { + TORCH_CHECK(num_bits == 4 || num_bits == 8, + "num_bits must be 4 or 8. Got = ", num_bits); + TORCH_CHECK(prob_m > 0 && prob_n > 0 && prob_k > 0, "Invalid MNK = [", + prob_m, ", ", prob_n, ", ", prob_k, "]"); + + int tot_m = prob_m; + int tot_m_blocks = div_ceil(tot_m, 16); + int pad = 16 * tot_m_blocks - tot_m; + + if (sms == -1) { + cudaDeviceGetAttribute(&sms, cudaDevAttrMultiProcessorCount, dev); + } + + int max_shared_mem = 0; + cudaDeviceGetAttribute(&max_shared_mem, + cudaDevAttrMaxSharedMemoryPerBlockOptin, dev); + TORCH_CHECK(max_shared_mem > 0); + + // Set thread config + exec_config_t exec_cfg; + if (thread_k != -1 && thread_n != -1) { + // User-defined config + exec_cfg = exec_config_t{ + 4, thread_config_t{thread_k, thread_n, default_threads} }; + } + else { + // Auto config + exec_cfg = determine_thread_config(prob_m, prob_n, prob_k, num_bits, + group_size, has_act_order, is_k_full, + max_shared_mem); + } + + TORCH_CHECK( + exec_cfg.max_m_blocks > 0 && + is_valid_config(exec_cfg.tb_cfg, exec_cfg.max_m_blocks, prob_m, + prob_n, prob_k, num_bits, group_size, has_act_order, + is_k_full, max_shared_mem), + "Invalid thread config: max_m_blocks = ", exec_cfg.max_m_blocks, + ", thread_k = ", exec_cfg.tb_cfg.thread_k, + ", thread_n = ", exec_cfg.tb_cfg.thread_n, + ", num_threads = ", exec_cfg.tb_cfg.num_threads, " for MKN = [", prob_m, + ", ", prob_k, ", ", prob_n, "] and num_bits = ", num_bits, + ", group_size = ", group_size, ", has_act_order = ", has_act_order, + ", is_k_full = ", is_k_full, ", max_shared_mem = ", max_shared_mem); + + int num_threads = exec_cfg.tb_cfg.num_threads; + thread_k = exec_cfg.tb_cfg.thread_k; + thread_n = exec_cfg.tb_cfg.thread_n; + + int thread_k_blocks = thread_k / 16; + int thread_n_blocks = thread_n / 16; + + int blocks = sms; + + TORCH_CHECK(prob_n % thread_n == 0, "prob_n = ", prob_n, + " is not divisible by thread_n = ", thread_n); + TORCH_CHECK(prob_k % thread_k == 0, "prob_k = ", prob_k, + " is not divisible by thread_k = ", thread_k); + + int group_blocks = 0; + if (has_act_order) { + if (is_k_full) { + TORCH_CHECK(group_size != -1); + group_blocks = group_size / 16; + TORCH_CHECK(prob_k % group_blocks == 0, "prob_k = ", prob_k, + " is not divisible by group_blocks = ", group_blocks); + } + else { + TORCH_CHECK(group_size == 0); + group_blocks = 0; + } + + } + else { + if (group_size == -1) { + group_blocks = -1; + } + else { + group_blocks = group_size / 16; + TORCH_CHECK(prob_k % group_blocks == 0, "prob_k = ", prob_k, + " is not divisible by group_blocks = ", group_blocks); + } + } + + const int4* A_ptr = (const int4*)A; + const int4* B_ptr = (const int4*)B; + int4* C_ptr = (int4*)C; + const int4* s_ptr = (const int4*)s; + const int* g_idx_ptr = (const int*)g_idx; + const int* perm_ptr = (const int*)perm; + int4* a_tmp_ptr = (int4*)a_tmp; + + int* locks = (int*)workspace; + + if (has_act_order) { + // Permute A columns + int block_rows = div_ceil(prob_m, blocks); + permute_cols_kernel << > > ( + A_ptr, perm_ptr, a_tmp_ptr, prob_m, prob_k, block_rows); + A_ptr = a_tmp_ptr; + } + + // If we have a full K, then we can run the non-act-order version of Marlin + // (since the weight rows are reordered by increasing group ids, and by + // having a full K, we have full original groups) + if (is_k_full) { + has_act_order = false; + } + + // Main loop + for (int i = 0; i < tot_m_blocks; i += exec_cfg.max_m_blocks) { + int thread_m_blocks = tot_m_blocks - i; + prob_m = tot_m - 16 * i; + int par = 1; + if (thread_m_blocks > exec_cfg.max_m_blocks) { + // Note that parallel > 1 currently only works for inputs without + // any padding + par = (16 * thread_m_blocks - pad) / (16 * exec_cfg.max_m_blocks); + if (par > max_par) + par = max_par; + prob_m = (16 * exec_cfg.max_m_blocks) * par; + i += exec_cfg.max_m_blocks * (par - 1); + thread_m_blocks = exec_cfg.max_m_blocks; + } + + // Define kernel configurations +#define undefined_error \ + TORCH_CHECK(false, "Unsupported shapes: MNK = [" + str(prob_m) + ", " + \ + str(prob_n) + ", " + str(prob_k) + "]" + \ + ", has_act_order = " + str(has_act_order) + \ + ", num_groups = " + str(num_groups) + \ + ", group_size = " + str(group_size) + \ + ", thread_m_blocks = " + str(thread_m_blocks) + \ + ", thread_n_blocks = " + str(thread_n_blocks) + \ + ", thread_k_blocks = " + str(thread_k_blocks)); + + /* std::cout << "MNK = [" + str(prob_m) + ", " + \ + str(prob_n) + ", " + str(prob_k) + "]" + \ + ", has_act_order = " + str(has_act_order) + \ + ", num_groups = " + str(num_groups) + \ + ", group_size = " + str(group_size) + \ + ", thread_m_blocks = " + str(thread_m_blocks) + \ + ", thread_n_blocks = " + str(thread_n_blocks) + \ + ", thread_k_blocks = " + str(thread_k_blocks) << std::endl;*/ + + /*if (false) { + } + // CALL_IF(4, 32, 2, 256) + // CALL_IF(4, 16, 4, 256) + __CALL_IF(4, 1, 16, 4, false, 4, 256) + __CALL_IF(4, 2, 16, 4, false, 4, 256) + // CALL_IF(4, 8, 8, 256) + __CALL_IF(4, 1, 8, 8, false, 4, 256) + __CALL_IF(4, 2, 8, 8, false, 4, 256) + // CALL_IF(4, 16, 4, 128) + __CALL_IF(4, 1, 16, 4, false, 4, 128) + __CALL_IF(4, 2, 16, 4, false, 4, 128) + // CALL_IF(4, 8, 8, 128) + __CALL_IF(4, 1, 8, 8, false, 4, 128) + __CALL_IF(4, 2, 8, 8, false, 4, 128) + else {undefined_error}*/ + + if (num_bits == 4 && num_threads == 256) + { + if (false) { + } + CALL_IF(4, 32, 2, 256) + CALL_IF(4, 16, 4, 256) + CALL_IF(4, 8, 8, 256) + else { + undefined_error + } + } + else if (num_bits == 4 && num_threads == 128) + { + if (false) { + } + CALL_IF(4, 8, 4, 128) + CALL_IF(4, 16, 4, 128) + CALL_IF(4, 4, 8, 128) + else { + undefined_error + } + } + // else if (num_bits == 8 && num_threads == 256) + // { + // if (false) { + // } + // CALL_IF(8, 32, 2, 256) + // CALL_IF(8, 16, 4, 256) + // CALL_IF(8, 8, 8, 256) + // else { + // undefined_error + // } + // } + // else if (num_bits == 8 && num_threads == 128) + // { + // if (false) { + // } + // CALL_IF(8, 8, 4, 128) + // CALL_IF(8, 16, 4, 128) + // CALL_IF(8, 4, 8, 128) + // else { + // undefined_error + // } + // } + else { + undefined_error + } + + A_ptr += 16 * thread_m_blocks * (prob_k / 8) * par; + C_ptr += 16 * thread_m_blocks * (prob_n / 8) * par; + } + } + +} // namespace gptq_marlin + +torch::Tensor gptq_marlin_gemm(torch::Tensor& a, torch::Tensor& b_q_weight, + torch::Tensor& b_scales, torch::Tensor& g_idx, + torch::Tensor& perm, torch::Tensor& workspace, + int64_t num_bits, torch::Tensor size_m_tensor, int64_t size_m, int64_t size_n, + int64_t size_k, int sms, bool is_k_full) { + const at::cuda::OptionalCUDAGuard device_guard(device_of(a)); + // Verify num_bits + TORCH_CHECK(num_bits == 4 || num_bits == 8, + "num_bits must be 4 or 8. Got = ", num_bits); + int pack_factor = 32 / num_bits; + + // Verify A + TORCH_CHECK(a.size(0) == size_m, "Shape mismatch: a.size(0) = ", a.size(0), + ", size_m = ", size_m); + TORCH_CHECK(a.size(1) == size_k, "Shape mismatch: a.size(1) = ", a.size(1), + ", size_k = ", size_k); + + // Verify B + TORCH_CHECK(size_k % gptq_marlin::tile_size == 0, "size_k = ", size_k, + " is not divisible by tile_size = ", gptq_marlin::tile_size); + TORCH_CHECK((size_k / gptq_marlin::tile_size) == b_q_weight.size(0), + "Shape mismatch: b_q_weight.size(0) = ", b_q_weight.size(0), + ", size_k = ", size_k, + ", tile_size = ", gptq_marlin::tile_size); + TORCH_CHECK(b_q_weight.size(1) % gptq_marlin::tile_size == 0, + "b_q_weight.size(1) = ", b_q_weight.size(1), + " is not divisible by tile_size = ", gptq_marlin::tile_size); + int actual_size_n = + (b_q_weight.size(1) / gptq_marlin::tile_size) * pack_factor; + TORCH_CHECK(size_n == actual_size_n, "size_n = ", size_n, + ", actual_size_n = ", actual_size_n); + + // Verify device and strides + TORCH_CHECK(a.device().is_cuda(), "A is not on GPU"); + TORCH_CHECK(a.is_contiguous(), "A is not contiguous"); + + TORCH_CHECK(b_q_weight.device().is_cuda(), "b_q_weight is not on GPU"); + TORCH_CHECK(b_q_weight.is_contiguous(), "b_q_weight is not contiguous"); + + TORCH_CHECK(b_scales.device().is_cuda(), "b_scales is not on GPU"); + TORCH_CHECK(b_scales.is_contiguous(), "b_scales is not contiguous"); + + TORCH_CHECK(g_idx.device().is_cuda(), "g_idx is not on GPU"); + TORCH_CHECK(g_idx.is_contiguous(), "g_idx is not contiguous"); + + TORCH_CHECK(perm.device().is_cuda(), "perm is not on GPU"); + TORCH_CHECK(perm.is_contiguous(), "perm is not contiguous"); + + // Alloc buffers + auto options = torch::TensorOptions().dtype(a.dtype()).device(a.device()); + torch::Tensor c = torch::empty({ size_m, size_n }, options); + torch::Tensor a_tmp = torch::empty({ size_m, size_k }, options); + + // thread_k: `k` size of a thread_tile in `weights` (can usually be left as + // auto -1) + int thread_k = -1; + // thread_n: `n` size of a thread_tile in `weights` (can usually be left as + // auto -1) + int thread_n = -1; + // sms: number of SMs to use for the kernel (can usually be left as auto -1) + // int sms = -1; //zbx + + // Verify g_idx and perm + TORCH_CHECK((g_idx.size(0) == 0 && perm.size(0) == 0) || + (g_idx.size(0) == size_k && perm.size(0) == size_k), + "Unexpected g_idx.size(0) = ", g_idx.size(0), + " and perm.size(0) = ", perm.size(0), + ", where size_k = ", size_k); + + // Detect groupsize and act_order + int num_groups = -1; + int group_size = -1; + bool has_act_order = g_idx.size(0) != 0; + + int b_rank = b_scales.sizes().size(); + TORCH_CHECK(b_rank == 2, "b_scales rank = ", b_rank, " is not 2"); + TORCH_CHECK(b_scales.size(1) == size_n, + "b_scales dim 1 = ", b_scales.size(1), + " is not size_n = ", size_n); + num_groups = b_scales.size(0); + + if (has_act_order) { + if (is_k_full) { + TORCH_CHECK(num_groups > 1, + "For act_order, num_groups must be > 1"); + TORCH_CHECK(size_k % num_groups == 0, "size_k = ", size_k, + ", is not divisible by num_groups = ", num_groups); + group_size = size_k / num_groups; + } + else { + group_size = 0; + } + + } + else { + if (num_groups > 1) { + TORCH_CHECK( + size_k % num_groups == 0, "size_k = ", size_k, + ", is not divisible by b_scales.size(0) = ", b_scales.size(0)); + group_size = size_k / num_groups; + } + else { + group_size = -1; + } + } + + // Verify workspace size + TORCH_CHECK( + size_n % gptq_marlin::min_thread_n == 0, "size_n = ", size_n, + ", is not divisible by min_thread_n = ", gptq_marlin::min_thread_n); + int min_workspace_size = + (size_n / gptq_marlin::min_thread_n) * gptq_marlin::max_par; + TORCH_CHECK(workspace.numel() >= min_workspace_size, + "workspace.numel = ", workspace.numel(), + " is below min_workspace_size = ", min_workspace_size); + + int dev = a.get_device(); + if (a.scalar_type() == at::ScalarType::Half) { + gptq_marlin::marlin_mm_f16i4( + a.data_ptr(), b_q_weight.data_ptr(), + c.data_ptr(), b_scales.data_ptr(), + g_idx.data_ptr(), perm.data_ptr(), a_tmp.data_ptr(), + size_m_tensor.data_ptr(), + size_m, size_n, size_k, workspace.data_ptr(), num_bits, + has_act_order, is_k_full, num_groups, group_size, dev, + at::cuda::getCurrentCUDAStream(dev), thread_k, thread_n, sms, + gptq_marlin::max_par); + } + else if (a.scalar_type() == at::ScalarType::BFloat16) { + gptq_marlin::marlin_mm_f16i4( + a.data_ptr(), b_q_weight.data_ptr(), + c.data_ptr(), b_scales.data_ptr(), + g_idx.data_ptr(), perm.data_ptr(), a_tmp.data_ptr(), + size_m_tensor.data_ptr(), + size_m, size_n, size_k, workspace.data_ptr(), num_bits, + has_act_order, is_k_full, num_groups, group_size, dev, + at::cuda::getCurrentCUDAStream(dev), thread_k, thread_n, sms, + gptq_marlin::max_par); + } + else { + TORCH_CHECK(false, + "gpt_marlin_gemm only supports bfloat16 and float16"); + } + + return c; +} + +#endif \ No newline at end of file diff --git a/archive/kt-sft/csrc/custom_marlin/gptq_marlin/gptq_marlin.cuh b/archive/kt-sft/csrc/custom_marlin/gptq_marlin/gptq_marlin.cuh new file mode 100644 index 0000000..5b4b059 --- /dev/null +++ b/archive/kt-sft/csrc/custom_marlin/gptq_marlin/gptq_marlin.cuh @@ -0,0 +1,76 @@ +// Adapted from +// https://github.com/vllm-project/vllm/tree/main/csrc/quantization/gptq_marlin +// Copyrigth 2024 The vLLM team. +// Copyright (c) 2024 by KVCache.AI, All Rights Reserved. +#pragma once + +#include + +#include +#include +#include +#include +#include +#include + +namespace gptq_marlin { + +// 8 warps are a good choice since every SM has 4 schedulers and having more +// than 1 warp per schedule allows some more latency hiding. At the same time, +// we want relatively few warps to have many registers per warp and small tiles. +static constexpr int default_threads = 256; + +static constexpr int pipe_stages = + 4; // 4 pipeline stages fit into shared memory + +static constexpr int min_thread_n = 64; +static constexpr int min_thread_k = 64; + +static constexpr int tile_size = 16; +static constexpr int max_par = 16; + +template struct Vec { + T elems[n]; + __device__ T &operator[](int i) { return elems[i]; } +}; + +using I4 = Vec; + +constexpr int div_ceil(int a, int b) { return (a + b - 1) / b; } + +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800 +// No support for async +#else + +__device__ inline void cp_async4_pred(void *smem_ptr, const void *glob_ptr, + bool pred = true) { + const int BYTES = 16; + uint32_t smem = static_cast(__cvta_generic_to_shared(smem_ptr)); + asm volatile("{\n" + " .reg .pred p;\n" + " setp.ne.b32 p, %0, 0;\n" + " @p cp.async.cg.shared.global [%1], [%2], %3;\n" + "}\n" ::"r"((int)pred), + "r"(smem), "l"(glob_ptr), "n"(BYTES)); +} + +__device__ inline void cp_async4(void *smem_ptr, const void *glob_ptr) { + const int BYTES = 16; + uint32_t smem = static_cast(__cvta_generic_to_shared(smem_ptr)); + asm volatile("{\n" + " cp.async.cg.shared.global [%0], [%1], %2;\n" + "}\n" ::"r"(smem), + "l"(glob_ptr), "n"(BYTES)); +} + +__device__ inline void cp_async_fence() { + asm volatile("cp.async.commit_group;\n" ::); +} + +template __device__ inline void cp_async_wait() { + asm volatile("cp.async.wait_group %0;\n" ::"n"(n)); +} + +#endif + +} // namespace gptq_marlin \ No newline at end of file diff --git a/archive/kt-sft/csrc/custom_marlin/gptq_marlin/gptq_marlin_dtypes.cuh b/archive/kt-sft/csrc/custom_marlin/gptq_marlin/gptq_marlin_dtypes.cuh new file mode 100644 index 0000000..3e8c3ca --- /dev/null +++ b/archive/kt-sft/csrc/custom_marlin/gptq_marlin/gptq_marlin_dtypes.cuh @@ -0,0 +1,77 @@ +// Adapted from +// https://github.com/vllm-project/vllm/tree/main/csrc/quantization/gptq_marlin +// Copyrigth 2024 The vLLM team. +// Copyright (c) 2024 by KVCache.AI, All Rights Reserved. +#ifndef _data_types_cuh +#define _data_types_cuh +#include "gptq_marlin.cuh" +#include +#include + +namespace gptq_marlin { + +template class ScalarType {}; + +template <> class ScalarType { + public: + using scalar_t = half; + using scalar_t2 = half2; + + // Matrix fragments for tensor core instructions; their precise layout is + // documented here: + // https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#matrix-fragments-for-mma-m16n8k16-with-floating-point-type + using FragA = Vec; + using FragB = Vec; + using FragC = Vec; + using FragS = Vec; + + static __device__ float inline num2float(const half x) { + return __half2float(x); + } + + static __device__ half2 inline num2num2(const half x) { + return __half2half2(x); + } + + static __device__ half2 inline nums2num2(const half x1, const half x2) { + return __halves2half2(x1, x2); + } + + static __host__ __device__ half inline float2num(const float x) { + return __float2half(x); + } +}; + +template <> class ScalarType { + public: + using scalar_t = nv_bfloat16; + using scalar_t2 = nv_bfloat162; + + using FragA = Vec; + using FragB = Vec; + using FragC = Vec; + using FragS = Vec; + +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800 + static __device__ float inline num2float(const nv_bfloat16 x) { + return __bfloat162float(x); + } + + static __device__ nv_bfloat162 inline num2num2(const nv_bfloat16 x) { + return __bfloat162bfloat162(x); + } + + static __device__ nv_bfloat162 inline nums2num2(const nv_bfloat16 x1, + const nv_bfloat16 x2) { + return __halves2bfloat162(x1, x2); + } + + static __host__ __device__ nv_bfloat16 inline float2num(const float x) { + return __float2bfloat16(x); + } +#endif +}; + +} // namespace gptq_marlin + +#endif \ No newline at end of file diff --git a/archive/kt-sft/csrc/custom_marlin/gptq_marlin/gptq_marlin_repack.cu b/archive/kt-sft/csrc/custom_marlin/gptq_marlin/gptq_marlin_repack.cu new file mode 100644 index 0000000..4adcbd5 --- /dev/null +++ b/archive/kt-sft/csrc/custom_marlin/gptq_marlin/gptq_marlin_repack.cu @@ -0,0 +1,350 @@ +#include "gptq_marlin.cuh" + +namespace gptq_marlin { + +static constexpr int repack_stages = 8; + +static constexpr int repack_threads = 256; + +static constexpr int tile_k_size = tile_size; +static constexpr int tile_n_size = tile_k_size * 4; + +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800 + +template +__global__ void marlin_repack_kernel( + uint32_t const* __restrict__ b_q_weight_ptr, + uint32_t const* __restrict__ perm_ptr, uint32_t* __restrict__ out_ptr, + int size_k, int size_n) {} + +} // namespace gptq_marlin + +torch::Tensor gptq_marlin_repack(torch::Tensor& b_q_weight, torch::Tensor& perm, + int64_t size_k, int64_t size_n, + int64_t num_bits) { + TORCH_CHECK_NOT_IMPLEMENTED( + false, "marlin_repack_from_gptq(..) requires CUDA_ARCH >= 8.0"); + return torch::empty({1, 1}); +} + +#else + +template +__global__ void marlin_repack_kernel( + uint32_t const* __restrict__ b_q_weight_ptr, + uint32_t const* __restrict__ perm_ptr, uint32_t* __restrict__ out_ptr, + int size_k, int size_n) { + constexpr int pack_factor = 32 / num_bits; + + int k_tiles = size_k / tile_k_size; + int n_tiles = size_n / tile_n_size; + int block_k_tiles = div_ceil(k_tiles, gridDim.x); + + int start_k_tile = blockIdx.x * block_k_tiles; + if (start_k_tile >= k_tiles) { + return; + } + + int finish_k_tile = min(start_k_tile + block_k_tiles, k_tiles); + + // Wait until the next thread tile has been loaded to shared memory. + auto wait_for_stage = [&]() { + // We only have `stages - 2` active fetches since we are double buffering + // and can only issue the next fetch when it is guaranteed that the previous + // shared memory load is fully complete (as it may otherwise be + // overwritten). + cp_async_wait(); + __syncthreads(); + }; + + extern __shared__ int4 sh[]; + + constexpr int perm_size = tile_k_size / 4; + + int4* sh_perm_ptr = sh; + int4* sh_pipe_ptr = sh_perm_ptr; + if constexpr (has_perm) { + sh_pipe_ptr += perm_size; + } + + constexpr int tile_ints = tile_k_size / pack_factor; + + constexpr int stage_n_threads = tile_n_size / 4; + constexpr int stage_k_threads = has_perm ? tile_k_size : tile_ints; + constexpr int stage_size = stage_k_threads * stage_n_threads; + + auto load_perm_to_shared = [&](int k_tile_id) { + int first_k_int4 = (k_tile_id * tile_k_size) / 4; + + int4 const* perm_int4_ptr = reinterpret_cast(perm_ptr); + + if (threadIdx.x < perm_size) { + sh_perm_ptr[threadIdx.x] = perm_int4_ptr[first_k_int4 + threadIdx.x]; + } + __syncthreads(); + }; + + auto fetch_to_shared = [&](int pipe, int k_tile_id, int n_tile_id) { + if (n_tile_id >= n_tiles) { + cp_async_fence(); + return; + } + + int first_n = n_tile_id * tile_n_size; + + int4* sh_ptr = sh_pipe_ptr + stage_size * pipe; + + if constexpr (has_perm) { + if (threadIdx.x < stage_size) { + int k_id = threadIdx.x / stage_n_threads; + int n_id = threadIdx.x % stage_n_threads; + + uint32_t const* sh_perm_int_ptr = + reinterpret_cast(sh_perm_ptr); + + int src_k = sh_perm_int_ptr[k_id]; + int src_k_packed = src_k / pack_factor; + + cp_async4( + &sh_ptr[k_id * stage_n_threads + n_id], + reinterpret_cast(&( + b_q_weight_ptr[src_k_packed * size_n + first_n + (n_id * 4)]))); + } + + } else { + if (threadIdx.x < stage_size) { + int k_id = threadIdx.x / stage_n_threads; + int n_id = threadIdx.x % stage_n_threads; + + int first_k = k_tile_id * tile_k_size; + int first_k_packed = first_k / pack_factor; + + cp_async4(&sh_ptr[k_id * stage_n_threads + n_id], + reinterpret_cast( + &(b_q_weight_ptr[(first_k_packed + k_id) * size_n + + first_n + (n_id * 4)]))); + } + } + + cp_async_fence(); + }; + + auto repack_tile = [&](int pipe, int k_tile_id, int n_tile_id) { + if (n_tile_id >= n_tiles) { + return; + } + + int warp_id = threadIdx.x / 32; + int th_id = threadIdx.x % 32; + + if (warp_id >= 4) { + return; + } + + int tc_col = th_id / 4; + int tc_row = (th_id % 4) * 2; + + constexpr int tc_offsets[4] = {0, 1, 8, 9}; + + int cur_n = warp_id * 16 + tc_col; + + constexpr int sh_stride = 64; + constexpr uint32_t mask = (1 << num_bits) - 1; + + int4* sh_stage_ptr = sh_pipe_ptr + stage_size * pipe; + uint32_t* sh_stage_int_ptr = reinterpret_cast(sh_stage_ptr); + + uint32_t* sh_perm_int_ptr = reinterpret_cast(sh_perm_ptr); + + uint32_t vals[8]; + + if constexpr (has_perm) { + for (int i = 0; i < 4; i++) { + int k_idx = tc_row + tc_offsets[i]; + + uint32_t src_k = sh_perm_int_ptr[k_idx]; + uint32_t src_k_pos = src_k % pack_factor; + + uint32_t b1_val = sh_stage_int_ptr[k_idx * sh_stride + cur_n]; + uint32_t b1_cur_val = (b1_val >> (src_k_pos * num_bits)) & mask; + + uint32_t b2_val = sh_stage_int_ptr[k_idx * sh_stride + cur_n + 8]; + uint32_t b2_cur_val = (b2_val >> (src_k_pos * num_bits)) & mask; + + vals[i] = b1_cur_val; + vals[4 + i] = b2_cur_val; + } + + } else { + uint32_t b1_vals[tile_ints]; + uint32_t b2_vals[tile_ints]; + + #pragma unroll + for (int i = 0; i < tile_ints; i++) { + b1_vals[i] = sh_stage_int_ptr[cur_n + sh_stride * i]; + b2_vals[i] = sh_stage_int_ptr[cur_n + 8 + sh_stride * i]; + } + + #pragma unroll + for (int i = 0; i < 4; i++) { + int cur_elem = tc_row + tc_offsets[i]; + int cur_int = cur_elem / pack_factor; + int cur_pos = cur_elem % pack_factor; + + vals[i] = (b1_vals[cur_int] >> (cur_pos * num_bits)) & mask; + vals[4 + i] = (b2_vals[cur_int] >> (cur_pos * num_bits)) & mask; + } + } + + constexpr int tile_size = tile_k_size * tile_n_size / pack_factor; + int out_offset = (k_tile_id * n_tiles + n_tile_id) * tile_size; + + // Result of: + // https://github.com/NVIDIA/FasterTransformer/blob/main/src/fastertransformer/cutlass_extensions/include/cutlass_extensions/interleaved_numeric_conversion.h + if constexpr (num_bits == 4) { + constexpr int pack_idx[8] = {0, 2, 4, 6, 1, 3, 5, 7}; + + uint32_t res = 0; + #pragma unroll + for (int i = 0; i < 8; i++) { + res |= vals[pack_idx[i]] << (i * 4); + } + + out_ptr[out_offset + th_id * 4 + warp_id] = res; + + } else { + constexpr int pack_idx[4] = {0, 2, 1, 3}; + + uint32_t res1 = 0; + uint32_t res2 = 0; + #pragma unroll + for (int i = 0; i < 4; i++) { + res1 |= vals[pack_idx[i]] << (i * 8); + res2 |= vals[4 + pack_idx[i]] << (i * 8); + } + + out_ptr[out_offset + th_id * 8 + (warp_id * 2) + 0] = res1; + out_ptr[out_offset + th_id * 8 + (warp_id * 2) + 1] = res2; + } + }; + + auto start_pipes = [&](int k_tile_id, int n_tile_id) { + #pragma unroll + for (int pipe = 0; pipe < repack_stages - 1; pipe++) { + fetch_to_shared(pipe, k_tile_id, n_tile_id + pipe); + } + + wait_for_stage(); + }; + #pragma unroll + for (int k_tile_id = start_k_tile; k_tile_id < finish_k_tile; k_tile_id++) { + int n_tile_id = 0; + + if constexpr (has_perm) { + load_perm_to_shared(k_tile_id); + } + + start_pipes(k_tile_id, n_tile_id); + + while (n_tile_id < n_tiles) { + #pragma unroll + for (int pipe = 0; pipe < repack_stages; pipe++) { + fetch_to_shared((pipe + repack_stages - 1) % repack_stages, k_tile_id, + n_tile_id + pipe + repack_stages - 1); + repack_tile(pipe, k_tile_id, n_tile_id + pipe); + wait_for_stage(); + } + n_tile_id += repack_stages; + } + } +} + +} // namespace gptq_marlin + + #define CALL_IF(NUM_BITS, HAS_PERM) \ + else if (num_bits == NUM_BITS && has_perm == HAS_PERM) { \ + cudaFuncSetAttribute( \ + gptq_marlin::marlin_repack_kernel, \ + cudaFuncAttributeMaxDynamicSharedMemorySize, max_shared_mem); \ + gptq_marlin::marlin_repack_kernel \ + <<>>( \ + b_q_weight_ptr, perm_ptr, out_ptr, size_k, size_n); \ + } + +torch::Tensor gptq_marlin_repack(torch::Tensor& b_q_weight, torch::Tensor& perm, + int64_t size_k, int64_t size_n, + int64_t num_bits) { + // Verify compatibility with marlin tile of 16x64 + TORCH_CHECK(size_k % gptq_marlin::tile_k_size == 0, "size_k = ", size_k, + " is not divisible by tile_k_size = ", gptq_marlin::tile_k_size); + TORCH_CHECK(size_n % gptq_marlin::tile_n_size == 0, "size_n = ", size_n, + " is not divisible by tile_n_size = ", gptq_marlin::tile_n_size); + + TORCH_CHECK(num_bits == 4 || num_bits == 8, + "num_bits must be 4 or 8. Got = ", num_bits); + int const pack_factor = 32 / num_bits; + + // Verify B + TORCH_CHECK((size_k / pack_factor) == b_q_weight.size(0), + "Shape mismatch: b_q_weight.size(0) = ", b_q_weight.size(0), + ", size_k = ", size_k, ", pack_factor = ", pack_factor); + TORCH_CHECK(b_q_weight.size(1) == size_n, + "b_q_weight.size(1) = ", b_q_weight.size(1), + " is not size_n = ", size_n); + + // Verify device and strides + TORCH_CHECK(b_q_weight.device().is_cuda(), "b_q_weight is not on GPU"); + TORCH_CHECK(b_q_weight.is_contiguous(), "b_q_weight is not contiguous"); + TORCH_CHECK(b_q_weight.dtype() == at::kInt, "b_q_weight type is not kInt"); + + TORCH_CHECK(perm.device().is_cuda(), "perm is not on GPU"); + TORCH_CHECK(perm.is_contiguous(), "perm is not contiguous"); + TORCH_CHECK(perm.dtype() == at::kInt, "perm type is not at::kInt"); + + // Alloc buffers + const at::cuda::OptionalCUDAGuard device_guard(device_of(b_q_weight)); + auto options = torch::TensorOptions() + .dtype(b_q_weight.dtype()) + .device(b_q_weight.device()); + torch::Tensor out = + torch::empty({size_k / gptq_marlin::tile_size, + size_n * gptq_marlin::tile_size / pack_factor}, + options); + + // Detect if there is act_order + bool has_perm = perm.size(0) != 0; + + // Get ptrs + uint32_t const* b_q_weight_ptr = + reinterpret_cast(b_q_weight.data_ptr()); + uint32_t const* perm_ptr = reinterpret_cast(perm.data_ptr()); + uint32_t* out_ptr = reinterpret_cast(out.data_ptr()); + + // Get dev info + int dev = b_q_weight.get_device(); + cudaStream_t stream = at::cuda::getCurrentCUDAStream(dev); + int blocks; + cudaDeviceGetAttribute(&blocks, cudaDevAttrMultiProcessorCount, dev); + + int max_shared_mem = 0; + cudaDeviceGetAttribute(&max_shared_mem, + cudaDevAttrMaxSharedMemoryPerBlockOptin, dev); + TORCH_CHECK(max_shared_mem > 0); + + if (false) { + } + CALL_IF(4, false) + CALL_IF(4, true) + CALL_IF(8, false) + CALL_IF(8, true) + else { + TORCH_CHECK(false, "Unsupported repack config: num_bits = ", num_bits, + ", has_perm = ", has_perm); + } + + return out; +} + +#endif \ No newline at end of file diff --git a/archive/kt-sft/csrc/custom_marlin/gptq_marlin/ops.h b/archive/kt-sft/csrc/custom_marlin/gptq_marlin/ops.h new file mode 100644 index 0000000..b3327d3 --- /dev/null +++ b/archive/kt-sft/csrc/custom_marlin/gptq_marlin/ops.h @@ -0,0 +1,24 @@ +/** + * @Description : + * @Author : Azure + * @Date : 2024-07-22 09:27:55 + * @Version : 1.0.0 + * @LastEditors : Azure + * @LastEditTime : 2024-07-26 08:35:00 + * @Copyright (c) 2024 by KVCache.AI, All Rights Reserved. + **/ +#pragma once + +#include +#include +#include + +torch::Tensor gptq_marlin_gemm(torch::Tensor &a, torch::Tensor &b_q_weight, + torch::Tensor &b_scales, torch::Tensor &g_idx, + torch::Tensor &perm, torch::Tensor &workspace, + int64_t num_bits, torch::Tensor size_m_tensor, int64_t size_m, int64_t size_n, + int64_t size_k, int sms, bool is_k_full); + +torch::Tensor gptq_marlin_repack(torch::Tensor& b_q_weight, torch::Tensor&perm, + int64_t size_k, int64_t size_n, + int64_t num_bits); \ No newline at end of file diff --git a/archive/kt-sft/csrc/custom_marlin/setup.py b/archive/kt-sft/csrc/custom_marlin/setup.py new file mode 100644 index 0000000..ec710ca --- /dev/null +++ b/archive/kt-sft/csrc/custom_marlin/setup.py @@ -0,0 +1,25 @@ +from setuptools import setup, Extension +from torch.utils import cpp_extension +from torch.utils.cpp_extension import BuildExtension, CUDAExtension +setup( + name='vLLMMarlin', + ext_modules=[ + CUDAExtension( + 'vLLMMarlin', [ + #'custom_gguf/dequant.cu', + 'binding.cpp', + 'gptq_marlin/gptq_marlin.cu', + 'gptq_marlin/gptq_marlin_repack.cu', + ], + extra_compile_args={ + 'cxx': ['-O3'], + 'nvcc': [ + '-O3', + '--use_fast_math', + '-Xcompiler', '-fPIC', + ] + }, + ) + ], + cmdclass={'build_ext': BuildExtension} +) \ No newline at end of file diff --git a/archive/kt-sft/csrc/custom_marlin/test_cuda_graph.py b/archive/kt-sft/csrc/custom_marlin/test_cuda_graph.py new file mode 100644 index 0000000..0024082 --- /dev/null +++ b/archive/kt-sft/csrc/custom_marlin/test_cuda_graph.py @@ -0,0 +1,335 @@ +import csv +import torch +import torch.nn as nn +import vLLMMarlin +torch.set_grad_enabled(False) +from utils.marlin_utils import ( + MarlinWorkspace, + marlin_quantize, + GPTQ_MARLIN_MIN_THREAD_N, + GPTQ_MARLIN_MIN_THREAD_K, + GPTQ_MARLIN_MAX_PARALLEL, +) + +def setup_seed(seed): + torch.manual_seed(seed) + torch.cuda.manual_seed_all(seed) + +setup_seed(20241223) + +torch.set_grad_enabled(False) +torch.set_default_dtype(torch.bfloat16) +global_dtype=torch.bfloat16 +global_device=torch.device("cuda",0) +global_num_cases:int=int(50) +torch.cuda.set_device(0) +torch.backends.cudnn.enabled =True +torch.backends.cudnn.benchmark = True + +max_batch_size = 512 +max_tp = 8 +L2_size = 73728 * 1024 + +def get_usable_mem(): + properties = torch.cuda.get_device_properties(global_device) + #print(f"Total memory: {properties.total_memory / (1024 ** 3):.2f} GB") + allocated_memory = torch.cuda.memory_allocated(global_device) + #print(f"Currently allocated memory: {allocated_memory / (1024 ** 2):.2f} MB") + reserved_memory = torch.cuda.memory_reserved(global_device) + #print(f"Currently reserved memory: {reserved_memory / (1024 ** 2):.2f} MB") + return properties.total_memory - 512 * 1024 ** 2 - allocated_memory# - reserved_memory + +def exp_range(start, stop, step = 2): + now = start + while now <= stop: + yield now + now *= step + +def timing(func, iters, epochs=100): + #warmup + for idx in range(iters): + func(idx) + + torch.cuda.synchronize() + cuda_graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(cuda_graph): + for idx in range(iters): + func(idx) + + for _ in range(2000): + cuda_graph.replay() + + start_event = torch.cuda.Event(enable_timing=True) + end_event = torch.cuda.Event(enable_timing=True) + stream = torch.cuda.Stream() + torch.cuda.synchronize() + #with torch.cuda.stream(stream): + start_event.record() + for _ in range(10): + cuda_graph.replay() + end_event.record() + torch.cuda.synchronize() + elapsed_time_ms0 = start_event.elapsed_time(end_event) + + start_event = torch.cuda.Event(enable_timing=True) + end_event = torch.cuda.Event(enable_timing=True) + torch.cuda.synchronize() + #with torch.cuda.stream(stream): + start_event.record() + for _ in range(epochs+10): + cuda_graph.replay() + end_event.record() + torch.cuda.synchronize() + elapsed_time_ms = start_event.elapsed_time(end_event) - elapsed_time_ms0 + + #print(elapsed_time_ms0, elapsed_time_ms) + return elapsed_time_ms/iters/epochs + +class LinearMarlin(nn.Linear): + marlin_q_w: torch.Tensor + marlin_s: torch.Tensor + g_idx: torch.Tensor + sort_indices: torch.Tensor + has_bias: bool + def __init__( + self, + in_features, + out_features, + bias = False, + device: str = "cuda", + num_bits: int = 4, # 4-bit/8-bit is supported + group_size: int = 64, # -1, 32, 64, 128 + act_order: bool = False, + is_k_full=True, + sms = -1, # sms in GPU + **kwargs, + ): + self.padding = False + assert device.lower() != "cpu", "Marlin quantized linear only supports GPU device" + if in_features%GPTQ_MARLIN_MIN_THREAD_K!=0 or out_features%GPTQ_MARLIN_MIN_THREAD_K!=0: + #print(f"warning!, in_features={in_features} or out_features={out_features} is undivisible by GPTQ_MARLIN_MIN_THREAD_K={GPTQ_MARLIN_MIN_THREAD_K} and GPTQ_MARLIN_MIN_THREAD_N={GPTQ_MARLIN_MIN_THREAD_N}, padding") + self.padding = True + self.orin_in_features = in_features + self.orin_out_features = out_features + in_features = (in_features+GPTQ_MARLIN_MIN_THREAD_K-1)//GPTQ_MARLIN_MIN_THREAD_K*GPTQ_MARLIN_MIN_THREAD_K + out_features = (out_features+GPTQ_MARLIN_MIN_THREAD_N-1)//GPTQ_MARLIN_MIN_THREAD_N*GPTQ_MARLIN_MIN_THREAD_N + #print(f"After padding: in_features={in_features}, out_features={out_features}") + + + super().__init__(in_features, out_features, bias, device) + self.has_bias = bias + self.device = device + self.num_bits = num_bits + self.group_size = group_size + self.act_order = act_order + # TODO: optimize every shape GEMM + + blocks_k, blocks_n = in_features//128, out_features//128 + + self.sms = sms + + self.is_k_full = is_k_full + + self.weight.requires_grad = False + self.weight.t_() + # Pack Marlin linear + #w_ref, marlin_q_w, marlin_s, g_idx, sort_indices, _ = marlin_quantize( + # self.weight, self.num_bits, self.group_size, self.act_order + #) + marlin_q_w = torch.randint(int(-1e9), int(1e9), (in_features//16, out_features*2), device=device, dtype=torch.int) + marlin_s = torch.randn((in_features//64, out_features), device=device) + self.workspace = MarlinWorkspace( + self.out_features, GPTQ_MARLIN_MIN_THREAD_N, GPTQ_MARLIN_MAX_PARALLEL, self.device + ) + self.marlin_q_w = marlin_q_w + self.marlin_s = marlin_s + self.g_idx = torch.empty((0), dtype=torch.int32, device=self.device) + self.sort_indices = torch.empty((0), dtype=torch.int32, device=self.device) + self.k = self.weight.shape[0] + self.n = self.weight.shape[1] + self.weight = None + """ + print(in_features, out_features) + print(marlin_q_w.shape) + print(marlin_q_w.dtype) + print(marlin_s.shape) + print(marlin_s.dtype) + print(self.workspace.scratch.shape) + print(self.workspace.scratch.dtype) + print(self.g_idx.shape) + print(self.g_idx.dtype) + print(self.sort_indices.shape) + print(self.sort_indices.dtype) + #print(w_ref.shape) + #print(w_ref.dtype) + """ + #w_ref = None + + def forward(self, x: torch.Tensor, bsz_tensor: torch.Tensor) -> torch.Tensor: + # Only support input x as BF16 and FP16 + x = x.to(self.device) + orig_shape = list(x.shape) + orig_dtype = x.dtype + x = x.reshape(-1, x.shape[-1]) + if self.padding: + padding_input=torch.empty(x.shape[0], self.in_features, device=x.device, dtype=x.dtype) + padding_input[:,:self.orin_in_features] = x + x = padding_input + marlin_s = self.marlin_s.to(x.dtype) + #print(self.sms * ((orig_shape[0]+63)//64)) + + sms = self.sms + + x = vLLMMarlin.gptq_marlin_gemm( + x, + self.marlin_q_w, + marlin_s, + self.g_idx, + self.sort_indices, + self.workspace.scratch, + self.num_bits, + bsz_tensor, + x.shape[0], + self.n, + x.shape[-1], + sms, + self.is_k_full, + ) + # TODO: don't padding bias + if self.has_bias: + x = x + self.bias + if self.padding: + x = x[:,:self.orin_out_features] + orig_shape[-1] = self.orin_out_features + else: + orig_shape[-1] = self.out_features + return x.reshape(orig_shape).to(orig_dtype) + +def benchLinearMarlin(input_dim, output_dim):#, out_file + print("benchmarking MLP Marlin") + print("-----------------------------------------------------------") + headers = ["batch_size", "tp", "used_time", "bandwidth GB/s", "TFLOPS", "cases", "padding", "sms"] + print(" | ".join(headers) + "\n") + rows = [] + for batch_size in exp_range(1, 64): + for tp in exp_range(1, max_tp): + torch.cuda.empty_cache() + if output_dim % tp != 0: + continue + cur_output_dim = output_dim // tp + modules = [] + inputs = [] + data_size = int(0.53125*input_dim*cur_output_dim) + input_size = int(2*batch_size*input_dim) + output_size = int(2*batch_size*cur_output_dim) + usable_mem = get_usable_mem() - 2 * input_dim * cur_output_dim + min_cases = max(global_num_cases, (2*L2_size) // (data_size+input_size)) + cases = int(min(min_cases, (usable_mem * 0.8) // (data_size+input_size))) + #print(usable_mem, data_size, input_size, cases) + + bsz_tensor = torch.tensor([batch_size], device=global_device, dtype=torch.int32) + + if cases == 0: + row = [f"{batch_size}", "OOM", "OOM", "OOM", "0", "False"] + rows.append(row) + break + for _ in range(cases): + modules.append(LinearMarlin(input_dim, cur_output_dim, sms=56, non_equal_division=False).to(device=global_device).eval()) + inputs.append(torch.randn(batch_size, 1, input_dim, device=global_device)) + + def forward(case_id): + modules[case_id](inputs[case_id], bsz_tensor) + + used_time = timing(forward, iters=cases) + bandwidth = (data_size+input_size+output_size)/used_time/1e6 + flops = 2*batch_size*input_dim*cur_output_dim + tflops = flops/used_time/1e9 + cur_sms = modules[0].sms + row = [f"{batch_size}", f"{tp}", f"{used_time}", f"{bandwidth}", f"{tflops}", f"{cases}", modules[0].padding, cur_sms] + rows.append(row) + print(f"{batch_size}", f"{tp}", f"{used_time}", f"{bandwidth}", f"{tflops}", f"{cases}", modules[0].padding, cur_sms) + + """ + with open(out_file, 'w', newline='') as csvfile: + csvwriter = csv.writer(csvfile) + csvwriter.writerow(headers) + for row in rows: + csvwriter.writerow(row) + """ + + """ + markdown_table = " | ".join(headers) + "\n" + markdown_table += " | ".join(["---"] * len(headers)) + "\n" + for row in rows: + markdown_table += " | ".join(row) + "\n" + + print(markdown_table) + """ + #print("finish write file", out_file) + #print("-------------------------------------------------------------") + +if __name__ == "__main__": + + benchLinearMarlin(5120, 3584) + exit(0) + + max_batch = 1 + cur_batch = 1 + + + marlin_linear = LinearMarlin(5120, 3584) + + input_tensor = torch.randn(max_batch, 1, 5120, device="cuda", dtype=torch.bfloat16) + bsz_tensor = torch.tensor([max_batch], device="cuda", dtype=torch.int32) + + out_truth = marlin_linear(input_tensor, bsz_tensor) + + print(out_truth) + + g = torch.cuda.CUDAGraph() + with torch.cuda.graph(g): + out_buf = marlin_linear(input_tensor, bsz_tensor) + + for i in range(10000): + g.replay() + + #torch.testing.assert_close(out_buf, out_truth, rtol=1e-3, atol=1e-3) + + marlin_linear = LinearMarlin(5120, 3584) + g = torch.cuda.CUDAGraph() + with torch.cuda.graph(g): + out_buf = marlin_linear(input_tensor, bsz_tensor) + + new_input = torch.randn(cur_batch, 1, 5120, device="cuda", dtype=torch.bfloat16) + bsz_tensor.copy_(torch.tensor([cur_batch], device="cuda", dtype=torch.int32)) + + new_out_truth = marlin_linear(new_input, bsz_tensor) + input_tensor[:cur_batch].copy_(new_input) + input_tensor[cur_batch:] = 0 + + g.replay() + + torch.cuda.synchronize() + + def printMinMax(tensor): + abs_tensor = torch.abs(tensor) + + min_val = torch.min(abs_tensor) + max_val = torch.max(abs_tensor) + + min_indices = (abs_tensor == min_val).nonzero(as_tuple=True) + max_indices = (abs_tensor == max_val).nonzero(as_tuple=True) + + print(f"min: {min_val.item()}") + print(f"min idx: {min_indices}") + print(f"max: {max_val.item()}") + print(f"max idx: {max_indices}") + + print(out_buf[:cur_batch].shape) + print(new_out_truth.shape) + + + printMinMax(out_buf[:cur_batch]) + printMinMax(new_out_truth) + + #torch.testing.assert_close(out_buf[:cur_batch, 0, :], new_out_truth[:cur_batch, 0, :], rtol=1e-3, atol=1e-3) diff --git a/archive/kt-sft/csrc/custom_marlin/utils/__init__.py b/archive/kt-sft/csrc/custom_marlin/utils/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/archive/kt-sft/csrc/custom_marlin/utils/format24.py b/archive/kt-sft/csrc/custom_marlin/utils/format24.py new file mode 100644 index 0000000..2434e79 --- /dev/null +++ b/archive/kt-sft/csrc/custom_marlin/utils/format24.py @@ -0,0 +1,308 @@ +# +# Modified by Roberto Lopez Castro (roberto.lopez.castro@udc.es). +# + +import torch + + +# This is PyTorch implementation of main part of reorder_meta() +# function, from tools/util/include/cutlass/util/host_reorder.h file +# of CUTLASS source tree. Furthermore, CUTLASS template for sparse +# GEMM decides upon layout of this matrix, and at the moment for the +# sparse GEMM executed on tensor cores, this is layout described by +# ColumnMajorInterleaved<2> data structure, in +# include/cutlass/layout/matrix.h of CUTLASS source tree. The +# reordering of meta matrix into meta_reordered matrix calculated +# according to these segments of CUTLASS code is re-implemented here. +# Note that this calculation produces offsets for scattering metadata +# matrix elements into reordered metadata matrix elements (or, +# equivalently, for gathering reordered metadata matrix element back +# into metadata matrix elements). +def _calculate_meta_reordering_scatter_offsets(m, meta_ncols, meta_dtype, + device): + dst_rows = torch.arange(0, m, device=device)[:, None].repeat(1, meta_ncols) + dst_cols = torch.arange(0, meta_ncols, device=device).repeat(m, 1) + + # Reorder the rows, then swizzle the 2x2 blocks. + group_x = 64 + group_y = 32 if meta_dtype.itemsize == 2 else 16 + + dst_rows = (dst_rows // group_x * group_x + (dst_rows % 2) * 2 + + (dst_rows % 8) // 4 + ((dst_rows % group_y) % 4) // 2 * 32 + + ((dst_rows % group_x) // 8) * 4) + + topright = ((dst_rows % 2 == 0) & (dst_cols % 2 == 1)).to(torch.int8) + bottomleft = ((dst_rows % 2 == 1) & (dst_cols % 2 == 0)).to(torch.int8) + dst_rows += topright - bottomleft + dst_cols -= topright - bottomleft + + # Assumed that meta tensor is to be stored in CUTLASS + # InterleavedColumnMajor layout, and reverse engineered + # corresponding code to store values into this tensor. + interleave = 2 + cols_maj = dst_cols // interleave + cols_min = dst_cols % interleave + return (cols_maj * m * interleave + dst_rows * interleave + + cols_min).view(-1) + + +# This function converts dense matrix into sparse semi-structured +# representation, producing "compressed" matrix, in the layout used by +# CUTLASS backend, and corresponding metadata matrix. +def sparse_semi_structured_from_dense_cutlass(dense): + if dense.dim() != 2: + raise RuntimeError( + f"Expected 2-dimensional dense tensor, got {dense.dim()}-dimensional tensor" # noqa: E501 + ) + + m, k = dense.shape + device = dense.device + + meta_dtype = torch.int8 + if dense.dtype == torch.int8: + meta_dtype = torch.int32 + elif dense.dtype in [torch.half, torch.bfloat16, torch.float, torch.int32]: + meta_dtype = torch.int16 + else: + raise RuntimeError(f"Invalid datatype {dense.dtype} of dense matrix") + quadbits_per_meta_elem = meta_dtype.itemsize * 8 // 4 + if quadbits_per_meta_elem not in (4, 8): + raise RuntimeError( + "Invalid number of elements per meta element calculated") + + if meta_dtype == torch.int32: + if m % 16 != 0: + raise RuntimeError( + f"Number of rows of dense matrix {m} must be divisible by 16") + else: + if m % 32 != 0: + raise RuntimeError( + f"Number of rows of dense matrix {m} must be divisible by 32") + if k % (4 * quadbits_per_meta_elem) != 0: + raise RuntimeError( + f"Number of columns of dense matrix {k} must be divisible by {4 * quadbits_per_meta_elem}" # noqa: E501 + ) + + if dense.dtype != torch.float: + ksparse = 4 + dense_4 = dense.view(-1, k // ksparse, ksparse) + m0, m1, m2, m3 = (dense_4 != 0).unbind(-1) + else: + ksparse = 2 + dense_2 = dense.view(-1, k // ksparse, ksparse) + m0, m2 = m1, m3 = (dense_2 != 0).unbind(-1) + meta_ncols = k // (ksparse * quadbits_per_meta_elem) + + # Encoding quadruples of True/False values as follows: + # [True, True, False, False] -> 0b0100 + # [True, False, True, False] -> 0b1000 + # [False, True, True, False] -> 0b1001 + # [True, False, False, True ] -> 0b1100 + # [False, True, False, True ] -> 0b1101 + # [False, False, True, True ] -> 0b1110 + # Thus, lower two bits in the encoding are index of the True value + # at the lowest index in the quadruple, and the higher two bits in + # the encoding are index of the other True value in the quadruple. + # In case there are less than two True values, than False value or + # values at some index or indices are considered True for the + # encoding. In case there are more than two True values, then the + # excess True value(s) at some indices are considered False for + # the encoding. The exact encodings used for these cases are as + # follows: + # [False, False, False, False] -> 0b1110 + # [False, False, False, True ] -> 0b1110 + # [False, False, True, False] -> 0b1110 + # [False, True, False, False] -> 0b1001 + # [False, True, True, True ] -> 0b1101 + # [True, False, False, False] -> 0b1000 + # [True, False, True, True ] -> 0b1100 + # [True, True, False, True ] -> 0b0100 + # [True, True, True, False] -> 0b0100 + # [True, True, True, True ] -> 0b0100 + # These particular encodings are chosen, with the help of Espresso + # logic minimizer software, for the purpose of minimization of + # corresponding Boolean functions, that translate non-zero flags + # into encoding bits. Note also possible choices for the first + # and last of these encodings were limited only to (0b0100, + # 0b1110), in order to produce valid encodings for 1:2 sparsity + # case. + + expr0 = m0 & m1 + expr1 = ~m0 & m1 + expr2 = ~m0 & ~m1 + bit0 = expr1 + bit1 = expr2 + bit2 = expr0 | expr2 | m3 + bit3 = expr1 | ~m1 + idxs0 = bit0 | (bit1.to(torch.int64) << 1) + idxs1 = bit2 | (bit3.to(torch.int64) << 1) + + if dense.dtype != torch.float: + sparse0 = dense_4.gather( + -1, idxs0.unsqueeze(-1)) # type: ignore[possibly-undefined] + sparse1 = dense_4.gather(-1, idxs1.unsqueeze(-1)) + sparse = torch.stack((sparse0, sparse1), dim=-1).view(m, k // 2) + else: + sparse = dense_2.gather(-1, + idxs0.unsqueeze(-1) // 2).view( + m, + k // 2) # type: ignore[possibly-undefined] + + meta_4 = idxs0 | (idxs1 << 2) + meta_n = meta_4.view( + (-1, meta_ncols, quadbits_per_meta_elem)).to(meta_dtype) + + if quadbits_per_meta_elem == 4: + meta = (meta_n[:, :, 0] + | (meta_n[:, :, 1] << 4) + | (meta_n[:, :, 2] << 8) + | (meta_n[:, :, 3] << 12)) + elif quadbits_per_meta_elem == 8: + meta = (meta_n[:, :, 0] + | (meta_n[:, :, 1] << 4) + | (meta_n[:, :, 2] << 8) + | (meta_n[:, :, 3] << 12) + | (meta_n[:, :, 4] << 16) + | (meta_n[:, :, 5] << 20) + | (meta_n[:, :, 6] << 24) + | (meta_n[:, :, 7] << 28)) + + # Reorder meta tensor elements. + meta_reordered = meta.new_empty( + (m * meta_ncols, )) # type: ignore[possibly-undefined] + meta_offsets = _calculate_meta_reordering_scatter_offsets( + m, meta_ncols, meta_dtype, device) + meta_reordered.scatter_(0, meta_offsets, meta.view(-1)) + + return (sparse, meta_reordered.view(m, meta_ncols)) + + +# This function performs reverse of the function above - it +# reconstructs dense matrix from a pair of "compressed" matrix, given +# in the layout used by CUTLASS backend, and accompanying metadata +# matrix. +def sparse_semi_structured_to_dense_cutlass(sparse, meta_reordered): + if sparse.dim() != 2: + raise RuntimeError( + f"Expected 2-dimensional sparse tensor, got {sparse.dim()}-dimensional tensor" # noqa: E501 + ) + + m, k = sparse.shape + device = sparse.device + + if meta_reordered.dim() != 2: + raise RuntimeError( + f"Expected 2-dimensional meta tensor, got {meta_reordered.dim()}-dimensional tensor" # noqa: E501 + ) + if meta_reordered.device != device: + raise RuntimeError( + f"Expected meta matrix to be on {device} device, got matrix on {meta_reordered.device} device" # noqa: E501 + ) + + meta_dtype = meta_reordered.dtype + if meta_dtype not in (torch.int16, torch.int32): + raise RuntimeError(f"Invalid datatype {meta_dtype} of meta matrix") + quadbits_per_meta_elem = meta_dtype.itemsize * 8 // 4 + + ksparse = 4 if sparse.dtype != torch.float else 2 + + meta_nrows, meta_ncols = meta_reordered.shape + if meta_nrows != m: + raise RuntimeError( + f"Number of rows of meta matrix {meta_nrows} must be equal to number of columns of spase matrix {m}" # noqa: E501 + ) + if meta_ncols * ksparse * quadbits_per_meta_elem != 2 * k: + raise RuntimeError( + f"Number of columns of sparse matrix {k} different from the {meta_ncols * ksparse * quadbits_per_meta_elem // 2}, " # noqa: E501 + "expected according to the number of columns of meta matrix") + + # Undo meta tensor elements reordering. + meta_offsets = _calculate_meta_reordering_scatter_offsets( + m, meta_ncols, meta_dtype, device) + meta = torch.gather(meta_reordered.view(-1), 0, + meta_offsets).view(m, meta_ncols) + + # Unpack sparse tensor back to original dense tensor, using + # information provided by meta tensor. Note that torch.float + # datatype is handled pretty much the same as + # torch.half/torch.bfloat16, as metadata for a pair of torch.float + # value is encoded as if underlying 8 bytes contain four + # torch.half/torch.bfloat16 values, where either first two or last + # two are zeros. + meta_2 = torch.empty( + (m, meta_ncols, 2 * quadbits_per_meta_elem), + dtype=meta_dtype, + device=device, + ) + if quadbits_per_meta_elem == 4: + meta_2[:, :, 0] = meta & 0b11 + meta_2[:, :, 1] = (meta >> 2) & 0b11 + meta_2[:, :, 2] = (meta >> 4) & 0b11 + meta_2[:, :, 3] = (meta >> 6) & 0b11 + meta_2[:, :, 4] = (meta >> 8) & 0b11 + meta_2[:, :, 5] = (meta >> 10) & 0b11 + meta_2[:, :, 6] = (meta >> 12) & 0b11 + meta_2[:, :, 7] = (meta >> 14) & 0b11 + elif quadbits_per_meta_elem == 8: + meta_2[:, :, 0] = meta & 0b11 + meta_2[:, :, 1] = (meta >> 2) & 0b11 + meta_2[:, :, 2] = (meta >> 4) & 0b11 + meta_2[:, :, 3] = (meta >> 6) & 0b11 + meta_2[:, :, 4] = (meta >> 8) & 0b11 + meta_2[:, :, 5] = (meta >> 10) & 0b11 + meta_2[:, :, 6] = (meta >> 12) & 0b11 + meta_2[:, :, 7] = (meta >> 14) & 0b11 + meta_2[:, :, 8] = (meta >> 16) & 0b11 + meta_2[:, :, 9] = (meta >> 18) & 0b11 + meta_2[:, :, 10] = (meta >> 20) & 0b11 + meta_2[:, :, 11] = (meta >> 22) & 0b11 + meta_2[:, :, 12] = (meta >> 24) & 0b11 + meta_2[:, :, 13] = (meta >> 26) & 0b11 + meta_2[:, :, 14] = (meta >> 28) & 0b11 + meta_2[:, :, 15] = (meta >> 30) & 0b11 + + dense_offsets = meta_2.view(-1) + ( + torch.arange(0, 2 * m * k // ksparse, device=device) * 4).view( + -1, 1).repeat(1, 2).view(-1) + + dense = torch.zeros((m * 2 * k, ), dtype=sparse.dtype, device=device) + if sparse.dtype != torch.float: + # dense.scatter_(0, dense_offsets, sparse.view(-1)) + dense.scatter_(0, dense_offsets, sparse.reshape(-1)) + else: + dense.view(torch.half).scatter_(0, dense_offsets, + sparse.view(torch.half).view(-1)) + + return dense.view(m, 2 * k) + + +def mask_creator(tensor): + """ + Class for creating N:M sparsity masks. + Masks will be created using the N:M ratio, where for every block of + M weights, N will be pruned based on ranked weight value. Each mask + will correspond to the given tensor. + + :param N: The number of weights in a group to keep + :param M: The size of a weight group + """ + N = 2 + M = 4 + + mask = None + # for i, tensor in enumerate(tensors): + if tensor.numel() % M != 0: + raise ValueError( + f"Tensor of size {tensor.shape} can't be evenly divided into " + f"{M} groups") + + num_groups = tensor.numel() // M + + # N:M sparsity for linear layers + tensor_temp = tensor.detach().abs().reshape(num_groups, M) + index = torch.argsort(tensor_temp, dim=1)[:, :int(M - N)] + + w_b = torch.ones(tensor_temp.shape, device=tensor_temp.device) + mask = w_b.scatter_(dim=1, index=index, value=0).reshape(tensor.shape) + + return mask \ No newline at end of file diff --git a/archive/kt-sft/csrc/custom_marlin/utils/marlin_24_perms.py b/archive/kt-sft/csrc/custom_marlin/utils/marlin_24_perms.py new file mode 100644 index 0000000..d79926f --- /dev/null +++ b/archive/kt-sft/csrc/custom_marlin/utils/marlin_24_perms.py @@ -0,0 +1,65 @@ +''' +Date: 2024-11-08 02:46:07 +LastEditors: djw +LastEditTime: 2024-11-08 02:46:41 +''' +"""This file is used for /tests and /benchmarks""" +from typing import Dict, List + +import numpy +import torch + + +# Precompute permutations for Marlin24 weight and scale shuffling # noqa: E501 +# +# Marlin works on [16*2,64] tiles. The goal of the permutations is to reorder the weight data so that it is compatible noqa: # noqa: E501 +# with the tensor-core format that is described here: +# https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#matrix-fragments-for-mma-m16n8k16-with-floating-point-type # noqa: E501 +# +# As a result of this reordering, the vector loads inside the kernel will get the data as it is needed for tensor-core # noqa: E501 +# (without the need to use ldmatrix instructions) # noqa: E501 +def get_perms_24(num_bits: int): + perm_list: List[int] = [] + for i in range(32): + perm1: List[int] = [] + col = i // 4 + col_o = col // 2 + for block in [0, 1]: + for row in [ + 2 * (i % 4), + 2 * (i % 4) + 1, + 2 * (i % 4 + 4), + 2 * (i % 4 + 4) + 1, + ]: + perm1.append(16 * row + col_o * 256 + 8 * (col % 2) + + 4 * block) + for j in range(4): + perm_list.extend([p + 1 * j for p in perm1]) + perm = numpy.array(perm_list) + + if num_bits == 4: + interleave = numpy.array([0, 2, 4, 6, 1, 3, 5, 7]) + elif num_bits == 8: + interleave = numpy.array([0, 2, 1, 3]) + else: + raise ValueError("num_bits must be 4 or 8, got {}".format(num_bits)) + + perm = perm.reshape((-1, len(interleave)))[:, interleave].ravel() + perm = torch.from_numpy(perm) + scale_perm: List[int] = [] + for i in range(8): + scale_perm.extend([i * 8 + j for j in [0, 4, 1, 5, 2, 6, 3, 7]]) + scale_perm_single: List[int] = [] + for i in range(8): + scale_perm_single.extend([8 * i + j for j in [0, 1, 2, 3, 4, 5, 6, 7]]) + return perm, scale_perm, scale_perm_single + + +marlin_24_perm: Dict[int, torch.Tensor] = {} +marlin_24_scale_perm: Dict[int, List[int]] = {} +marlin_24_scale_perm_single: Dict[int, List[int]] = {} +for num_bits in [4, 8]: + perm_24, scale_perm_24, scale_perm_single_24 = get_perms_24(num_bits) + marlin_24_perm[num_bits] = perm_24 + marlin_24_scale_perm[num_bits] = scale_perm_24 + marlin_24_scale_perm_single[num_bits] = scale_perm_single_24 \ No newline at end of file diff --git a/archive/kt-sft/csrc/custom_marlin/utils/marlin_perms.py b/archive/kt-sft/csrc/custom_marlin/utils/marlin_perms.py new file mode 100644 index 0000000..62255ec --- /dev/null +++ b/archive/kt-sft/csrc/custom_marlin/utils/marlin_perms.py @@ -0,0 +1,65 @@ +''' +Date: 2024-11-08 02:46:47 +LastEditors: djw +LastEditTime: 2024-11-08 02:46:55 +''' +"""This file is used for /tests and /benchmarks""" +from typing import Dict, List + +import numpy +import torch + + +# Precompute permutations for Marlin weight and scale shuffling # noqa: E501 +# +# Marlin works on [16,64] tiles. The goal of the permutations is to reorder the weight data so that it is compatible noqa: # noqa: E501 +# with the tensor-core format that is described here: +# https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#matrix-fragments-for-mma-m16n8k16-with-floating-point-type # noqa: E501 +# +# As a result of this reordering, the vector loads inside the kernel will get the data as it is needed for tensor-core # noqa: E501 +# (without the need to use ldmatrix instructions) # noqa: E501 +def get_perms(num_bits: int): + perm_list: List[int] = [] + for i in range(32): + perm1: List[int] = [] + col = i // 4 + for block in [0, 1]: + for row in [ + 2 * (i % 4), + 2 * (i % 4) + 1, + 2 * (i % 4 + 4), + 2 * (i % 4 + 4) + 1, + ]: + perm1.append(16 * row + col + 8 * block) + for j in range(4): + perm_list.extend([p + 256 * j for p in perm1]) + + perm = numpy.array(perm_list) + + if num_bits == 4: + interleave = numpy.array([0, 2, 4, 6, 1, 3, 5, 7]) + elif num_bits == 8: + interleave = numpy.array([0, 2, 1, 3]) + else: + raise Exception("num_bits must be 4 or 8, got {}".format(num_bits)) + + perm = perm.reshape((-1, len(interleave)))[:, interleave].ravel() + perm = torch.from_numpy(perm) + 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 perm, scale_perm, scale_perm_single + + +marlin_perm: Dict[int, torch.Tensor] = {} +marlin_scale_perm: Dict[int, List[int]] = {} +marlin_scale_perm_single: Dict[int, List[int]] = {} +for num_bits in [4, 8]: + perm, scale_perm, scale_perm_single = get_perms(num_bits) + marlin_perm[num_bits] = perm + marlin_scale_perm[num_bits] = scale_perm + marlin_scale_perm_single[num_bits] = scale_perm_single \ No newline at end of file diff --git a/archive/kt-sft/csrc/custom_marlin/utils/marlin_utils.py b/archive/kt-sft/csrc/custom_marlin/utils/marlin_utils.py new file mode 100644 index 0000000..ccecdc3 --- /dev/null +++ b/archive/kt-sft/csrc/custom_marlin/utils/marlin_utils.py @@ -0,0 +1,234 @@ +"""This file is used for /tests and /benchmarks""" +import random + +import numpy +import torch + +from .format24 import ( + mask_creator, sparse_semi_structured_from_dense_cutlass) +from .marlin_24_perms import ( + marlin_24_perm, marlin_24_scale_perm, marlin_24_scale_perm_single) +from .marlin_perms import ( + marlin_perm, marlin_scale_perm, marlin_scale_perm_single) +from .quant_utils import ( + get_pack_factor, quantize_weights, sort_weights, dequantize_weights) + + + +__cuda_arch = torch.cuda.get_device_capability() + +MARLIN_TILE = 16 + +GPTQ_MARLIN_TILE = 16 +GPTQ_MARLIN_MIN_THREAD_N = 64 +GPTQ_MARLIN_MIN_THREAD_K = 128 +GPTQ_MARLIN_MAX_PARALLEL = 16 + +GPTQ_MARLIN_SUPPORTED_NUM_BITS = [4, 8] +GPTQ_MARLIN_SUPPORTED_GROUP_SIZES = [-1, 32, 64, 128] +GPTQ_MARLIN_SUPPORTED_SYM = [True] + +def is_marlin_supported(): + return __cuda_arch[0] >= 8 + + +def marlin_permute_weights(q_w, size_k, size_n, perm, tile=MARLIN_TILE): + assert q_w.shape == (size_k, size_n) + assert size_k % tile == 0, f"size_k = {size_k}, tile = {tile}" + assert size_n % tile == 0, f"size_k = {size_n}, tile = {tile}" + + # Permute weights to 16x64 marlin tiles + q_w = q_w.reshape((size_k // tile, tile, size_n // tile, tile)) + q_w = q_w.permute((0, 2, 1, 3)) + q_w = q_w.reshape((size_k // tile, size_n * tile)) + + q_w = q_w.reshape((-1, perm.numel()))[:, perm].reshape(q_w.shape) + + return q_w + + +def marlin_weights(q_w, size_k, size_n, num_bits, perm): + # Permute + q_w = marlin_permute_weights(q_w, size_k, size_n, perm) + + # Pack + pack_factor = get_pack_factor(num_bits) + orig_device = q_w.device + + q_w = q_w.cpu().numpy().astype(numpy.uint32) + + q_packed = numpy.zeros((q_w.shape[0], q_w.shape[1] // pack_factor), + dtype=numpy.uint32) + for i in range(pack_factor): + q_packed |= q_w[:, i::pack_factor] << num_bits * i + + q_packed = torch.from_numpy(q_packed.astype(numpy.int32)).to(orig_device) + + return q_packed + + +def marlin_permute_scales(s, size_k, size_n, group_size, scale_perm, + scale_perm_single): + 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_quantize( + w: torch.Tensor, + num_bits: int, + group_size: int, + act_order: bool, +): + size_k, size_n = w.shape + + # Normalize group_size + if group_size == -1: + group_size = size_k + assert group_size <= size_k + + # Quantize (and apply act_order if provided) + w_ref, q_w, s, g_idx, rand_perm = quantize_weights(w, num_bits, group_size, + act_order) + + # For act_order, sort the "weights" and "g_idx" so that group ids are + # increasing + sort_indices = torch.empty(0, dtype=torch.int, device=w.device) + if act_order: + q_w, g_idx, sort_indices = sort_weights(q_w, g_idx) + + # Reformat to marlin + marlin_q_w = marlin_weights(q_w, size_k, size_n, num_bits, + marlin_perm[num_bits]) + marlin_s = marlin_permute_scales(s, size_k, size_n, group_size, + marlin_scale_perm[num_bits], + marlin_scale_perm_single[num_bits]) + + # Create result + res_list = [w_ref, marlin_q_w, marlin_s, g_idx, sort_indices, rand_perm] + for i in range(len(res_list)): + res_list[i] = res_list[i].to(w.device) + + return res_list + + +def inject_24(w, size_k, size_n): + assert w.shape == (size_k, size_n) + + mask = mask_creator(w.t()).t().cuda().bool() + + return (mask * w).contiguous(), mask.contiguous() + + +def check_24(w, num_rows_to_sample=50, _verbose=False): + BLOCK_SIZE = 4 + MAX_NON_ZEROS = 2 + + w = w.t().contiguous() + + print("check_24: w.shape = {}".format(w.shape)) + + num_rows, num_cols = w.shape + sampled_row_idxs = random.choices(range(num_rows), k=num_rows_to_sample) + if _verbose: + print(f"Sampled row idxs = {sampled_row_idxs}") + + total_segments = 0 + non_24_segments = 0 + for i in sampled_row_idxs: + for j in range(0, num_cols - BLOCK_SIZE, BLOCK_SIZE): + total_segments += 1 + block = w[i, j:j + BLOCK_SIZE] + num_nonzero = torch.count_nonzero(block) + if num_nonzero > MAX_NON_ZEROS: + print("i = {} j = {} block = {}".format(i, j, block)) + non_24_segments += 1 + + print(f"{non_24_segments} / {total_segments} do not have 2:4 structure.") + + +def compress_quantized_24_weight(q_24, size_k, size_n, num_bits): + assert q_24.shape == (size_k, size_n) + + # Remove zp to normalize over 0 + max_q_val = (1 << num_bits) - 1 + zp = (max_q_val + 1) // 2 + q_24_no_zp = q_24 - zp + + # Compress + q_24_no_zp = q_24_no_zp.t().contiguous() + q_24_no_zp_comp, meta = sparse_semi_structured_from_dense_cutlass( + q_24_no_zp) + q_24_no_zp_comp = q_24_no_zp_comp.t().contiguous() + + # Restore zp + q_24_comp = q_24_no_zp_comp + zp + + # Resize meta to its actual shape (without moving any data) + meta = meta.resize_(meta.shape[1] // 2, meta.shape[0] * 2) + + return q_24_comp, meta + + +def marlin_24_quantize( + w: torch.Tensor, + num_bits: int, + group_size: int, +): + size_k, size_n = w.shape + + # Normalize group_size + if group_size == -1: + group_size = size_k + assert group_size <= size_k + + # Inject 2:4 sparsity + w_24, mask_24 = inject_24(w, size_k, size_n) + + # Quantize + w_24_ref, q_w_24, s, g_idx, rand_perm = quantize_weights(w_24, + num_bits, + group_size, + act_order=False) + + # Compress quantized weight + q_w_24_comp, meta = compress_quantized_24_weight(q_w_24, size_k, size_n, + num_bits) + size_k_comp = size_k // 2 + + # Reformat to marlin + marlin_24_q_w_comp = marlin_weights(q_w_24_comp, size_k_comp, size_n, + num_bits, marlin_24_perm[num_bits]) + marlin_24_s = marlin_permute_scales(s, size_k, size_n, group_size, + marlin_24_scale_perm[num_bits], + marlin_24_scale_perm_single[num_bits]) + + # Create result + res_list = [w_24_ref, marlin_24_q_w_comp, meta, marlin_24_s] + for i in range(len(res_list)): + res_list[i] = res_list[i].to(w.device) + + return res_list + + +def compute_max_diff(output, output_ref): + return torch.mean(torch.abs(output - output_ref)) / torch.mean( + torch.abs(output_ref)) + + +class MarlinWorkspace: + + def __init__(self, out_features, min_thread_n, max_parallel, device): + assert (out_features % min_thread_n == 0), ( + "out_features = {} is undivisible by min_thread_n = {}".format( + out_features, min_thread_n)) + + max_workspace_size = ((out_features // min_thread_n) * max_parallel) + + self.scratch = torch.zeros(max_workspace_size, + dtype=torch.int, + device=device) \ No newline at end of file diff --git a/archive/kt-sft/csrc/custom_marlin/utils/quant_utils.py b/archive/kt-sft/csrc/custom_marlin/utils/quant_utils.py new file mode 100644 index 0000000..077d22f --- /dev/null +++ b/archive/kt-sft/csrc/custom_marlin/utils/quant_utils.py @@ -0,0 +1,195 @@ +"""This file is used for /tests and /benchmarks""" +import numpy +import torch + +SUPPORTED_NUM_BITS = [4, 8] +SUPPORTED_GROUP_SIZES = [-1, 32, 64, 128] + + +def get_pack_factor(num_bits): + assert num_bits in SUPPORTED_NUM_BITS, f"Unsupported num_bits = {num_bits}" + return 32 // num_bits + + +def permute_rows(q_w: torch.Tensor, w_ref: torch.Tensor, group_size: int): + assert q_w.shape == w_ref.shape + + orig_device = q_w.device + k_size, _ = q_w.shape + + g_idx = torch.zeros((k_size, ), dtype=torch.int32) + for i in range(k_size): + g_idx[i] = i // group_size + + # Simulate act_order by doing a random permutation on K + rand_perm = torch.randperm(k_size) + + g_idx = g_idx[rand_perm].contiguous() + q_w = q_w[rand_perm, :].contiguous() + w_ref = w_ref[rand_perm, :].contiguous() + + return ( + w_ref.to(device=orig_device), + q_w.to(device=orig_device), + g_idx.to(device=orig_device), + rand_perm.to(device=orig_device), + ) + + +# Function: Dequantize quantized weights +def dequantize_weights(qweight, qzeros, scales, g_idx, bits=4, group_size=128, device='cuda:0'): + # Create a tensor for bitwise right shift operation + wf = torch.tensor(list(range(0, 32, bits)), dtype=torch.int32, device=device).unsqueeze(0) + + # Apply bitwise right shift and convert qzeros to the appropriate type + zeros = torch.bitwise_right_shift(torch.unsqueeze(qzeros, 2).expand(-1, -1, 32 // bits), wf.unsqueeze(0)).to(torch.int16 if bits == 8 else torch.int8) + torch.bitwise_and(zeros, (2 ** bits) - 1, out=zeros) + + # Reshape the zeros tensor + zeros = zeros + 1 + zeros = zeros.reshape(-1, 1, zeros.shape[1] * zeros.shape[2]) + + # Reshape the scales tensor + scales = scales.reshape(-1, 1, scales.shape[-1]) + + # Similar bitwise right shift operation for qweight and reshape + weight = torch.bitwise_right_shift(torch.unsqueeze(qweight, 1).expand(-1, 32 // bits, -1), wf.unsqueeze(-1)).to(torch.int16 if bits == 8 else torch.int8) + torch.bitwise_and(weight, (2 ** bits) - 1, out=weight) + weight = weight.reshape(-1, group_size, weight.shape[2]) + + # Apply dequantization formula and reshape the final weight + weight = (scales * (weight - zeros)) + weight = weight.reshape(weight.shape[0] * weight.shape[1], weight.shape[2]) + + # Return the transposed weight + return weight.transpose(0, 1) + +def quantize_weights(w: torch.Tensor, num_bits: int, group_size: int, + act_order: bool): + orig_device = w.device + size_k, size_n = w.shape + + assert w.is_floating_point(), "w must be float" + assert num_bits in SUPPORTED_NUM_BITS, f"Unsupported num_bits = {num_bits}" + assert group_size in SUPPORTED_GROUP_SIZES + [ + size_k + ], f"Unsupported groupsize = {group_size}" + + if group_size == -1: + group_size = size_k + assert group_size <= size_k + + max_q_val = 2**num_bits - 1 + half_q_val = (max_q_val + 1) // 2 + + # Reshape to [groupsize, -1] + if group_size < size_k: + w = w.view((-1, group_size, size_n)) + w = w.permute(1, 0, 2) + w = w.reshape((group_size, -1)) + + # Compute scale for each group + s = torch.max(torch.abs(w), 0, keepdim=True)[0] + s *= 2 / max_q_val # 2 => symmetric + + # Quantize + q_w = torch.round(w / s).int() + q_w += half_q_val + q_w = torch.clamp(q_w, 0, max_q_val) + + # Compute ref (dequantized) + w_ref = (q_w - half_q_val).half() * s + + # Restore original shapes + if group_size < size_k: + + def reshape_w(w): + w = w.reshape((group_size, -1, size_n)) + w = w.permute(1, 0, 2) + w = w.reshape((size_k, size_n)).contiguous() + return w + + q_w = reshape_w(q_w) + w_ref = reshape_w(w_ref) + + s = s.reshape((-1, size_n)).contiguous() + + # Apply act_order + g_idx = torch.empty(0, dtype=torch.int, device=w.device) + rand_perm = torch.empty(0, dtype=torch.int, device=w.device) + if act_order: + assert ( + group_size < size_k + ), "For act_order, groupsize = {} must be less than size_k = {}".format( + group_size, size_k) + + w_ref, q_w, g_idx, rand_perm = permute_rows(q_w, w_ref, group_size) + + return ( + w_ref.to(device=orig_device), + q_w.to(device=orig_device), + s.to(device=orig_device), + g_idx.to(device=orig_device), + rand_perm.to(device=orig_device), + ) + + +def sort_weights(q_w: torch.Tensor, g_idx: torch.Tensor): + orig_device = q_w.device + + sort_indices = torch.argsort(g_idx).to( + dtype=torch.int32) # Sort based on g_idx + + g_idx = g_idx[sort_indices].contiguous() + q_w = q_w[sort_indices, :].contiguous() + + return ( + q_w.to(device=orig_device), + g_idx.to(device=orig_device), + sort_indices.to(device=orig_device), + ) + + +def gptq_pack( + q_w: torch.Tensor, + num_bits: int, + size_k: int, + size_n: int, +): + assert q_w.shape == (size_k, size_n) + + pack_factor = get_pack_factor(num_bits) + assert size_k % pack_factor == 0 + + orig_device = q_w.device + + q_w = q_w.cpu().numpy().astype(numpy.uint32) + + q_res = numpy.zeros((size_k // pack_factor, size_n), dtype=numpy.uint32) + + for i in range(pack_factor): + q_res |= q_w[i::pack_factor, :] << num_bits * i + + q_res = torch.from_numpy(q_res.astype(numpy.int32)).to(orig_device) + return q_res + +def gptq_unpack( + q_res: torch.Tensor, + num_bits: int, + size_k: int, + size_n: int, +): + pack_factor = 32 // num_bits + assert size_k % pack_factor == 0 + + orig_device = q_res.device + + q_res = q_res.cpu().numpy() + + q_w = numpy.zeros((size_k, size_n), dtype=numpy.uint32) + + for i in range(pack_factor): + q_w[i::pack_factor, :] = (q_res >> (num_bits * i)) & ((1 << num_bits) - 1) + + q_w = torch.from_numpy(q_w.astype(numpy.int32)).to(orig_device) + return q_w \ No newline at end of file diff --git a/archive/kt-sft/csrc/ktransformers_ext/CMakeLists.txt b/archive/kt-sft/csrc/ktransformers_ext/CMakeLists.txt new file mode 100644 index 0000000..b215653 --- /dev/null +++ b/archive/kt-sft/csrc/ktransformers_ext/CMakeLists.txt @@ -0,0 +1,415 @@ +cmake_minimum_required(VERSION 3.16) +project(cpuinfer_ext VERSION 0.1.0) + + +set(CMAKE_CXX_STANDARD 17) + + +set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -O3 -ffast-math -fopenmp") +add_compile_definitions(_GLIBCXX_USE_CXX11_ABI=${_GLIBCXX_USE_CXX11_ABI}) +set(CMAKE_BUILD_TYPE "Release") +# set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -g -ffast-math -fopenmp") +# set(CMAKE_BUILD_TYPE "Debug") +set(CMAKE_EXPORT_COMPILE_COMMANDS ON) + + +include(CheckCXXCompilerFlag) +set(CMAKE_POSITION_INDEPENDENT_CODE ON) + + +option(LLAMA_NATIVE "llama: enable -march=native flag" ON) + +# instruction set specific +if (LLAMA_NATIVE) + set(INS_ENB OFF) +else() + set(INS_ENB ON) +endif() + +option(LLAMA_AVX "llama: enable AVX" OFF) +option(LLAMA_AVX2 "llama: enable AVX2" OFF) +option(LLAMA_AVX512 "llama: enable AVX512" OFF) +option(LLAMA_AVX512_VBMI "llama: enable AVX512-VBMI" OFF) +option(LLAMA_AVX512_VNNI "llama: enable AVX512-VNNI" OFF) +option(LLAMA_AVX512_BF16 "llama: enable AVX512-BF16" OFF) +option(LLAMA_FMA "llama: enable FMA" OFF) +# in MSVC F16C is implied with AVX2/AVX512 +if (NOT MSVC) + option(LLAMA_F16C "llama: enable F16C" OFF) +endif() +option(LLAMA_AVX512_FANCY_SIMD "llama: enable AVX512-VL, AVX512-BW, AVX512-DQ, AVX512-VNNI" OFF) +option(KTRANSFORMERS_USE_CUDA "ktransformers: use CUDA" ON) +option(KTRANSFORMERS_USE_MUSA "ktransformers: use MUSA" OFF) +option(KTRANSFORMERS_USE_ROCM "ktransformers: use ROCM" OFF) +option(KTRANSFORMERS_USE_XPU "ktransformers: use XPU" OFF) +option(KTRANSFORMERS_USE_NPU "ktransformers: use NPU" OFF) + +if(KTRANSFORMERS_USE_NPU) + add_definitions(-DKTRANSFORMERS_USE_NPU=1) +endif() + +# Architecture specific +# TODO: probably these flags need to be tweaked on some architectures +# feel free to update the Makefile for your architecture and send a pull request or issue +message(STATUS "CMAKE_SYSTEM_PROCESSOR: ${CMAKE_SYSTEM_PROCESSOR}") +if (MSVC) + string(TOLOWER "${CMAKE_GENERATOR_PLATFORM}" CMAKE_GENERATOR_PLATFORM_LWR) + message(STATUS "CMAKE_GENERATOR_PLATFORM: ${CMAKE_GENERATOR_PLATFORM}") +else () + set(CMAKE_GENERATOR_PLATFORM_LWR "") +endif () + +if (NOT MSVC) + if (LLAMA_STATIC) + add_link_options(-static) + if (MINGW) + add_link_options(-static-libgcc -static-libstdc++) + endif() + endif() + if (LLAMA_GPROF) + add_compile_options(-pg) + endif() +endif() + +set(ARCH_FLAGS "") + +if (CMAKE_OSX_ARCHITECTURES STREQUAL "arm64" OR CMAKE_GENERATOR_PLATFORM_LWR STREQUAL "arm64" OR + (NOT CMAKE_OSX_ARCHITECTURES AND NOT CMAKE_GENERATOR_PLATFORM_LWR AND + CMAKE_SYSTEM_PROCESSOR MATCHES "^(aarch64|arm.*|ARM64)$")) + message(STATUS "ARM detected") + if (MSVC) + add_compile_definitions(__aarch64__) # MSVC defines _M_ARM64 instead + add_compile_definitions(__ARM_NEON) + add_compile_definitions(__ARM_FEATURE_FMA) + + set(CMAKE_REQUIRED_FLAGS_PREV ${CMAKE_REQUIRED_FLAGS}) + string(JOIN " " CMAKE_REQUIRED_FLAGS ${CMAKE_REQUIRED_FLAGS} "/arch:armv8.2") + check_cxx_source_compiles("#include \nint main() { int8x16_t _a, _b; int32x4_t _s = vdotq_s32(_s, _a, _b); return 0; }" GGML_COMPILER_SUPPORT_DOTPROD) + if (GGML_COMPILER_SUPPORT_DOTPROD) + add_compile_definitions(__ARM_FEATURE_DOTPROD) + endif () + check_cxx_source_compiles("#include \nint main() { float16_t _a; float16x8_t _s = vdupq_n_f16(_a); return 0; }" GGML_COMPILER_SUPPORT_FP16_VECTOR_ARITHMETIC) + if (GGML_COMPILER_SUPPORT_FP16_VECTOR_ARITHMETIC) + add_compile_definitions(__ARM_FEATURE_FP16_VECTOR_ARITHMETIC) + endif () + set(CMAKE_REQUIRED_FLAGS ${CMAKE_REQUIRED_FLAGS_PREV}) + else() + if(KTRANSFORMERS_USE_NPU) + list(APPEND ARCH_FLAGS -march=armv8.2-a+fp16+fp16fml+dotprod -lnuma) + endif() + check_cxx_compiler_flag(-mfp16-format=ieee COMPILER_SUPPORTS_FP16_FORMAT_I3E) + if (NOT "${COMPILER_SUPPORTS_FP16_FORMAT_I3E}" STREQUAL "") + list(APPEND ARCH_FLAGS -mfp16-format=ieee) + endif() + if (${CMAKE_SYSTEM_PROCESSOR} MATCHES "armv6") + # Raspberry Pi 1, Zero + list(APPEND ARCH_FLAGS -mfpu=neon-fp-armv8 -mno-unaligned-access) + endif() + if (${CMAKE_SYSTEM_PROCESSOR} MATCHES "armv7") + if ("${CMAKE_SYSTEM_NAME}" STREQUAL "Android") + # Android armeabi-v7a + list(APPEND ARCH_FLAGS -mfpu=neon-vfpv4 -mno-unaligned-access -funsafe-math-optimizations) + else() + # Raspberry Pi 2 + list(APPEND ARCH_FLAGS -mfpu=neon-fp-armv8 -mno-unaligned-access -funsafe-math-optimizations) + endif() + endif() + if (${CMAKE_SYSTEM_PROCESSOR} MATCHES "armv8") + # Android arm64-v8a + # Raspberry Pi 3, 4, Zero 2 (32-bit) + list(APPEND ARCH_FLAGS -mno-unaligned-access) + endif() + endif() +elseif (CMAKE_OSX_ARCHITECTURES STREQUAL "x86_64" OR CMAKE_GENERATOR_PLATFORM_LWR MATCHES "^(x86_64|i686|amd64|x64|win32)$" OR + (NOT CMAKE_OSX_ARCHITECTURES AND NOT CMAKE_GENERATOR_PLATFORM_LWR AND + CMAKE_SYSTEM_PROCESSOR MATCHES "^(x86_64|i686|AMD64)$")) + message(STATUS "x86 detected") + if(NOT KTRANSFORMERS_USE_NPU) + set(HOST_IS_X86 TRUE) + set(HAS_AVX512 TRUE) + set(__HAS_AMX__ TRUE) + add_compile_definitions(__x86_64__) + # check AVX512 + execute_process( + COMMAND lscpu + OUTPUT_VARIABLE LSCPU_OUTPUT + OUTPUT_STRIP_TRAILING_WHITESPACE + ) + # message(STATUS "LSCPU_OUTPUT: ${LSCPU_OUTPUT}") + + string(FIND "${LSCPU_OUTPUT}" "avx512" COMPILER_SUPPORTS_AVX512F) + + if (COMPILER_SUPPORTS_AVX512F GREATER -1) + message(STATUS "Compiler and CPU support AVX512F (tested by compiling a program)") + add_compile_definitions(__HAS_AVX512F__) + else() + message(STATUS "Compiler and/or CPU do NOT support AVX512F") + set(HAS_AVX512 False) + endif() + + # check AMX + string(FIND "${LSCPU_OUTPUT}" "amx" COMPILER_SUPPORTS_AMX) + + if(COMPILER_SUPPORTS_AMX GREATER -1) + message(STATUS "Compiler supports AMX") + add_compile_definitions(__HAS_AMX__) + else() + message(STATUS "Compiler does NOT support AMX") + endif() + endif() + if (MSVC) + # instruction set detection for MSVC only + if (LLAMA_NATIVE) + include(cmake/FindSIMD.cmake) + endif () + if (LLAMA_AVX512) + list(APPEND ARCH_FLAGS /arch:AVX512) + # MSVC has no compile-time flags enabling specific + # AVX512 extensions, neither it defines the + # macros corresponding to the extensions. + # Do it manually. + if (LLAMA_AVX512_VBMI) + add_compile_definitions($<$:__AVX512VBMI__>) + add_compile_definitions($<$:__AVX512VBMI__>) + endif() + if (LLAMA_AVX512_VNNI) + add_compile_definitions($<$:__AVX512VNNI__>) + add_compile_definitions($<$:__AVX512VNNI__>) + endif() + if (LLAMA_AVX512_FANCY_SIMD) + add_compile_definitions($<$:__AVX512VL__>) + add_compile_definitions($<$:__AVX512VL__>) + add_compile_definitions($<$:__AVX512BW__>) + add_compile_definitions($<$:__AVX512BW__>) + add_compile_definitions($<$:__AVX512DQ__>) + add_compile_definitions($<$:__AVX512DQ__>) + add_compile_definitions($<$:__AVX512VNNI__>) + add_compile_definitions($<$:__AVX512VNNI__>) + endif() + if (LLAMA_AVX512_BF16) + add_compile_definitions($<$:__AVX512BF16__>) + add_compile_definitions($<$:__AVX512BF16__>) + endif() + elseif (LLAMA_AVX2) + list(APPEND ARCH_FLAGS /arch:AVX2) + elseif (LLAMA_AVX) + list(APPEND ARCH_FLAGS /arch:AVX) + endif() + else() + if (LLAMA_NATIVE) + list(APPEND ARCH_FLAGS -mfma -mavx -mavx2) + list(APPEND ARCH_FLAGS -march=native) + endif() + if (LLAMA_F16C) + list(APPEND ARCH_FLAGS -mf16c) + endif() + if (LLAMA_FMA) + list(APPEND ARCH_FLAGS -mfma) + endif() + if (LLAMA_AVX) + list(APPEND ARCH_FLAGS -mavx) + endif() + if (LLAMA_AVX2) + list(APPEND ARCH_FLAGS -mavx2) + endif() + if (LLAMA_AVX512) + list(APPEND ARCH_FLAGS -mavx512f) + list(APPEND ARCH_FLAGS -mavx512bw) + endif() + if (LLAMA_AVX512_VBMI) + list(APPEND ARCH_FLAGS -mavx512vbmi) + endif() + if (LLAMA_AVX512_VNNI) + list(APPEND ARCH_FLAGS -mavx512vnni) + endif() + if (LLAMA_AVX512_FANCY_SIMD) + message(STATUS "AVX512-VL, AVX512-BW, AVX512-DQ, AVX512-VNNI enabled") + list(APPEND ARCH_FLAGS -mavx512vl) + list(APPEND ARCH_FLAGS -mavx512bw) + list(APPEND ARCH_FLAGS -mavx512dq) + list(APPEND ARCH_FLAGS -mavx512vnni) + list(APPEND ARCH_FLAGS -mavx512vpopcntdq) + endif() + if (LLAMA_AVX512_BF16) + list(APPEND ARCH_FLAGS -mavx512bf16) + endif() + endif() +elseif (${CMAKE_SYSTEM_PROCESSOR} MATCHES "ppc64") + message(STATUS "PowerPC detected") + if (${CMAKE_SYSTEM_PROCESSOR} MATCHES "ppc64le") + list(APPEND ARCH_FLAGS -mcpu=powerpc64le) + else() + list(APPEND ARCH_FLAGS -mcpu=native -mtune=native) + #TODO: Add targets for Power8/Power9 (Altivec/VSX) and Power10(MMA) and query for big endian systems (ppc64/le/be) + endif() +else() + message(STATUS "Unknown architecture") +endif() + +# message(STATUS "CUDAToolkit_ROOT:${CUDAToolkit_ROOT}") +# find_package(FindCUDAToolkit REQUIRED) +# if(CUDAToolkit_FOUND) +# message(STATUS "Found CUDA cudart lib at:${CUDAToolkit_LIBRARY_DIR}") +# else() +# message(STATUS "Can't found CUDA lib") +# endif() + +if (NOT EXISTS $ENV{ROCM_PATH}) + if (NOT EXISTS /opt/rocm) + set(ROCM_PATH /usr) + else() + set(ROCM_PATH /opt/rocm) + endif() +else() + set(ROCM_PATH $ENV{ROCM_PATH}) +endif() + +list(APPEND CMAKE_PREFIX_PATH ${ROCM_PATH}) +list(APPEND CMAKE_PREFIX_PATH "${ROCM_PATH}/lib64/cmake") + +if (NOT EXISTS $ENV{MUSA_PATH}) + if (NOT EXISTS /opt/musa) + set(MUSA_PATH /usr/local/musa) + else() + set(MUSA_PATH /opt/musa) + endif() +else() + set(MUSA_PATH $ENV{MUSA_PATH}) +endif() + +list(APPEND CMAKE_MODULE_PATH "${MUSA_PATH}/cmake") + +add_compile_options("$<$:${ARCH_FLAGS}>") +add_compile_options("$<$:${ARCH_FLAGS}>") + +add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/../../../third_party/pybind11 ${CMAKE_CURRENT_BINARY_DIR}/third_party/pybind11) +add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/../../../third_party/llama.cpp ${CMAKE_CURRENT_BINARY_DIR}/third_party/llama.cpp) + +include_directories(${CMAKE_CURRENT_SOURCE_DIR}/../../../third_party) +if (WIN32) + include_directories("$ENV{CUDA_PATH}/include") + add_compile_definitions(KTRANSFORMERS_USE_CUDA=1) +elseif (UNIX) + if (KTRANSFORMERS_USE_ROCM) + find_package(HIP REQUIRED) + if(HIP_FOUND) + include_directories("${HIP_INCLUDE_DIRS}") + add_compile_definitions(KTRANSFORMERS_USE_ROCM=1) + endif() + elseif (KTRANSFORMERS_USE_MUSA) + if (NOT EXISTS $ENV{MUSA_PATH}) + if (NOT EXISTS /opt/musa) + set(MUSA_PATH /usr/local/musa) + else() + set(MUSA_PATH /opt/musa) + endif() + else() + set(MUSA_PATH $ENV{MUSA_PATH}) + endif() + + list(APPEND CMAKE_MODULE_PATH "${MUSA_PATH}/cmake") + + find_package(MUSAToolkit) + if (MUSAToolkit_FOUND) + message(STATUS "MUSA Toolkit found") + add_compile_definitions(KTRANSFORMERS_USE_MUSA=1) + endif() + elseif (KTRANSFORMERS_USE_XPU) + add_compile_definitions(KTRANSFORMERS_USE_XPU=1) + elseif (KTRANSFORMERS_USE_CUDA) + find_package(CUDA REQUIRED) + include_directories("${CUDA_INCLUDE_DIRS}") + include(CheckLanguage) + check_language(CUDA) + if(CMAKE_CUDA_COMPILER) + message(STATUS "CUDA detected") + find_package(CUDAToolkit REQUIRED) + include_directories(${CUDAToolkit_INCLUDE_DIRS}) + endif() + message(STATUS "enabling CUDA") + enable_language(CUDA) + add_compile_definitions(KTRANSFORMERS_USE_CUDA=1) + endif() +endif() + +aux_source_directory(${CMAKE_CURRENT_SOURCE_DIR} SOURCE_DIR1) +aux_source_directory(${CMAKE_CURRENT_SOURCE_DIR}/cpu_backend SOURCE_DIR2) +aux_source_directory(${CMAKE_CURRENT_SOURCE_DIR}/operators/llamafile SOURCE_DIR3) +# aux_source_directory(${CMAKE_CURRENT_SOURCE_DIR}/../../../third_party/llamafile SOURCE_DIR4) +file(GLOB LLAMAFILE_SOURCES "${CMAKE_CURRENT_SOURCE_DIR}/../../../third_party/llamafile/*.cpp") +list(REMOVE_ITEM LLAMAFILE_SOURCES + "${CMAKE_CURRENT_SOURCE_DIR}/../../../third_party/llamafile/sgemm_arm.cpp" + "${CMAKE_CURRENT_SOURCE_DIR}/../../../third_party/llamafile/sgemm_x86.cpp" +) +set(SOURCE_DIR4 ${LLAMAFILE_SOURCES}) +aux_source_directory(${CMAKE_CURRENT_SOURCE_DIR}/operators/kvcache SOURCE_DIR5) + +if (HOST_IS_X86 AND HAS_AVX512 AND __HAS_AMX__) + aux_source_directory(${CMAKE_CURRENT_SOURCE_DIR}/operators/amx SOURCE_DIR6) +endif() + + +set(ALL_SOURCES ${SOURCE_DIR1} ${SOURCE_DIR2} ${SOURCE_DIR3} ${SOURCE_DIR4} ${SOURCE_DIR5} ${SOURCE_DIR6}) + +file(GLOB_RECURSE FMT_SOURCES "${CMAKE_CURRENT_SOURCE_DIR}/*.cpp" "${CMAKE_CURRENT_SOURCE_DIR}/*.hpp" "${CMAKE_CURRENT_SOURCE_DIR}/*.h") + +add_custom_target( + format + COMMAND clang-format + -i + -style=file + ${FMT_SOURCES} + COMMENT "Running clang-format on all source files" +) + + +add_library(llamafile STATIC ${SOURCE_DIR4}) + +message(STATUS "CMAKE_CXX_FLAGS: ${CMAKE_CXX_FLAGS}") +message(STATUS "ARCH_FLAGS: ${ARCH_FLAGS}") +pybind11_add_module(${PROJECT_NAME} MODULE ${ALL_SOURCES}) +target_link_libraries(${PROJECT_NAME} PRIVATE llama) + + +if(WIN32) + target_link_libraries(${PROJECT_NAME} PRIVATE "$ENV{CUDA_PATH}/lib/x64/cudart.lib")#CUDA::cudart +elseif(UNIX) + if (KTRANSFORMERS_USE_ROCM) + add_compile_definitions(USE_HIP=1) + target_link_libraries(${PROJECT_NAME} PRIVATE "${ROCM_PATH}/lib/libamdhip64.so") + message(STATUS "Building for HIP") + elseif(KTRANSFORMERS_USE_MUSA) + target_link_libraries(${PROJECT_NAME} PRIVATE MUSA::musart) + elseif(KTRANSFORMERS_USE_XPU) + elseif(KTRANSFORMERS_USE_CUDA AND NOT KTRANSFORMERS_USE_MUSA) + target_link_libraries(${PROJECT_NAME} PRIVATE "${CUDAToolkit_LIBRARY_DIR}/libcudart.so") + endif() +endif() + +# Define the USE_NUMA option +option(USE_NUMA "Disable NUMA support" OFF) + +# Check if the USE_NUMA environment variable is set +if(DEFINED ENV{USE_NUMA}) + set(USE_NUMA ON) +endif() + +if(USE_NUMA) + message(STATUS "NUMA support is enabled") +else() + message(STATUS "NUMA support is disabled") +endif() + +find_library(NUMA_LIBRARY NAMES numa) + +if(NUMA_LIBRARY AND USE_NUMA) + message(STATUS "NUMA library found: ${NUMA_LIBRARY} - enabling NUMA support") + target_link_libraries(${PROJECT_NAME} PRIVATE ${NUMA_LIBRARY}) + target_compile_definitions(${PROJECT_NAME} PRIVATE USE_NUMA) +else() + if(USE_NUMA) + message(FATAL_ERROR "NUMA library not found - maybe sudo apt install libnuma-dev") + else() + message(STATUS "NUMA library not found or user not set USE_NUMA - disabling NUMA support") + endif() +endif() diff --git a/archive/kt-sft/csrc/ktransformers_ext/bench/bench_attention.py b/archive/kt-sft/csrc/ktransformers_ext/bench/bench_attention.py new file mode 100644 index 0000000..1b8b9b8 --- /dev/null +++ b/archive/kt-sft/csrc/ktransformers_ext/bench/bench_attention.py @@ -0,0 +1,178 @@ +#!/usr/bin/env python +# coding=utf-8 +""" +Description : +Author : Jianwei Dong +Date : 2024-08-28 10:32:05 +Version : 1.0.0 +LastEditors : Jianwei Dong +LastEditTime : 2024-08-28 10:32:05 +Copyright (c) 2024 by KVCache.AI, All Rights Reserved. +""" +import os, sys +import time + +sys.path.append(os.path.dirname(__file__) + "/../build") +import cpuinfer_ext +import torch + +layer_num = 10 +kv_head_num = 8 +q_head_num = 32 +head_dim = 128 +block_len = 128 +anchor_num = 1 + +anchor_type = cpuinfer_ext.kvcache.AnchorType.DYNAMIC +kv_type = cpuinfer_ext.kvcache.ggml_type.FP16 +retrieval_type = cpuinfer_ext.kvcache.RetrievalType.LAYER +layer_step: int = 1 +token_step: int = 1 +layer_offset: int = 0 +max_thread_num: int = 64 +max_batch_size: int = 1 +max_block_num: int = 1024 +CPUInfer = cpuinfer_ext.CPUInfer(max_thread_num) + +warm_up_iter = 1000 +test_iter = 10000 + + +def bench_linear(cache_seqlen: int): + with torch.inference_mode(mode=True): + cache_seqlens = torch.tensor([cache_seqlen], dtype=torch.int32, device="cpu") + seqlens_zero = torch.zeros((1,), dtype=torch.int32, device="cpu") + + config = cpuinfer_ext.kvcache.KVCacheConfig( + layer_num, + kv_head_num, + q_head_num, + head_dim, + block_len, + anchor_num, + anchor_type, + kv_type, + retrieval_type, + layer_step, + token_step, + layer_offset, + max_block_num, + max_batch_size, + max_thread_num, + ) + local_kvcache = cpuinfer_ext.kvcache.KVCache(config) + block_table = ( + torch.arange(max_block_num, dtype=torch.int32, device="cpu") + .contiguous() + .view(1, -1) + ) + + for layer_idx in range(layer_num): + k_cache = torch.randn( + (1, cache_seqlen, kv_head_num, head_dim), + dtype=torch.float16, + device="cpu", + ).contiguous() + v_cache = torch.randn( + (1, cache_seqlen, kv_head_num, head_dim), + dtype=torch.float16, + device="cpu", + ).contiguous() + + CPUInfer.submit( + local_kvcache.update_kvcache_fp16( + k_cache.data_ptr(), + v_cache.data_ptr(), + layer_idx, + block_table.data_ptr(), + 1, + max_block_num, + seqlens_zero.data_ptr(), + cache_seqlen, + ) + ) + CPUInfer.sync() + + input = torch.randn( + (1, 1, q_head_num, head_dim), dtype=torch.float16, device="cpu" + ).contiguous() + output = torch.empty( + (1, 1, q_head_num, head_dim), dtype=torch.float16, device="cpu" + ).contiguous() + + # attn_lse: (bsz, q_len, q_head_num) + attn_lse = torch.empty( + (1, 1, q_head_num), dtype=torch.float32, device="cpu" + ).contiguous() + input = input / 100 + + # warm up + for i in range(warm_up_iter): + CPUInfer.submit( + local_kvcache.attn( + input.data_ptr(), + output.data_ptr(), + attn_lse.data_ptr(), + i % layer_num, + 0, + 1, + 1, + max_block_num, + block_table.data_ptr(), + cache_seqlens.data_ptr(), + -1, + -1, + -1, + ) + ) + CPUInfer.sync() + + # test + start = time.perf_counter() + for i in range(test_iter): + CPUInfer.submit( + local_kvcache.attn( + input.data_ptr(), + output.data_ptr(), + attn_lse.data_ptr(), + i % layer_num, + 0, + 1, + 1, + max_block_num, + block_table.data_ptr(), + cache_seqlens.data_ptr(), + -1, + -1, + -1, + ) + ) + CPUInfer.sync() + end = time.perf_counter() + total_time = end - start + print("cache sequence length: ", cache_seqlen) + print("Time(s): ", total_time) + print("Iteration: ", test_iter) + print("Time(us) per iteration: ", total_time / test_iter * 1000000) + print( + "Bandwidth: ", + cache_seqlen + * kv_head_num + * head_dim + * 2 + * 2 + * test_iter + / total_time + / 1000 + / 1000 + / 1000, + "GB/s", + ) + print("") + + +bench_linear(1024) +bench_linear(4096) +bench_linear(16384) +bench_linear(32768) +bench_linear(65536) diff --git a/archive/kt-sft/csrc/ktransformers_ext/bench/bench_attention_torch.py b/archive/kt-sft/csrc/ktransformers_ext/bench/bench_attention_torch.py new file mode 100644 index 0000000..25f20e7 --- /dev/null +++ b/archive/kt-sft/csrc/ktransformers_ext/bench/bench_attention_torch.py @@ -0,0 +1,94 @@ +#!/usr/bin/env python +# coding=utf-8 +""" +Description : +Author : Jianwei Dong +Date : 2024-08-28 10:32:05 +Version : 1.0.0 +LastEditors : Jianwei Dong +LastEditTime : 2024-08-28 10:32:05 +Copyright (c) 2024 by KVCache.AI, All Rights Reserved. +""" +import os, sys +import time + +sys.path.append(os.path.dirname(__file__) + "/../build") +import cpuinfer_ext +import torch + +layer_num = 10 +kv_head_num = 8 +q_head_num = 32 +head_dim = 128 +block_len = 128 +anchor_num = 1 +warm_up_iter = 1000 +test_iter = 10000 + + +def bench_linear(cache_seqlen: int, device): + with torch.inference_mode(mode=True): + + kvcaches = [] + + for layer_idx in range(layer_num): + k_cache = torch.randn( + (1, 32, cache_seqlen, head_dim), + dtype=torch.float16, + device=device, + ).contiguous() + v_cache = torch.randn( + (1, 32, cache_seqlen, head_dim), + dtype=torch.float16, + device=device, + ).contiguous() + + kvcaches.append((k_cache, v_cache)) + + input = torch.randn( + (1, q_head_num, 1, head_dim), dtype=torch.float16, device=device + ).contiguous() + input = input / 100 + + # warm up + for i in range(warm_up_iter): + k_cache = kvcaches[i % layer_num][0] + v_cache = kvcaches[i % layer_num][1] + torch.nn.functional.scaled_dot_product_attention(input, k_cache, v_cache) + + # test + start = time.perf_counter() + for i in range(test_iter): + k_cache = kvcaches[i % layer_num][0] + v_cache = kvcaches[i % layer_num][1] + torch.nn.functional.scaled_dot_product_attention(input, k_cache, v_cache) + end = time.perf_counter() + total_time = end - start + print("cache sequence length: ", cache_seqlen) + print("Time(s): ", total_time) + print("Iteration: ", test_iter) + print("Time(us) per iteration: ", total_time / test_iter * 1000000) + print( + "Bandwidth: ", + cache_seqlen + * q_head_num + * head_dim + * 2 + * 2 + * test_iter + / total_time + / 1000 + / 1000 + / 1000, + "GB/s", + ) + print("") + + +bench_linear(1024, "cpu") +bench_linear(4096, "cpu") +bench_linear(1024, "cuda") +bench_linear(4096, "cuda") +bench_linear(16384, "cuda") +bench_linear(32768, "cuda") +bench_linear(65536, "cuda") diff --git a/archive/kt-sft/csrc/ktransformers_ext/bench/bench_linear.py b/archive/kt-sft/csrc/ktransformers_ext/bench/bench_linear.py new file mode 100644 index 0000000..3189afd --- /dev/null +++ b/archive/kt-sft/csrc/ktransformers_ext/bench/bench_linear.py @@ -0,0 +1,121 @@ +#!/usr/bin/env python +# coding=utf-8 +''' +Description : +Author : chenht2022 +Date : 2024-07-25 10:31:59 +Version : 1.0.0 +LastEditors : chenht2022 +LastEditTime : 2024-08-06 10:35:35 +Copyright (c) 2024 by KVCache.AI, All Rights Reserved. +''' +import os, sys +import time +sys.path.append(os.path.dirname(__file__) + '/../build') +import cpuinfer_ext +import torch + +input_size = 16384 +output_size = 5120 +stride = 16 +group_max_len = 1024 +layer_num = 10 +qlen = 1 +CPUInfer = cpuinfer_ext.CPUInfer(64) +warm_up_iter = 1000 +test_iter = 10000 + +def bench_linear(quant_mode: str): + with torch.inference_mode(mode=True): + + hidden_type = 30 # ggml_type::GGML_TYPE_BF16 + if quant_mode == "fp32": + proj_type = 0 # ggml_type::GGML_TYPE_F32 + bytes_per_elem = 4.000000 + elif quant_mode == "fp16": + proj_type = 1 # ggml_type::GGML_TYPE_F16 + bytes_per_elem = 2.000000 + elif quant_mode == "bf16": + proj_type = 30 # ggml_type::GGML_TYPE_BF16 + bytes_per_elem = 2.000000 + elif quant_mode == "q8_0": + proj_type = 8 # ggml_type::GGML_TYPE_Q8_0 + bytes_per_elem = 1.062500 + elif quant_mode == "q6_k": + proj_type = 14 # ggml_type::GGML_TYPE_Q6_K + bytes_per_elem = 0.820312 + elif quant_mode == "q5_k_m": + proj_type = 13 # ggml_type::GGML_TYPE_Q5_K + bytes_per_elem = 0.687500 + elif quant_mode == "q4_k_m": + proj_type = 12 # ggml_type::GGML_TYPE_Q4_K + bytes_per_elem = 0.562500 + elif quant_mode == "q3_k_m": + proj_type = 11 # ggml_type::GGML_TYPE_Q3_K + bytes_per_elem = 0.429688 + elif quant_mode == "q2_k": + proj_type = 10 # ggml_type::GGML_TYPE_Q2_K + bytes_per_elem = 0.328125 + elif quant_mode == "iq3_xs": + proj_type = 21 # ggml_type::GGML_TYPE_IQ3_S + bytes_per_elem = 0.429688 + elif quant_mode == "iq2_xxs": + proj_type = 16 # ggml_type::GGML_TYPE_IQ2_XXS + bytes_per_elem = 0.257812 + else: + assert(False) + + linears = [] + projs = [] + for _ in range(layer_num): + proj = torch.randn((output_size, input_size), dtype=torch.float32, device = "cuda").to("cpu").contiguous() + config = cpuinfer_ext.linear.LinearConfig(input_size, output_size, stride, group_max_len, proj.data_ptr(), proj_type, hidden_type) + linear = cpuinfer_ext.linear.Linear(config) + projs.append(proj) + linears.append(linear) + input = torch.randn((layer_num, qlen, input_size), dtype=torch.bfloat16, device = "cuda").to("cpu").contiguous() + output = torch.empty((layer_num, qlen, output_size), dtype=torch.bfloat16, device = "cuda").to("cpu").contiguous() + + # warm up + for i in range(warm_up_iter): + CPUInfer.submit( + linears[i % layer_num].forward( + qlen, + input[i % layer_num].data_ptr(), + output[i % layer_num].data_ptr() + ) + ) + CPUInfer.sync() + + # test + start = time.perf_counter() + for i in range(test_iter): + CPUInfer.submit( + linears[i % layer_num].forward( + qlen, + input[i % layer_num].data_ptr(), + output[i % layer_num].data_ptr() + ) + ) + CPUInfer.sync() + end = time.perf_counter() + total_time = end - start + print('Quant mode: ', quant_mode) + print('Time(s): ', total_time) + print('Iteration: ', test_iter) + print('Time(us) per iteration: ', total_time / test_iter * 1000000) + print('Bandwidth: ', input_size * output_size * bytes_per_elem * test_iter / total_time / 1000 / 1000 / 1000, 'GB/s') + print('') + +bench_linear("fp32") +bench_linear("fp16") +bench_linear("bf16") +bench_linear("q8_0") +bench_linear("q6_k") +bench_linear("q5_k_m") +bench_linear("q4_k_m") +bench_linear("q3_k_m") +bench_linear("q2_k") +# Not supported on __x86_64__ +# bench_linear("iq3_xs") +# bench_linear("iq2_xxs") diff --git a/archive/kt-sft/csrc/ktransformers_ext/bench/bench_linear_torch.py b/archive/kt-sft/csrc/ktransformers_ext/bench/bench_linear_torch.py new file mode 100644 index 0000000..72e0e75 --- /dev/null +++ b/archive/kt-sft/csrc/ktransformers_ext/bench/bench_linear_torch.py @@ -0,0 +1,83 @@ +#!/usr/bin/env python +# coding=utf-8 +''' +Description : +Author : chenht2022 +Date : 2024-07-25 10:31:59 +Version : 1.0.0 +LastEditors : chenht2022 +LastEditTime : 2024-07-25 10:32:48 +Copyright (c) 2024 by KVCache.AI, All Rights Reserved. +''' +import os, sys +import time +import torch +import torch.nn.quantized as nnq + +scale, zero_point = 0.1, 0 # Adjust scale and zero_point based on your dataset + +input_size = 16384 +output_size = 5120 +layer_num = 10 +qlen = 1 +warm_up_iter = 1000 +test_iter = 10000 + +def bench_linear(quant_mode: str): + with torch.inference_mode(mode=True): + if quant_mode == "fp32": + proj_type = torch.float32 + bytes_per_elem = 4.000000 + elif quant_mode == "fp16": + proj_type = torch.float16 + bytes_per_elem = 2.000000 + elif quant_mode == "bf16": + proj_type = torch.bfloat16 + bytes_per_elem = 2.000000 + elif quant_mode == "qint8": + proj_type = torch.qint8 + bytes_per_elem = 1.000000 + else: + assert(False) + + projs = [] + for _ in range(layer_num): + proj = torch.randn((output_size, input_size), dtype = torch.float32, device = "cuda").to("cpu").contiguous() + if quant_mode == "qint8": + proj_q = torch.quantize_per_tensor(proj, scale, zero_point, torch.qint8) + quantized_layer = nnq.Linear(input_size, output_size) + quantized_layer.set_weight_bias(proj_q, None) + projs.append(quantized_layer) + else: + projs.append(proj.to(proj_type)) + input = torch.randn((layer_num, qlen, input_size), dtype=torch.bfloat16, device = "cuda").to("cpu").contiguous() + + # warm up + for i in range(warm_up_iter): + if isinstance(projs[i % layer_num], nnq.Linear): + input_q = torch.quantize_per_tensor(input[i % layer_num].to(torch.float32), scale, zero_point, torch.quint8) + t_output = projs[i % layer_num](input_q) + else: + t_output = torch.mm(input[i % layer_num].to(proj_type), projs[i % layer_num].t()) + + # test + start = time.perf_counter() + for i in range(test_iter): + if isinstance(projs[i % layer_num], nnq.Linear): + input_q = torch.quantize_per_tensor(input[i % layer_num].to(torch.float32), scale, zero_point, torch.quint8) + t_output = projs[i % layer_num](input_q) + else: + t_output = torch.mm(input[i % layer_num].to(proj_type), projs[i % layer_num].t()) + end = time.perf_counter() + total_time = end - start + print('Quant mode: ', quant_mode) + print('Time(s): ', total_time) + print('Iteration: ', test_iter) + print('Time(us) per iteration: ', total_time / test_iter * 1000000) + print('Bandwidth: ', input_size * output_size * bytes_per_elem * test_iter / total_time / 1000 / 1000 / 1000, 'GB/s') + print('') + +bench_linear("fp32") +bench_linear("fp16") +bench_linear("bf16") +bench_linear("qint8") diff --git a/archive/kt-sft/csrc/ktransformers_ext/bench/bench_mlp.py b/archive/kt-sft/csrc/ktransformers_ext/bench/bench_mlp.py new file mode 100644 index 0000000..690f9d9 --- /dev/null +++ b/archive/kt-sft/csrc/ktransformers_ext/bench/bench_mlp.py @@ -0,0 +1,150 @@ +#!/usr/bin/env python +# coding=utf-8 +''' +Description : +Author : chenht2022 +Date : 2024-07-16 10:43:18 +Version : 1.0.0 +LastEditors : chenht2022 +LastEditTime : 2024-08-06 10:36:04 +Copyright (c) 2024 by KVCache.AI, All Rights Reserved. +''' +import os, sys +import time +sys.path.append(os.path.dirname(__file__) + '/../build') +import cpuinfer_ext +import torch + +hidden_size = 5120 +intermediate_size = 3072 +stride = 16 +group_max_len = 1024 +layer_num = 10 +qlen = 1 +CPUInfer = cpuinfer_ext.CPUInfer(64) +warm_up_iter = 1000 +test_iter = 10000 + +def bench_mlp(quant_mode: str): + with torch.inference_mode(mode=True): + + hidden_type = 30 # ggml_type::GGML_TYPE_BF16 + if quant_mode == "fp32": + gate_type = 0 # ggml_type::GGML_TYPE_F32 + up_type = 0 # ggml_type::GGML_TYPE_F32 + down_type = 0 # ggml_type::GGML_TYPE_F32 + bytes_per_elem = 4.000000 + elif quant_mode == "fp16": + gate_type = 1 # ggml_type::GGML_TYPE_F16 + up_type = 1 # ggml_type::GGML_TYPE_F16 + down_type = 1 # ggml_type::GGML_TYPE_F16 + bytes_per_elem = 2.000000 + elif quant_mode == "bf16": + gate_type = 30 # ggml_type::GGML_TYPE_BF16 + up_type = 30 # ggml_type::GGML_TYPE_BF16 + down_type = 30 # ggml_type::GGML_TYPE_BF16 + bytes_per_elem = 2.000000 + elif quant_mode == "q8_0": + gate_type = 8 # ggml_type::GGML_TYPE_Q8_0 + up_type = 8 # ggml_type::GGML_TYPE_Q8_0 + down_type = 8 # ggml_type::GGML_TYPE_Q8_0 + bytes_per_elem = 1.062500 + elif quant_mode == "q6_k": + gate_type = 14 # ggml_type::GGML_TYPE_Q6_K + up_type = 14 # ggml_type::GGML_TYPE_Q6_K + down_type = 14 # ggml_type::GGML_TYPE_Q6_K + bytes_per_elem = 0.820312 + elif quant_mode == "q5_k_m": + gate_type = 13 # ggml_type::GGML_TYPE_Q5_K + up_type = 13 # ggml_type::GGML_TYPE_Q5_K + down_type = 14 # ggml_type::GGML_TYPE_Q6_K + bytes_per_elem = 0.731771 + elif quant_mode == "q4_k_m": + gate_type = 12 # ggml_type::GGML_TYPE_Q4_K + up_type = 12 # ggml_type::GGML_TYPE_Q4_K + down_type = 14 # ggml_type::GGML_TYPE_Q6_K + bytes_per_elem = 0.648437 + elif quant_mode == "q3_k_m": + gate_type = 11 # ggml_type::GGML_TYPE_Q3_K + up_type = 11 # ggml_type::GGML_TYPE_Q3_K + down_type = 13 # ggml_type::GGML_TYPE_Q5_K + bytes_per_elem = 0.515625 + elif quant_mode == "q2_k": + gate_type = 10 # ggml_type::GGML_TYPE_Q2_K + up_type = 10 # ggml_type::GGML_TYPE_Q2_K + down_type = 11 # ggml_type::GGML_TYPE_Q3_K + bytes_per_elem = 0.328125 + elif quant_mode == "iq3_xs": + gate_type = 21 # ggml_type::GGML_TYPE_IQ3_S + up_type = 21 # ggml_type::GGML_TYPE_IQ3_S + down_type = 21 # ggml_type::GGML_TYPE_IQ3_S + bytes_per_elem = 0.429688 + elif quant_mode == "iq2_xxs": + gate_type = 16 # ggml_type::GGML_TYPE_IQ2_XXS + up_type = 16 # ggml_type::GGML_TYPE_IQ2_XXS + down_type = 16 # ggml_type::GGML_TYPE_IQ2_XXS + bytes_per_elem = 0.257812 + else: + assert(False) + + + mlps = [] + gate_projs = [] + up_projs = [] + down_projs = [] + for _ in range(layer_num): + gate_proj = torch.randn((intermediate_size, hidden_size), dtype=torch.float32, device = "cuda").to("cpu").contiguous() + up_proj = torch.randn((intermediate_size, hidden_size), dtype=torch.float32, device = "cuda").to("cpu").contiguous() + down_proj = torch.randn((hidden_size, intermediate_size), dtype=torch.float32, device = "cuda").to("cpu").contiguous() + config = cpuinfer_ext.mlp.MLPConfig(hidden_size, intermediate_size, stride, group_max_len, gate_proj.data_ptr(), up_proj.data_ptr(), down_proj.data_ptr(), gate_type, up_type, down_type, hidden_type) + mlp = cpuinfer_ext.mlp.MLP(config) + gate_projs.append(gate_proj) + up_projs.append(up_proj) + down_projs.append(down_proj) + mlps.append(mlp) + input = torch.randn((layer_num, qlen, hidden_size), dtype=torch.bfloat16, device = "cuda").to("cpu").contiguous() + output = torch.empty((layer_num, qlen, hidden_size), dtype=torch.bfloat16, device = "cuda").to("cpu").contiguous() + + # warm up + for i in range(warm_up_iter): + CPUInfer.submit( + mlps[i % layer_num].forward( + qlen, + input[i % layer_num].data_ptr(), + output[i % layer_num].data_ptr() + ) + ) + CPUInfer.sync() + + # test + start = time.perf_counter() + for i in range(test_iter): + CPUInfer.submit( + mlps[i % layer_num].forward( + qlen, + input[i % layer_num].data_ptr(), + output[i % layer_num].data_ptr() + ) + ) + CPUInfer.sync() + end = time.perf_counter() + total_time = end - start + print('Quant mode: ', quant_mode) + print('Time(s): ', total_time) + print('Iteration: ', test_iter) + print('Time(us) per iteration: ', total_time / test_iter * 1000000) + print('Bandwidth: ', hidden_size * intermediate_size * 3 * bytes_per_elem * test_iter / total_time / 1000 / 1000 / 1000, 'GB/s') + print('') + +bench_mlp("fp32") +bench_mlp("fp16") +bench_mlp("bf16") +bench_mlp("q8_0") +bench_mlp("q6_k") +bench_mlp("q5_k_m") +bench_mlp("q4_k_m") +bench_mlp("q3_k_m") +bench_mlp("q2_k") +# Not supported on __x86_64__ +# bench_linear("iq3_xs") +# bench_linear("iq2_xxs") diff --git a/archive/kt-sft/csrc/ktransformers_ext/bench/bench_mlp_torch.py b/archive/kt-sft/csrc/ktransformers_ext/bench/bench_mlp_torch.py new file mode 100644 index 0000000..7b811d8 --- /dev/null +++ b/archive/kt-sft/csrc/ktransformers_ext/bench/bench_mlp_torch.py @@ -0,0 +1,110 @@ +#!/usr/bin/env python +# coding=utf-8 +''' +Description : +Author : chenht2022 +Date : 2024-07-16 10:43:18 +Version : 1.0.0 +LastEditors : chenht2022 +LastEditTime : 2024-07-25 10:32:53 +Copyright (c) 2024 by KVCache.AI, All Rights Reserved. +''' +import os, sys +import time +import torch +import torch.nn.quantized as nnq + +scale, zero_point = 0.1, 0 # Adjust scale and zero_point based on your dataset + +hidden_size = 5120 +intermediate_size = 3072 +layer_num = 10 +qlen = 1 +warm_up_iter = 1000 +test_iter = 10000 + +def act_fn(x): + return x / (1.0 + torch.exp(-x)) + +def mlp_torch(input, gate_proj, up_proj, down_proj): + if isinstance(gate_proj, nnq.Linear): + input_q = torch.quantize_per_tensor(input.to(torch.float32), scale, zero_point, torch.quint8) + gate_buf = gate_proj(input_q) + up_buf = up_proj(input_q) + gate_buf = gate_buf.dequantize() + up_buf = up_buf.dequantize() + intermediate = act_fn(gate_buf) * up_buf + intermediate_q = torch.quantize_per_tensor(intermediate, scale, zero_point, torch.quint8) + expert_output = down_proj(intermediate_q) + ret = expert_output.dequantize() + else: + gate_buf = torch.mm(input.to(gate_proj.dtype), gate_proj.t()) + up_buf = torch.mm(input.to(up_proj.dtype), up_proj.t()) + intermediate = act_fn(gate_buf) * up_buf + ret = torch.mm(intermediate.to(down_proj.dtype), down_proj.t()) + return ret + +def bench_mlp(quant_mode: str): + with torch.inference_mode(mode=True): + if quant_mode == "fp32": + proj_type = torch.float32 + bytes_per_elem = 4.000000 + elif quant_mode == "fp16": + proj_type = torch.float16 + bytes_per_elem = 2.000000 + elif quant_mode == "bf16": + proj_type = torch.bfloat16 + bytes_per_elem = 2.000000 + elif quant_mode == "qint8": + proj_type = torch.qint8 + bytes_per_elem = 1.000000 + else: + assert(False) + + gate_projs = [] + up_projs = [] + down_projs = [] + for _ in range(layer_num): + gate_proj = torch.randn((intermediate_size, hidden_size), dtype=torch.float32, device = "cuda").to("cpu").contiguous() + up_proj = torch.randn((intermediate_size, hidden_size), dtype=torch.float32, device = "cuda").to("cpu").contiguous() + down_proj = torch.randn((hidden_size, intermediate_size), dtype=torch.float32, device = "cuda").to("cpu").contiguous() + if quant_mode == "qint8": + gate_proj_q = torch.quantize_per_tensor(gate_proj, scale, zero_point, torch.qint8) + quantized_gate = nnq.Linear(hidden_size, intermediate_size) + quantized_gate.set_weight_bias(gate_proj_q, None) + up_proj_q = torch.quantize_per_tensor(up_proj, scale, zero_point, torch.qint8) + quantized_up = nnq.Linear(hidden_size, intermediate_size) + quantized_up.set_weight_bias(up_proj_q, None) + down_proj_q = torch.quantize_per_tensor(down_proj, scale, zero_point, torch.qint8) + quantized_down = nnq.Linear(intermediate_size, hidden_size) + quantized_down.set_weight_bias(down_proj_q, None) + gate_projs.append(quantized_gate) + up_projs.append(quantized_up) + down_projs.append(quantized_down) + else: + gate_projs.append(gate_proj.to(proj_type)) + up_projs.append(up_proj.to(proj_type)) + down_projs.append(down_proj.to(proj_type)) + input = torch.randn((layer_num, qlen, hidden_size), dtype=torch.bfloat16, device = "cuda").to("cpu").contiguous() + + # warm up + for i in range(warm_up_iter): + mlp_torch(input[i % layer_num], gate_projs[i % layer_num], up_projs[i % layer_num], down_projs[i % layer_num]) + + # test + start = time.perf_counter() + for i in range(test_iter): + mlp_torch(input[i % layer_num], gate_projs[i % layer_num], up_projs[i % layer_num], down_projs[i % layer_num]) + end = time.perf_counter() + total_time = end - start + print('Quant mode: ', quant_mode) + print('Time(s): ', total_time) + print('Iteration: ', test_iter) + print('Time(us) per iteration: ', total_time / test_iter * 1000000) + print('Bandwidth: ', hidden_size * intermediate_size * 3 * bytes_per_elem * test_iter / total_time / 1000 / 1000 / 1000, 'GB/s') + print('') + +bench_mlp("fp32") +bench_mlp("fp16") +bench_mlp("bf16") +bench_mlp("qint8") diff --git a/archive/kt-sft/csrc/ktransformers_ext/bench/bench_moe.py b/archive/kt-sft/csrc/ktransformers_ext/bench/bench_moe.py new file mode 100644 index 0000000..6d617b7 --- /dev/null +++ b/archive/kt-sft/csrc/ktransformers_ext/bench/bench_moe.py @@ -0,0 +1,160 @@ +#!/usr/bin/env python +# coding=utf-8 +''' +Description : +Author : chenht2022 +Date : 2024-07-25 10:32:05 +Version : 1.0.0 +LastEditors : chenht2022 +LastEditTime : 2024-08-06 10:41:28 +Copyright (c) 2024 by KVCache.AI, All Rights Reserved. +''' +import os, sys +import time +sys.path.append(os.path.dirname(__file__) + '/../build') +import cpuinfer_ext +import torch + +expert_num = 160 +hidden_size = 5120 +intermediate_size = 1536 +stride = 16 +group_min_len = 10 +group_max_len = 1024 +n_routed_experts = 6 +layer_num = 10 +qlen = 1 +CPUInfer = cpuinfer_ext.CPUInfer(64) +warm_up_iter = 1000 +test_iter = 10000 + +def bench_moe(quant_mode: str): + with torch.inference_mode(mode=True): + hidden_type = 30 # ggml_type::GGML_TYPE_BF16 + if quant_mode == "fp32": + gate_type = 0 # ggml_type::GGML_TYPE_F32 + up_type = 0 # ggml_type::GGML_TYPE_F32 + down_type = 0 # ggml_type::GGML_TYPE_F32 + bytes_per_elem = 4.000000 + elif quant_mode == "fp16": + gate_type = 1 # ggml_type::GGML_TYPE_F16 + up_type = 1 # ggml_type::GGML_TYPE_F16 + down_type = 1 # ggml_type::GGML_TYPE_F16 + bytes_per_elem = 2.000000 + elif quant_mode == "bf16": + gate_type = 30 # ggml_type::GGML_TYPE_BF16 + up_type = 30 # ggml_type::GGML_TYPE_BF16 + down_type = 30 # ggml_type::GGML_TYPE_BF16 + bytes_per_elem = 2.000000 + elif quant_mode == "q8_0": + gate_type = 8 # ggml_type::GGML_TYPE_Q8_0 + up_type = 8 # ggml_type::GGML_TYPE_Q8_0 + down_type = 8 # ggml_type::GGML_TYPE_Q8_0 + bytes_per_elem = 1.062500 + elif quant_mode == "q6_k": + gate_type = 14 # ggml_type::GGML_TYPE_Q6_K + up_type = 14 # ggml_type::GGML_TYPE_Q6_K + down_type = 14 # ggml_type::GGML_TYPE_Q6_K + bytes_per_elem = 0.820312 + elif quant_mode == "q5_k_m": + gate_type = 13 # ggml_type::GGML_TYPE_Q5_K + up_type = 13 # ggml_type::GGML_TYPE_Q5_K + down_type = 14 # ggml_type::GGML_TYPE_Q6_K + bytes_per_elem = 0.731771 + elif quant_mode == "q4_k_m": + gate_type = 12 # ggml_type::GGML_TYPE_Q4_K + up_type = 12 # ggml_type::GGML_TYPE_Q4_K + down_type = 14 # ggml_type::GGML_TYPE_Q6_K + bytes_per_elem = 0.648437 + elif quant_mode == "q3_k_m": + gate_type = 11 # ggml_type::GGML_TYPE_Q3_K + up_type = 11 # ggml_type::GGML_TYPE_Q3_K + down_type = 13 # ggml_type::GGML_TYPE_Q5_K + bytes_per_elem = 0.515625 + elif quant_mode == "q2_k": + gate_type = 10 # ggml_type::GGML_TYPE_Q2_K + up_type = 10 # ggml_type::GGML_TYPE_Q2_K + down_type = 11 # ggml_type::GGML_TYPE_Q3_K + bytes_per_elem = 0.328125 + elif quant_mode == "iq3_xs": + gate_type = 21 # ggml_type::GGML_TYPE_IQ3_S + up_type = 21 # ggml_type::GGML_TYPE_IQ3_S + down_type = 21 # ggml_type::GGML_TYPE_IQ3_S + bytes_per_elem = 0.429688 + elif quant_mode == "iq2_xxs": + gate_type = 16 # ggml_type::GGML_TYPE_IQ2_XXS + up_type = 16 # ggml_type::GGML_TYPE_IQ2_XXS + down_type = 16 # ggml_type::GGML_TYPE_IQ2_XXS + bytes_per_elem = 0.257812 + else: + assert(False) + + + moes = [] + gate_projs = [] + up_projs = [] + down_projs = [] + for _ in range(layer_num): + gate_proj = torch.randn((expert_num, intermediate_size, hidden_size), dtype=torch.float32, device = "cuda").to("cpu").contiguous() + up_proj = torch.randn((expert_num, intermediate_size, hidden_size), dtype=torch.float32, device = "cuda").to("cpu").contiguous() + down_proj = torch.randn((expert_num, hidden_size, intermediate_size), dtype=torch.float32, device = "cuda").to("cpu").contiguous() + config = cpuinfer_ext.moe.MOEConfig(expert_num, n_routed_experts, hidden_size, intermediate_size, stride, group_min_len, group_max_len, gate_proj.data_ptr(), up_proj.data_ptr(), down_proj.data_ptr(), gate_type, up_type, down_type, hidden_type) + moe = cpuinfer_ext.moe.MOE(config) + gate_projs.append(gate_proj) + up_projs.append(up_proj) + down_projs.append(down_proj) + moes.append(moe) + expert_ids = torch.stack([torch.stack([torch.randperm(expert_num, dtype=torch.int64, device = "cuda")[:n_routed_experts] for _ in range(qlen)]) for _ in range(layer_num)]).to("cpu").contiguous() + weights = torch.rand((layer_num, qlen, n_routed_experts), dtype=torch.float32, device = "cuda").to("cpu").contiguous() + input = torch.randn((layer_num, qlen, hidden_size), dtype=torch.bfloat16, device = "cuda").to("cpu").contiguous() + output = torch.empty((layer_num, qlen, hidden_size), dtype=torch.bfloat16, device = "cuda").to("cpu").contiguous() + + # warm up + for i in range(warm_up_iter): + CPUInfer.submit( + moes[i % layer_num].forward( + qlen, + n_routed_experts, + expert_ids[i % layer_num].data_ptr(), + weights[i % layer_num].data_ptr(), + input[i % layer_num].data_ptr(), + output[i % layer_num].data_ptr() + ) + ) + CPUInfer.sync() + + # test + start = time.perf_counter() + for i in range(test_iter): + CPUInfer.submit( + moes[i % layer_num].forward( + qlen, + n_routed_experts, + expert_ids[i % layer_num].data_ptr(), + weights[i % layer_num].data_ptr(), + input[i % layer_num].data_ptr(), + output[i % layer_num].data_ptr() + ) + ) + CPUInfer.sync() + end = time.perf_counter() + total_time = end - start + print('Quant mode: ', quant_mode) + print('Time(s): ', total_time) + print('Iteration: ', test_iter) + print('Time(us) per iteration: ', total_time / test_iter * 1000000) + print('Bandwidth: ', hidden_size * intermediate_size * 3 * n_routed_experts * bytes_per_elem * test_iter / total_time / 1000 / 1000 / 1000, 'GB/s') + print('') + +bench_moe("fp32") +bench_moe("fp16") +bench_moe("bf16") +bench_moe("q8_0") +bench_moe("q6_k") +bench_moe("q5_k_m") +bench_moe("q4_k_m") +bench_moe("q3_k_m") +bench_moe("q2_k") +# Not supported on __x86_64__ +# bench_linear("iq3_xs") +# bench_linear("iq2_xxs") diff --git a/archive/kt-sft/csrc/ktransformers_ext/bench/bench_moe_amx.py b/archive/kt-sft/csrc/ktransformers_ext/bench/bench_moe_amx.py new file mode 100644 index 0000000..1101ff8 --- /dev/null +++ b/archive/kt-sft/csrc/ktransformers_ext/bench/bench_moe_amx.py @@ -0,0 +1,107 @@ +#!/usr/bin/env python +# coding=utf-8 +''' +Description : +Author : chenht2022 +Date : 2025-04-25 18:28:12 +Version : 1.0.0 +LastEditors : chenht2022 +LastEditTime : 2025-04-25 18:28:12 +Copyright (c) 2024 by KVCache.AI, All Rights Reserved. +''' +import os, sys +import time +sys.path.append(os.path.dirname(__file__) + '/../build') +import cpuinfer_ext +import torch + +expert_num = 8 +hidden_size = 7168 +intermediate_size = 2048 +max_len = 25600 +n_routed_experts = 8 +layer_num = 10 +qlen = 1024 +CPUInfer = cpuinfer_ext.CPUInfer(65) +warm_up_iter = 100 +test_iter = 100 + +def bench_moe(quant_mode: str): + with torch.inference_mode(mode=True): + if quant_mode == "bf16": + bytes_per_elem = 2.000000 + elif quant_mode == "int8": + bytes_per_elem = 1.000000 + else: + assert(False) + + + moes = [] + gate_projs = [] + up_projs = [] + down_projs = [] + for _ in range(layer_num): + gate_proj = torch.randn((expert_num, intermediate_size, hidden_size), dtype=torch.float32, device = "cuda").to("cpu").contiguous() + up_proj = torch.randn((expert_num, intermediate_size, hidden_size), dtype=torch.float32, device = "cuda").to("cpu").contiguous() + down_proj = torch.randn((expert_num, hidden_size, intermediate_size), dtype=torch.float32, device = "cuda").to("cpu").contiguous() + config = cpuinfer_ext.moe.AMX_MOEConfig(expert_num, n_routed_experts, hidden_size, intermediate_size, max_len, gate_proj.data_ptr(), up_proj.data_ptr(), down_proj.data_ptr()) + if quant_mode == "bf16": + moe = cpuinfer_ext.moe.AMXBF16_MOE(config) + CPUInfer.submit(moe.load_weights()) + CPUInfer.sync() + elif quant_mode == "int8": + moe = cpuinfer_ext.moe.AMXInt8_MOE(config) + CPUInfer.submit(moe.load_weights()) + CPUInfer.sync() + gate_projs.append(gate_proj) + up_projs.append(up_proj) + down_projs.append(down_proj) + moes.append(moe) + expert_ids = torch.stack([torch.stack([torch.randperm(expert_num, dtype=torch.int64, device = "cuda")[:n_routed_experts] for _ in range(qlen)]) for _ in range(layer_num)]).to("cpu").contiguous() + weights = torch.rand((layer_num, qlen, n_routed_experts), dtype=torch.float32, device = "cuda").to("cpu").contiguous() + input = torch.randn((layer_num, qlen, hidden_size), dtype=torch.bfloat16, device = "cuda").to("cpu").contiguous() + output = torch.empty((layer_num, qlen, hidden_size), dtype=torch.bfloat16, device = "cuda").to("cpu").contiguous() + qlen_tensor = torch.tensor([qlen], dtype=torch.int32) + + # warm up + for i in range(warm_up_iter): + CPUInfer.submit( + moes[i % layer_num].forward( + qlen, + n_routed_experts, + expert_ids[i % layer_num].data_ptr(), + weights[i % layer_num].data_ptr(), + input[i % layer_num].data_ptr(), + output[i % layer_num].data_ptr(), + qlen_tensor.data_ptr() + ) + ) + CPUInfer.sync() + + # test + start = time.perf_counter() + for i in range(test_iter): + CPUInfer.submit( + moes[i % layer_num].forward( + qlen, + n_routed_experts, + expert_ids[i % layer_num].data_ptr(), + weights[i % layer_num].data_ptr(), + input[i % layer_num].data_ptr(), + output[i % layer_num].data_ptr(), + qlen_tensor.data_ptr() + ) + ) + CPUInfer.sync() + end = time.perf_counter() + total_time = end - start + print('Quant mode: ', quant_mode) + print('Time(s): ', total_time) + print('Iteration: ', test_iter) + print('Time(us) per iteration: ', total_time / test_iter * 1000000) + print('Bandwidth: ', hidden_size * intermediate_size * 3 * n_routed_experts * bytes_per_elem * test_iter / total_time / 1000 / 1000 / 1000, 'GB/s') + print('Flops: ', hidden_size * intermediate_size * qlen * 3 * n_routed_experts * 2 * test_iter / total_time / 1000 / 1000 / 1000, 'GFLOPS') + print('') + +bench_moe("bf16") +bench_moe("int8") diff --git a/archive/kt-sft/csrc/ktransformers_ext/bench/bench_moe_torch.py b/archive/kt-sft/csrc/ktransformers_ext/bench/bench_moe_torch.py new file mode 100644 index 0000000..1aecf40 --- /dev/null +++ b/archive/kt-sft/csrc/ktransformers_ext/bench/bench_moe_torch.py @@ -0,0 +1,152 @@ +#!/usr/bin/env python +# coding=utf-8 +''' +Description : +Author : chenht2022 +Date : 2024-07-25 10:32:05 +Version : 1.0.0 +LastEditors : chenht2022 +LastEditTime : 2024-07-25 10:32:57 +Copyright (c) 2024 by KVCache.AI, All Rights Reserved. +''' +import os, sys +import time +import torch +import torch.nn.quantized as nnq + +scale, zero_point = 0.1, 0 # Adjust scale and zero_point based on your dataset + +expert_num = 160 +hidden_size = 5120 +intermediate_size = 1536 +n_routed_experts = 6 +layer_num = 10 +qlen = 1 +warm_up_iter = 1000 +test_iter = 10000 + +def act_fn(x): + return x / (1.0 + torch.exp(-x)) + +def mlp_torch(input, gate_proj, up_proj, down_proj): + if isinstance(gate_proj, nnq.Linear): + input_q = torch.quantize_per_tensor(input.to(torch.float32), scale, zero_point, torch.quint8) + gate_buf = gate_proj(input_q) + up_buf = up_proj(input_q) + gate_buf = gate_buf.dequantize() + up_buf = up_buf.dequantize() + intermediate = act_fn(gate_buf) * up_buf + intermediate_q = torch.quantize_per_tensor(intermediate, scale, zero_point, torch.quint8) + expert_output = down_proj(intermediate_q) + ret = expert_output.dequantize() + else: + gate_buf = torch.mm(input.to(gate_proj.dtype), gate_proj.t()) + up_buf = torch.mm(input.to(up_proj.dtype), up_proj.t()) + intermediate = act_fn(gate_buf) * up_buf + ret = torch.mm(intermediate.to(down_proj.dtype), down_proj.t()) + return ret + +def moe_torch(input, expert_ids, weights, gate_proj, up_proj, down_proj): + cnts = expert_ids.new_zeros((expert_ids.shape[0], expert_num)) + cnts.scatter_(1, expert_ids, 1) + tokens_per_expert = cnts.sum(dim=0) + idxs = expert_ids.view(-1).argsort() + sorted_tokens = input[idxs // expert_ids.shape[1]] + + outputs = [] + start_idx = 0 + for i, num_tokens in enumerate(tokens_per_expert): + end_idx = start_idx + num_tokens + if num_tokens == 0: + continue + tokens_for_this_expert = sorted_tokens[start_idx:end_idx] + expert_out = mlp_torch(tokens_for_this_expert, gate_proj[i], up_proj[i], down_proj[i]) + outputs.append(expert_out) + start_idx = end_idx + + outs = torch.cat(outputs, dim=0) if len(outputs) else sorted_tokens.new_empty(0) + + new_x = torch.empty_like(outs) + new_x[idxs] = outs + t_output = ( + new_x.view(*expert_ids.shape, -1) + .type(weights.dtype) + .mul_(weights.unsqueeze(dim=-1)) + .sum(dim=1) + .type(new_x.dtype) + ) + return t_output + +def bench_moe(quant_mode: str): + with torch.inference_mode(mode=True): + if quant_mode == "fp32": + proj_type = torch.float32 + bytes_per_elem = 4.000000 + elif quant_mode == "fp16": + proj_type = torch.float16 + bytes_per_elem = 2.000000 + elif quant_mode == "bf16": + proj_type = torch.bfloat16 + bytes_per_elem = 2.000000 + elif quant_mode == "qint8": + proj_type = torch.qint8 + bytes_per_elem = 1.000000 + else: + assert(False) + + gate_projs = [] + up_projs = [] + down_projs = [] + for _ in range(layer_num): + gate_proj = torch.randn((expert_num, intermediate_size, hidden_size), dtype=torch.float32, device = "cuda").to("cpu").contiguous() + up_proj = torch.randn((expert_num, intermediate_size, hidden_size), dtype=torch.float32, device = "cuda").to("cpu").contiguous() + down_proj = torch.randn((expert_num, hidden_size, intermediate_size), dtype=torch.float32, device = "cuda").to("cpu").contiguous() + if quant_mode == "qint8": + quantized_gate_proj = [] + quantized_up_proj = [] + quantized_down_proj = [] + for i in range(expert_num): + gate_proj_q = torch.quantize_per_tensor(gate_proj[i], scale, zero_point, torch.qint8) + quantized_gate = nnq.Linear(hidden_size, intermediate_size) + quantized_gate.set_weight_bias(gate_proj_q, None) + quantized_gate_proj.append(quantized_gate) + up_proj_q = torch.quantize_per_tensor(up_proj[i], scale, zero_point, torch.qint8) + quantized_up = nnq.Linear(hidden_size, intermediate_size) + quantized_up.set_weight_bias(up_proj_q, None) + quantized_up_proj.append(quantized_up) + down_proj_q = torch.quantize_per_tensor(down_proj[i], scale, zero_point, torch.qint8) + quantized_down = nnq.Linear(intermediate_size, hidden_size) + quantized_down.set_weight_bias(down_proj_q, None) + quantized_down_proj.append(quantized_down) + gate_projs.append(quantized_gate_proj) + up_projs.append(quantized_up_proj) + down_projs.append(quantized_down_proj) + else: + gate_projs.append(gate_proj.to(proj_type)) + up_projs.append(up_proj.to(proj_type)) + down_projs.append(down_proj.to(proj_type)) + expert_ids = torch.stack([torch.stack([torch.randperm(expert_num, dtype=torch.int64, device = "cuda")[:n_routed_experts] for _ in range(qlen)]) for _ in range(layer_num)]).to("cpu").contiguous() + weights = torch.rand((layer_num, qlen, n_routed_experts), dtype=torch.float32, device = "cuda").to("cpu").contiguous() + input = torch.randn((layer_num, qlen, hidden_size), dtype=torch.bfloat16, device = "cuda").to("cpu").contiguous() + + # warm up + for i in range(warm_up_iter): + moe_torch(input[i % layer_num], expert_ids[i % layer_num], weights[i % layer_num], gate_projs[i % layer_num], up_projs[i % layer_num], down_projs[i % layer_num]) + + # test + start = time.perf_counter() + for i in range(test_iter): + moe_torch(input[i % layer_num], expert_ids[i % layer_num], weights[i % layer_num], gate_projs[i % layer_num], up_projs[i % layer_num], down_projs[i % layer_num]) + end = time.perf_counter() + total_time = end - start + print('Quant mode: ', quant_mode) + print('Time(s): ', total_time) + print('Iteration: ', test_iter) + print('Time(us) per iteration: ', total_time / test_iter * 1000000) + print('Bandwidth: ', hidden_size * intermediate_size * 3 * n_routed_experts * bytes_per_elem * test_iter / total_time / 1000 / 1000 / 1000, 'GB/s') + print('') + +bench_moe("fp32") +bench_moe("fp16") +bench_moe("bf16") +bench_moe("qint8") diff --git a/archive/kt-sft/csrc/ktransformers_ext/cmake/FindSIMD.cmake b/archive/kt-sft/csrc/ktransformers_ext/cmake/FindSIMD.cmake new file mode 100644 index 0000000..33377ec --- /dev/null +++ b/archive/kt-sft/csrc/ktransformers_ext/cmake/FindSIMD.cmake @@ -0,0 +1,100 @@ +include(CheckCSourceRuns) + +set(AVX_CODE " + #include + int main() + { + __m256 a; + a = _mm256_set1_ps(0); + return 0; + } +") + +set(AVX512_CODE " + #include + int main() + { + __m512i a = _mm512_set_epi8(0, 0, 0, 0, 0, 0, 0, 0, + 0, 0, 0, 0, 0, 0, 0, 0, + 0, 0, 0, 0, 0, 0, 0, 0, + 0, 0, 0, 0, 0, 0, 0, 0, + 0, 0, 0, 0, 0, 0, 0, 0, + 0, 0, 0, 0, 0, 0, 0, 0, + 0, 0, 0, 0, 0, 0, 0, 0, + 0, 0, 0, 0, 0, 0, 0, 0); + __m512i b = a; + __mmask64 equality_mask = _mm512_cmp_epi8_mask(a, b, _MM_CMPINT_EQ); + return 0; + } +") + +set(AVX2_CODE " + #include + int main() + { + __m256i a = {0}; + a = _mm256_abs_epi16(a); + __m256i x; + _mm256_extract_epi64(x, 0); // we rely on this in our AVX2 code + return 0; + } +") + +set(FMA_CODE " + #include + int main() + { + __m256 acc = _mm256_setzero_ps(); + const __m256 d = _mm256_setzero_ps(); + const __m256 p = _mm256_setzero_ps(); + acc = _mm256_fmadd_ps( d, p, acc ); + return 0; + } +") + +macro(check_sse type flags) + set(__FLAG_I 1) + set(CMAKE_REQUIRED_FLAGS_SAVE ${CMAKE_REQUIRED_FLAGS}) + foreach (__FLAG ${flags}) + if (NOT ${type}_FOUND) + set(CMAKE_REQUIRED_FLAGS ${__FLAG}) + check_c_source_runs("${${type}_CODE}" HAS_${type}_${__FLAG_I}) + if (HAS_${type}_${__FLAG_I}) + set(${type}_FOUND TRUE CACHE BOOL "${type} support") + set(${type}_FLAGS "${__FLAG}" CACHE STRING "${type} flags") + endif() + math(EXPR __FLAG_I "${__FLAG_I}+1") + endif() + endforeach() + set(CMAKE_REQUIRED_FLAGS ${CMAKE_REQUIRED_FLAGS_SAVE}) + + if (NOT ${type}_FOUND) + set(${type}_FOUND FALSE CACHE BOOL "${type} support") + set(${type}_FLAGS "" CACHE STRING "${type} flags") + endif() + + mark_as_advanced(${type}_FOUND ${type}_FLAGS) +endmacro() + +# flags are for MSVC only! +check_sse("AVX" " ;/arch:AVX") +if (NOT ${AVX_FOUND}) + set(LLAMA_AVX OFF) +else() + set(LLAMA_AVX ON) +endif() + +check_sse("AVX2" " ;/arch:AVX2") +check_sse("FMA" " ;/arch:AVX2") +if ((NOT ${AVX2_FOUND}) OR (NOT ${FMA_FOUND})) + set(LLAMA_AVX2 OFF) +else() + set(LLAMA_AVX2 ON) +endif() + +check_sse("AVX512" " ;/arch:AVX512") +if (NOT ${AVX512_FOUND}) + set(LLAMA_AVX512 OFF) +else() + set(LLAMA_AVX512 ON) +endif() diff --git a/archive/kt-sft/csrc/ktransformers_ext/cpu_backend/backend.cpp b/archive/kt-sft/csrc/ktransformers_ext/cpu_backend/backend.cpp new file mode 100644 index 0000000..7478d5c --- /dev/null +++ b/archive/kt-sft/csrc/ktransformers_ext/cpu_backend/backend.cpp @@ -0,0 +1,154 @@ +/** + * @Description : + * @Author : chenht2022 + * @Date : 2024-07-22 02:03:05 + * @Version : 1.0.0 + * @LastEditors : chenht2022 + * @LastEditTime : 2024-07-25 10:33:34 + * @Copyright (c) 2024 by KVCache.AI, All Rights Reserved. + **/ + +#include "backend.h" + +#ifdef USE_NUMA +#include +#include + +thread_local int Backend::numa_node = -1; +#endif + +thread_local int Backend::thread_local_id = -1; + +Backend::Backend(int max_thread_num) { + max_thread_num_ = max_thread_num; + thread_state_.resize(max_thread_num_); + for (int i = 0; i < max_thread_num_; i++) { + thread_state_[i].curr = std::make_unique>(); + thread_state_[i].status = + std::make_unique>(ThreadStatus::WAITING); + } + workers_.resize(max_thread_num_); + for (int i = 1; i < max_thread_num_; i++) { + workers_[i] = std::thread(&Backend::worker_thread, this, i); + } +} + +Backend::~Backend() { + for (int i = 0; i < max_thread_num_; i++) { + thread_state_[i].status->store(ThreadStatus::EXIT, + std::memory_order_release); + } + for (int i = 1; i < max_thread_num_; i++) { + if (workers_[i].joinable()) { + workers_[i].join(); + } + } +} + +int Backend::get_thread_num() { return max_thread_num_; } + +void Backend::do_work_stealing_job(int task_num, + std::function init_func, + std::function compute_func, + std::function finalize_func) { + init_func_ = init_func; + compute_func_ = compute_func; + finalize_func_ = finalize_func; +#ifdef USE_NUMA + // numa node location will be calculated based on the number of threads + thread_num_ = max_thread_num_; +#else + thread_num_ = std::min(max_thread_num_, task_num); +#endif + int base = task_num / thread_num_; + int remain = task_num % thread_num_; + thread_state_[0].end = base + (0 < remain); + + // 为主线程设置 thread_local_id + thread_local_id = 0; + + for (int i = 1; i < thread_num_; i++) { + thread_state_[i].curr->store(thread_state_[i - 1].end, + std::memory_order_relaxed); + thread_state_[i].end = thread_state_[i - 1].end + base + (i < remain); + thread_state_[i].status->store(ThreadStatus::WORKING, + std::memory_order_release); + } + thread_state_[0].curr->store(0, std::memory_order_relaxed); + thread_state_[0].status->store(ThreadStatus::WORKING, + std::memory_order_release); + process_tasks(0); + for (int i = 1; i < thread_num_; i++) { + while (thread_state_[i].status->load(std::memory_order_acquire) == + ThreadStatus::WORKING) { + } + } +} + +void Backend::process_tasks(int thread_id) { + + #ifdef USE_NUMA + if(numa_node == -1){ + numa_node = thread_id * numa_num_configured_nodes() / thread_num_; + struct bitmask* mask = numa_bitmask_alloc(numa_num_configured_nodes()); + numa_bitmask_setbit(mask, numa_node); + numa_bind(mask); + } + #endif + + if (init_func_ != nullptr) { + init_func_(thread_id); + } + while (true) { + int task_id = thread_state_[thread_id].curr->fetch_add( + 1, std::memory_order_acq_rel); + if (task_id >= thread_state_[thread_id].end) { + break; + } + compute_func_(task_id); + } + for (int t_offset = 1; t_offset < thread_num_; t_offset++) { + int t_i = (thread_id + t_offset) % thread_num_; + if (thread_state_[t_i].status->load(std::memory_order_acquire) != + ThreadStatus::WORKING) { + continue; + } + while (true) { + int task_id = thread_state_[t_i].curr->fetch_add( + 1, std::memory_order_acq_rel); + if (task_id >= thread_state_[t_i].end) { + break; + } + compute_func_(task_id); + } + } + if (finalize_func_ != nullptr) { + finalize_func_(thread_id); + } + thread_state_[thread_id].status->store(ThreadStatus::WAITING, + std::memory_order_release); +} + +void Backend::worker_thread(int thread_id) { + auto start = std::chrono::steady_clock::now(); + thread_local_id = thread_id; // 设置线程本地变量 + while (true) { + ThreadStatus status = + thread_state_[thread_id].status->load(std::memory_order_acquire); + if (status == ThreadStatus::WORKING) { + process_tasks(thread_id); + start = std::chrono::steady_clock::now(); + } else if (status == ThreadStatus::WAITING) { + auto now = std::chrono::steady_clock::now(); + auto duration = + std::chrono::duration_cast(now - + start) + .count(); + if (duration > 50) { + std::this_thread::sleep_for(std::chrono::milliseconds(1)); + } + } else if (status == ThreadStatus::EXIT) { + return; + } + } +} \ No newline at end of file diff --git a/archive/kt-sft/csrc/ktransformers_ext/cpu_backend/backend.h b/archive/kt-sft/csrc/ktransformers_ext/cpu_backend/backend.h new file mode 100644 index 0000000..7a95f27 --- /dev/null +++ b/archive/kt-sft/csrc/ktransformers_ext/cpu_backend/backend.h @@ -0,0 +1,58 @@ +/** + * @Description : + * @Author : chenht2022 + * @Date : 2024-07-22 02:03:05 + * @Version : 1.0.0 + * @LastEditors : chenht2022 + * @LastEditTime : 2024-07-25 10:33:38 + * @Copyright (c) 2024 by KVCache.AI, All Rights Reserved. + **/ +#ifndef CPUINFER_BACKEND_H +#define CPUINFER_BACKEND_H + +#include +#include +#include +#include +#include +#include +#include + +enum ThreadStatus { + WORKING, + WAITING, + EXIT, +}; + +struct ThreadState { + std::unique_ptr> status; + std::unique_ptr> curr; + int end; +}; + +class Backend { + public: + Backend(int); + ~Backend(); + int get_thread_num(); + void do_work_stealing_job(int, std::function, + std::function, + std::function); + #ifdef USE_NUMA + static thread_local int numa_node; + #endif + static thread_local int thread_local_id; + + private: + int thread_num_; + int max_thread_num_; + std::vector thread_state_; // [thread_num] + std::function init_func_; + std::function compute_func_; + std::function finalize_func_; + std::vector workers_; + + void process_tasks(int); + void worker_thread(int); +}; +#endif \ No newline at end of file diff --git a/archive/kt-sft/csrc/ktransformers_ext/cpu_backend/cpuinfer.h b/archive/kt-sft/csrc/ktransformers_ext/cpu_backend/cpuinfer.h new file mode 100644 index 0000000..7b1d898 --- /dev/null +++ b/archive/kt-sft/csrc/ktransformers_ext/cpu_backend/cpuinfer.h @@ -0,0 +1,98 @@ +/** + * @Description : + * @Author : chenht2022 + * @Date : 2024-07-16 10:43:18 + * @Version : 1.0.0 + * @LastEditors : chenht2022 + * @LastEditTime : 2024-08-07 09:47:43 + * @Copyright (c) 2024 by KVCache.AI, All Rights Reserved. + **/ + #ifndef CPUINFER_CPUINFER_H + #define CPUINFER_CPUINFER_H + + #include + #include + #include + #include + #include + #include + #include + #include + #ifdef KTRANSFORMERS_USE_CUDA + #include "vendors/cuda.h" + #elif KTRANSFORMERS_USE_MUSA + #include "vendors/musa.h" + #elif KTRANSFORMERS_USE_ROCM + #define __HIP_PLATFORM_AMD__ + #include "vendors/hip.h" + #endif + + #include "backend.h" + #include "task_queue.h" + #include "./vendors/vendor.h" + + #include "llama.cpp/ggml-impl.h" + + class CPUInfer { + public: + CPUInfer(int thread_num) { + backend_ = new Backend(thread_num - 1); + task_queue_ = new TaskQueue(); + for (int i = 0; i < (1 << 16); ++i) { + ggml_table_f32_f16[i] = GGML_COMPUTE_FP16_TO_FP32(i); + } + } + + ~CPUInfer() { + delete backend_; + delete task_queue_; + } + + template + void enqueue(Func f, Obj* obj, Args... args) { + task_queue_->enqueue([=]() { + std::invoke(f, *obj, args..., backend_); + }); + } + + void submit(std::pair params) { + void (*func)(void*) = (void (*)(void*))params.first; + void* args = (void*)params.second; + *((CPUInfer**)args) = this; + func(args); + } + + void sync() { + task_queue_->sync(); + } + + void submit_with_cuda_stream(intptr_t user_cuda_stream, std::pair params) { + #if defined(KTRANSFORMERS_USE_CUDA) || defined(KTRANSFORMERS_USE_MUSA) || defined(KTRANSFORMERS_USE_ROCM) + void (*func)(void*) = (void (*)(void*))params.first; + void* args = (void*)params.second; + *((CPUInfer**)args) = this; + cudaLaunchHostFunc((cudaStream_t)user_cuda_stream, (cudaHostFn_t)func, args); + #else + throw std::runtime_error("submit_with_cuda_stream is not supported on this platforma"); + #endif + } + + static void sync_(void* cpu_infer_ptr) { + CPUInfer* cpuinfer = (CPUInfer*)cpu_infer_ptr; + cpuinfer->sync(); + } + + void sync_with_cuda_stream(intptr_t user_cuda_stream) { + #if defined(KTRANSFORMERS_USE_CUDA) || defined(KTRANSFORMERS_USE_MUSA) || defined(KTRANSFORMERS_USE_ROCM) + cudaLaunchHostFunc((cudaStream_t)user_cuda_stream, (cudaHostFn_t)&sync_, (void*)this); + #else + throw std::runtime_error("sync_with_cuda_stream is not supported on this platforma"); + #endif + } + + public: + Backend* backend_; + TaskQueue* task_queue_; + }; + + #endif \ No newline at end of file diff --git a/archive/kt-sft/csrc/ktransformers_ext/cpu_backend/shared_mem_buffer.cpp b/archive/kt-sft/csrc/ktransformers_ext/cpu_backend/shared_mem_buffer.cpp new file mode 100644 index 0000000..0fe7275 --- /dev/null +++ b/archive/kt-sft/csrc/ktransformers_ext/cpu_backend/shared_mem_buffer.cpp @@ -0,0 +1,56 @@ +/** + * @Description : + * @Author : chenht2022 + * @Date : 2024-08-05 04:49:08 + * @Version : 1.0.0 + * @LastEditors : chenht2022 + * @LastEditTime : 2024-08-05 09:21:29 + * @Copyright (c) 2024 by KVCache.AI, All Rights Reserved. + **/ +#include "shared_mem_buffer.h" +#include + +SharedMemBuffer::SharedMemBuffer() { + buffer_ = nullptr; + size_ = 0; +} + +SharedMemBuffer::~SharedMemBuffer() { + if (buffer_) { + free(buffer_); + } +} + +void SharedMemBuffer::alloc(void* object, std::vector> requests) { + uint64_t size = 0; + for (auto& request : requests) { + size += request.second; + } + if (size > size_) { + if (buffer_) { + free(buffer_); + } + buffer_ = std::aligned_alloc(64, size); + + size_ = size; + for (auto& obj_requests : hist_requests_) { + for (auto& requests : obj_requests.second) { + arrange(requests); + } + } + } + arrange(requests); + hist_requests_[object].push_back(requests); +} + +void SharedMemBuffer::dealloc(void* object) { + hist_requests_.erase(object); +} + +void SharedMemBuffer::arrange(std::vector> requests) { + uint64_t offset = 0; + for (auto& request : requests) { + *(request.first) = (uint8_t*)buffer_ + offset; + offset += request.second; + } +} \ No newline at end of file diff --git a/archive/kt-sft/csrc/ktransformers_ext/cpu_backend/shared_mem_buffer.h b/archive/kt-sft/csrc/ktransformers_ext/cpu_backend/shared_mem_buffer.h new file mode 100644 index 0000000..58e7aa6 --- /dev/null +++ b/archive/kt-sft/csrc/ktransformers_ext/cpu_backend/shared_mem_buffer.h @@ -0,0 +1,37 @@ +/** + * @Description : + * @Author : chenht2022 + * @Date : 2024-08-05 04:49:08 + * @Version : 1.0.0 + * @LastEditors : chenht2022 + * @LastEditTime : 2024-08-05 06:36:41 + * @Copyright (c) 2024 by KVCache.AI, All Rights Reserved. + **/ + + #ifndef CPUINFER_SHAREDMEMBUFFER_H + #define CPUINFER_SHAREDMEMBUFFER_H + + #include + #include + #include + #include + + class SharedMemBuffer { + public: + SharedMemBuffer(); + ~SharedMemBuffer(); + + void alloc(void* object, std::vector> requests); + void dealloc(void* object); + + private: + void* buffer_; + uint64_t size_; + std::map>>> hist_requests_; + + void arrange(std::vector> requests); + }; + + static SharedMemBuffer shared_mem_buffer; + + #endif \ No newline at end of file diff --git a/archive/kt-sft/csrc/ktransformers_ext/cpu_backend/task_queue.cpp b/archive/kt-sft/csrc/ktransformers_ext/cpu_backend/task_queue.cpp new file mode 100644 index 0000000..297a4ea --- /dev/null +++ b/archive/kt-sft/csrc/ktransformers_ext/cpu_backend/task_queue.cpp @@ -0,0 +1,67 @@ +/** + * @Description : + * @Author : chenht2022 + * @Date : 2024-07-17 12:25:51 + * @Version : 1.0.0 + * @LastEditors : chenht2022 + * @LastEditTime : 2024-10-09 11:08:10 + * @Copyright (c) 2024 by KVCache.AI, All Rights Reserved. + **/ +#include "task_queue.h" + +TaskQueue::TaskQueue() { + worker = std::thread(&TaskQueue::processTasks, this); + sync_flag.store(true, std::memory_order_seq_cst); + exit_flag.store(false, std::memory_order_seq_cst); +} + +TaskQueue::~TaskQueue() { + { + mutex.lock(); + exit_flag.store(true, std::memory_order_seq_cst); + mutex.unlock(); + } + cv.notify_all(); + if (worker.joinable()) { + worker.join(); + } +} + +void TaskQueue::enqueue(std::function task) { + { + mutex.lock(); + tasks.push(task); + sync_flag.store(false, std::memory_order_seq_cst); + mutex.unlock(); + } + cv.notify_one(); +} + +void TaskQueue::sync() { + while (!sync_flag.load(std::memory_order_seq_cst)) + ; +} + +void TaskQueue::processTasks() { + while (true) { + std::function task; + { + mutex.lock(); + cv.wait(mutex, [this]() { return !tasks.empty() || exit_flag.load(std::memory_order_seq_cst); }); + if (exit_flag.load(std::memory_order_seq_cst) && tasks.empty()) { + return; + } + task = tasks.front(); + tasks.pop(); + mutex.unlock(); + } + task(); + { + mutex.lock(); + if (tasks.empty()) { + sync_flag.store(true, std::memory_order_seq_cst); + } + mutex.unlock(); + } + } +} \ No newline at end of file diff --git a/archive/kt-sft/csrc/ktransformers_ext/cpu_backend/task_queue.h b/archive/kt-sft/csrc/ktransformers_ext/cpu_backend/task_queue.h new file mode 100644 index 0000000..4f9f112 --- /dev/null +++ b/archive/kt-sft/csrc/ktransformers_ext/cpu_backend/task_queue.h @@ -0,0 +1,138 @@ +/** + * @Description : + * @Author : chenht2022 + * @Date : 2024-07-16 10:43:18 + * @Version : 1.0.0 + * @LastEditors : chenht + * @LastEditTime : 2024-10-09 11:08:07 + * @Copyright (c) 2024 by KVCache.AI, All Rights Reserved. + **/ +#ifndef CPUINFER_TASKQUEUE_H +#define CPUINFER_TASKQUEUE_H + +#include +#include +#include +#include +#include +#include +#include +#ifdef _WIN32 +#include +#endif + +class custom_mutex { + private: +#ifdef _WIN32 + CRITICAL_SECTION cs; +#else + std::mutex mtx; +#endif + + public: + custom_mutex() { +#ifdef _WIN32 + InitializeCriticalSection(&cs); +#else + // No initialization required for std::mutex +#endif + } + + ~custom_mutex() { +#ifdef _WIN32 + DeleteCriticalSection(&cs); +#endif + } + + void lock() { +#ifdef _WIN32 + EnterCriticalSection(&cs); +#else + mtx.lock(); +#endif + } + + void unlock() { +#ifdef _WIN32 + LeaveCriticalSection(&cs); +#else + mtx.unlock(); +#endif + } + +#ifdef _WIN32 + CRITICAL_SECTION* get_handle() { + return &cs; + } +#else + std::mutex* get_handle() { + return &mtx; + } +#endif +}; + +class custom_condition_variable { + private: +#ifdef _WIN32 + CONDITION_VARIABLE cond_var; +#else + std::condition_variable cond_var; +#endif + + public: + custom_condition_variable() { +#ifdef _WIN32 + InitializeConditionVariable(&cond_var); +#endif + } + + template + void wait(custom_mutex& mutex, Predicate pred) { +#ifdef _WIN32 + while (!pred()) { + SleepConditionVariableCS(&cond_var, mutex.get_handle(), INFINITE); + } +#else + std::unique_lock lock(*mutex.get_handle(), std::adopt_lock); + cond_var.wait(lock, pred); + lock.release(); +#endif + } + + void notify_one() { +#ifdef _WIN32 + WakeConditionVariable(&cond_var); +#else + cond_var.notify_one(); +#endif + } + + void notify_all() { +#ifdef _WIN32 + WakeAllConditionVariable(&cond_var); +#else + cond_var.notify_all(); +#endif + } +}; + +class TaskQueue { + public: + TaskQueue(); + ~TaskQueue(); + + void enqueue(std::function); + + void sync(); + + private: + void processTasks(); + + std::queue> tasks; + custom_mutex mutex; + custom_condition_variable cv; + std::thread worker; + std::atomic sync_flag; + std::atomic exit_flag; +}; +#endif \ No newline at end of file diff --git a/archive/kt-sft/csrc/ktransformers_ext/cpu_backend/vendors/README.md b/archive/kt-sft/csrc/ktransformers_ext/cpu_backend/vendors/README.md new file mode 100644 index 0000000..d179f66 --- /dev/null +++ b/archive/kt-sft/csrc/ktransformers_ext/cpu_backend/vendors/README.md @@ -0,0 +1,3 @@ +## TODO + +This directory can be removed after updating the version of `llama.cpp`. \ No newline at end of file diff --git a/archive/kt-sft/csrc/ktransformers_ext/cpu_backend/vendors/cuda.h b/archive/kt-sft/csrc/ktransformers_ext/cpu_backend/vendors/cuda.h new file mode 100644 index 0000000..1746b07 --- /dev/null +++ b/archive/kt-sft/csrc/ktransformers_ext/cpu_backend/vendors/cuda.h @@ -0,0 +1,15 @@ +#pragma once + +#include +#include +#include +#include +#include + +#if CUDART_VERSION < 11020 +#define CU_DEVICE_ATTRIBUTE_VIRTUAL_MEMORY_MANAGEMENT_SUPPORTED CU_DEVICE_ATTRIBUTE_VIRTUAL_ADDRESS_MANAGEMENT_SUPPORTED +#define CUBLAS_TF32_TENSOR_OP_MATH CUBLAS_TENSOR_OP_MATH +#define CUBLAS_COMPUTE_16F CUDA_R_16F +#define CUBLAS_COMPUTE_32F CUDA_R_32F +#define cublasComputeType_t cudaDataType_t +#endif // CUDART_VERSION < 11020 diff --git a/archive/kt-sft/csrc/ktransformers_ext/cpu_backend/vendors/hip.h b/archive/kt-sft/csrc/ktransformers_ext/cpu_backend/vendors/hip.h new file mode 100644 index 0000000..abbc1e8 --- /dev/null +++ b/archive/kt-sft/csrc/ktransformers_ext/cpu_backend/vendors/hip.h @@ -0,0 +1,172 @@ +#pragma once + +#define HIP_ENABLE_WARP_SYNC_BUILTINS 1 +#include +#include +#include +#include +#ifdef __HIP_PLATFORM_AMD__ +// for rocblas_initialize() +#include "rocblas/rocblas.h" +#endif // __HIP_PLATFORM_AMD__ + +#define CUBLAS_COMPUTE_16F HIPBLAS_R_16F +#define CUBLAS_COMPUTE_32F HIPBLAS_R_32F +#define CUBLAS_COMPUTE_32F_FAST_16F HIPBLAS_R_32F +#define CUBLAS_GEMM_DEFAULT HIPBLAS_GEMM_DEFAULT +#define CUBLAS_GEMM_DEFAULT_TENSOR_OP HIPBLAS_GEMM_DEFAULT +#define CUBLAS_OP_N HIPBLAS_OP_N +#define CUBLAS_OP_T HIPBLAS_OP_T +#define CUBLAS_STATUS_SUCCESS HIPBLAS_STATUS_SUCCESS +#define CUBLAS_TF32_TENSOR_OP_MATH 0 +#define CUDA_R_16F HIPBLAS_R_16F +#define CUDA_R_32F HIPBLAS_R_32F +#define CU_DEVICE_ATTRIBUTE_VIRTUAL_MEMORY_MANAGEMENT_SUPPORTED hipDeviceAttributeVirtualMemoryManagementSupported +#define CU_MEM_ALLOC_GRANULARITY_RECOMMENDED hipMemAllocationGranularityRecommended +#define CU_MEM_ALLOCATION_TYPE_PINNED hipMemAllocationTypePinned +#define CU_MEM_LOCATION_TYPE_DEVICE hipMemLocationTypeDevice +#define CU_MEM_ACCESS_FLAGS_PROT_READWRITE hipMemAccessFlagsProtReadWrite +#define CU_CHECK(fn) {hipError_t err = fn; if(err != hipSuccess) { GGML_ABORT("HipVMM Failure: %s\n", hipGetErrorString(err)); }} +#define __shfl_sync(mask, var, laneMask, width) __shfl(var, laneMask, width) +#define __shfl_xor_sync(mask, var, laneMask, width) __shfl_xor(var, laneMask, width) +#define cublasComputeType_t hipblasDatatype_t //deprecated, new hipblasComputeType_t not in 5.6 +#define cublasCreate hipblasCreate +#define cublasDestroy hipblasDestroy +#define cublasGemmEx hipblasGemmEx +#define cublasGemmBatchedEx hipblasGemmBatchedEx +#define cublasGemmStridedBatchedEx hipblasGemmStridedBatchedEx +#define cublasHandle_t hipblasHandle_t +#define cublasSetMathMode(handle, mode) CUBLAS_STATUS_SUCCESS +#define cublasSetStream hipblasSetStream +#define cublasSgemm hipblasSgemm +#define cublasStatus_t hipblasStatus_t +#define cublasOperation_t hipblasOperation_t +#define cudaDataType_t hipblasDatatype_t //deprecated, new hipblasDatatype not in 5.6 +#define cudaDeviceCanAccessPeer hipDeviceCanAccessPeer +#define cudaDeviceDisablePeerAccess hipDeviceDisablePeerAccess +#define cudaDeviceEnablePeerAccess hipDeviceEnablePeerAccess +#define cudaDeviceProp hipDeviceProp_t +#define cudaDeviceSynchronize hipDeviceSynchronize +#define cudaError_t hipError_t +#define cudaErrorPeerAccessAlreadyEnabled hipErrorPeerAccessAlreadyEnabled +#define cudaErrorPeerAccessNotEnabled hipErrorPeerAccessNotEnabled +#define cudaEventCreateWithFlags hipEventCreateWithFlags +#define cudaEventDisableTiming hipEventDisableTiming +#define cudaEventRecord hipEventRecord +#define cudaEventSynchronize hipEventSynchronize +#define cudaEvent_t hipEvent_t +#define cudaEventDestroy hipEventDestroy +#define cudaFree hipFree +#define cudaFreeHost hipHostFree +#define cudaGetDevice hipGetDevice +#define cudaGetDeviceCount hipGetDeviceCount +#define cudaGetDeviceProperties hipGetDeviceProperties +#define cudaGetErrorString hipGetErrorString +#define cudaGetLastError hipGetLastError +#define cudaHostRegister hipHostRegister +#define cudaHostRegisterPortable hipHostRegisterPortable +#define cudaHostRegisterReadOnly hipHostRegisterReadOnly +#define cudaHostUnregister hipHostUnregister +#define cudaLaunchHostFunc hipLaunchHostFunc +#define cudaMalloc hipMalloc +#define cudaMallocHost(ptr, size) hipHostMalloc(ptr, size, hipHostMallocDefault) +#define cudaMemcpy hipMemcpy +#define cudaMemcpyAsync hipMemcpyAsync +#define cudaMemcpyPeerAsync hipMemcpyPeerAsync +#define cudaMemcpy2DAsync hipMemcpy2DAsync +#define cudaMemcpyDeviceToDevice hipMemcpyDeviceToDevice +#define cudaMemcpyDeviceToHost hipMemcpyDeviceToHost +#define cudaMemcpyHostToDevice hipMemcpyHostToDevice +#define cudaMemcpyKind hipMemcpyKind +#define cudaMemset hipMemset +#define cudaMemsetAsync hipMemsetAsync +#define cudaMemGetInfo hipMemGetInfo +#define cudaOccupancyMaxPotentialBlockSize hipOccupancyMaxPotentialBlockSize +#define cudaSetDevice hipSetDevice +#define cuDeviceGet hipDeviceGet +#define CUdevice hipDevice_t +#define CUdeviceptr hipDeviceptr_t +#define cuMemUnmap hipMemUnmap +#define CUmemAccessDesc hipMemAccessDesc +#define cuMemAddressFree hipMemAddressFree +#define cuMemRelease hipMemRelease +#define CUmemGenericAllocationHandle hipMemGenericAllocationHandle_t +#define cuMemCreate hipMemCreate +#define cuMemAddressReserve hipMemAddressReserve +#define cuMemMap hipMemMap +#define cuMemSetAccess hipMemSetAccess +#define cuMemGetAllocationGranularity hipMemGetAllocationGranularity +#define CUmemAllocationProp hipMemAllocationProp +#define cuDeviceGetAttribute hipDeviceGetAttribute +#define cudaStreamCreateWithFlags hipStreamCreateWithFlags +#define cudaStreamDestroy hipStreamDestroy +#define cudaStreamFireAndForget hipStreamFireAndForget +#define cudaStreamNonBlocking hipStreamNonBlocking +#define cudaStreamPerThread hipStreamPerThread +#define cudaStreamSynchronize hipStreamSynchronize +#define cudaStreamWaitEvent(stream, event, flags) hipStreamWaitEvent(stream, event, flags) +#define cudaGraphExec_t hipGraphExec_t +#define cudaGraphNode_t hipGraphNode_t +#define cudaKernelNodeParams hipKernelNodeParams +#define cudaKernelNodeParams hipKernelNodeParams +#define cudaGraphExecDestroy hipGraphExecDestroy +#define cudaGraphLaunch hipGraphLaunch +#define cudaErrorGraphExecUpdateFailure hipErrorGraphExecUpdateFailure +#define cudaGraphExecUpdateResultInfo hipGraphExecUpdateResult +#define cudaGraphNodeType hipGraphNodeType +#define cudaGraphNodeTypeKernel hipGraphNodeTypeKernel +#define cudaGraphInstantiate hipGraphInstantiate +#define cudaStreamEndCapture hipStreamEndCapture +#define cudaGraphDestroy hipGraphDestroy +#define cudaGraphKernelNodeSetParams hipGraphKernelNodeSetParams +#define cudaErrorInvalidDeviceFunction hipErrorInvalidDeviceFunction +#define cudaGraphKernelNodeGetParams hipGraphKernelNodeGetParams +#define cudaGraphNodeGetType hipGraphNodeGetType +#define cudaGraphGetNodes hipGraphGetNodes +#define cudaGraphExecUpdate hipGraphExecUpdate +#define cudaStreamCaptureModeRelaxed hipStreamCaptureModeRelaxed +#define cudaStreamBeginCapture hipStreamBeginCapture +#define cudaGraph_t hipGraph_t +#define cudaStream_t hipStream_t +#define cudaSuccess hipSuccess +#define cudaHostFn_t hipHostFn_t +#define __trap() do { abort(); __builtin_unreachable(); } while(0) +#define CUBLAS_STATUS_SUCCESS HIPBLAS_STATUS_SUCCESS +#define CUBLAS_STATUS_NOT_INITIALIZED HIPBLAS_STATUS_NOT_INITIALIZED +#define CUBLAS_STATUS_ALLOC_FAILED HIPBLAS_STATUS_ALLOC_FAILED +#define CUBLAS_STATUS_INVALID_VALUE HIPBLAS_STATUS_INVALID_VALUE +#define CUBLAS_STATUS_ARCH_MISMATCH HIPBLAS_STATUS_ARCH_MISMATCH +#define CUBLAS_STATUS_MAPPING_ERROR HIPBLAS_STATUS_MAPPING_ERROR +#define CUBLAS_STATUS_EXECUTION_FAILED HIPBLAS_STATUS_EXECUTION_FAILED +#define CUBLAS_STATUS_INTERNAL_ERROR HIPBLAS_STATUS_INTERNAL_ERROR +#define CUBLAS_STATUS_NOT_SUPPORTED HIPBLAS_STATUS_NOT_SUPPORTED + +#define __CUDA_ARCH__ 1300 + +#if defined(__gfx803__) || defined(__gfx900__) || defined(__gfx906__) +#define GCN +#endif + +#if defined(__gfx908__) || defined(__gfx90a__) || defined(__gfx942__) +#define CDNA +#endif + +#if defined(__gfx1100__) || defined(__gfx1101__) || defined(__gfx1102__) || defined(__gfx1103__) || \ + defined(__gfx1150__) || defined(__gfx1151__) +#define RDNA3 +#endif + +#if defined(__gfx1030__) || defined(__gfx1031__) || defined(__gfx1032__) || defined(__gfx1033__) || \ + defined(__gfx1034__) || defined(__gfx1035__) || defined(__gfx1036__) || defined(__gfx1037__) +#define RDNA2 +#endif + +#if defined(__gfx1010__) || defined(__gfx1012__) +#define RDNA1 +#endif + +#ifndef __has_builtin + #define __has_builtin(x) 0 +#endif + +typedef hip_bfloat16 nv_bfloat16; diff --git a/archive/kt-sft/csrc/ktransformers_ext/cpu_backend/vendors/musa.h b/archive/kt-sft/csrc/ktransformers_ext/cpu_backend/vendors/musa.h new file mode 100644 index 0000000..6cc1b69 --- /dev/null +++ b/archive/kt-sft/csrc/ktransformers_ext/cpu_backend/vendors/musa.h @@ -0,0 +1,137 @@ +#pragma once + +#include +#include +#include +#include +#include +#define CUBLAS_COMPUTE_16F CUDA_R_16F +#define CUBLAS_COMPUTE_32F CUDA_R_32F +#define CUBLAS_COMPUTE_32F_FAST_16F MUBLAS_COMPUTE_32F_FAST_16F +#define CUBLAS_GEMM_DEFAULT MUBLAS_GEMM_DEFAULT +#define CUBLAS_GEMM_DEFAULT_TENSOR_OP MUBLAS_GEMM_DEFAULT +#define CUBLAS_OP_N MUBLAS_OP_N +#define CUBLAS_OP_T MUBLAS_OP_T +#define CUBLAS_STATUS_SUCCESS MUBLAS_STATUS_SUCCESS +#define CUBLAS_TF32_TENSOR_OP_MATH MUBLAS_MATH_MODE_DEFAULT +#define CUDA_R_16F MUSA_R_16F +#define CUDA_R_32F MUSA_R_32F +#define cublasComputeType_t cudaDataType_t +#define cublasCreate mublasCreate +#define cublasDestroy mublasDestroy +#define cublasGemmEx mublasGemmEx +#define cublasGemmBatchedEx mublasGemmBatchedEx +#define cublasGemmStridedBatchedEx mublasGemmStridedBatchedEx +#define cublasHandle_t mublasHandle_t +#define cublasSetMathMode mublasSetMathMode +#define cublasSetStream mublasSetStream +#define cublasSgemm mublasSgemm +#define cublasStatus_t mublasStatus_t +#define cublasOperation_t mublasOperation_t +#define cublasGetStatusString mublasStatus_to_string +#define cudaDataType_t musaDataType_t +#define cudaDeviceCanAccessPeer musaDeviceCanAccessPeer +#define cudaDeviceDisablePeerAccess musaDeviceDisablePeerAccess +#define cudaDeviceEnablePeerAccess musaDeviceEnablePeerAccess +#define cudaDeviceProp musaDeviceProp +#define cudaDeviceSynchronize musaDeviceSynchronize +#define cudaError_t musaError_t +#define cudaErrorPeerAccessAlreadyEnabled musaErrorPeerAccessAlreadyEnabled +#define cudaErrorPeerAccessNotEnabled musaErrorPeerAccessNotEnabled +#define cudaEventCreateWithFlags musaEventCreateWithFlags +#define cudaEventDisableTiming musaEventDisableTiming +#define cudaEventRecord musaEventRecord +#define cudaEventSynchronize musaEventSynchronize +#define cudaEvent_t musaEvent_t +#define cudaEventDestroy musaEventDestroy +#define cudaFree musaFree +#define cudaFreeHost musaFreeHost +#define cudaGetDevice musaGetDevice +#define cudaGetDeviceCount musaGetDeviceCount +#define cudaGetDeviceProperties musaGetDeviceProperties +#define cudaGetErrorString musaGetErrorString +#define cudaGetLastError musaGetLastError +#define cudaHostRegister musaHostRegister +#define cudaHostRegisterPortable musaHostRegisterPortable +#define cudaHostRegisterReadOnly musaHostRegisterReadOnly +#define cudaHostUnregister musaHostUnregister +#define cudaLaunchHostFunc musaLaunchHostFunc +#define cudaMalloc musaMalloc +#define cudaMallocHost musaMallocHost +#define cudaMallocManaged musaMallocManaged +#define cudaMemcpy musaMemcpy +#define cudaMemcpyAsync musaMemcpyAsync +#define cudaMemcpyPeerAsync musaMemcpyPeerAsync +#define cudaMemcpy2DAsync musaMemcpy2DAsync +#define cudaMemcpyDeviceToDevice musaMemcpyDeviceToDevice +#define cudaMemcpyDeviceToHost musaMemcpyDeviceToHost +#define cudaMemcpyHostToDevice musaMemcpyHostToDevice +#define cudaMemcpyKind musaMemcpyKind +#define cudaMemset musaMemset +#define cudaMemsetAsync musaMemsetAsync +#define cudaMemGetInfo musaMemGetInfo +#define cudaOccupancyMaxPotentialBlockSize musaOccupancyMaxPotentialBlockSize +#define cudaSetDevice musaSetDevice +#define cudaStreamCreateWithFlags musaStreamCreateWithFlags +#define cudaStreamDestroy musaStreamDestroy +#define cudaStreamFireAndForget musaStreamFireAndForget +#define cudaStreamNonBlocking musaStreamNonBlocking +#define cudaStreamPerThread musaStreamPerThread +#define cudaStreamSynchronize musaStreamSynchronize +#define cudaStreamWaitEvent musaStreamWaitEvent +#define cudaStream_t musaStream_t +#define cudaSuccess musaSuccess + +// Additional mappings for MUSA virtual memory pool +#define CU_DEVICE_ATTRIBUTE_VIRTUAL_MEMORY_MANAGEMENT_SUPPORTED MU_DEVICE_ATTRIBUTE_VIRTUAL_ADDRESS_MANAGEMENT_SUPPORTED +#define CU_MEM_ACCESS_FLAGS_PROT_READWRITE MU_MEM_ACCESS_FLAGS_PROT_READWRITE +#define CU_MEM_ALLOC_GRANULARITY_RECOMMENDED MU_MEM_ALLOC_GRANULARITY_RECOMMENDED +#define CU_MEM_ALLOCATION_TYPE_PINNED MU_MEM_ALLOCATION_TYPE_PINNED +#define CU_MEM_LOCATION_TYPE_DEVICE MU_MEM_LOCATION_TYPE_DEVICE +#define CUdevice MUdevice +#define CUdeviceptr MUdeviceptr +#define CUmemAccessDesc MUmemAccessDesc +#define CUmemAllocationProp MUmemAllocationProp +#define CUmemGenericAllocationHandle MUmemGenericAllocationHandle +#define cuDeviceGet muDeviceGet +#define cuDeviceGetAttribute muDeviceGetAttribute +#define cuMemAddressFree muMemAddressFree +#define cuMemAddressReserve muMemAddressReserve +#define cuMemCreate muMemCreate +#define cuMemGetAllocationGranularity muMemGetAllocationGranularity +#define cuMemMap muMemMap +#define cuMemRelease muMemRelease +#define cuMemSetAccess muMemSetAccess +#define cuMemUnmap muMemUnmap +#define cudaFuncAttributeMaxDynamicSharedMemorySize musaFuncAttributeMaxDynamicSharedMemorySize +#define cudaFuncSetAttribute musaFuncSetAttribute +#define cudaMemcpy3DPeerParms musaMemcpy3DPeerParms +#define make_cudaExtent make_musaExtent +#define make_cudaPitchedPtr make_musaPitchedPtr + +// Additional mappings for MUSA graphs +#define CUDA_SUCCESS MUSA_SUCCESS +#define CUresult MUresult +#define cuGetErrorString muGetErrorString +#define cudaErrorGraphExecUpdateFailure musaErrorGraphExecUpdateFailure +#define cudaErrorInvalidDeviceFunction musaErrorInvalidDeviceFunction +#define cudaGraphDestroy musaGraphDestroy +#define cudaGraphExecDestroy musaGraphExecDestroy +#define cudaGraphExec_t musaGraphExec_t +#define cudaGraphExecUpdate musaGraphExecUpdate +#define cudaGraphExecUpdateResultInfo musaGraphExecUpdateResult +#define cudaGraphGetNodes musaGraphGetNodes +#define cudaGraphInstantiate musaGraphInstantiate +#define cudaGraphKernelNodeGetParams musaGraphKernelNodeGetParams +#define cudaGraphKernelNodeSetParams musaGraphKernelNodeSetParams +#define cudaGraphLaunch musaGraphLaunch +#define cudaGraphNodeGetType musaGraphNodeGetType +#define cudaGraphNode_t musaGraphNode_t +#define cudaGraphNodeType musaGraphNodeType +#define cudaGraphNodeTypeKernel musaGraphNodeTypeKernel +#define cudaGraph_t musaGraph_t +#define cudaKernelNodeParams musaKernelNodeParams +#define cudaStreamCaptureModeRelaxed musaStreamCaptureModeRelaxed +#define cudaStreamEndCapture musaStreamEndCapture + +typedef mt_bfloat16 nv_bfloat16; diff --git a/archive/kt-sft/csrc/ktransformers_ext/cpu_backend/vendors/vendor.h b/archive/kt-sft/csrc/ktransformers_ext/cpu_backend/vendors/vendor.h new file mode 100644 index 0000000..8470438 --- /dev/null +++ b/archive/kt-sft/csrc/ktransformers_ext/cpu_backend/vendors/vendor.h @@ -0,0 +1,13 @@ +#ifndef CPUINFER_VENDOR_VENDOR_H +#define CPUINFER_VENDOR_VENDOR_H + +#ifdef USE_CUDA +#include "cuda.h" +#elif USE_HIP +#define __HIP_PLATFORM_AMD__ +#include "hip.h" +#elif USE_MUSA +#include "musa.h" +#endif + +#endif // CPUINFER_VENDOR_VENDOR_H \ No newline at end of file diff --git a/archive/kt-sft/csrc/ktransformers_ext/cuda/binding.cpp b/archive/kt-sft/csrc/ktransformers_ext/cuda/binding.cpp new file mode 100644 index 0000000..4aa00c6 --- /dev/null +++ b/archive/kt-sft/csrc/ktransformers_ext/cuda/binding.cpp @@ -0,0 +1,71 @@ +/** + * @Description : + * @Author : Azure-Tang, Boxin Zhang + * @Date : 2024-07-25 13:38:30 + * @Version : 0.2.2 + * @Copyright (c) 2024 by KVCache.AI, All Rights Reserved. +**/ + +#include "custom_gguf/ops.h" +#ifdef KTRANSFORMERS_USE_CUDA +#include "gptq_marlin/ops.h" +#endif +// Python bindings +#include +#include +#include +#include +#include +// namespace py = pybind11; + +PYBIND11_MODULE(KTransformersOps, m) { + + m.def("dequantize_q8_0", [](const intptr_t data, int num_bytes, int blk_size, const int ele_per_blk, torch::Device device, py::object target_dtype) { + torch::Dtype dtype = torch::python::detail::py_object_to_dtype(target_dtype); + return dequantize_q8_0((int8_t*)data, num_bytes, blk_size, ele_per_blk, device, dtype); + }, "Function to dequantize q8_0 data.", + py::arg("data"), py::arg("num_bytes"), py::arg("blk_size"), py::arg("ele_per_blk"), py::arg("device"), py::arg("target_dtype")); + + m.def("dequantize_q6_k", [](const intptr_t data, int num_bytes, int blk_size, const int ele_per_blk, torch::Device device, py::object target_dtype) { + torch::Dtype dtype = torch::python::detail::py_object_to_dtype(target_dtype); + return dequantize_q6_k((int8_t*)data, num_bytes, blk_size, ele_per_blk, device, dtype); + }, "Function to dequantize q6_k data.", + py::arg("data"), py::arg("num_bytes"), py::arg("blk_size"), py::arg("ele_per_blk"), py::arg("device"), py::arg("target_dtype")); + + m.def("dequantize_q5_k", [](const intptr_t data, int num_bytes, int blk_size, const int ele_per_blk, torch::Device device, py::object target_dtype) { + torch::Dtype dtype = torch::python::detail::py_object_to_dtype(target_dtype); + return dequantize_q5_k((int8_t*)data, num_bytes, blk_size, ele_per_blk, device, dtype); + }, "Function to dequantize q5_k data.", + py::arg("data"), py::arg("num_bytes"), py::arg("blk_size"), py::arg("ele_per_blk"), py::arg("device"), py::arg("target_dtype")); + + m.def("dequantize_q4_k", [](const intptr_t data, int num_bytes, int blk_size, const int ele_per_blk, torch::Device device, py::object target_dtype) { + torch::Dtype dtype = torch::python::detail::py_object_to_dtype(target_dtype); + return dequantize_q4_k((int8_t*)data, num_bytes, blk_size, ele_per_blk, device, dtype); + }, "Function to dequantize q4_k data.", + py::arg("data"), py::arg("num_bytes"), py::arg("blk_size"), py::arg("ele_per_blk"), py::arg("device"), py::arg("target_dtype")); + + m.def("dequantize_q3_k", [](const intptr_t data, int num_bytes, int blk_size, const int ele_per_blk, torch::Device device, py::object target_dtype) { + torch::Dtype dtype = torch::python::detail::py_object_to_dtype(target_dtype); + return dequantize_q3_k((int8_t*)data, num_bytes, blk_size, ele_per_blk, device, dtype); + }, "Function to dequantize q3_k data.", + py::arg("data"), py::arg("num_bytes"), py::arg("blk_size"), py::arg("ele_per_blk"), py::arg("device"), py::arg("target_dtype")); + + m.def("dequantize_q2_k", [](const intptr_t data, int num_bytes, int blk_size, const int ele_per_blk, torch::Device device, py::object target_dtype) { + torch::Dtype dtype = torch::python::detail::py_object_to_dtype(target_dtype); + return dequantize_q2_k((int8_t*)data, num_bytes, blk_size, ele_per_blk, device, dtype); + }, "Function to dequantize q2_k data.", + py::arg("data"), py::arg("num_bytes"), py::arg("blk_size"), py::arg("ele_per_blk"), py::arg("device"), py::arg("target_dtype")); + + m.def("dequantize_iq4_xs", [](const intptr_t data, int num_bytes, int blk_size, const int ele_per_blk, torch::Device device, py::object target_dtype) { + torch::Dtype dtype = torch::python::detail::py_object_to_dtype(target_dtype); + return dequantize_iq4_xs((int8_t*)data, num_bytes, blk_size, ele_per_blk, device, dtype); + }, "Function to dequantize iq4_xs data.", + py::arg("data"), py::arg("num_bytes"), py::arg("blk_size"), py::arg("ele_per_blk"), py::arg("device"), py::arg("target_dtype")); + +#ifdef KTRANSFORMERS_USE_CUDA + m.def("gptq_marlin_gemm", &gptq_marlin_gemm, "Function to perform GEMM using Marlin quantization.", + py::arg("a"), py::arg("b_q_weight"), py::arg("b_scales"), py::arg("g_idx"), + py::arg("perm"), py::arg("workspace"), py::arg("num_bits"), py::arg("size_m"), + py::arg("size_n"), py::arg("size_k"), py::arg("is_k_full")); +#endif +} \ No newline at end of file diff --git a/archive/kt-sft/csrc/ktransformers_ext/cuda/custom_gguf/dequant.cu b/archive/kt-sft/csrc/ktransformers_ext/cuda/custom_gguf/dequant.cu new file mode 100644 index 0000000..e77e007 --- /dev/null +++ b/archive/kt-sft/csrc/ktransformers_ext/cuda/custom_gguf/dequant.cu @@ -0,0 +1,882 @@ +/* + * @Description : + * @Author : Azure-Tang, Boxin Zhang + * @Date : 2024-07-25 13:38:30 + * @Version : 0.2.2 + * Adapted from https://github.com/ggerganov/ggml/blob/fca1caafea7de9fbd7efc733b9818f9cf2da3050/src/ggml-quants.c + * Copyright (c) 2023-2024 The ggml authors + * Copyright (c) 2024 by KVCache.AI, All Rights Reserved. + */ +#include +#include +#include +#include +#include +#include +#include +#include + +#ifdef __HIP_PLATFORM_AMD__ +typedef __hip_bfloat16 nv_bfloat16; +#endif + +__global__ void dequantize_q8_0_fp32_kernel(const int8_t* data, float* output, const int blk_size, const int ele_per_blk, const int num_blocks) { + long long global_idx = blockIdx.x * blockDim.x + threadIdx.x; + for (long long block_id = global_idx; block_id < num_blocks; block_id += blockDim.x * gridDim.x){ + float* __restrict__ output_blk = (float*)(output + block_id * ele_per_blk); + const int8_t* cur_block = data + block_id * blk_size; + float scale = __half2float(*((half*)cur_block)); + cur_block += 2; + for (int i = 0; i < ele_per_blk; i++){ + output_blk[i] = scale * cur_block[i]; + } + } +} + +__global__ void dequantize_q8_0_fp16_kernel(const int8_t* data, __half* output, const int blk_size, const int ele_per_blk, const int num_blocks) { + long long global_idx = blockIdx.x * blockDim.x + threadIdx.x; + for (long long block_id = global_idx; block_id < num_blocks; block_id += blockDim.x * gridDim.x) { + __half* __restrict__ output_blk = (__half*)(output + block_id * ele_per_blk); + const int8_t* cur_block = data + block_id * blk_size; + float scale = __half2float(*((half*)cur_block)); + cur_block += 2; + for (int i = 0; i < ele_per_blk; i++) { + output_blk[i] = __float2half(scale * cur_block[i]); + } + } +} + +__global__ void dequantize_q8_0_bf16_kernel(const int8_t* data, nv_bfloat16* output, const int blk_size, const int ele_per_blk, const int num_blocks) { + long long global_idx = blockIdx.x * blockDim.x + threadIdx.x; + for (long long block_id = global_idx; block_id < num_blocks; block_id += blockDim.x * gridDim.x) { + nv_bfloat16* __restrict__ output_blk = (nv_bfloat16*)(output + block_id * ele_per_blk); + const int8_t* cur_block = data + block_id * blk_size; + float scale = __half2float(*((half*)cur_block)); + cur_block += 2; + for (int i = 0; i < ele_per_blk; i++) { + output_blk[i] = __float2bfloat16(scale * cur_block[i]); + } + } +} + +// __device__ void get_scale_min_k4(int j, const uint8_t * __restrict__ q, uint8_t * __restrict__ d, uint8_t * __restrict__ m) { +__device__ void get_scale_min_k4(int j, const uint8_t * q, uint8_t * __restrict__ d, uint8_t * __restrict__ m) { + if (j < 4) { + *d = q[j] & 63; *m = q[j + 4] & 63; + } else { + *d = (q[j+4] & 0xF) | ((q[j-4] >> 6) << 4); + *m = (q[j+4] >> 4) | ((q[j-0] >> 6) << 4); + } +} + +__global__ void dequantize_q2_k_fp32_kernel(const int8_t* data, float* output, const int blk_size, const int ele_per_blk, const int num_blocks) { + long long global_idx = blockIdx.x * blockDim.x + threadIdx.x; + for (long long block_id=global_idx; block_id(data + block_id * blk_size + 80))); + const float min = __half2float(*(reinterpret_cast(data + block_id * blk_size + 82))); + + const uint8_t * __restrict__ q = (uint8_t*)(data + block_id * blk_size + 16); + + int is = 0; + float dl, ml; + + for (int n = 0; n < 256; n += 128) { + int shift = 0; + for (int j = 0; j < 4; ++j) { + uint8_t* scales = (uint8_t*)(data + block_id * blk_size + (is++)); + uint8_t sc = *scales; + dl = d * (sc & 0xF); ml = min * (sc >> 4); + for (int l = 0; l < 16; ++l) *output_blk++ = dl * ((int8_t)((q[l] >> shift) & 3)) - ml; + + scales = (uint8_t*)(data + block_id * blk_size + (is++)); + sc = *scales; + + dl = d * (sc & 0xF); ml = min * (sc >> 4); + for (int l = 0; l < 16; ++l) *output_blk++ = dl * ((int8_t)((q[l+16] >> shift) & 3)) - ml; + + shift += 2; + } + q += 32; + } + } +} + +__global__ void dequantize_q2_k_fp16_kernel(const int8_t* data, __half* output, const int blk_size, const int ele_per_blk, const int num_blocks) { + long long global_idx = blockIdx.x * blockDim.x + threadIdx.x; + for (long long block_id=global_idx; block_id(data + block_id * blk_size + 80))); + const float min = __half2float(*(reinterpret_cast(data + block_id * blk_size + 82))); + + const uint8_t * __restrict__ q = (uint8_t*)(data + block_id * blk_size + 16); + + int is = 0; + float dl, ml; + + for (int n = 0; n < 256; n += 128) { + int shift = 0; + for (int j = 0; j < 4; ++j) { + uint8_t* scales = (uint8_t*)(data + block_id * blk_size + (is++)); + uint8_t sc = *scales; + dl = d * (sc & 0xF); ml = min * (sc >> 4); + for (int l = 0; l < 16; ++l) *output_blk++ = __float2half(dl * ((int8_t)((q[l] >> shift) & 3)) - ml); + + scales = (uint8_t*)(data + block_id * blk_size + (is++)); + sc = *scales; + + dl = d * (sc & 0xF); ml = min * (sc >> 4); + for (int l = 0; l < 16; ++l) *output_blk++ = __float2half(dl * ((int8_t)((q[l+16] >> shift) & 3)) - ml); + + shift += 2; + } + q += 32; + } + } +} + +__global__ void dequantize_q2_k_bf16_kernel(const int8_t* data, nv_bfloat16* output, const int blk_size, const int ele_per_blk, const int num_blocks) { + long long global_idx = blockIdx.x * blockDim.x + threadIdx.x; + for (long long block_id=global_idx; block_id(data + block_id * blk_size + 80))); + const float min = __half2float(*(reinterpret_cast(data + block_id * blk_size + 82))); + + const uint8_t * __restrict__ q = (uint8_t*)(data + block_id * blk_size + 16); + + int is = 0; + float dl, ml; + + for (int n = 0; n < 256; n += 128) { + int shift = 0; + for (int j = 0; j < 4; ++j) { + uint8_t* scales = (uint8_t*)(data + block_id * blk_size + (is++)); + uint8_t sc = *scales; + dl = d * (sc & 0xF); ml = min * (sc >> 4); + for (int l = 0; l < 16; ++l) *output_blk++ = __float2bfloat16(dl * ((int8_t)((q[l] >> shift) & 3)) - ml); + + scales = (uint8_t*)(data + block_id * blk_size + (is++)); + sc = *scales; + + dl = d * (sc & 0xF); ml = min * (sc >> 4); + for (int l = 0; l < 16; ++l) *output_blk++ = __float2bfloat16(dl * ((int8_t)((q[l+16] >> shift) & 3)) - ml); + + shift += 2; + } + q += 32; + } + } +} + +__global__ void dequantize_q3_k_fp32_kernel(const int8_t* data, float* output, const int blk_size, const int ele_per_blk, const int num_blocks) { + + long long global_idx = blockIdx.x * blockDim.x + threadIdx.x; + const uint32_t kmask1 = 0x03030303; + const uint32_t kmask2 = 0x0f0f0f0f; + for (long long block_id=global_idx; block_id(data + block_id * blk_size + 108))); + + const uint8_t * __restrict__ q = (uint8_t*)(data + block_id * blk_size + 32); + const uint8_t * __restrict__ hm = (uint8_t*)(data + block_id * blk_size + 0); + uint8_t m = 1; + + + uint8_t* block_scales = (uint8_t*)(data + block_id * blk_size + 96); + + for (int i = 0; i < 3; i++) { + aux[i] = 0; + for (int j = 0; j < 4; j++) { + aux[i] |= ((uint32_t)block_scales[i * 4 + j]) << (j * 8); + } + } + + uint32_t tmp = aux[2]; + aux[2] = ((aux[0] >> 4) & kmask2) | (((tmp >> 4) & kmask1) << 4); + aux[3] = ((aux[1] >> 4) & kmask2) | (((tmp >> 6) & kmask1) << 4); + aux[0] = (aux[0] & kmask2) | (((tmp >> 0) & kmask1) << 4); + aux[1] = (aux[1] & kmask2) | (((tmp >> 2) & kmask1) << 4); + + int is = 0; + float dl; + for (int n = 0; n < 256; n += 128) { + int shift = 0; + for (int j = 0; j < 4; ++j) { + + dl = d_all * (scales[is++] - 32); + for (int l = 0; l < 16; ++l) { + *output_blk++ = dl * ((int8_t)((q[l+ 0] >> shift) & 3) - ((hm[l+ 0] & m) ? 0 : 4)); + } + + dl = d_all * (scales[is++] - 32); + for (int l = 0; l < 16; ++l) { + *output_blk++ = dl * ((int8_t)((q[l+16] >> shift) & 3) - ((hm[l+16] & m) ? 0 : 4)); + } + + shift += 2; + m <<= 1; + } + q += 32; + } + } +} + +__global__ void dequantize_q3_k_fp16_kernel(const int8_t* data, __half* output, const int blk_size, const int ele_per_blk, const int num_blocks) { + + long long global_idx = blockIdx.x * blockDim.x + threadIdx.x; + const uint32_t kmask1 = 0x03030303; + const uint32_t kmask2 = 0x0f0f0f0f; + for (long long block_id=global_idx; block_id(data + block_id * blk_size + 108))); + + const uint8_t * __restrict__ q = (uint8_t*)(data + block_id * blk_size + 32); + const uint8_t * __restrict__ hm = (uint8_t*)(data + block_id * blk_size + 0); + uint8_t m = 1; + + + uint8_t* block_scales = (uint8_t*)(data + block_id * blk_size + 96); + + for (int i = 0; i < 3; i++) { + aux[i] = 0; + for (int j = 0; j < 4; j++) { + aux[i] |= ((uint32_t)block_scales[i * 4 + j]) << (j * 8); + } + } + + uint32_t tmp = aux[2]; + aux[2] = ((aux[0] >> 4) & kmask2) | (((tmp >> 4) & kmask1) << 4); + aux[3] = ((aux[1] >> 4) & kmask2) | (((tmp >> 6) & kmask1) << 4); + aux[0] = (aux[0] & kmask2) | (((tmp >> 0) & kmask1) << 4); + aux[1] = (aux[1] & kmask2) | (((tmp >> 2) & kmask1) << 4); + + int is = 0; + float dl; + for (int n = 0; n < 256; n += 128) { + int shift = 0; + for (int j = 0; j < 4; ++j) { + + dl = d_all * (scales[is++] - 32); + for (int l = 0; l < 16; ++l) { + *output_blk++ = __float2half(dl * ((int8_t)((q[l+ 0] >> shift) & 3) - ((hm[l+ 0] & m) ? 0 : 4))); + } + + dl = d_all * (scales[is++] - 32); + for (int l = 0; l < 16; ++l) { + *output_blk++ = __float2half(dl * ((int8_t)((q[l+16] >> shift) & 3) - ((hm[l+16] & m) ? 0 : 4))); + } + + shift += 2; + m <<= 1; + } + q += 32; + } + } +} + +__global__ void dequantize_q3_k_bf16_kernel(const int8_t* data, nv_bfloat16* output, const int blk_size, const int ele_per_blk, const int num_blocks) { + + long long global_idx = blockIdx.x * blockDim.x + threadIdx.x; + const uint32_t kmask1 = 0x03030303; + const uint32_t kmask2 = 0x0f0f0f0f; + for (long long block_id=global_idx; block_id(data + block_id * blk_size + 108))); + + const uint8_t * __restrict__ q = (uint8_t*)(data + block_id * blk_size + 32); + const uint8_t * __restrict__ hm = (uint8_t*)(data + block_id * blk_size + 0); + uint8_t m = 1; + + + uint8_t* block_scales = (uint8_t*)(data + block_id * blk_size + 96); + + for (int i = 0; i < 3; i++) { + aux[i] = 0; + for (int j = 0; j < 4; j++) { + aux[i] |= ((uint32_t)block_scales[i * 4 + j]) << (j * 8); + } + } + + uint32_t tmp = aux[2]; + aux[2] = ((aux[0] >> 4) & kmask2) | (((tmp >> 4) & kmask1) << 4); + aux[3] = ((aux[1] >> 4) & kmask2) | (((tmp >> 6) & kmask1) << 4); + aux[0] = (aux[0] & kmask2) | (((tmp >> 0) & kmask1) << 4); + aux[1] = (aux[1] & kmask2) | (((tmp >> 2) & kmask1) << 4); + + int is = 0; + float dl; + for (int n = 0; n < 256; n += 128) { + int shift = 0; + for (int j = 0; j < 4; ++j) { + + dl = d_all * (scales[is++] - 32); + for (int l = 0; l < 16; ++l) { + *output_blk++ = __float2bfloat16(dl * ((int8_t)((q[l+ 0] >> shift) & 3) - ((hm[l+ 0] & m) ? 0 : 4))); + } + + dl = d_all * (scales[is++] - 32); + for (int l = 0; l < 16; ++l) { + *output_blk++ = __float2bfloat16(dl * ((int8_t)((q[l+16] >> shift) & 3) - ((hm[l+16] & m) ? 0 : 4))); + } + + shift += 2; + m <<= 1; + } + q += 32; + } + } +} + + +__global__ void dequantize_q4_k_fp32_kernel(const int8_t* data, float* output, const int blk_size, const int ele_per_blk, const int num_blocks) { + long long global_idx = blockIdx.x * blockDim.x + threadIdx.x; + for (long long block_id=global_idx; block_id(data + block_id * 144 + 0))); + const float min = __half2float(*(reinterpret_cast(data + block_id * 144 + 2))); + int is = 0; + uint8_t sc, m; + for (int j = 0; j < ele_per_blk; j += 64) { + uint8_t* scales = (uint8_t*)(data + block_id * 144 + 4); + get_scale_min_k4(is + 0, scales, &sc, &m); + const float d1 = d * sc; const float m1 = min * m; + get_scale_min_k4(is + 1, scales, &sc, &m); + const float d2 = d * sc; const float m2 = min * m; + for (int l = 0; l < 32; ++l) *output_blk++ = d1 * (q[l] & 0xF) - m1; + for (int l = 0; l < 32; ++l) *output_blk++ = d2 * (q[l] >> 4) - m2; + q += 32; is += 2; + } + } +} + +__global__ void dequantize_q4_k_fp16_kernel(const int8_t* data, __half* output, const int blk_size, const int ele_per_blk, const int num_blocks) { + long long global_idx = blockIdx.x * blockDim.x + threadIdx.x; + for (long long block_id=global_idx; block_id(data + block_id * 144 + 0))); + const float min = __half2float(*(reinterpret_cast(data + block_id * 144 + 2))); + int is = 0; + uint8_t sc, m; + for (int j = 0; j < ele_per_blk; j += 64) { + uint8_t* scales = (uint8_t*)(data + block_id * 144 + 4); + get_scale_min_k4(is + 0, scales, &sc, &m); + const float d1 = d * sc; const float m1 = min * m; + get_scale_min_k4(is + 1, scales, &sc, &m); + const float d2 = d * sc; const float m2 = min * m; + for (int l = 0; l < 32; ++l) *output_blk++ = __float2half(d1 * (q[l] & 0xF) - m1); + for (int l = 0; l < 32; ++l) *output_blk++ = __float2half(d2 * (q[l] >> 4) - m2); + q += 32; is += 2; + } + } +} + +__global__ void dequantize_q4_k_bf16_kernel(const int8_t* data, nv_bfloat16* output, const int blk_size, const int ele_per_blk, const int num_blocks) { + long long global_idx = blockIdx.x * blockDim.x + threadIdx.x; + for (long long block_id=global_idx; block_id(data + block_id * 144 + 0))); + const float min = __half2float(*(reinterpret_cast(data + block_id * 144 + 2))); + int is = 0; + uint8_t sc, m; + for (int j = 0; j < ele_per_blk; j += 64) { + uint8_t* scales = (uint8_t*)(data + block_id * 144 + 4); + get_scale_min_k4(is + 0, scales, &sc, &m); + const float d1 = d * sc; const float m1 = min * m; + get_scale_min_k4(is + 1, scales, &sc, &m); + const float d2 = d * sc; const float m2 = min * m; + for (int l = 0; l < 32; ++l) *output_blk++ = __float2bfloat16(d1 * (q[l] & 0xF) - m1); + for (int l = 0; l < 32; ++l) *output_blk++ = __float2bfloat16(d2 * (q[l] >> 4) - m2); + q += 32; is += 2; + } + } +} + +__global__ void dequantize_q5_k_fp32_kernel(const int8_t* data, float* output, const int blk_size, const int ele_per_blk, const int num_blocks) { + long long global_idx = blockIdx.x * blockDim.x + threadIdx.x; + for (long long block_id = global_idx; block_id < num_blocks; block_id += blockDim.x * gridDim.x){ + float* __restrict__ output_blk = (float*)(output + block_id * ele_per_blk); + + const float d = __half2float(*(reinterpret_cast(data + block_id * blk_size + 0))); + const float min = __half2float(*(reinterpret_cast(data + block_id * blk_size + 2))); + + const uint8_t * __restrict__ qh = (uint8_t*)(data + block_id * blk_size + 16); + const uint8_t * __restrict__ ql = (uint8_t*)(data + block_id * blk_size + 48); + + int is = 0; + uint8_t sc, m; + uint8_t u1 = 1, u2 = 2; + uint8_t* scales = (uint8_t*)(data + block_id * blk_size + 4); + + for (int j = 0; j < 256; j += 64) { + get_scale_min_k4(is + 0, scales, &sc, &m); + const float d1 = d * sc; const float m1 = min * m; + get_scale_min_k4(is + 1, scales, &sc, &m); + const float d2 = d * sc; const float m2 = min * m; + for (int l = 0; l < 32; ++l) *output_blk++ = d1 * ((ql[l] & 0xF) + (qh[l] & u1 ? 16 : 0)) - m1; + for (int l = 0; l < 32; ++l) *output_blk++ = d2 * ((ql[l] >> 4) + (qh[l] & u2 ? 16 : 0)) - m2; + ql += 32; is += 2; + u1 <<= 2; u2 <<= 2; + } + } +} + +__global__ void dequantize_q5_k_fp16_kernel(const int8_t* data, __half* output, const int blk_size, const int ele_per_blk, const int num_blocks) { + long long global_idx = blockIdx.x * blockDim.x + threadIdx.x; + for (long long block_id = global_idx; block_id < num_blocks; block_id += blockDim.x * gridDim.x){ + __half* __restrict__ output_blk = (__half*)(output + block_id * ele_per_blk); + + const float d = __half2float(*(reinterpret_cast(data + block_id * blk_size + 0))); + const float min = __half2float(*(reinterpret_cast(data + block_id * blk_size + 2))); + + const uint8_t * __restrict__ qh = (uint8_t*)(data + block_id * blk_size + 16); + const uint8_t * __restrict__ ql = (uint8_t*)(data + block_id * blk_size + 48); + + int is = 0; + uint8_t sc, m; + uint8_t u1 = 1, u2 = 2; + uint8_t* scales = (uint8_t*)(data + block_id * blk_size + 4); + + for (int j = 0; j < 256; j += 64) { + get_scale_min_k4(is + 0, scales, &sc, &m); + const float d1 = d * sc; const float m1 = min * m; + get_scale_min_k4(is + 1, scales, &sc, &m); + const float d2 = d * sc; const float m2 = min * m; + for (int l = 0; l < 32; ++l) *output_blk++ = __float2half(d1 * ((ql[l] & 0xF) + (qh[l] & u1 ? 16 : 0)) - m1); + for (int l = 0; l < 32; ++l) *output_blk++ = __float2half(d2 * ((ql[l] >> 4) + (qh[l] & u2 ? 16 : 0)) - m2); + ql += 32; is += 2; + u1 <<= 2; u2 <<= 2; + } + } +} + +__global__ void dequantize_q5_k_bf16_kernel(const int8_t* data, nv_bfloat16* output, const int blk_size, const int ele_per_blk, const int num_blocks) { + long long global_idx = blockIdx.x * blockDim.x + threadIdx.x; + for (long long block_id = global_idx; block_id < num_blocks; block_id += blockDim.x * gridDim.x){ + nv_bfloat16* __restrict__ output_blk = (nv_bfloat16*)(output + block_id * ele_per_blk); + + const float d = __half2float(*(reinterpret_cast(data + block_id * blk_size + 0))); + const float min = __half2float(*(reinterpret_cast(data + block_id * blk_size + 2))); + + const uint8_t * __restrict__ qh = (uint8_t*)(data + block_id * blk_size + 16); + const uint8_t * __restrict__ ql = (uint8_t*)(data + block_id * blk_size + 48); + + int is = 0; + uint8_t sc, m; + uint8_t u1 = 1, u2 = 2; + uint8_t* scales = (uint8_t*)(data + block_id * blk_size + 4); + + for (int j = 0; j < 256; j += 64) { + get_scale_min_k4(is + 0, scales, &sc, &m); + const float d1 = d * sc; const float m1 = min * m; + get_scale_min_k4(is + 1, scales, &sc, &m); + const float d2 = d * sc; const float m2 = min * m; + for (int l = 0; l < 32; ++l) *output_blk++ = __float2bfloat16(d1 * ((ql[l] & 0xF) + (qh[l] & u1 ? 16 : 0)) - m1); + for (int l = 0; l < 32; ++l) *output_blk++ = __float2bfloat16(d2 * ((ql[l] >> 4) + (qh[l] & u2 ? 16 : 0)) - m2); + ql += 32; is += 2; + u1 <<= 2; u2 <<= 2; + } + } +} + +__global__ void dequantize_q6_k_fp32_kernel(const int8_t* data, float* output, const int blk_size, const int ele_per_blk, const int num_blocks) { + long long global_idx = blockIdx.x * blockDim.x + threadIdx.x; + for (long long block_id=global_idx; block_id(data + block_id * blk_size + 208))); + + const uint8_t * __restrict__ ql = (uint8_t*)(data + block_id * blk_size); + const uint8_t * __restrict__ qh = (uint8_t*)(data + block_id * blk_size + 128); + const int8_t * __restrict__ sc = (int8_t*)(data + block_id * blk_size + 192); + + + for (int n = 0; n < ele_per_blk; n += 128) { + for (int l = 0; l < 32; ++l) { + int is = l/16; + const int8_t q1 = (int8_t)((ql[l + 0] & 0xF) | (((qh[l] >> 0) & 3) << 4)) - 32; + const int8_t q2 = (int8_t)((ql[l + 32] & 0xF) | (((qh[l] >> 2) & 3) << 4)) - 32; + const int8_t q3 = (int8_t)((ql[l + 0] >> 4) | (((qh[l] >> 4) & 3) << 4)) - 32; + const int8_t q4 = (int8_t)((ql[l + 32] >> 4) | (((qh[l] >> 6) & 3) << 4)) - 32; + output_blk[l + 0] = d * sc[is + 0] * q1; + output_blk[l + 32] = d * sc[is + 2] * q2; + output_blk[l + 64] = d * sc[is + 4] * q3; + output_blk[l + 96] = d * sc[is + 6] * q4; + } + output_blk += 128; + ql += 64; + qh += 32; + sc += 8; + } + } +} + +__global__ void dequantize_q6_k_fp16_kernel(const int8_t* data, __half* output, const int blk_size, const int ele_per_blk, const int num_blocks) { + long long global_idx = blockIdx.x * blockDim.x + threadIdx.x; + for (long long block_id=global_idx; block_id(data + block_id * blk_size + 208))); + + const uint8_t * __restrict__ ql = (uint8_t*)(data + block_id * blk_size); + const uint8_t * __restrict__ qh = (uint8_t*)(data + block_id * blk_size + 128); + const int8_t * __restrict__ sc = (int8_t*)(data + block_id * blk_size + 192); + + + for (int n = 0; n < ele_per_blk; n += 128) { + for (int l = 0; l < 32; ++l) { + int is = l/16; + const int8_t q1 = (int8_t)((ql[l + 0] & 0xF) | (((qh[l] >> 0) & 3) << 4)) - 32; + const int8_t q2 = (int8_t)((ql[l + 32] & 0xF) | (((qh[l] >> 2) & 3) << 4)) - 32; + const int8_t q3 = (int8_t)((ql[l + 0] >> 4) | (((qh[l] >> 4) & 3) << 4)) - 32; + const int8_t q4 = (int8_t)((ql[l + 32] >> 4) | (((qh[l] >> 6) & 3) << 4)) - 32; + output_blk[l + 0] = __float2half(d * sc[is + 0] * q1); + output_blk[l + 32] = __float2half(d * sc[is + 2] * q2); + output_blk[l + 64] = __float2half(d * sc[is + 4] * q3); + output_blk[l + 96] = __float2half(d * sc[is + 6] * q4); + } + output_blk += 128; + ql += 64; + qh += 32; + sc += 8; + } + } +} + +__global__ void dequantize_q6_k_bf16_kernel(const int8_t* data, nv_bfloat16* output, const int blk_size, const int ele_per_blk, const int num_blocks) { + long long global_idx = blockIdx.x * blockDim.x + threadIdx.x; + for (long long block_id=global_idx; block_id(data + block_id * blk_size + 208))); + + const uint8_t * __restrict__ ql = (uint8_t*)(data + block_id * blk_size); + const uint8_t * __restrict__ qh = (uint8_t*)(data + block_id * blk_size + 128); + const int8_t * __restrict__ sc = (int8_t*)(data + block_id * blk_size + 192); + + + for (int n = 0; n < ele_per_blk; n += 128) { + for (int l = 0; l < 32; ++l) { + int is = l/16; + const int8_t q1 = (int8_t)((ql[l + 0] & 0xF) | (((qh[l] >> 0) & 3) << 4)) - 32; + const int8_t q2 = (int8_t)((ql[l + 32] & 0xF) | (((qh[l] >> 2) & 3) << 4)) - 32; + const int8_t q3 = (int8_t)((ql[l + 0] >> 4) | (((qh[l] >> 4) & 3) << 4)) - 32; + const int8_t q4 = (int8_t)((ql[l + 32] >> 4) | (((qh[l] >> 6) & 3) << 4)) - 32; + output_blk[l + 0] = __float2bfloat16(d * sc[is + 0] * q1); + output_blk[l + 32] = __float2bfloat16(d * sc[is + 2] * q2); + output_blk[l + 64] = __float2bfloat16(d * sc[is + 4] * q3); + output_blk[l + 96] = __float2bfloat16(d * sc[is + 6] * q4); + } + output_blk += 128; + ql += 64; + qh += 32; + sc += 8; + } + } +} + +static constexpr __device__ int8_t kvalues_iq4nl[16] = {-127, -104, -83, -65, -49, -35, -22, -10, 1, 13, 25, 38, 53, 69, 89, 113}; + +__global__ void dequantize_iq4_xs_fp32_kernel(const int8_t* data, float* output, const int blk_size, const int ele_per_blk, const int num_blocks) { + long long global_idx = blockIdx.x * blockDim.x + threadIdx.x; + for (long long block_id=global_idx; block_id(data + block_id * blk_size))); + const uint16_t scales_h = *(reinterpret_cast(data + block_id * blk_size + 2)); + const uint8_t* scales_l = (uint8_t*)(data + block_id * blk_size + 2 + 2); + const uint8_t* qs = (uint8_t*)(data + block_id * blk_size + 2 + 2 + 4); + + for (int ib = 0; ib < 8; ++ib) { + const int ls = ((scales_l[ib / 2] >> 4 * (ib % 2)) & 0xf) | (((scales_h >> 2 * ib) & 3) << 4); + const float dl = d * (ls - 32); + for (int j = 0; j < 16; ++j) { + output_blk[j + 0] = dl * kvalues_iq4nl[qs[j] & 0xf]; + output_blk[j + 16] = dl * kvalues_iq4nl[qs[j] >> 4]; + } + output_blk += 32; + qs += 16; + } + } +} + +__global__ void dequantize_iq4_xs_fp16_kernel(const int8_t* data, __half* output, const int blk_size, const int ele_per_blk, const int num_blocks) { + long long global_idx = blockIdx.x * blockDim.x + threadIdx.x; + for (long long block_id=global_idx; block_id(data + block_id * blk_size))); + const uint16_t scales_h = *(reinterpret_cast(data + block_id * blk_size + 2)); + const uint8_t* scales_l = (uint8_t*)(data + block_id * blk_size + 2 + 2); + const uint8_t* qs = (uint8_t*)(data + block_id * blk_size + 2 + 2 + 4); + + for (int ib = 0; ib < 8; ++ib) { + const int ls = ((scales_l[ib / 2] >> 4 * (ib % 2)) & 0xf) | (((scales_h >> 2 * ib) & 3) << 4); + const float dl = d * (ls - 32); + for (int j = 0; j < 16; ++j) { + output_blk[j + 0] = __float2half(dl * kvalues_iq4nl[qs[j] & 0xf]); + output_blk[j + 16] = __float2half(dl * kvalues_iq4nl[qs[j] >> 4]); + } + output_blk += 32; + qs += 16; + } + } +} + +__global__ void dequantize_iq4_xs_bf16_kernel(const int8_t* data, nv_bfloat16* output, const int blk_size, const int ele_per_blk, const int num_blocks) { + long long global_idx = blockIdx.x * blockDim.x + threadIdx.x; + for (long long block_id=global_idx; block_id(data + block_id * blk_size))); + const uint16_t scales_h = *(reinterpret_cast(data + block_id * blk_size + 2)); + const uint8_t* scales_l = (uint8_t*)(data + block_id * blk_size + 2 + 2); + const uint8_t* qs = (uint8_t*)(data + block_id * blk_size + 2 + 2 + 4); + + for (int ib = 0; ib < 8; ++ib) { + const int ls = ((scales_l[ib / 2] >> 4 * (ib % 2)) & 0xf) | (((scales_h >> 2 * ib) & 3) << 4); + const float dl = d * (ls - 32); + for (int j = 0; j < 16; ++j) { + output_blk[j + 0] = __float2bfloat16(dl * kvalues_iq4nl[qs[j] & 0xf]); + output_blk[j + 16] = __float2bfloat16(dl * kvalues_iq4nl[qs[j] >> 4]); + } + output_blk += 32; + qs += 16; + } + } +} + +torch::Tensor dequantize_q8_0(const int8_t* data, const int num_bytes, const int blk_size, const int ele_per_blk, const torch::Device device, const torch::Dtype target_dtype) { + int num_blocks = num_bytes / blk_size; + const at::cuda::OptionalCUDAGuard device_guard(device); + + auto options = torch::TensorOptions().dtype(torch::kInt8).device(device).memory_format(torch::MemoryFormat::Contiguous); + auto data_gpu = torch::empty({ num_bytes }, options); + + cudaMemcpy(data_gpu.data_ptr(), data, num_bytes, cudaMemcpyHostToDevice); + //data_gpu.copy_(data, false); + + // Create output tensor + auto output = torch::zeros({ num_blocks, ele_per_blk }, torch::dtype(target_dtype).device(device)); + + switch (target_dtype) { + case torch::kFloat16: + dequantize_q8_0_fp16_kernel<<<512, 256>>>(data_gpu.data_ptr(), (__half*)output.data_ptr(), blk_size, ele_per_blk, num_blocks); + break; + case torch::kBFloat16: + dequantize_q8_0_bf16_kernel<<<512, 256>>>(data_gpu.data_ptr(), (nv_bfloat16*)output.data_ptr(), blk_size, ele_per_blk, num_blocks); + break; + case torch::kFloat32: + dequantize_q8_0_fp32_kernel<<<512, 256>>>(data_gpu.data_ptr(), output.data_ptr(), blk_size, ele_per_blk, num_blocks); + break; + default: + printf("target type not support\n"); + exit(0); + } + + cudaDeviceSynchronize(); + return output; +} + + +torch::Tensor dequantize_q6_k(const int8_t* data, const int num_bytes, const int blk_size, const int ele_per_blk, const torch::Device device, const torch::Dtype target_dtype) { + // data.numel%blk_size should be 0, else raise err + int num_blocks = num_bytes / blk_size; + + const at::cuda::OptionalCUDAGuard device_guard(device); + auto options = torch::TensorOptions().dtype(torch::kInt8).device(device).memory_format(torch::MemoryFormat::Contiguous); + auto data_gpu = torch::empty({num_bytes}, options); + + cudaMemcpy(data_gpu.data_ptr(), data, num_bytes, cudaMemcpyHostToDevice); + //data_gpu.copy_(data, false); + + // Create output tensor + auto output = torch::zeros({num_blocks, ele_per_blk}, torch::dtype(target_dtype).device(device)); + + switch (target_dtype) { + case torch::kFloat16: + dequantize_q6_k_fp16_kernel<<<512, 256>>>(data_gpu.data_ptr(), (__half*)output.data_ptr(), blk_size, ele_per_blk, num_blocks); + break; + case torch::kBFloat16: + dequantize_q6_k_bf16_kernel<<<512, 256>>>(data_gpu.data_ptr(), (nv_bfloat16*)output.data_ptr(), blk_size, ele_per_blk, num_blocks); + break; + case torch::kFloat32: + dequantize_q6_k_fp32_kernel<<<512, 256>>>(data_gpu.data_ptr(), output.data_ptr(), blk_size, ele_per_blk, num_blocks); + break; + default: + printf("target type not support\n"); + exit(0); + } + cudaDeviceSynchronize(); + return output; +} + +torch::Tensor dequantize_q5_k(const int8_t* data, const int num_bytes, const int blk_size, const int ele_per_blk, const torch::Device device, const torch::Dtype target_dtype) { + int num_blocks = num_bytes / blk_size; + const at::cuda::OptionalCUDAGuard device_guard(device); + + auto options = torch::TensorOptions().dtype(torch::kInt8).device(device).memory_format(torch::MemoryFormat::Contiguous); + auto data_gpu = torch::empty({num_bytes}, options); + + cudaMemcpy(data_gpu.data_ptr(), data, num_bytes, cudaMemcpyHostToDevice); + //data_gpu.copy_(data, false); + + // Create output tensor + auto output = torch::zeros({num_blocks, ele_per_blk}, torch::dtype(target_dtype).device(device)); + + switch (target_dtype) { + case torch::kFloat16: + dequantize_q5_k_fp16_kernel<<<512, 256>>>(data_gpu.data_ptr(), (__half*)output.data_ptr(), blk_size, ele_per_blk, num_blocks); + break; + case torch::kBFloat16: + dequantize_q5_k_bf16_kernel<<<512, 256>>>(data_gpu.data_ptr(), (nv_bfloat16*)output.data_ptr(), blk_size, ele_per_blk, num_blocks); + break; + case torch::kFloat32: + dequantize_q5_k_fp32_kernel<<<512, 256>>>(data_gpu.data_ptr(), output.data_ptr(), blk_size, ele_per_blk, num_blocks); + break; + default: + printf("target type not support\n"); + exit(0); + } + cudaDeviceSynchronize(); + return output; +} + +torch::Tensor dequantize_q4_k(const int8_t* data, const int num_bytes, const int blk_size, const int ele_per_blk, const torch::Device device, const torch::Dtype target_dtype) { + // data.numel%blk_size should be 0, else raise err + int num_blocks = num_bytes / blk_size; + const at::cuda::OptionalCUDAGuard device_guard(device); + + auto options = torch::TensorOptions().dtype(torch::kInt8).device(device).memory_format(torch::MemoryFormat::Contiguous); + auto data_gpu = torch::empty({num_bytes}, options); + + cudaMemcpy(data_gpu.data_ptr(), data, num_bytes, cudaMemcpyHostToDevice); + //data_gpu.copy_(data, false); + + // Create output tensor + auto output = torch::zeros({num_blocks, ele_per_blk}, torch::dtype(target_dtype).device(device)); + + switch (target_dtype) { + case torch::kFloat16: + dequantize_q4_k_fp16_kernel<<<512, 256>>>(data_gpu.data_ptr(), (__half*)output.data_ptr(), blk_size, ele_per_blk, num_blocks); + break; + case torch::kBFloat16: + dequantize_q4_k_bf16_kernel<<<512, 256>>>(data_gpu.data_ptr(), (nv_bfloat16*)output.data_ptr(), blk_size, ele_per_blk, num_blocks); + break; + case torch::kFloat32: + dequantize_q4_k_fp32_kernel<<<512, 256>>>(data_gpu.data_ptr(), output.data_ptr(), blk_size, ele_per_blk, num_blocks); + break; + default: + printf("target type not support\n"); + exit(0); + } + cudaDeviceSynchronize(); + return output; +} + +torch::Tensor dequantize_q3_k(const int8_t* data, const int num_bytes, const int blk_size, const int ele_per_blk, const torch::Device device, const torch::Dtype target_dtype) { + int num_blocks = num_bytes / blk_size; + const at::cuda::OptionalCUDAGuard device_guard(device); + + auto options = torch::TensorOptions().dtype(torch::kInt8).device(device).memory_format(torch::MemoryFormat::Contiguous); + auto data_gpu = torch::empty({num_bytes}, options); + + cudaMemcpy(data_gpu.data_ptr(), data, num_bytes, cudaMemcpyHostToDevice); + //data_gpu.copy_(data, false); + + // Create output tensor + auto output = torch::zeros({num_blocks, ele_per_blk}, torch::dtype(target_dtype).device(device)); + + switch (target_dtype) { + case torch::kFloat16: + dequantize_q3_k_fp16_kernel<<<512, 256>>>(data_gpu.data_ptr(), (__half*)output.data_ptr(), blk_size, ele_per_blk, num_blocks); + break; + case torch::kBFloat16: + dequantize_q3_k_bf16_kernel<<<512, 256>>>(data_gpu.data_ptr(), (nv_bfloat16*)output.data_ptr(), blk_size, ele_per_blk, num_blocks); + break; + case torch::kFloat32: + dequantize_q3_k_fp32_kernel<<<512, 256>>>(data_gpu.data_ptr(), output.data_ptr(), blk_size, ele_per_blk, num_blocks); + break; + default: + printf("target type not support\n"); + exit(0); + } + cudaDeviceSynchronize(); + return output; +} + +torch::Tensor dequantize_q2_k(const int8_t* data, const int num_bytes, const int blk_size, const int ele_per_blk, const torch::Device device, const torch::Dtype target_dtype) { + int num_blocks = num_bytes / blk_size; + const at::cuda::OptionalCUDAGuard device_guard(device); + + auto options = torch::TensorOptions().dtype(torch::kInt8).device(device).memory_format(torch::MemoryFormat::Contiguous); + auto data_gpu = torch::empty({num_bytes}, options); + + cudaMemcpy(data_gpu.data_ptr(), data, num_bytes, cudaMemcpyHostToDevice); + //data_gpu.copy_(data, false); + + // Create output tensor + auto output = torch::zeros({num_blocks, ele_per_blk}, torch::dtype(target_dtype).device(device)); + + switch (target_dtype) { + case torch::kFloat16: + dequantize_q2_k_fp16_kernel<<<512, 256>>>(data_gpu.data_ptr(), (__half*)output.data_ptr(), blk_size, ele_per_blk, num_blocks); + break; + case torch::kBFloat16: + dequantize_q2_k_bf16_kernel<<<512, 256>>>(data_gpu.data_ptr(), (nv_bfloat16*)output.data_ptr(), blk_size, ele_per_blk, num_blocks); + break; + case torch::kFloat32: + dequantize_q2_k_fp32_kernel<<<512, 256>>>(data_gpu.data_ptr(), output.data_ptr(), blk_size, ele_per_blk, num_blocks); + break; + default: + printf("target type not support\n"); + exit(0); + } + cudaDeviceSynchronize(); + return output; +} + +torch::Tensor dequantize_iq4_xs(const int8_t* data, const int num_bytes, const int blk_size, const int ele_per_blk, const torch::Device device, const torch::Dtype target_dtype) { + int num_blocks = num_bytes / blk_size; + const at::cuda::OptionalCUDAGuard device_guard(device); + + auto options = torch::TensorOptions().dtype(torch::kInt8).device(device).memory_format(torch::MemoryFormat::Contiguous); + auto data_gpu = torch::empty({num_bytes}, options); + + cudaMemcpy(data_gpu.data_ptr(), data, num_bytes, cudaMemcpyHostToDevice); + //data_gpu.copy_(data, false); + + // Create output tensor + auto output = torch::zeros({num_blocks, ele_per_blk}, torch::dtype(target_dtype).device(device)); + + switch (target_dtype) { + case torch::kFloat16: + dequantize_iq4_xs_fp16_kernel<<<512, 256>>>(data_gpu.data_ptr(), (__half*)output.data_ptr(), blk_size, ele_per_blk, num_blocks); + break; + case torch::kBFloat16: + dequantize_iq4_xs_bf16_kernel<<<512, 256>>>(data_gpu.data_ptr(), (nv_bfloat16*)output.data_ptr(), blk_size, ele_per_blk, num_blocks); + break; + case torch::kFloat32: + dequantize_iq4_xs_fp32_kernel<<<512, 256>>>(data_gpu.data_ptr(), output.data_ptr(), blk_size, ele_per_blk, num_blocks); + break; + default: + printf("target type not support\n"); + exit(0); + } + cudaDeviceSynchronize(); + return output; +} diff --git a/archive/kt-sft/csrc/ktransformers_ext/cuda/custom_gguf/ops.h b/archive/kt-sft/csrc/ktransformers_ext/cuda/custom_gguf/ops.h new file mode 100644 index 0000000..fe9161a --- /dev/null +++ b/archive/kt-sft/csrc/ktransformers_ext/cuda/custom_gguf/ops.h @@ -0,0 +1,22 @@ +/** + * @Description : + * @Author : Azure-Tang + * @Date : 2024-07-22 09:27:55 + * @Version : 1.0.0 + * @LastEditors : kkk1nak0 + * @LastEditTime : 2024-08-12 03:48:46 + * @Copyright (c) 2024 by KVCache.AI, All Rights Reserved. +**/ +#pragma once + +#include +#include +#include + +torch::Tensor dequantize_q8_0(const int8_t* data, const int num_bytes, const int blk_size, const int ele_per_blk, const torch::Device device, const torch::Dtype target_dtype); +torch::Tensor dequantize_q6_k(const int8_t* data, const int num_bytes, const int blk_size, const int ele_per_blk, const torch::Device device, const torch::Dtype target_dtype); +torch::Tensor dequantize_q5_k(const int8_t* data, const int num_bytes, const int blk_size, const int ele_per_blk, const torch::Device device, const torch::Dtype target_dtype); +torch::Tensor dequantize_q4_k(const int8_t* data, const int num_bytes, const int blk_size, const int ele_per_blk, const torch::Device device, const torch::Dtype target_dtype); +torch::Tensor dequantize_q3_k(const int8_t* data, const int num_bytes, const int blk_size, const int ele_per_blk, const torch::Device device, const torch::Dtype target_dtype); +torch::Tensor dequantize_q2_k(const int8_t* data, const int num_bytes, const int blk_size, const int ele_per_blk, const torch::Device device, const torch::Dtype target_dtype); +torch::Tensor dequantize_iq4_xs(const int8_t* data, const int num_bytes, const int blk_size, const int ele_per_blk, const torch::Device device, const torch::Dtype target_dtype); \ No newline at end of file diff --git a/archive/kt-sft/csrc/ktransformers_ext/cuda/gptq_marlin/gptq_marlin.cu b/archive/kt-sft/csrc/ktransformers_ext/cuda/gptq_marlin/gptq_marlin.cu new file mode 100644 index 0000000..87f4581 --- /dev/null +++ b/archive/kt-sft/csrc/ktransformers_ext/cuda/gptq_marlin/gptq_marlin.cu @@ -0,0 +1,1907 @@ +/* + * Modified by Neural Magic + * Copyright (C) Marlin.2024 Elias Frantar + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +/* + * Adapted from https://github.com/IST-DASLab/marlin + */ +/* + * Adapted from https://github.com/vllm-project/vllm/tree/main/csrc/quantization/gptq_marlin + */ +#include "gptq_marlin.cuh" +#include "gptq_marlin_dtypes.cuh" +#include +#define STATIC_ASSERT_SCALAR_TYPE_VALID(scalar_t) \ + static_assert(std::is_same::value || \ + std::is_same::value, \ + "only float16 and bfloat16 is supported"); + +template +inline std::string str(T x) { + return std::to_string(x); +} + +namespace gptq_marlin { + +#if (defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800) || defined(__HIP_PLATFORM_AMD__) + +__global__ void permute_cols_kernel(int4 const* __restrict__ a_int4_ptr, + int const* __restrict__ perm_int_ptr, + int4* __restrict__ out_int4_ptr, int size_m, + int size_k, int block_rows) {} + +template shared + // fetch pipeline + const bool has_act_order, // whether act_order is enabled + const int group_blocks = -1 // number of consecutive 16x16 blocks + // with a separate quantization scale + > +__global__ void Marlin( + const int4* __restrict__ A, // fp16 input matrix of shape mxk + const int4* __restrict__ B, // 4bit quantized weight matrix of shape kxn + int4* __restrict__ C, // fp16 output buffer of shape mxn + const int4* __restrict__ scales_ptr, // fp16 quantization scales of shape + // (k/groupsize)xn + const int* __restrict__ g_idx, // int32 group indices of shape k + int num_groups, // number of scale groups per output channel + int prob_m, // batch dimension m + int prob_n, // output dimension n + int prob_k, // reduction dimension k + int* locks // extra global storage for barrier synchronization +) {} + +} // namespace gptq_marlin + +torch::Tensor gptq_marlin_gemm(torch::Tensor& a, torch::Tensor& b_q_weight, + torch::Tensor& b_scales, torch::Tensor& g_idx, + torch::Tensor& perm, torch::Tensor& workspace, + int64_t num_bits, int64_t size_m, int64_t size_n, + int64_t size_k, bool is_k_full) { + TORCH_CHECK_NOT_IMPLEMENTED(false, + "marlin_gemm(..) requires CUDA_ARCH >= 8.0"); + return torch::empty({1, 1}); +} + +#else + +// m16n8k16 tensor core mma instruction with fp16 inputs and fp32 +// output/accumulation. +template +__device__ inline void mma(const typename ScalarType::FragA& a_frag, + const typename ScalarType::FragB& frag_b, + typename ScalarType::FragC& frag_c) { + const uint32_t* a = reinterpret_cast(&a_frag); + const uint32_t* b = reinterpret_cast(&frag_b); + float* c = reinterpret_cast(&frag_c); + if constexpr (std::is_same::value) { + asm volatile( + "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 " + "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};\n" + : "=f"(c[0]), "=f"(c[1]), "=f"(c[2]), "=f"(c[3]) + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1]), + "f"(c[0]), "f"(c[1]), "f"(c[2]), "f"(c[3])); + } else if constexpr (std::is_same::value) { + asm volatile( + "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 " + "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};\n" + : "=f"(c[0]), "=f"(c[1]), "=f"(c[2]), "=f"(c[3]) + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1]), + "f"(c[0]), "f"(c[1]), "f"(c[2]), "f"(c[3])); + } else { + STATIC_ASSERT_SCALAR_TYPE_VALID(scalar_t); + } +} + +// Instruction for loading a full 16x16 matrix fragment of operand A from shared +// memory, directly in tensor core layout. +template +__device__ inline void ldsm4(typename ScalarType::FragA& frag_a, + const void* smem_ptr) { + uint32_t* a = reinterpret_cast(&frag_a); + uint32_t smem = static_cast(__cvta_generic_to_shared(smem_ptr)); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0,%1,%2,%3}, [%4];\n" + : "=r"(a[0]), "=r"(a[1]), "=r"(a[2]), "=r"(a[3]) + : "r"(smem)); +} + +// Lookup-table based 3-input logical operation; explicitly used for +// dequantization as the compiler does not seem to automatically recognize it in +// all cases. +template +__device__ inline int lop3(int a, int b, int c) { + int res; + asm volatile("lop3.b32 %0, %1, %2, %3, %4;\n" + : "=r"(res) + : "r"(a), "r"(b), "r"(c), "n"(lut)); + return res; +} + +// Constructs destination register by taking bytes from 2 sources (based on +// mask) +template +__device__ inline uint32_t prmt(uint32_t a) { + uint32_t res; + asm volatile("prmt.b32 %0, %1, %2, %3;\n" + : "=r"(res) + : "r"(a), "n"(start_byte), "n"(mask)); + return res; +} + +// Efficiently dequantize an int32 value into a full B-fragment of 4 fp16 +// values. We mostly follow the strategy in the link below, with some small +// changes: +// - FP16: +// https://github.com/NVIDIA/FasterTransformer/blob/release/v5.3_tag/src/fastertransformer/cutlass_extensions/include/cutlass_extensions/interleaved_numeric_conversion.h#L215-L287 +// - BF16: +// https://github.com/NVIDIA/FasterTransformer/blob/release/v5.3_tag/src/fastertransformer/cutlass_extensions/include/cutlass_extensions/interleaved_numeric_conversion.h#L327-L385 +template +__device__ inline typename ScalarType::FragB dequant_4bit(int q) { + STATIC_ASSERT_SCALAR_TYPE_VALID(scalar_t); +} + +template <> +__device__ inline typename ScalarType::FragB dequant_4bit(int q) { + const int LO = 0x000f000f; + const int HI = 0x00f000f0; + const int EX = 0x64006400; + // Guarantee that the `(a & b) | c` operations are LOP3s. + int lo = lop3<(0xf0 & 0xcc) | 0xaa>(q, LO, EX); + int hi = lop3<(0xf0 & 0xcc) | 0xaa>(q, HI, EX); + // We want signed int4 outputs, hence we fuse the `-8` symmetric zero point + // directly into `SUB` and `ADD`. + const int SUB = 0x64086408; + const int MUL = 0x2c002c00; + const int ADD = 0xd480d480; + typename ScalarType::FragB frag_b; + frag_b[0] = __hsub2(*reinterpret_cast(&lo), + *reinterpret_cast(&SUB)); + frag_b[1] = __hfma2(*reinterpret_cast(&hi), + *reinterpret_cast(&MUL), + *reinterpret_cast(&ADD)); + return frag_b; +} + +template <> +__device__ inline typename ScalarType::FragB +dequant_4bit(int q) { + static constexpr uint32_t MASK = 0x000f000f; + static constexpr uint32_t EX = 0x43004300; + + // Guarantee that the `(a & b) | c` operations are LOP3s. + + int lo = lop3<(0xf0 & 0xcc) | 0xaa>(q, MASK, EX); + q >>= 4; + int hi = lop3<(0xf0 & 0xcc) | 0xaa>(q, MASK, EX); + + typename ScalarType::FragB frag_b; + static constexpr uint32_t MUL = 0x3F803F80; + static constexpr uint32_t ADD = 0xC308C308; + + frag_b[0] = __hfma2(*reinterpret_cast(&lo), + *reinterpret_cast(&MUL), + *reinterpret_cast(&ADD)); + frag_b[1] = __hfma2(*reinterpret_cast(&hi), + *reinterpret_cast(&MUL), + *reinterpret_cast(&ADD)); + return frag_b; +} + +// Fast Int8ToFp16/Int8ToBf16: Efficiently dequantize 8bit int values to fp16 or +// bf16 Reference: +// - FP16: +// https://github.com/NVIDIA/FasterTransformer/blob/release/v5.3_tag/src/fastertransformer/cutlass_extensions/include/cutlass_extensions/interleaved_numeric_conversion.h#L53-L85 +// - BF16: +// https://github.com/NVIDIA/FasterTransformer/blob/release/v5.3_tag/src/fastertransformer/cutlass_extensions/include/cutlass_extensions/interleaved_numeric_conversion.h#L125-L175 +template +__device__ inline typename ScalarType::FragB dequant_8bit(int q) { + STATIC_ASSERT_SCALAR_TYPE_VALID(scalar_t); +} + +template <> +__device__ inline typename ScalarType::FragB dequant_8bit(int q) { + static constexpr uint32_t mask_for_elt_01 = 0x5250; + static constexpr uint32_t mask_for_elt_23 = 0x5351; + static constexpr uint32_t start_byte_for_fp16 = 0x64646464; + + uint32_t lo = prmt(q); + uint32_t hi = prmt(q); + + static constexpr uint32_t I8s_TO_F16s_MAGIC_NUM = 0x64806480; + + typename ScalarType::FragB frag_b; + frag_b[0] = __hsub2(*reinterpret_cast(&lo), + *reinterpret_cast(&I8s_TO_F16s_MAGIC_NUM)); + frag_b[1] = __hsub2(*reinterpret_cast(&hi), + *reinterpret_cast(&I8s_TO_F16s_MAGIC_NUM)); + return frag_b; +} + +template <> +__device__ inline typename ScalarType::FragB +dequant_8bit(int q) { + typename ScalarType::FragB frag_b; + + float fp32_intermediates[4]; + uint32_t* fp32_intermediates_casted = + reinterpret_cast(fp32_intermediates); + + static constexpr uint32_t fp32_base = 0x4B000000; + fp32_intermediates_casted[0] = __byte_perm(q, fp32_base, 0x7650); + fp32_intermediates_casted[1] = __byte_perm(q, fp32_base, 0x7652); + fp32_intermediates_casted[2] = __byte_perm(q, fp32_base, 0x7651); + fp32_intermediates_casted[3] = __byte_perm(q, fp32_base, 0x7653); + + fp32_intermediates[0] -= 8388736.f; + fp32_intermediates[1] -= 8388736.f; + fp32_intermediates[2] -= 8388736.f; + fp32_intermediates[3] -= 8388736.f; + + uint32_t* bf16_result_ptr = reinterpret_cast(&frag_b); + bf16_result_ptr[0] = __byte_perm(fp32_intermediates_casted[0], + fp32_intermediates_casted[1], 0x7632); + bf16_result_ptr[1] = __byte_perm(fp32_intermediates_casted[2], + fp32_intermediates_casted[3], 0x7632); + + return frag_b; +} + +// Multiply dequantized values by the corresponding quantization scale; used +// only for grouped quantization. +template +__device__ inline void scale(typename ScalarType::FragB& frag_b, + typename ScalarType::FragS& frag_s, + int i) { + using scalar_t2 = typename ScalarType::scalar_t2; + scalar_t2 s = + ScalarType::num2num2(reinterpret_cast(&frag_s)[i]); + frag_b[0] = __hmul2(frag_b[0], s); + frag_b[1] = __hmul2(frag_b[1], s); +} + +// Same as above, but for act_order (each K is multiplied individually) +template +__device__ inline void scale4(typename ScalarType::FragB& frag_b, + typename ScalarType::FragS& frag_s_1, + typename ScalarType::FragS& frag_s_2, + typename ScalarType::FragS& frag_s_3, + typename ScalarType::FragS& frag_s_4, + int i) { + using scalar_t2 = typename ScalarType::scalar_t2; + scalar_t2 s_val_1_2; + s_val_1_2.x = reinterpret_cast(&frag_s_1)[i]; + s_val_1_2.y = reinterpret_cast(&frag_s_2)[i]; + + scalar_t2 s_val_3_4; + s_val_3_4.x = reinterpret_cast(&frag_s_3)[i]; + s_val_3_4.y = reinterpret_cast(&frag_s_4)[i]; + + frag_b[0] = __hmul2(frag_b[0], s_val_1_2); + frag_b[1] = __hmul2(frag_b[1], s_val_3_4); +} + +// Given 2 floats multiply by 2 scales (halves) +template +__device__ inline void scale_float(float* c, + typename ScalarType::FragS& s) { + scalar_t* s_ptr = reinterpret_cast(&s); + c[0] = __fmul_rn(c[0], ScalarType::num2float(s_ptr[0])); + c[1] = __fmul_rn(c[1], ScalarType::num2float(s_ptr[1])); +} + +// Wait until barrier reaches `count`, then lock for current threadblock. +__device__ inline void barrier_acquire(int* lock, int count) { + if (threadIdx.x == 0) { + int state = -1; + do + // Guarantee that subsequent writes by this threadblock will be visible + // globally. + asm volatile("ld.global.acquire.gpu.b32 %0, [%1];\n" + : "=r"(state) + : "l"(lock)); + while (state != count); + } + __syncthreads(); +} + +// Release barrier and increment visitation count. +__device__ inline void barrier_release(int* lock, bool reset = false) { + __syncthreads(); + if (threadIdx.x == 0) { + if (reset) { + lock[0] = 0; + return; + } + int val = 1; + // Make sure that all writes since acquiring this barrier are visible + // globally, while releasing the barrier. + asm volatile("fence.acq_rel.gpu;\n"); + asm volatile("red.relaxed.gpu.global.add.s32 [%0], %1;\n" + : + : "l"(lock), "r"(val)); + } +} + +// For a given "a" of size [M,K] performs a permutation of the K columns based +// on the given "perm" indices. +__global__ void permute_cols_kernel(int4 const* __restrict__ a_int4_ptr, + int const* __restrict__ perm_int_ptr, + int4* __restrict__ out_int4_ptr, int size_m, + int size_k, int block_rows) { + int start_row = block_rows * blockIdx.x; + int finish_row = start_row + block_rows; + if (finish_row > size_m) { + finish_row = size_m; + } + int cur_block_rows = finish_row - start_row; + + int row_stride = size_k * sizeof(half) / 16; + + auto permute_row = [&](int row) { + int iters = size_k / default_threads; + int rest = size_k % default_threads; + + int offset = row * row_stride; + + half const* a_row_half = reinterpret_cast(a_int4_ptr + offset); + half* out_half = reinterpret_cast(out_int4_ptr + offset); + + int base_k = 0; + + for (int i = 0; i < iters; i++) { + int cur_k = base_k + threadIdx.x; + int src_pos = perm_int_ptr[cur_k]; + + out_half[cur_k] = a_row_half[src_pos]; + + base_k += default_threads; + } + + if (rest) { + if (threadIdx.x < rest) { + int cur_k = base_k + threadIdx.x; + int src_pos = perm_int_ptr[cur_k]; + + out_half[cur_k] = a_row_half[src_pos]; + } + } + }; + + for (int i = 0; i < cur_block_rows; i++) { + int cur_row = start_row + i; + if (cur_row < size_m) { + permute_row(cur_row); + } + } +} + +template shared + // fetch pipeline + const bool has_act_order, // whether act_order is enabled + const int group_blocks = -1 // number of consecutive 16x16 blocks + // with a separate quantization scale + > +__global__ void Marlin( + const int4* __restrict__ A, // fp16 input matrix of shape mxk + const int4* __restrict__ B, // 4bit quantized weight matrix of shape kxn + int4* __restrict__ C, // fp16 output buffer of shape mxn + const int4* __restrict__ scales_ptr, // fp16 quantization scales of shape + // (k/groupsize)xn + const int* __restrict__ g_idx, // int32 group indices of shape k + int num_groups, // number of scale groups per output channel + int prob_m, // batch dimension m + int prob_n, // output dimension n + int prob_k, // reduction dimension k + int* locks // extra global storage for barrier synchronization +) { + // Each threadblock processes one "stripe" of the B matrix with (roughly) the + // same size, which might involve multiple column "slices" (of width 16 * + // `thread_n_blocks`). Stripes are defined as shown in the 3x3 matrix 5 SM + // example: + // 0 1 3 + // 0 2 3 + // 1 2 4 + // While this kind of partitioning makes things somewhat more complicated, it + // ensures good utilization of all SMs for many kinds of shape and GPU + // configurations, while requiring as few slow global cross-threadblock + // reductions as possible. + using Dtype = ScalarType; + using scalar_t2 = typename ScalarType::scalar_t2; + using FragA = typename ScalarType::FragA; + using FragB = typename ScalarType::FragB; + using FragC = typename ScalarType::FragC; + using FragS = typename ScalarType::FragS; + + constexpr int pack_factor = 32 / num_bits; + + // For larger GEMMs we run multiple batchsize 64 versions in parallel for a + // better partitioning with less reductions + int parallel = 1; + if (prob_m > 16 * thread_m_blocks) { + parallel = prob_m / (16 * thread_m_blocks); + prob_m = 16 * thread_m_blocks; + } + + int k_tiles = prob_k / 16 / thread_k_blocks; + int n_tiles = prob_n / 16 / thread_n_blocks; + int iters = div_ceil(k_tiles * n_tiles * parallel, gridDim.x); + + if constexpr (!has_act_order && group_blocks != -1) { + if (group_blocks >= thread_k_blocks) { + // Ensure that the number of tiles in each stripe is a multiple of the + // groupsize; this avoids an annoying special case where a stripe starts + // in the middle of group. + iters = (group_blocks / thread_k_blocks) * + div_ceil(iters, (group_blocks / thread_k_blocks)); + } + } + + int slice_row = (iters * blockIdx.x) % k_tiles; + int slice_col_par = (iters * blockIdx.x) / k_tiles; + int slice_col = slice_col_par; + int slice_iters; // number of threadblock tiles in the current slice + int slice_count = + 0; // total number of active threadblocks in the current slice + int slice_idx; // index of threadblock in current slice; numbered bottom to + // top + + // We can easily implement parallel problem execution by just remapping + // indices and advancing global pointers + if (slice_col_par >= n_tiles) { + A += (slice_col_par / n_tiles) * 16 * thread_m_blocks * prob_k / 8; + C += (slice_col_par / n_tiles) * 16 * thread_m_blocks * prob_n / 8; + locks += (slice_col_par / n_tiles) * n_tiles; + slice_col = slice_col_par % n_tiles; + } + + // Compute all information about the current slice which is required for + // synchronization. + auto init_slice = [&]() { + slice_iters = + iters * (blockIdx.x + 1) - (k_tiles * slice_col_par + slice_row); + if (slice_iters < 0 || slice_col_par >= n_tiles * parallel) slice_iters = 0; + if (slice_iters == 0) return; + if (slice_row + slice_iters > k_tiles) slice_iters = k_tiles - slice_row; + slice_count = 1; + slice_idx = 0; + int col_first = iters * div_ceil(k_tiles * slice_col_par, iters); + if (col_first <= k_tiles * (slice_col_par + 1)) { + int col_off = col_first - k_tiles * slice_col_par; + slice_count = div_ceil(k_tiles - col_off, iters); + if (col_off > 0) slice_count++; + int delta_first = iters * blockIdx.x - col_first; + if (delta_first < 0 || (col_off == 0 && delta_first == 0)) + slice_idx = slice_count - 1; + else { + slice_idx = slice_count - 1 - delta_first / iters; + if (col_off > 0) slice_idx--; + } + } + if (slice_col == n_tiles) { + A += 16 * thread_m_blocks * prob_k / 8; + C += 16 * thread_m_blocks * prob_n / 8; + locks += n_tiles; + slice_col = 0; + } + }; + init_slice(); + + // A sizes/strides + + // stride of the A matrix in global memory + int a_gl_stride = prob_k / 8; + // stride of an A matrix tile in shared memory + constexpr int a_sh_stride = 16 * thread_k_blocks / 8; + // delta between subsequent A tiles in global memory + constexpr int a_gl_rd_delta_o = 16 * thread_k_blocks / 8; + // between subsequent accesses within a tile + int a_gl_rd_delta_i = a_gl_stride * (threads / a_gl_rd_delta_o); + // between shared memory writes + constexpr int a_sh_wr_delta = a_sh_stride * (threads / a_gl_rd_delta_o); + // between shared memory tile reads + constexpr int a_sh_rd_delta_o = 2 * ((threads / 32) / (thread_n_blocks / 4)); + // within a shared memory tile + constexpr int a_sh_rd_delta_i = a_sh_stride * 16; + // overall size of a tile + constexpr int a_sh_stage = a_sh_stride * (16 * thread_m_blocks); + // number of shared write iterations for a tile + constexpr int a_sh_wr_iters = div_ceil(a_sh_stage, a_sh_wr_delta); + + // B sizes/strides + int b_gl_stride = 16 * prob_n / (pack_factor * 4); + constexpr int b_sh_stride = ((thread_n_blocks * 16) * 16 / pack_factor) / 4; + constexpr int b_thread_vecs = num_bits == 4 ? 1 : 2; + constexpr int b_sh_stride_threads = b_sh_stride / b_thread_vecs; + + int b_gl_rd_delta_o = b_gl_stride * thread_k_blocks; + int b_gl_rd_delta_i = b_gl_stride * (threads / b_sh_stride_threads); + constexpr int b_sh_wr_delta = threads * b_thread_vecs; + constexpr int b_sh_rd_delta = threads * b_thread_vecs; + constexpr int b_sh_stage = b_sh_stride * thread_k_blocks; + constexpr int b_sh_wr_iters = b_sh_stage / b_sh_wr_delta; + + // Scale sizes/strides without act_order + int s_gl_stride = prob_n / 8; + constexpr int s_sh_stride = 16 * thread_n_blocks / 8; + constexpr int s_tb_groups = + !has_act_order && group_blocks != -1 && group_blocks < thread_k_blocks + ? thread_k_blocks / group_blocks + : 1; + constexpr int s_sh_stage = s_tb_groups * s_sh_stride; + int s_gl_rd_delta = s_gl_stride; + + // Scale size/strides with act_order + constexpr int tb_k = 16 * thread_k_blocks; + constexpr int g_idx_stage = has_act_order ? (tb_k * sizeof(int)) / 16 : 0; + // constexpr int act_s_row_stride = 1; + // int act_s_col_stride = act_s_row_stride * num_groups; + int act_s_col_stride = 1; + int act_s_col_warp_stride = act_s_col_stride * 8; + int tb_n_warps = thread_n_blocks / 4; + int act_s_col_tb_stride = act_s_col_warp_stride * tb_n_warps; + + // Global A read index of current thread. + int a_gl_rd = a_gl_stride * (threadIdx.x / a_gl_rd_delta_o) + + (threadIdx.x % a_gl_rd_delta_o); + a_gl_rd += a_gl_rd_delta_o * slice_row; + // Shared write index of current thread. + int a_sh_wr = a_sh_stride * (threadIdx.x / a_gl_rd_delta_o) + + (threadIdx.x % a_gl_rd_delta_o); + // Shared read index. + int a_sh_rd = + a_sh_stride * ((threadIdx.x % 32) % 16) + (threadIdx.x % 32) / 16; + a_sh_rd += 2 * ((threadIdx.x / 32) / (thread_n_blocks / 4)); + + int b_gl_rd = b_gl_stride * (threadIdx.x / b_sh_stride_threads) + + (threadIdx.x % b_sh_stride_threads) * b_thread_vecs; + b_gl_rd += b_sh_stride * slice_col; + b_gl_rd += b_gl_rd_delta_o * slice_row; + int b_sh_wr = threadIdx.x * b_thread_vecs; + int b_sh_rd = threadIdx.x * b_thread_vecs; + + // For act_order + constexpr int k_iter_size = tb_k / b_sh_wr_iters; + int slice_k_start = tb_k * slice_row; + int slice_k_finish = slice_k_start + tb_k * slice_iters; + int slice_k_start_shared_fetch = slice_k_start; + int slice_n_offset = act_s_col_tb_stride * slice_col; + + // No act_order + int s_gl_rd; + if constexpr (!has_act_order) { + if constexpr (group_blocks == -1) { + s_gl_rd = s_sh_stride * slice_col + threadIdx.x; + } else { + s_gl_rd = s_gl_stride * ((thread_k_blocks * slice_row) / group_blocks) + + s_sh_stride * slice_col + threadIdx.x; + } + } + int s_sh_wr = threadIdx.x; + bool s_sh_wr_pred = threadIdx.x < s_sh_stride; + + // We use a different scale layout for grouped and column-wise quantization as + // we scale a `half2` tile in column-major layout in the former and in + // row-major in the latter case. + int s_sh_rd; + if constexpr (group_blocks != -1) + s_sh_rd = 8 * ((threadIdx.x / 32) % (thread_n_blocks / 4)) + + (threadIdx.x % 32) / 4; + else + s_sh_rd = 8 * ((threadIdx.x / 32) % (thread_n_blocks / 4)) + + (threadIdx.x % 32) % 4; + + // Precompute which thread should not read memory in which iterations; this is + // needed if there are more threads than required for a certain tilesize or + // when the batchsize is not a multiple of 16. + bool a_sh_wr_pred[a_sh_wr_iters]; + #pragma unroll + for (int i = 0; i < a_sh_wr_iters; i++) + a_sh_wr_pred[i] = a_sh_wr_delta * i + a_sh_wr < a_sh_stride * prob_m; + + // To ensure that writing and reading A tiles to/from shared memory, the + // latter in fragment format, is fully bank conflict free, we need to use a + // rather fancy XOR-based layout. The key here is that neither reads nor + // writes of the 16-byte `int4` blocks of 8 consecutive threads involve the + // same shared memory banks. Further, it seems (based on NSight-Compute) that + // each warp must also write a consecutive memory segment? + auto transform_a = [&](int i) { + int row = i / a_gl_rd_delta_o; + return a_gl_rd_delta_o * row + (i % a_gl_rd_delta_o) ^ row; + }; + // Since the computation of this remapping is non-trivial and, due to our main + // loop unrolls, all shared memory accesses are static, we simply precompute + // both transformed reads and writes. + int a_sh_wr_trans[a_sh_wr_iters]; + #pragma unroll + for (int i = 0; i < a_sh_wr_iters; i++) + a_sh_wr_trans[i] = transform_a(a_sh_wr_delta * i + a_sh_wr); + int a_sh_rd_trans[b_sh_wr_iters][thread_m_blocks]; + #pragma unroll + for (int i = 0; i < b_sh_wr_iters; i++) { + #pragma unroll + for (int j = 0; j < thread_m_blocks; j++) + a_sh_rd_trans[i][j] = + transform_a(a_sh_rd_delta_o * i + a_sh_rd_delta_i * j + a_sh_rd); + } + + // Since B-accesses have non-constant stride they have to be computed at + // runtime; we break dependencies between subsequent accesses with a tile by + // maintining multiple pointers (we have enough registers), a tiny + // optimization. + const int4* B_ptr[b_sh_wr_iters]; + #pragma unroll + for (int i = 0; i < b_sh_wr_iters; i++) + B_ptr[i] = B + b_gl_rd_delta_i * i + b_gl_rd; + + extern __shared__ int4 sh[]; + // Shared memory storage for global fetch pipelines. + int4* sh_a = sh; + int4* sh_b = sh_a + (stages * a_sh_stage); + int4* sh_g_idx = sh_b + (stages * b_sh_stage); + int4* sh_s = sh_g_idx + (stages * g_idx_stage); + + // Register storage for double buffer of shared memory reads. + FragA frag_a[2][thread_m_blocks]; + I4 frag_b_quant[2][b_thread_vecs]; + FragC frag_c[thread_m_blocks][4][2]; + FragS frag_s[2][4]; // No act-order + FragS act_frag_s[2][4][4]; // For act-order + + // Zero accumulators. + auto zero_accums = [&]() { + #pragma unroll + for (int i = 0; i < thread_m_blocks * 4 * 2 * 4; i++) + reinterpret_cast(frag_c)[i] = 0; + }; + + int sh_first_group_id = -1; + int sh_num_groups = -1; + constexpr int sh_max_num_groups = 32; + + auto fetch_scales_to_shared = [&](bool is_async, int first_group_id, + int last_group_id) { + sh_first_group_id = first_group_id; + sh_num_groups = last_group_id - first_group_id + 1; + + if (sh_num_groups < sh_max_num_groups) { + sh_num_groups = sh_max_num_groups; + } + + if (sh_first_group_id + sh_num_groups > num_groups) { + sh_num_groups = num_groups - sh_first_group_id; + } + + int row_offset = first_group_id * s_gl_stride; + + if (is_async) { + for (int i = 0; i < sh_num_groups; i++) { + if (threadIdx.x < s_sh_stride) { + cp_async4_pred(&sh_s[(i * s_sh_stride) + threadIdx.x], + &scales_ptr[row_offset + (i * s_gl_stride) + + slice_n_offset + threadIdx.x]); + } + } + } else { + for (int i = 0; i < sh_num_groups; i++) { + if (threadIdx.x < s_sh_stride) { + sh_s[(i * s_sh_stride) + threadIdx.x] = + scales_ptr[row_offset + (i * s_gl_stride) + slice_n_offset + + threadIdx.x]; + } + } + } + }; + // Asynchronously fetch the next A, B and s tile from global to the next + // shared memory pipeline location. + auto fetch_to_shared = [&](int pipe, int a_off, bool pred = true) { + if (pred) { + int4* sh_a_stage = sh_a + a_sh_stage * pipe; + #pragma unroll + for (int i = 0; i < a_sh_wr_iters; i++) { + cp_async4_pred( + &sh_a_stage[a_sh_wr_trans[i]], + &A[a_gl_rd_delta_i * i + a_gl_rd + a_gl_rd_delta_o * a_off], + a_sh_wr_pred[i]); + } + int4* sh_b_stage = sh_b + b_sh_stage * pipe; + #pragma unroll + for (int i = 0; i < b_sh_wr_iters; i++) { + #pragma unroll + for (int j = 0; j < b_thread_vecs; j++) { + cp_async4(&sh_b_stage[b_sh_wr_delta * i + b_sh_wr + j], B_ptr[i] + j); + } + + B_ptr[i] += b_gl_rd_delta_o; + } + + if constexpr (has_act_order) { + // Fetch g_idx thread-block portion + int full_pipe = a_off; + int cur_k = slice_k_start_shared_fetch + tb_k * full_pipe; + if (cur_k < prob_k && cur_k < slice_k_finish) { + int4* sh_g_idx_stage = sh_g_idx + g_idx_stage * pipe; + + int4 const* cur_g_idx_stage_ptr = + reinterpret_cast(&g_idx[cur_k]); + + if (threadIdx.x < g_idx_stage) { + cp_async4_pred(&sh_g_idx_stage[threadIdx.x], + &cur_g_idx_stage_ptr[threadIdx.x]); + } + } + } else { + if constexpr (group_blocks != -1) { + int4* sh_s_stage = sh_s + s_sh_stage * pipe; + + if constexpr (group_blocks >= thread_k_blocks) { + // Only fetch scales if this tile starts a new group + if (pipe % (group_blocks / thread_k_blocks) == 0) { + if (s_sh_wr_pred) { + cp_async4(&sh_s_stage[s_sh_wr], &scales_ptr[s_gl_rd]); + } + s_gl_rd += s_gl_rd_delta; + } + } else { + for (int i = 0; i < s_tb_groups; i++) { + if (s_sh_wr_pred) { + cp_async4(&sh_s_stage[i * s_sh_stride + s_sh_wr], + &scales_ptr[s_gl_rd]); + } + s_gl_rd += s_gl_rd_delta; + } + } + } + } + } + // Insert a fence even when we are winding down the pipeline to ensure that + // waiting is also correct at this point. + cp_async_fence(); + }; + + // Wait until the next thread tile has been loaded to shared memory. + auto wait_for_stage = [&]() { + // We only have `stages - 2` active fetches since we are double buffering + // and can only issue the next fetch when it is guaranteed that the previous + // shared memory load is fully complete (as it may otherwise be + // overwritten). + cp_async_wait(); + __syncthreads(); + }; + + // Load the next sub-tile from the current location in the shared memory pipe + // into the current register buffer. + auto fetch_to_registers = [&](int k, int pipe) { + int4* sh_a_stage = sh_a + a_sh_stage * pipe; + #pragma unroll + for (int i = 0; i < thread_m_blocks; i++) + ldsm4(frag_a[k % 2][i], + &sh_a_stage[a_sh_rd_trans[k % b_sh_wr_iters][i]]); + int4* sh_b_stage = sh_b + b_sh_stage * pipe; + + #pragma unroll + for (int i = 0; i < b_thread_vecs; i++) { + frag_b_quant[k % 2][i] = *reinterpret_cast( + &sh_b_stage[b_sh_rd_delta * (k % b_sh_wr_iters) + b_sh_rd + i]); + } + }; + + bool is_same_group[stages]; + int same_group_id[stages]; + + auto init_same_group = [&](int pipe) { + if constexpr (!has_act_order) { + is_same_group[pipe] = false; + same_group_id[pipe] = 0; + return; + } + + int4* sh_g_idx_stage = sh_g_idx + g_idx_stage * pipe; + int* sh_g_idx_int_ptr = reinterpret_cast(sh_g_idx_stage); + + int group_id_1 = sh_g_idx_int_ptr[0]; + int group_id_2 = sh_g_idx_int_ptr[tb_k - 1]; + + is_same_group[pipe] = group_id_1 == group_id_2; + same_group_id[pipe] = group_id_1; + }; + + auto fetch_scales_to_registers = [&](int k, int full_pipe) { + int pipe = full_pipe % stages; + + if constexpr (!has_act_order) { + // No act-order case + if constexpr (group_blocks != -1) { + if constexpr (group_blocks >= thread_k_blocks) { + int4* sh_s_stage = + sh_s + s_sh_stage * ((group_blocks / thread_k_blocks) * + (pipe / (group_blocks / thread_k_blocks))); + reinterpret_cast(&frag_s[k % 2])[0] = sh_s_stage[s_sh_rd]; + } else { + int warp_id = threadIdx.x / 32; + int n_warps = thread_n_blocks / 4; + + int warp_row = warp_id / n_warps; + + int cur_k = warp_row * 16; + cur_k += k_iter_size * (k % b_sh_wr_iters); + + int k_blocks = cur_k / 16; + int cur_group_id = k_blocks / group_blocks; + + int4* sh_s_stage = sh_s + s_sh_stage * pipe; + + reinterpret_cast(&frag_s[k % 2])[0] = + sh_s_stage[s_sh_rd + cur_group_id * s_sh_stride]; + } + } + + return; + } + + // Act-order case + + // Determine K of the "current" thread-block + int cur_k = slice_k_start + tb_k * full_pipe; + if (cur_k >= prob_k || cur_k >= slice_k_finish) { + return; + } + + // Reset (to current thread-block) since we read g_idx portion from the + // shared memory + cur_k = 0; + + // Progress to current iteration + cur_k += k_iter_size * (k % b_sh_wr_iters); + + // Determine "position" inside the thread-block (based on warp and + // thread-id) + int warp_id = threadIdx.x / 32; + int n_warps = + thread_n_blocks / 4; // Each warp processes 4 16-size tiles over N + + int warp_row = warp_id / n_warps; + int warp_col = warp_id % n_warps; + + cur_k += warp_row * 16; + + int th_id = threadIdx.x % 32; + cur_k += (th_id % 4) * 2; // Due to tensor-core layout for fp16 B matrix + + int s_col_shift = + /*slice_n_offset +*/ (act_s_col_warp_stride * warp_col) + + (th_id / 4) * act_s_col_stride; + + if (is_same_group[pipe]) { + if (k % 2 == 0) { + *(reinterpret_cast(&(act_frag_s[k % 2][0][0]))) = + sh_s[(same_group_id[pipe] - sh_first_group_id) * s_sh_stride + + s_col_shift]; + } else { + *(reinterpret_cast(&(act_frag_s[k % 2][0][0]))) = + *(reinterpret_cast(&(act_frag_s[(k - 1) % 2][0][0]))); + } + + for (int i = 1; i < 4; i++) { + *(reinterpret_cast(&(act_frag_s[k % 2][i][0]))) = + *(reinterpret_cast(&(act_frag_s[k % 2][0][0]))); + } + return; + } + + int4* sh_g_idx_stage = sh_g_idx + g_idx_stage * pipe; + int* sh_g_idx_int_ptr = reinterpret_cast(sh_g_idx_stage); + + constexpr int k_frag_offsets[4] = {0, 1, 8, + 9}; // Tensor core offsets per thread + + #pragma unroll + for (int i = 0; i < 4; i++) { + int actual_k = cur_k + k_frag_offsets[i]; + + int group_id = sh_g_idx_int_ptr[actual_k]; + int rel_group_id = group_id - sh_first_group_id; + + *(reinterpret_cast(&(act_frag_s[k % 2][i][0]))) = + sh_s[rel_group_id * s_sh_stride + s_col_shift]; + } + }; + + // Execute the actual tensor core matmul of a sub-tile. + auto matmul = [&](int k) { + // We have the m dimension as the inner loop in order to encourage overlapping + // dequantization and matmul operations. + #pragma unroll + for (int j = 0; j < 4; j++) { + FragB frag_b0; + FragB frag_b1; + if constexpr (num_bits == 4) { + int b_quant = frag_b_quant[k % 2][0][j]; + int b_quant_shift = b_quant >> 8; + + frag_b0 = dequant_4bit(b_quant); + frag_b1 = dequant_4bit(b_quant_shift); + + } else { + int* frag_b_quant_ptr = reinterpret_cast(frag_b_quant[k % 2]); + int b_quant_0 = frag_b_quant_ptr[j * 2 + 0]; + int b_quant_1 = frag_b_quant_ptr[j * 2 + 1]; + + frag_b0 = dequant_8bit(b_quant_0); + frag_b1 = dequant_8bit(b_quant_1); + } + + // Apply scale to frag_b0 + if constexpr (has_act_order) { + scale4(frag_b0, act_frag_s[k % 2][0][j], + act_frag_s[k % 2][1][j], act_frag_s[k % 2][2][j], + act_frag_s[k % 2][3][j], 0); + } else { + if constexpr (group_blocks != -1) { + scale(frag_b0, frag_s[k % 2][j], 0); + } + } + + // Apply scale to frag_b1 + if constexpr (has_act_order) { + scale4(frag_b1, act_frag_s[k % 2][0][j], + act_frag_s[k % 2][1][j], act_frag_s[k % 2][2][j], + act_frag_s[k % 2][3][j], 1); + + } else { + if constexpr (group_blocks != -1) { + scale(frag_b1, frag_s[k % 2][j], 1); + } + } + + #pragma unroll + for (int i = 0; i < thread_m_blocks; i++) { + mma(frag_a[k % 2][i], frag_b0, frag_c[i][j][0]); + mma(frag_a[k % 2][i], frag_b1, frag_c[i][j][1]); + } + } + }; + + // Since we slice across the k dimension of a tile in order to increase the + // number of warps while keeping the n dimension of a tile reasonable, we have + // multiple warps that accumulate their partial sums of the same output + // location; which we have to reduce over in the end. We do in shared memory. + auto thread_block_reduce = [&]() { + constexpr int red_off = threads / b_sh_stride_threads / 2; + if (red_off >= 1) { + int red_idx = threadIdx.x / b_sh_stride_threads; + constexpr int red_sh_stride = b_sh_stride_threads * 4 * 2; + constexpr int red_sh_delta = b_sh_stride_threads; + int red_sh_rd = red_sh_stride * (threadIdx.x / b_sh_stride_threads) + + (threadIdx.x % b_sh_stride_threads); + + // Parallel logarithmic shared memory reduction. We make sure to avoid any + // unnecessary read or write iterations, e.g., for two warps we write only + // once by warp 1 and read only once by warp 0. + + #pragma unroll + for (int m_block = 0; m_block < thread_m_blocks; m_block++) { + #pragma unroll + for (int i = red_off; i > 0; i /= 2) { + if (i <= red_idx && red_idx < 2 * i) { + #pragma unroll + for (int j = 0; j < 4 * 2; j++) { + int red_sh_wr = + red_sh_delta * j + (red_sh_rd - red_sh_stride * i); + if (i < red_off) { + float* c_rd = + reinterpret_cast(&sh[red_sh_delta * j + red_sh_rd]); + float* c_wr = reinterpret_cast(&sh[red_sh_wr]); + #pragma unroll + for (int k = 0; k < 4; k++) + reinterpret_cast(frag_c)[4 * 2 * m_block + j][k] += + c_rd[k] + c_wr[k]; + } + sh[red_sh_wr] = + reinterpret_cast(&frag_c)[4 * 2 * m_block + j]; + } + } + __syncthreads(); + } + if (red_idx == 0) { + #pragma unroll + for (int i = 0; i < 4 * 2; i++) { + float* c_rd = + reinterpret_cast(&sh[red_sh_delta * i + red_sh_rd]); + #pragma unroll + for (int j = 0; j < 4; j++) + reinterpret_cast(frag_c)[4 * 2 * m_block + i][j] += + c_rd[j]; + } + } + __syncthreads(); + } + } + }; + + // Since multiple threadblocks may process parts of the same column slice, we + // finally have to globally reduce over the results. As the striped + // partitioning minimizes the number of such reductions and our outputs are + // usually rather small, we perform this reduction serially in L2 cache. + auto global_reduce = [&](bool first = false, bool last = false) { + // We are very careful here to reduce directly in the output buffer to + // maximize L2 cache utilization in this step. To do this, we write out + // results in FP16 (but still reduce with FP32 compute). + constexpr int active_threads = 32 * thread_n_blocks / 4; + if (threadIdx.x < active_threads) { + int c_gl_stride = prob_n / 8; + int c_gl_wr_delta_o = 8 * c_gl_stride; + int c_gl_wr_delta_i = 4 * (active_threads / 32); + int c_gl_wr = c_gl_stride * ((threadIdx.x % 32) / 4) + + 4 * (threadIdx.x / 32) + threadIdx.x % 4; + c_gl_wr += (2 * thread_n_blocks) * slice_col; + constexpr int c_sh_wr_delta = active_threads; + int c_sh_wr = threadIdx.x; + + int row = (threadIdx.x % 32) / 4; + + if (!first) { + // Interestingly, doing direct global accesses here really seems to mess up + // the compiler and lead to slowdowns, hence we also use async-copies even + // though these fetches are not actually asynchronous. + #pragma unroll + for (int i = 0; i < thread_m_blocks * 4; i++) { + cp_async4_pred( + &sh[c_sh_wr + c_sh_wr_delta * i], + &C[c_gl_wr + c_gl_wr_delta_o * (i / 2) + + c_gl_wr_delta_i * (i % 2)], + i < (thread_m_blocks - 1) * 4 || 8 * (i / 2) + row < prob_m); + } + cp_async_fence(); + cp_async_wait<0>(); + } + + #pragma unroll + for (int i = 0; i < thread_m_blocks * 4; i++) { + if (i < (thread_m_blocks - 1) * 4 || 8 * (i / 2) + row < prob_m) { + if (!first) { + int4 c_red = sh[c_sh_wr + i * c_sh_wr_delta]; + #pragma unroll + for (int j = 0; j < 2 * 4; j++) { + reinterpret_cast( + &frag_c)[4 * 2 * 4 * (i / 4) + 4 * j + (i % 4)] += + Dtype::num2float(reinterpret_cast(&c_red)[j]); + } + } + if (!last) { + int4 c; + #pragma unroll + for (int j = 0; j < 2 * 4; j++) { + reinterpret_cast(&c)[j] = + Dtype::float2num(reinterpret_cast( + &frag_c)[4 * 2 * 4 * (i / 4) + 4 * j + (i % 4)]); + } + C[c_gl_wr + c_gl_wr_delta_o * (i / 2) + c_gl_wr_delta_i * (i % 2)] = + c; + } + } + } + } + }; + + // Write out the reduce final result in the correct layout. We only actually + // reshuffle matrix fragments in this step, the reduction above is performed + // in fragment layout. + auto write_result = [&]() { + int c_gl_stride = prob_n / 8; + constexpr int c_sh_stride = 2 * thread_n_blocks + 1; + int c_gl_wr_delta = c_gl_stride * (threads / (2 * thread_n_blocks)); + constexpr int c_sh_rd_delta = + c_sh_stride * (threads / (2 * thread_n_blocks)); + + int c_gl_wr = c_gl_stride * (threadIdx.x / (2 * thread_n_blocks)) + + (threadIdx.x % (2 * thread_n_blocks)); + c_gl_wr += (2 * thread_n_blocks) * slice_col; + int c_sh_wr = + (4 * c_sh_stride) * ((threadIdx.x % 32) / 4) + (threadIdx.x % 32) % 4; + c_sh_wr += 32 * (threadIdx.x / 32); + int c_sh_rd = c_sh_stride * (threadIdx.x / (2 * thread_n_blocks)) + + (threadIdx.x % (2 * thread_n_blocks)); + + int c_gl_wr_end = c_gl_stride * prob_m; + + // We first reorder in shared memory to guarantee the most efficient final + // global write patterns + auto write = [&](int idx, float c0, float c1, FragS& s) { + scalar_t2 res = + Dtype::nums2num2(Dtype::float2num(c0), Dtype::float2num(c1)); + + // For per-column quantization we finally apply the scale here (only for + // 4-bit) + if constexpr (!has_act_order && group_blocks == -1 && num_bits == 4) { + res = __hmul2(res, s[0]); + } + + ((scalar_t2*)sh)[idx] = res; + }; + + if (threadIdx.x / 32 < thread_n_blocks / 4) { + #pragma unroll + for (int i = 0; i < thread_m_blocks; i++) { + #pragma unroll + for (int j = 0; j < 4; j++) { + int wr = c_sh_wr + 8 * j; + write(wr + (4 * c_sh_stride) * 0 + 0, frag_c[i][j][0][0], + frag_c[i][j][0][1], frag_s[j / 2][2 * (j % 2) + 0]); + write(wr + (4 * c_sh_stride) * 8 + 0, frag_c[i][j][0][2], + frag_c[i][j][0][3], frag_s[j / 2][2 * (j % 2) + 0]); + write(wr + (4 * c_sh_stride) * 0 + 4, frag_c[i][j][1][0], + frag_c[i][j][1][1], frag_s[j / 2][2 * (j % 2) + 1]); + write(wr + (4 * c_sh_stride) * 8 + 4, frag_c[i][j][1][2], + frag_c[i][j][1][3], frag_s[j / 2][2 * (j % 2) + 1]); + } + c_sh_wr += 16 * (4 * c_sh_stride); + } + } + __syncthreads(); + + #pragma unroll + for (int i = 0; + i < div_ceil(16 * thread_m_blocks, threads / (2 * thread_n_blocks)); + i++) { + if (c_gl_wr < c_gl_wr_end) { + C[c_gl_wr] = sh[c_sh_rd]; + c_gl_wr += c_gl_wr_delta; + c_sh_rd += c_sh_rd_delta; + } + } + }; + + // Start global fetch and register load pipelines. + auto start_pipes = [&]() { + + #pragma unroll + for (int i = 0; i < stages - 1; i++) { + if (has_act_order && i == 0) { + int last_g_idx = slice_k_start + stages * tb_k * 2; + if (last_g_idx >= prob_k) { + last_g_idx = prob_k - 1; + } + fetch_scales_to_shared(true, g_idx[slice_k_start], g_idx[last_g_idx]); + } + fetch_to_shared(i, i, i < slice_iters); + } + + zero_accums(); + wait_for_stage(); + init_same_group(0); + fetch_to_registers(0, 0); + fetch_scales_to_registers(0, 0); + a_gl_rd += a_gl_rd_delta_o * (stages - 1); + slice_k_start_shared_fetch += tb_k * (stages - 1); + }; + if (slice_iters) { + start_pipes(); + } + + // Main loop. + while (slice_iters) { + // We unroll over both the global fetch and the register load pipeline to + // ensure all shared memory accesses are static. Note that both pipelines + // have even length meaning that the next iteration will always start at + // index 0. + + #pragma unroll + for (int pipe = 0; pipe < stages;) { + #pragma unroll + for (int k = 0; k < b_sh_wr_iters; k++) { + fetch_to_registers(k + 1, pipe % stages); + fetch_scales_to_registers(k + 1, pipe); + if (k == b_sh_wr_iters - 2) { + fetch_to_shared((pipe + stages - 1) % stages, pipe, + slice_iters >= stages); + pipe++; + wait_for_stage(); + init_same_group(pipe % stages); + } + matmul(k); + } + slice_iters--; + if (slice_iters == 0) { + break; + } + } + + a_gl_rd += a_gl_rd_delta_o * stages; + slice_k_start += tb_k * stages; + slice_k_start_shared_fetch += tb_k * stages; + + if constexpr (has_act_order) { + int first_group_id = g_idx[slice_k_start]; + int last_g_idx = slice_k_start + stages * tb_k * 2; + if (last_g_idx >= prob_k) { + last_g_idx = prob_k - 1; + } + int last_group_id = g_idx[last_g_idx]; + if (last_group_id >= sh_first_group_id + sh_num_groups) { + fetch_scales_to_shared(false, first_group_id, last_group_id); + __syncthreads(); + } + } + + // Process results and, if necessary, proceed to the next column slice. + // While this pattern may not be the most readable, other ways of writing + // the loop seemed to noticeably worse performance after compilation. + if (slice_iters == 0) { + cp_async_wait<0>(); + bool last = slice_idx == slice_count - 1; + // For per-column scales, we only fetch them here in the final step before + // write-out + if constexpr (!has_act_order && group_blocks == -1) { + if constexpr (num_bits == 8) { + if (s_sh_wr_pred) { + cp_async4(&sh_s[s_sh_wr], &scales_ptr[s_gl_rd]); + } + cp_async_fence(); + } else { + if (last) { + if (s_sh_wr_pred) { + cp_async4(&sh_s[s_sh_wr], &scales_ptr[s_gl_rd]); + } + cp_async_fence(); + } + } + } + + thread_block_reduce(); + if constexpr (!has_act_order && group_blocks == -1) { + if constexpr (num_bits == 8) { + cp_async_wait<0>(); + __syncthreads(); + if (threadIdx.x / 32 < thread_n_blocks / 4) { + reinterpret_cast(&frag_s)[0] = sh_s[s_sh_rd + 0]; + reinterpret_cast(&frag_s)[1] = sh_s[s_sh_rd + 4]; + } + + } else { + if (last) { + cp_async_wait<0>(); + __syncthreads(); + if (threadIdx.x / 32 < thread_n_blocks / 4) { + reinterpret_cast(&frag_s)[0] = sh_s[s_sh_rd + 0]; + reinterpret_cast(&frag_s)[1] = sh_s[s_sh_rd + 4]; + } + } + } + } + + // For 8-bit channelwise, we apply the scale before the global reduction + // that converts the fp32 results to fp16 (so that we avoid possible + // overflow in fp16) + if constexpr (!has_act_order && group_blocks == -1 && num_bits == 8) { + if (threadIdx.x / 32 < thread_n_blocks / 4) { + #pragma unroll + for (int i = 0; i < thread_m_blocks; i++) { + #pragma unroll + for (int j = 0; j < 4; j++) { + scale_float( + reinterpret_cast(&frag_c[i][j][0][0]), + frag_s[j / 2][2 * (j % 2) + 0]); + scale_float( + reinterpret_cast(&frag_c[i][j][0][2]), + frag_s[j / 2][2 * (j % 2) + 0]); + + scale_float( + reinterpret_cast(&frag_c[i][j][1][0]), + frag_s[j / 2][2 * (j % 2) + 1]); + scale_float( + reinterpret_cast(&frag_c[i][j][1][2]), + frag_s[j / 2][2 * (j % 2) + 1]); + } + } + } + } + + if (slice_count > 1) { // only globally reduce if there is more than one + // block in a slice + barrier_acquire(&locks[slice_col], slice_idx); + global_reduce(slice_idx == 0, last); + barrier_release(&locks[slice_col], last); + } + if (last) // only the last block in a slice actually writes the result + write_result(); + slice_row = 0; + slice_col_par++; + slice_col++; + init_slice(); + if (slice_iters) { + a_gl_rd = a_gl_stride * (threadIdx.x / a_gl_rd_delta_o) + + (threadIdx.x % a_gl_rd_delta_o); + #pragma unroll + for (int i = 0; i < b_sh_wr_iters; i++) + B_ptr[i] += b_sh_stride - b_gl_rd_delta_o * k_tiles; + if (slice_col == 0) { + #pragma unroll + for (int i = 0; i < b_sh_wr_iters; i++) B_ptr[i] -= b_gl_stride; + } + + // Update slice k/n for scales loading + if constexpr (has_act_order) { + slice_k_start = tb_k * slice_row; + slice_k_finish = slice_k_start + tb_k * slice_iters; + slice_k_start_shared_fetch = slice_k_start; + slice_n_offset = act_s_col_tb_stride * slice_col; + + } else { + s_gl_rd = s_sh_stride * slice_col + threadIdx.x; + } + + start_pipes(); + } + } + } +} + + #define __CALL_IF(NUM_BITS, THREAD_M_BLOCKS, THREAD_N_BLOCKS, \ + THREAD_K_BLOCKS, HAS_ACT_ORDER, GROUP_BLOCKS, NUM_THREADS) \ + else if (num_bits == NUM_BITS && thread_m_blocks == THREAD_M_BLOCKS && \ + thread_n_blocks == THREAD_N_BLOCKS && \ + thread_k_blocks == THREAD_K_BLOCKS && \ + has_act_order == HAS_ACT_ORDER && group_blocks == GROUP_BLOCKS && \ + num_threads == NUM_THREADS) { \ + cudaFuncSetAttribute( \ + Marlin, \ + cudaFuncAttributeMaxDynamicSharedMemorySize, max_shared_mem); \ + Marlin<<>>( \ + A_ptr, B_ptr, C_ptr, s_ptr, g_idx_ptr, num_groups, prob_m, prob_n, \ + prob_k, locks); \ + } + +typedef struct { + int thread_k; + int thread_n; + int num_threads; +} thread_config_t; + +typedef struct { + int max_m_blocks; + thread_config_t tb_cfg; +} exec_config_t; + +thread_config_t small_batch_thread_configs[] = { + // Ordered by priority + + // thread_k, thread_n, num_threads + {128, 128, 256}, + {64, 128, 128}, + {128, 64, 128}, +}; + +thread_config_t large_batch_thread_configs[] = { + // Ordered by priority + + // thread_k, thread_n, num_threads + {64, 256, 256}, + {64, 128, 128}, + {128, 64, 128}, + +}; + +int get_scales_cache_size(thread_config_t const& th_config, int prob_m, + int prob_n, int prob_k, int num_bits, int group_size, + bool has_act_order, bool is_k_full) { + bool cache_scales_chunk = has_act_order && !is_k_full; + + int tb_n = th_config.thread_n; + int tb_k = th_config.thread_k; + + // Get max scale groups per thread-block + int tb_groups; + if (group_size == -1) { + tb_groups = 1; + } else if (group_size == 0) { + tb_groups = div_ceil(tb_k, 32); // Worst case is 32 group size + } else { + tb_groups = div_ceil(tb_k, group_size); + } + + if (cache_scales_chunk) { + int load_groups = + tb_groups * pipe_stages * 2; // Chunk size is 2x pipeline over dim K + load_groups = max(load_groups, 32); // We load at least 32 scale groups + return load_groups * tb_n * 2; + + } else { + int tb_scales = tb_groups * tb_n * 2; + + return tb_scales * pipe_stages; + } +} + +bool is_valid_cache_size(thread_config_t const& th_config, int max_m_blocks, + int prob_m, int prob_n, int prob_k, int num_bits, + int scales_cache_size, int max_shared_mem) { + int pack_factor = 32 / num_bits; + + // Get B size + int tb_k = th_config.thread_k; + int tb_n = th_config.thread_n; + + int b_size = (tb_k * tb_n / pack_factor) * 4; + + // Get A size + int m_blocks = div_ceil(prob_m, 16); + int tb_max_m = 16; + + while (true) { + if (m_blocks >= max_m_blocks) { + tb_max_m *= max_m_blocks; + break; + } + + max_m_blocks--; + if (max_m_blocks == 0) { + TORCH_CHECK(false, "Unexpected m_blocks = ", m_blocks); + } + } + + int a_size = (tb_max_m * tb_k) * 2; + + float pipe_size = (a_size + b_size) * pipe_stages; + + TORCH_CHECK(max_shared_mem / 2 > scales_cache_size); // Sanity + + return pipe_size < 0.95f * (max_shared_mem - scales_cache_size); +} + +bool is_valid_config(thread_config_t const& th_config, int max_m_blocks, + int prob_m, int prob_n, int prob_k, int num_bits, + int group_size, bool has_act_order, bool is_k_full, + int max_shared_mem) { + // Sanity + if (th_config.thread_k == -1 || th_config.thread_n == -1 || + th_config.num_threads == -1) { + return false; + } + + // Verify K/N are divisible by thread K/N + if (prob_k % th_config.thread_k != 0 || prob_n % th_config.thread_n != 0) { + return false; + } + + // Verify min for thread K/N + if (th_config.thread_n < min_thread_n || th_config.thread_k < min_thread_k) { + return false; + } + + // num_threads must be at least 128 (= 4 warps) + if (th_config.num_threads < 128) { + return false; + } + + // Determine cache for scales + int scales_cache_size = + get_scales_cache_size(th_config, prob_m, prob_n, prob_k, num_bits, + group_size, has_act_order, is_k_full); + + // Check that pipeline fits into cache + if (!is_valid_cache_size(th_config, max_m_blocks, prob_m, prob_n, prob_k, + num_bits, scales_cache_size, max_shared_mem)) { + return false; + } + + return true; +} + +exec_config_t determine_thread_config(int prob_m, int prob_n, int prob_k, + int num_bits, int group_size, + bool has_act_order, bool is_k_full, + int max_shared_mem) { + int max_m_blocks = 4; + while (max_m_blocks > 0) { + if (prob_m <= 16) { + for (auto th_config : small_batch_thread_configs) { + if (is_valid_config(th_config, max_m_blocks, prob_m, prob_n, prob_k, + num_bits, group_size, has_act_order, is_k_full, + max_shared_mem)) { + return exec_config_t{max_m_blocks, th_config}; + } + } + } else { + for (auto th_config : large_batch_thread_configs) { + if (is_valid_config(th_config, max_m_blocks, prob_m, prob_n, prob_k, + num_bits, group_size, has_act_order, is_k_full, + max_shared_mem)) { + return exec_config_t{max_m_blocks, th_config}; + } + } + } + + max_m_blocks--; // Process less M blocks per invocation to reduce cache + // usage + } + + return exec_config_t{0, {-1, -1, -1}}; +} + + #define CALL_IF(NUM_BITS, N_BLOCKS, K_BLOCKS, NUM_THREADS) \ + __CALL_IF(NUM_BITS, 1, N_BLOCKS, K_BLOCKS, true, 0, NUM_THREADS) \ + __CALL_IF(NUM_BITS, 2, N_BLOCKS, K_BLOCKS, true, 0, NUM_THREADS) \ + __CALL_IF(NUM_BITS, 3, N_BLOCKS, K_BLOCKS, true, 0, NUM_THREADS) \ + __CALL_IF(NUM_BITS, 4, N_BLOCKS, K_BLOCKS, true, 0, NUM_THREADS) \ + \ + __CALL_IF(NUM_BITS, 1, N_BLOCKS, K_BLOCKS, false, -1, NUM_THREADS) \ + __CALL_IF(NUM_BITS, 1, N_BLOCKS, K_BLOCKS, false, 2, NUM_THREADS) \ + __CALL_IF(NUM_BITS, 1, N_BLOCKS, K_BLOCKS, false, 4, NUM_THREADS) \ + __CALL_IF(NUM_BITS, 1, N_BLOCKS, K_BLOCKS, false, 8, NUM_THREADS) \ + \ + __CALL_IF(NUM_BITS, 2, N_BLOCKS, K_BLOCKS, false, -1, NUM_THREADS) \ + __CALL_IF(NUM_BITS, 2, N_BLOCKS, K_BLOCKS, false, 2, NUM_THREADS) \ + __CALL_IF(NUM_BITS, 2, N_BLOCKS, K_BLOCKS, false, 4, NUM_THREADS) \ + __CALL_IF(NUM_BITS, 2, N_BLOCKS, K_BLOCKS, false, 8, NUM_THREADS) \ + \ + __CALL_IF(NUM_BITS, 3, N_BLOCKS, K_BLOCKS, false, -1, NUM_THREADS) \ + __CALL_IF(NUM_BITS, 3, N_BLOCKS, K_BLOCKS, false, 2, NUM_THREADS) \ + __CALL_IF(NUM_BITS, 3, N_BLOCKS, K_BLOCKS, false, 4, NUM_THREADS) \ + __CALL_IF(NUM_BITS, 3, N_BLOCKS, K_BLOCKS, false, 8, NUM_THREADS) \ + \ + __CALL_IF(NUM_BITS, 4, N_BLOCKS, K_BLOCKS, false, -1, NUM_THREADS) \ + __CALL_IF(NUM_BITS, 4, N_BLOCKS, K_BLOCKS, false, 2, NUM_THREADS) \ + __CALL_IF(NUM_BITS, 4, N_BLOCKS, K_BLOCKS, false, 4, NUM_THREADS) \ + __CALL_IF(NUM_BITS, 4, N_BLOCKS, K_BLOCKS, false, 8, NUM_THREADS) + +template +void marlin_mm_f16i4(const void* A, const void* B, void* C, void* s, + void* g_idx, void* perm, void* a_tmp, int prob_m, + int prob_n, int prob_k, void* workspace, int num_bits, + bool has_act_order, bool is_k_full, int num_groups, + int group_size, int dev, cudaStream_t stream, int thread_k, + int thread_n, int sms, int max_par) { + TORCH_CHECK(num_bits == 4 || num_bits == 8, + "num_bits must be 4 or 8. Got = ", num_bits); + TORCH_CHECK(prob_m > 0 && prob_n > 0 && prob_k > 0, "Invalid MNK = [", prob_m, + ", ", prob_n, ", ", prob_k, "]"); + + int tot_m = prob_m; + int tot_m_blocks = div_ceil(tot_m, 16); + int pad = 16 * tot_m_blocks - tot_m; + + if (sms == -1) { + cudaDeviceGetAttribute(&sms, cudaDevAttrMultiProcessorCount, dev); + } + + int max_shared_mem = 0; + cudaDeviceGetAttribute(&max_shared_mem, + cudaDevAttrMaxSharedMemoryPerBlockOptin, dev); + TORCH_CHECK(max_shared_mem > 0); + + // Set thread config + exec_config_t exec_cfg; + if (thread_k != -1 && thread_n != -1) { + // User-defined config + exec_cfg = + exec_config_t{4, thread_config_t{thread_k, thread_n, default_threads}}; + } else { + // Auto config + exec_cfg = + determine_thread_config(prob_m, prob_n, prob_k, num_bits, group_size, + has_act_order, is_k_full, max_shared_mem); + } + + TORCH_CHECK(exec_cfg.max_m_blocks > 0 && + is_valid_config(exec_cfg.tb_cfg, exec_cfg.max_m_blocks, + prob_m, prob_n, prob_k, num_bits, group_size, + has_act_order, is_k_full, max_shared_mem), + "Invalid thread config: max_m_blocks = ", exec_cfg.max_m_blocks, + ", thread_k = ", exec_cfg.tb_cfg.thread_k, + ", thread_n = ", exec_cfg.tb_cfg.thread_n, + ", num_threads = ", exec_cfg.tb_cfg.num_threads, " for MKN = [", + prob_m, ", ", prob_k, ", ", prob_n, "] and num_bits = ", num_bits, + ", group_size = ", group_size, + ", has_act_order = ", has_act_order, ", is_k_full = ", is_k_full, + ", max_shared_mem = ", max_shared_mem); + + int num_threads = exec_cfg.tb_cfg.num_threads; + thread_k = exec_cfg.tb_cfg.thread_k; + thread_n = exec_cfg.tb_cfg.thread_n; + + int thread_k_blocks = thread_k / 16; + int thread_n_blocks = thread_n / 16; + + int blocks = sms; + + TORCH_CHECK(prob_n % thread_n == 0, "prob_n = ", prob_n, + " is not divisible by thread_n = ", thread_n); + TORCH_CHECK(prob_k % thread_k == 0, "prob_k = ", prob_k, + " is not divisible by thread_k = ", thread_k); + + int group_blocks = 0; + if (has_act_order) { + if (is_k_full) { + TORCH_CHECK(group_size != -1); + group_blocks = group_size / 16; + TORCH_CHECK(prob_k % group_blocks == 0, "prob_k = ", prob_k, + " is not divisible by group_blocks = ", group_blocks); + } else { + TORCH_CHECK(group_size == 0); + group_blocks = 0; + } + + } else { + if (group_size == -1) { + group_blocks = -1; + } else { + group_blocks = group_size / 16; + TORCH_CHECK(prob_k % group_blocks == 0, "prob_k = ", prob_k, + " is not divisible by group_blocks = ", group_blocks); + } + } + + const int4* A_ptr = (const int4*)A; + const int4* B_ptr = (const int4*)B; + int4* C_ptr = (int4*)C; + const int4* s_ptr = (const int4*)s; + const int* g_idx_ptr = (const int*)g_idx; + const int* perm_ptr = (const int*)perm; + int4* a_tmp_ptr = (int4*)a_tmp; + + int* locks = (int*)workspace; + + if (has_act_order) { + // Permute A columns + int block_rows = div_ceil(prob_m, blocks); + permute_cols_kernel<<>>( + A_ptr, perm_ptr, a_tmp_ptr, prob_m, prob_k, block_rows); + A_ptr = a_tmp_ptr; + } + + // If we have a full K, then we can run the non-act-order version of Marlin + // (since the weight rows are reordered by increasing group ids, and by having + // a full K, we have full original groups) + if (is_k_full) { + has_act_order = false; + } + + // Main loop + for (int i = 0; i < tot_m_blocks; i += exec_cfg.max_m_blocks) { + int thread_m_blocks = tot_m_blocks - i; + prob_m = tot_m - 16 * i; + int par = 1; + if (thread_m_blocks > exec_cfg.max_m_blocks) { + // Note that parallel > 1 currently only works for inputs without any + // padding + par = (16 * thread_m_blocks - pad) / (16 * exec_cfg.max_m_blocks); + if (par > max_par) par = max_par; + prob_m = (16 * exec_cfg.max_m_blocks) * par; + i += exec_cfg.max_m_blocks * (par - 1); + thread_m_blocks = exec_cfg.max_m_blocks; + } + + + + // Define kernel configurations +#define undefined_error TORCH_CHECK(false, "Unsupported shapes: MNK = [" + str(prob_m) + ", " + \ + str(prob_n) + ", " + str(prob_k) + "]" + \ + ", has_act_order = " + str(has_act_order) + \ + ", num_groups = " + str(num_groups) + \ + ", group_size = " + str(group_size) + \ + ", thread_m_blocks = " + str(thread_m_blocks) + \ + ", thread_n_blocks = " + str(thread_n_blocks) + \ + ", thread_k_blocks = " + str(thread_k_blocks)); + + + if (num_bits == 4 && num_threads == 256) + { + if (false) { + } + CALL_IF(4, 32, 2, 256) + CALL_IF(4, 16, 4, 256) + CALL_IF(4, 8, 8, 256) + else { + undefined_error + } + } + else if (num_bits == 4 && num_threads == 128) + { + if (false) { + } + CALL_IF(4, 8, 4, 128) + CALL_IF(4, 4, 8, 128) + else { + undefined_error + } + } + else if (num_bits == 8 && num_threads == 256) + { + if (false) { + } + CALL_IF(8, 32, 2, 256) + CALL_IF(8, 16, 4, 256) + CALL_IF(8, 8, 8, 256) + else { + undefined_error + } + } + else if (num_bits == 8 && num_threads == 128) + { + if (false) { + } + CALL_IF(8, 8, 4, 128) + CALL_IF(8, 4, 8, 128) + else { + undefined_error + } + } + else { + undefined_error + } + + A_ptr += 16 * thread_m_blocks * (prob_k / 8) * par; + C_ptr += 16 * thread_m_blocks * (prob_n / 8) * par; + } +} + +} // namespace gptq_marlin + +torch::Tensor gptq_marlin_gemm(torch::Tensor& a, torch::Tensor& b_q_weight, + torch::Tensor& b_scales, torch::Tensor& g_idx, + torch::Tensor& perm, torch::Tensor& workspace, + int64_t num_bits, int64_t size_m, int64_t size_n, + int64_t size_k, bool is_k_full) { + const at::cuda::OptionalCUDAGuard device_guard(device_of(a)); + // Verify num_bits + TORCH_CHECK(num_bits == 4 || num_bits == 8, + "num_bits must be 4 or 8. Got = ", num_bits); + int pack_factor = 32 / num_bits; + + // Verify A + TORCH_CHECK(a.size(0) == size_m, "Shape mismatch: a.size(0) = ", a.size(0), + ", size_m = ", size_m); + TORCH_CHECK(a.size(1) == size_k, "Shape mismatch: a.size(1) = ", a.size(1), + ", size_k = ", size_k); + + // Verify B + TORCH_CHECK(size_k % gptq_marlin::tile_size == 0, "size_k = ", size_k, + " is not divisible by tile_size = ", gptq_marlin::tile_size); + TORCH_CHECK((size_k / gptq_marlin::tile_size) == b_q_weight.size(0), + "Shape mismatch: b_q_weight.size(0) = ", b_q_weight.size(0), + ", size_k = ", size_k, ", tile_size = ", gptq_marlin::tile_size); + TORCH_CHECK(b_q_weight.size(1) % gptq_marlin::tile_size == 0, + "b_q_weight.size(1) = ", b_q_weight.size(1), + " is not divisible by tile_size = ", gptq_marlin::tile_size); + int actual_size_n = + (b_q_weight.size(1) / gptq_marlin::tile_size) * pack_factor; + TORCH_CHECK(size_n == actual_size_n, "size_n = ", size_n, + ", actual_size_n = ", actual_size_n); + + // Verify device and strides + TORCH_CHECK(a.device().is_cuda(), "A is not on GPU"); + TORCH_CHECK(a.is_contiguous(), "A is not contiguous"); + + TORCH_CHECK(b_q_weight.device().is_cuda(), "b_q_weight is not on GPU"); + TORCH_CHECK(b_q_weight.is_contiguous(), "b_q_weight is not contiguous"); + + TORCH_CHECK(b_scales.device().is_cuda(), "b_scales is not on GPU"); + TORCH_CHECK(b_scales.is_contiguous(), "b_scales is not contiguous"); + + TORCH_CHECK(g_idx.device().is_cuda(), "g_idx is not on GPU"); + TORCH_CHECK(g_idx.is_contiguous(), "g_idx is not contiguous"); + + TORCH_CHECK(perm.device().is_cuda(), "perm is not on GPU"); + TORCH_CHECK(perm.is_contiguous(), "perm is not contiguous"); + + // Alloc buffers + auto options = torch::TensorOptions().dtype(a.dtype()).device(a.device()); + torch::Tensor c = torch::empty({size_m, size_n}, options); + torch::Tensor a_tmp = torch::empty({size_m, size_k}, options); + + // thread_k: `k` size of a thread_tile in `weights` (can usually be left as + // auto -1) + int thread_k = -1; + // thread_n: `n` size of a thread_tile in `weights` (can usually be left as + // auto -1) + int thread_n = -1; + // sms: number of SMs to use for the kernel (can usually be left as auto -1) + int sms = -1; + + // Verify g_idx and perm + TORCH_CHECK((g_idx.size(0) == 0 && perm.size(0) == 0) || + (g_idx.size(0) == size_k && perm.size(0) == size_k), + "Unexpected g_idx.size(0) = ", g_idx.size(0), + " and perm.size(0) = ", perm.size(0), + ", where size_k = ", size_k); + + // Detect groupsize and act_order + int num_groups = -1; + int group_size = -1; + bool has_act_order = g_idx.size(0) != 0; + + int b_rank = b_scales.sizes().size(); + TORCH_CHECK(b_rank == 2, "b_scales rank = ", b_rank, " is not 2"); + TORCH_CHECK(b_scales.size(1) == size_n, "b_scales dim 1 = ", b_scales.size(1), + " is not size_n = ", size_n); + num_groups = b_scales.size(0); + + if (has_act_order) { + if (is_k_full) { + TORCH_CHECK(num_groups > 1, "For act_order, num_groups must be > 1"); + TORCH_CHECK(size_k % num_groups == 0, "size_k = ", size_k, + ", is not divisible by num_groups = ", num_groups); + group_size = size_k / num_groups; + } else { + group_size = 0; + } + + } else { + if (num_groups > 1) { + TORCH_CHECK( + size_k % num_groups == 0, "size_k = ", size_k, + ", is not divisible by b_scales.size(0) = ", b_scales.size(0)); + group_size = size_k / num_groups; + } else { + group_size = -1; + } + } + + // Verify workspace size + TORCH_CHECK( + size_n % gptq_marlin::min_thread_n == 0, "size_n = ", size_n, + ", is not divisible by min_thread_n = ", gptq_marlin::min_thread_n); + int min_workspace_size = + (size_n / gptq_marlin::min_thread_n) * gptq_marlin::max_par; + TORCH_CHECK(workspace.numel() >= min_workspace_size, + "workspace.numel = ", workspace.numel(), + " is below min_workspace_size = ", min_workspace_size); + + int dev = a.get_device(); + if (a.scalar_type() == at::ScalarType::Half) { + gptq_marlin::marlin_mm_f16i4( + a.data_ptr(), b_q_weight.data_ptr(), c.data_ptr(), + b_scales.data_ptr(), g_idx.data_ptr(), perm.data_ptr(), + a_tmp.data_ptr(), size_m, size_n, size_k, + workspace.data_ptr(), num_bits, has_act_order, is_k_full, num_groups, + group_size, dev, at::cuda::getCurrentCUDAStream(dev), thread_k, + thread_n, sms, gptq_marlin::max_par); + } else if (a.scalar_type() == at::ScalarType::BFloat16) { + gptq_marlin::marlin_mm_f16i4( + a.data_ptr(), b_q_weight.data_ptr(), + c.data_ptr(), b_scales.data_ptr(), + g_idx.data_ptr(), perm.data_ptr(), a_tmp.data_ptr(), + size_m, size_n, size_k, workspace.data_ptr(), num_bits, has_act_order, + is_k_full, num_groups, group_size, dev, + at::cuda::getCurrentCUDAStream(dev), thread_k, thread_n, sms, + gptq_marlin::max_par); + } else { + TORCH_CHECK(false, "gpt_marlin_gemm only supports bfloat16 and float16"); + } + + return c; +} + +#endif diff --git a/archive/kt-sft/csrc/ktransformers_ext/cuda/gptq_marlin/gptq_marlin.cuh b/archive/kt-sft/csrc/ktransformers_ext/cuda/gptq_marlin/gptq_marlin.cuh new file mode 100644 index 0000000..ccf9cfd --- /dev/null +++ b/archive/kt-sft/csrc/ktransformers_ext/cuda/gptq_marlin/gptq_marlin.cuh @@ -0,0 +1,80 @@ +// Adapted from +// https://github.com/vllm-project/vllm/tree/main/csrc/quantization/gptq_marlin +// Copyrigth 2024 The vLLM team. +// Copyright (c) 2024 by KVCache.AI, All Rights Reserved. +#pragma once + +#include + +#include +#include +#include +#include +#include +#include + +namespace gptq_marlin { + +// 8 warps are a good choice since every SM has 4 schedulers and having more +// than 1 warp per schedule allows some more latency hiding. At the same time, +// we want relatively few warps to have many registers per warp and small tiles. +static constexpr int default_threads = 256; + +static constexpr int pipe_stages = + 4; // 4 pipeline stages fit into shared memory + +static constexpr int min_thread_n = 64; +static constexpr int min_thread_k = 64; + +static constexpr int tile_size = 16; +static constexpr int max_par = 16; + +template +struct Vec { + T elems[n]; + __device__ T& operator[](int i) { return elems[i]; } +}; + +using I4 = Vec; + +constexpr int div_ceil(int a, int b) { return (a + b - 1) / b; } + +#if (defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800) || defined (__HIP_PLATFORM_AMD__) +// No support for async +#else + +__device__ inline void cp_async4_pred(void* smem_ptr, const void* glob_ptr, + bool pred = true) { + const int BYTES = 16; + uint32_t smem = static_cast(__cvta_generic_to_shared(smem_ptr)); + asm volatile( + "{\n" + " .reg .pred p;\n" + " setp.ne.b32 p, %0, 0;\n" + " @p cp.async.cg.shared.global [%1], [%2], %3;\n" + "}\n" ::"r"((int)pred), + "r"(smem), "l"(glob_ptr), "n"(BYTES)); +} + +__device__ inline void cp_async4(void* smem_ptr, const void* glob_ptr) { + const int BYTES = 16; + uint32_t smem = static_cast(__cvta_generic_to_shared(smem_ptr)); + asm volatile( + "{\n" + " cp.async.cg.shared.global [%0], [%1], %2;\n" + "}\n" ::"r"(smem), + "l"(glob_ptr), "n"(BYTES)); +} + +__device__ inline void cp_async_fence() { + asm volatile("cp.async.commit_group;\n" ::); +} + +template +__device__ inline void cp_async_wait() { + asm volatile("cp.async.wait_group %0;\n" ::"n"(n)); +} + +#endif + +} // namespace gptq_marlin diff --git a/archive/kt-sft/csrc/ktransformers_ext/cuda/gptq_marlin/gptq_marlin_dtypes.cuh b/archive/kt-sft/csrc/ktransformers_ext/cuda/gptq_marlin/gptq_marlin_dtypes.cuh new file mode 100644 index 0000000..80f6ea4 --- /dev/null +++ b/archive/kt-sft/csrc/ktransformers_ext/cuda/gptq_marlin/gptq_marlin_dtypes.cuh @@ -0,0 +1,85 @@ +// Adapted from +// https://github.com/vllm-project/vllm/tree/main/csrc/quantization/gptq_marlin +// Copyrigth 2024 The vLLM team. +// Copyright (c) 2024 by KVCache.AI, All Rights Reserved. +#ifndef _data_types_cuh +#define _data_types_cuh +#include "gptq_marlin.cuh" +#include +#include + +#ifdef __HIP_PLATFORM_AMD__ +typedef __hip_bfloat16 nv_bfloat16; +typedef __hip_bfloat162 nv_bfloat162; +#endif + +namespace gptq_marlin { + +template +class ScalarType {}; + +template <> +class ScalarType { + public: + using scalar_t = half; + using scalar_t2 = half2; + + // Matrix fragments for tensor core instructions; their precise layout is + // documented here: + // https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#matrix-fragments-for-mma-m16n8k16-with-floating-point-type + using FragA = Vec; + using FragB = Vec; + using FragC = Vec; + using FragS = Vec; + + static __device__ float inline num2float(const half x) { + return __half2float(x); + } + + static __device__ half2 inline num2num2(const half x) { + return __half2half2(x); + } + + static __device__ half2 inline nums2num2(const half x1, const half x2) { + return __halves2half2(x1, x2); + } + + static __host__ __device__ half inline float2num(const float x) { + return __float2half(x); + } +}; + +template <> +class ScalarType { + public: + using scalar_t = nv_bfloat16; + using scalar_t2 = nv_bfloat162; + + using FragA = Vec; + using FragB = Vec; + using FragC = Vec; + using FragS = Vec; + +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800 + static __device__ float inline num2float(const nv_bfloat16 x) { + return __bfloat162float(x); + } + + static __device__ nv_bfloat162 inline num2num2(const nv_bfloat16 x) { + return __bfloat162bfloat162(x); + } + + static __device__ nv_bfloat162 inline nums2num2(const nv_bfloat16 x1, + const nv_bfloat16 x2) { + return __halves2bfloat162(x1, x2); + } + + static __host__ __device__ nv_bfloat16 inline float2num(const float x) { + return __float2bfloat16(x); + } +#endif +}; + +} // namespace gptq_marlin + +#endif diff --git a/archive/kt-sft/csrc/ktransformers_ext/cuda/gptq_marlin/ops.h b/archive/kt-sft/csrc/ktransformers_ext/cuda/gptq_marlin/ops.h new file mode 100644 index 0000000..47cefea --- /dev/null +++ b/archive/kt-sft/csrc/ktransformers_ext/cuda/gptq_marlin/ops.h @@ -0,0 +1,24 @@ +/** + * @Description : + * @Author : Azure + * @Date : 2024-07-22 09:27:55 + * @Version : 1.0.0 + * @LastEditors : Azure + * @LastEditTime : 2024-07-26 08:35:00 + * @Copyright (c) 2024 by KVCache.AI, All Rights Reserved. +**/ +#pragma once + +#include +#include +#include + +torch::Tensor gptq_marlin_gemm(torch::Tensor& a, torch::Tensor& b_q_weight, + torch::Tensor& b_scales, torch::Tensor& g_idx, + torch::Tensor& perm, torch::Tensor& workspace, + int64_t num_bits, int64_t size_m, int64_t size_n, + int64_t size_k, bool is_k_full); + +// torch::Tensor gptq_marlin_repack(torch::Tensor& b_q_weight, torch::Tensor& perm, +// int64_t size_k, int64_t size_n, +// int64_t num_bits); \ No newline at end of file diff --git a/archive/kt-sft/csrc/ktransformers_ext/cuda/setup.py b/archive/kt-sft/csrc/ktransformers_ext/cuda/setup.py new file mode 100644 index 0000000..156bb0e --- /dev/null +++ b/archive/kt-sft/csrc/ktransformers_ext/cuda/setup.py @@ -0,0 +1,26 @@ + +from setuptools import setup, Extension +from torch.utils import cpp_extension +from torch.utils.cpp_extension import BuildExtension, CUDAExtension +setup( + name='KTransformersOps', + ext_modules=[ + CUDAExtension( + 'KTransformersOps', [ + 'custom_gguf/dequant.cu', + 'binding.cpp', + 'gptq_marlin/gptq_marlin.cu', + # 'gptq_marlin_repack.cu', + ], + extra_compile_args={ + 'cxx': ['-O3'], + 'nvcc': [ + '-O3', + '--use_fast_math', + '-Xcompiler', '-fPIC', + ] + }, + ) + ], + cmdclass={'build_ext': BuildExtension} +) \ No newline at end of file diff --git a/archive/kt-sft/csrc/ktransformers_ext/cuda/test_dequant.py b/archive/kt-sft/csrc/ktransformers_ext/cuda/test_dequant.py new file mode 100644 index 0000000..c39d6c7 --- /dev/null +++ b/archive/kt-sft/csrc/ktransformers_ext/cuda/test_dequant.py @@ -0,0 +1,16 @@ +import os +import sys +sys.path.insert(0,"/home/zbx/ktransformers") +from ktransformers.util.custom_loader import GGUFLoader +import torch + +gguf_loader_1 = GGUFLoader("/mnt/data/model/DeepseekV3-q4km-gguf") +gguf_loader_2 = GGUFLoader("/mnt/data/chenht/model/gguf_for_ktransformers/DeepSeek-V3-bf16/") + +torch.set_default_dtype(torch.bfloat16) + +tensor_1 = gguf_loader_1.load_gguf_tensor("blk.0.attn_kv_a_mqa.weight", "cuda") +tensor_2 = gguf_loader_2.load_gguf_tensor("blk.0.attn_kv_a_mqa.weight", "cuda") + +print(tensor_1[0, -64:]) +print(tensor_2[0, -64:]) \ No newline at end of file diff --git a/archive/kt-sft/csrc/ktransformers_ext/examples/test_attention.py b/archive/kt-sft/csrc/ktransformers_ext/examples/test_attention.py new file mode 100644 index 0000000..5627a0e --- /dev/null +++ b/archive/kt-sft/csrc/ktransformers_ext/examples/test_attention.py @@ -0,0 +1,142 @@ +#!/usr/bin/env python +# coding=utf-8 +""" +Description : +Author : Jianwei Dong +Date : 2024-08-28 10:32:05 +Version : 1.0.0 +LastEditors : chenht2022 +LastEditTime : 2024-08-28 10:32:05 +Copyright (c) 2024 by KVCache.AI, All Rights Reserved. +""" +import os, sys +import time + +sys.path.append(os.path.dirname(__file__) + "/../build") +import cpuinfer_ext +from flash_attn import flash_attn_with_kvcache +import torch + +layer_num = 10 +kv_head_num = 8 +q_head_num = 32 +head_dim = 128 +block_len = 128 +anchor_num = 1 +cache_seqlen = 8192 +cache_seqlens = torch.tensor([cache_seqlen], dtype=torch.int32, device="cpu") +seqlens_zero = torch.zeros((1,), dtype=torch.int32, device="cpu") +anchor_type = cpuinfer_ext.kvcache.AnchorType.DYNAMIC +kv_type = cpuinfer_ext.kvcache.ggml_type.FP16 +retrieval_type = cpuinfer_ext.kvcache.RetrievalType.LAYER +layer_step: int = 1 +token_step: int = 1 +layer_offset: int = 0 +max_thread_num: int = 2 +max_batch_size: int = 1 +max_block_num: int = 512 +CPUInfer = cpuinfer_ext.CPUInfer(max_thread_num) +validation_iter = 100 + +with torch.inference_mode(mode=True): + config = cpuinfer_ext.kvcache.KVCacheConfig( + layer_num, + kv_head_num, + q_head_num, + head_dim, + block_len, + anchor_num, + anchor_type, + kv_type, + retrieval_type, + layer_step, + token_step, + layer_offset, + max_block_num, + max_batch_size, + max_thread_num, + ) + local_kvcache = cpuinfer_ext.kvcache.KVCache(config) + + kvcaches = [] + block_table = ( + torch.arange(max_block_num, dtype=torch.int32, device="cpu") + .contiguous() + .view(1, -1) + ) + + for layer_idx in range(layer_num): + k_cache = torch.randn( + (1, cache_seqlen, kv_head_num, head_dim), dtype=torch.float16, device="cpu" + ).contiguous() + v_cache = torch.randn( + (1, cache_seqlen, kv_head_num, head_dim), dtype=torch.float16, device="cpu" + ).contiguous() + + CPUInfer.submit( + local_kvcache.update_kvcache_fp16( + k_cache.data_ptr(), + v_cache.data_ptr(), + layer_idx, + block_table.data_ptr(), + 1, + max_block_num, + seqlens_zero.data_ptr(), + cache_seqlen, + ) + ) + CPUInfer.sync() + + kvcaches.append((k_cache.to("cuda"), v_cache.to("cuda"))) + + # validation + for i in range(validation_iter): + + k_cache = kvcaches[i % layer_num][0] + v_cache = kvcaches[i % layer_num][1] + input = torch.randn( + (1, 1, q_head_num, head_dim), dtype=torch.float16, device="cpu" + ).contiguous() + output = torch.empty( + (1, 1, q_head_num, head_dim), dtype=torch.float16, device="cpu" + ).contiguous() + + # attn_lse: (bsz, q_len, q_head_num) + attn_lse = torch.empty( + (1, 1, q_head_num), dtype=torch.float32, device="cpu" + ).contiguous() + input = input / 100 + + CPUInfer.submit( + local_kvcache.attn( + input.data_ptr(), + output.data_ptr(), + attn_lse.data_ptr(), + i % layer_num, + 0, + 1, + 1, + max_block_num, + block_table.data_ptr(), + cache_seqlens.data_ptr(), + -1, + -1, + -1, + ) + ) + CPUInfer.sync() + # print("cpuinfer output", output) + + t_output = flash_attn_with_kvcache( + q=input.to("cuda"), + k_cache=k_cache, + v_cache=v_cache, + cache_seqlens=cache_seqlens.to("cuda"), + ) + # print("torch output", t_output) + + diff = torch.mean(torch.abs(output.to("cuda") - t_output)) / torch.mean( + torch.abs(t_output) + ) + print("diff = ", diff) + assert diff < 0.001 diff --git a/archive/kt-sft/csrc/ktransformers_ext/examples/test_linear.py b/archive/kt-sft/csrc/ktransformers_ext/examples/test_linear.py new file mode 100644 index 0000000..7a331db --- /dev/null +++ b/archive/kt-sft/csrc/ktransformers_ext/examples/test_linear.py @@ -0,0 +1,62 @@ +#!/usr/bin/env python +# coding=utf-8 +''' +Description : +Author : chenht2022 +Date : 2024-07-25 10:32:05 +Version : 1.0.0 +LastEditors : chenht2022 +LastEditTime : 2024-08-06 10:36:59 +Copyright (c) 2024 by KVCache.AI, All Rights Reserved. +''' +import os, sys +import time +sys.path.append(os.path.dirname(__file__) + '/../build') +import cpuinfer_ext +import torch + +input_size = 16384 +output_size = 5120 +stride = 32 +group_max_len = 1024 +proj_type = 1 # ggml_type::GGML_TYPE_F16 +hidden_type = 1 # ggml_type::GGML_TYPE_F16 +qlen = 30 +layer_num = 10 +CPUInfer = cpuinfer_ext.CPUInfer(48) +validation_iter = 100 + +with torch.inference_mode(mode=True): + linears = [] + projs = [] + for _ in range(layer_num): + proj = torch.randn((output_size, input_size), dtype=torch.float16, device = "cuda").to("cpu").contiguous() + config = cpuinfer_ext.linear.LinearConfig(input_size, output_size, stride, group_max_len, proj.data_ptr(), proj_type, hidden_type) + linear = cpuinfer_ext.linear.Linear(config) + projs.append(proj) + linears.append(linear) + + # validation + for i in range(validation_iter): + linear = linears[i % layer_num] + input = torch.randn((qlen, input_size), dtype=torch.float16).contiguous() + output = torch.empty((qlen, output_size), dtype=torch.float16).contiguous() + input = input / 100 + + CPUInfer.submit( + linear.forward( + qlen, + input.data_ptr(), + output.data_ptr() + ) + ) + CPUInfer.sync() + # print('cpuinfer output', output) + + proj = projs[i%layer_num] + t_output = torch.mm(input, proj.t()) + # print('torch output', t_output) + + diff = torch.mean(torch.abs(output - t_output)) / torch.mean(torch.abs(t_output)) + print('diff = ', diff) + assert(diff < 0.001) diff --git a/archive/kt-sft/csrc/ktransformers_ext/examples/test_mlp.py b/archive/kt-sft/csrc/ktransformers_ext/examples/test_mlp.py new file mode 100644 index 0000000..9805e72 --- /dev/null +++ b/archive/kt-sft/csrc/ktransformers_ext/examples/test_mlp.py @@ -0,0 +1,82 @@ +#!/usr/bin/env python +# coding=utf-8 +''' +Description : +Author : chenht2022 +Date : 2024-07-25 10:32:05 +Version : 1.0.0 +LastEditors : chenht2022 +LastEditTime : 2024-08-06 10:37:28 +Copyright (c) 2024 by KVCache.AI, All Rights Reserved. +''' +import os, sys +import time +sys.path.append(os.path.dirname(__file__) + '/../build') +import cpuinfer_ext +import torch + +hidden_size = 5120 +intermediate_size = 3072 +stride = 32 +group_max_len = 1024 +gate_type = 1 # ggml_type::GGML_TYPE_F16 +up_type = 1 # ggml_type::GGML_TYPE_F16 +down_type = 1 # ggml_type::GGML_TYPE_F16 +hidden_type = 1 # ggml_type::GGML_TYPE_F16 +qlen = 30 +layer_num = 10 +CPUInfer = cpuinfer_ext.CPUInfer(48) +validation_iter = 100 + +def act_fn(x): + return x / (1.0 + torch.exp(-x)) + +def mlp_torch(input, gate_proj, up_proj, down_proj): + gate_buf = torch.mm(input, gate_proj.t()) + up_buf = torch.mm(input, up_proj.t()) + intermediate = act_fn(gate_buf) * up_buf + ret = torch.mm(intermediate, down_proj.t()) + return ret + +with torch.inference_mode(mode=True): + mlps = [] + gate_projs = [] + up_projs = [] + down_projs = [] + for _ in range(layer_num): + gate_proj = torch.randn((intermediate_size, hidden_size), dtype=torch.float16, device = "cuda").to("cpu").contiguous() + up_proj = torch.randn((intermediate_size, hidden_size), dtype=torch.float16, device = "cuda").to("cpu").contiguous() + down_proj = torch.randn((hidden_size, intermediate_size), dtype=torch.float16, device = "cuda").to("cpu").contiguous() + config = cpuinfer_ext.mlp.MLPConfig(hidden_size, intermediate_size, stride, group_max_len, gate_proj.data_ptr(), up_proj.data_ptr(), down_proj.data_ptr(), gate_type, up_type, down_type, hidden_type) + mlp = cpuinfer_ext.mlp.MLP(config) + gate_projs.append(gate_proj) + up_projs.append(up_proj) + down_projs.append(down_proj) + mlps.append(mlp) + + # validation + for i in range(validation_iter): + mlp = mlps[i % layer_num] + input = torch.randn((qlen, hidden_size), dtype=torch.float16).contiguous() + output = torch.empty((qlen, hidden_size), dtype=torch.float16).contiguous() + input = input / 100 + + CPUInfer.submit( + mlp.forward( + qlen, + input.data_ptr(), + output.data_ptr() + ) + ) + CPUInfer.sync() + # print('cpuinfer output', output) + + gate_proj = gate_projs[i%layer_num] + up_proj = up_projs[i%layer_num] + down_proj = down_projs[i%layer_num] + t_output = mlp_torch(input, gate_proj, up_proj, down_proj) + # print('torch output', t_output) + + diff = torch.mean(torch.abs(output - t_output)) / torch.mean(torch.abs(t_output)) + print('diff = ', diff) + assert(diff < 0.001) diff --git a/archive/kt-sft/csrc/ktransformers_ext/examples/test_moe.py b/archive/kt-sft/csrc/ktransformers_ext/examples/test_moe.py new file mode 100644 index 0000000..3fa4dbd --- /dev/null +++ b/archive/kt-sft/csrc/ktransformers_ext/examples/test_moe.py @@ -0,0 +1,121 @@ +#!/usr/bin/env python +# coding=utf-8 +''' +Description : +Author : chenht2022 +Date : 2024-07-25 10:32:05 +Version : 1.0.0 +LastEditors : chenht2022 +LastEditTime : 2024-08-06 10:38:05 +Copyright (c) 2024 by KVCache.AI, All Rights Reserved. +''' +import os, sys +import time +sys.path.append(os.path.dirname(__file__) + '/../build') +import cpuinfer_ext +import torch + +expert_num = 160 +hidden_size = 5120 +intermediate_size = 1536 +stride = 32 +group_min_len = 10 +group_max_len = 1024 +gate_type = 1 # ggml_type::GGML_TYPE_F16 +up_type = 1 # ggml_type::GGML_TYPE_F16 +down_type = 1 # ggml_type::GGML_TYPE_F16 +hidden_type = 1 # ggml_type::GGML_TYPE_F16 +n_routed_experts = 6 +qlen = 30 +layer_num = 10 +CPUInfer = cpuinfer_ext.CPUInfer(48) +validation_iter = 100 + +def act_fn(x): + return x / (1.0 + torch.exp(-x)) + +def mlp_torch(input, gate_proj, up_proj, down_proj): + gate_buf = torch.mm(input, gate_proj.t()) + up_buf = torch.mm(input, up_proj.t()) + intermediate = act_fn(gate_buf) * up_buf + ret = torch.mm(intermediate, down_proj.t()) + return ret + +def moe_torch(input, expert_ids, weights, gate_proj, up_proj, down_proj): + cnts = expert_ids.new_zeros((expert_ids.shape[0], expert_num)) + cnts.scatter_(1, expert_ids, 1) + tokens_per_expert = cnts.sum(dim=0) + idxs = expert_ids.view(-1).argsort() + sorted_tokens = input[idxs // expert_ids.shape[1]] + + outputs = [] + start_idx = 0 + for i, num_tokens in enumerate(tokens_per_expert): + end_idx = start_idx + num_tokens + if num_tokens == 0: + continue + tokens_for_this_expert = sorted_tokens[start_idx:end_idx] + expert_out = mlp_torch(tokens_for_this_expert, gate_proj[i], up_proj[i], down_proj[i]) + outputs.append(expert_out) + start_idx = end_idx + + outs = torch.cat(outputs, dim=0) if len(outputs) else sorted_tokens.new_empty(0) + + new_x = torch.empty_like(outs) + new_x[idxs] = outs + t_output = ( + new_x.view(*expert_ids.shape, -1) + .type(weights.dtype) + .mul_(weights.unsqueeze(dim=-1)) + .sum(dim=1) + .type(new_x.dtype) + ) + return t_output + +with torch.inference_mode(mode=True): + moes = [] + gate_projs = [] + up_projs = [] + down_projs = [] + for _ in range(layer_num): + gate_proj = torch.randn((expert_num, intermediate_size, hidden_size), dtype=torch.float16, device = "cuda").to("cpu").contiguous() + up_proj = torch.randn((expert_num, intermediate_size, hidden_size), dtype=torch.float16, device = "cuda").to("cpu").contiguous() + down_proj = torch.randn((expert_num, hidden_size, intermediate_size), dtype=torch.float16, device = "cuda").to("cpu").contiguous() + config = cpuinfer_ext.moe.MOEConfig(expert_num, n_routed_experts, hidden_size, intermediate_size, stride, group_min_len, group_max_len, gate_proj.data_ptr(), up_proj.data_ptr(), down_proj.data_ptr(), gate_type, up_type, down_type, hidden_type) + moe = cpuinfer_ext.moe.MOE(config) + gate_projs.append(gate_proj) + up_projs.append(up_proj) + down_projs.append(down_proj) + moes.append(moe) + + # validation + for i in range(validation_iter): + expert_ids = torch.stack([torch.randperm(expert_num)[:n_routed_experts] for _ in range(qlen)]).contiguous() + weights = torch.rand((qlen, n_routed_experts), dtype=torch.float32).contiguous() + input = torch.randn((qlen, hidden_size), dtype=torch.float16).contiguous() + output = torch.empty((qlen, hidden_size), dtype=torch.float16).contiguous() + input = input / 100 + + moe = moes[i % layer_num] + CPUInfer.submit( + moe.forward( + qlen, + n_routed_experts, + expert_ids.data_ptr(), + weights.data_ptr(), + input.data_ptr(), + output.data_ptr() + ) + ) + CPUInfer.sync() + # print('cpuinfer output', output) + + gate_proj = gate_projs[i%layer_num] + up_proj = up_projs[i%layer_num] + down_proj = down_projs[i%layer_num] + t_output = moe_torch(input, expert_ids, weights, gate_proj, up_proj, down_proj) + # print('torch output', t_output) + + diff = torch.mean(torch.abs(output - t_output)) / torch.mean(torch.abs(t_output)) + print('diff = ', diff) + assert(diff < 0.001) diff --git a/archive/kt-sft/csrc/ktransformers_ext/examples/test_sft_amx_moe.py b/archive/kt-sft/csrc/ktransformers_ext/examples/test_sft_amx_moe.py new file mode 100644 index 0000000..d59be12 --- /dev/null +++ b/archive/kt-sft/csrc/ktransformers_ext/examples/test_sft_amx_moe.py @@ -0,0 +1,693 @@ +#!/usr/bin/env python +# coding=utf-8 +''' +Description : +Author : chenht2022 +Date : 2024-07-25 10:32:05 +Version : 1.0.0 +LastEditors : chenht2022 +LastEditTime : 2024-08-06 10:38:05 +Copyright (c) 2024 by KVCache.AI, All Rights Reserved. +''' +import os, sys +import time +sys.path.append(os.path.dirname(__file__) + '/../build') +import cpuinfer_ext +import torch +from pathlib import Path +import numpy as np + +expert_num = 10 +hidden_size = 5120 +intermediate_size = 1536 +max_len = 1024 + +n_routed_experts = 2 +qlen = 600 +layer_num = 10 +num_threads = 112 +validation_iter = 1 +LAYER_IDX = 0 +DUMP_DIR = Path(os.getenv("SFT_DEBUG_PATH", "debug")) + +dtype = torch.bfloat16 +gradtype = torch.bfloat16 +# torch.backends.cuda.matmul.allow_tf32 = False + +import shutil +folder_path = "/home/lpl/kt-sft/debug" +if os.path.exists(folder_path): + shutil.rmtree(folder_path) +os.makedirs(folder_path) + +def act_fn(x): + return x / (1.0 + torch.exp(-x)) + +def silu_fwd(x: torch.Tensor) -> torch.Tensor: + return x / (1. + torch.exp(-x)) + +def silu_grad(x: torch.Tensor) -> torch.Tensor: + """SiLU激活函数的梯度""" + sigmoid_x = torch.sigmoid(x) + return sigmoid_x * (1. + x * (1. - sigmoid_x)) + +class SiLU(torch.autograd.Function): + @staticmethod + def forward(ctx, inp): + ctx.save_for_backward(inp) + return silu_fwd(inp) + + @staticmethod + def backward(ctx, grad_out): + (inp,) = ctx.saved_tensors + sig = torch.sigmoid(inp) + return grad_out * (sig + inp * sig * (1. - sig)) + +silu = SiLU.apply # 可求导版本 + +# -------------------- Torch MLP / MoE 参考实现 -------------------- +def mlp_torch(x, gate, up, down, req_grad=False): + g = torch.mm(x, gate.t()) + u = torch.mm(x, up.t()) + if req_grad: + inter = silu_fwd(g) * u + else: + inter = silu_fwd(g) * u + return torch.mm(inter, down.t()) + +def moe_torch(x, eid, w, gate, up, down, req_grad=False): + """eid: [T,k] int64, w: [T,k] float""" + T, k = eid.shape + tok_cnt = torch.zeros(expert_num, dtype=torch.int64) + for e in eid.view(-1): + tok_cnt[e] += 1 + # 打包 token + order = eid.view(-1).argsort() + packed = x[order // k] + + outputs, start = [], 0 + for e in range(expert_num): + num = tok_cnt[e].item() + if not num: + continue + end = start + num + o = mlp_torch(packed[start:end], gate[e], up[e], down[e], req_grad) + outputs.append(o) + start = end + if outputs: + out_all = torch.cat(outputs, 0) + else: + out_all = packed.new_empty(0, hidden_size) + + # 还原顺序并做加权 + out_restore = torch.empty_like(out_all) + out_restore[order] = out_all + out_restore = out_restore.view(T, k, hidden_size) + out = (out_restore * w.unsqueeze(-1)).sum(1) + return out + +def moe_backward_python(x, eid, w, gate, up, down, grad_output, gate_u_cache, up_v_cache): + """ + Python模拟C++的MoE backward计算 - 完全仿照sft_moe.hpp的实现 + 参数: + x: 输入 [T, hidden_size] + eid: expert_ids [T, k] + w: weights [T, k] + gate, up, down: 权重矩阵 + grad_output: 输出梯度 [T, hidden_size] + gate_u_cache, up_v_cache: forward时缓存的中间结果 + 返回: + grad_input: 输入梯度 [T, hidden_size] + """ + T, k = eid.shape + expert_num = gate.shape[0] + hidden_size = gate.shape[2] + intermediate_size = gate.shape[1] + + print("\n========== Python Backward详细对拍 ==========") + print(f"输入形状: T={T}, k={k}, hidden_size={hidden_size}, intermediate_size={intermediate_size}") + print(f"\n--- Python Token 0 ---") + print(f" Expert 0: weight={w[0, 0].item():.6f}") + + # 初始化梯度 + grad_input = torch.zeros_like(x, dtype=torch.float32) + + # print(f"grad_output:{grad_output}") + # print(f"gate_u_cache:{gate_u_cache}") + # print(f"up_v_cache:{up_v_cache}") + + # print(f"grad_output:{grad_output}") + # print(f"gate_u_cache:{gate_u_cache}") + # print(f"up_v_cache:{up_v_cache}") + + # 按C++的方式组织数据：按expert分组 + # 1. 统计每个expert处理的token数量 + expert_token_counts = torch.zeros(expert_num, dtype=torch.int64) + for i in range(T): + for j in range(k): + expert_token_counts[eid[i, j]] += 1 + + # 2. 构建expert到token的映射 + expert_token_indices = [[] for _ in range(expert_num)] + expert_token_positions = [[] for _ in range(expert_num)] + + for i in range(T): + for j in range(k): + expert_id = int(eid[i, j].item()) + expert_token_indices[expert_id].append(i) + expert_token_positions[expert_id].append(j) + + # 3. 为每个expert分配本地存储空间 + max_tokens_per_expert = int(expert_token_counts.max().item()) if expert_token_counts.max() > 0 else 0 + + # 本地存储空间（模拟C++中的m_local_*_ptr_） + local_input = torch.zeros(expert_num, max_tokens_per_expert, hidden_size, dtype=torch.float32) + local_gate_output = torch.zeros(expert_num, max_tokens_per_expert, intermediate_size, dtype=torch.float32) + local_up_output = torch.zeros(expert_num, max_tokens_per_expert, intermediate_size, dtype=torch.float32) + local_down_output_grad = torch.zeros(expert_num, max_tokens_per_expert, hidden_size, dtype=torch.float32) + local_down_input_grad = torch.zeros(expert_num, max_tokens_per_expert, intermediate_size, dtype=torch.float32) + local_gate_output_grad = torch.zeros(expert_num, max_tokens_per_expert, intermediate_size, dtype=torch.float32) + local_up_output_grad = torch.zeros(expert_num, max_tokens_per_expert, intermediate_size, dtype=torch.float32) + local_gate_input_grad = torch.zeros(expert_num, max_tokens_per_expert, hidden_size, dtype=torch.float32) + local_up_input_grad = torch.zeros(expert_num, max_tokens_per_expert, hidden_size, dtype=torch.float32) + + # 4. 复制输入数据和梯度到本地存储 + for expert_id in range(expert_num): + for local_idx, (token_idx, expert_pos) in enumerate(zip(expert_token_indices[expert_id], expert_token_positions[expert_id])): + local_input[expert_id, local_idx] = x[token_idx].to(torch.float32) + local_down_output_grad[expert_id, local_idx] = grad_output[token_idx].to(torch.float32) + + # 5. 重新计算forward的中间结果（模拟C++中的forward计算） + for expert_id in range(expert_num): + num_tokens = expert_token_counts[expert_id] + if num_tokens == 0: + continue + + # 计算gate和up的输出 + local_input_expert = local_input[expert_id, :num_tokens] # [num_tokens, hidden_size] + gate_output = torch.mm(local_input_expert, gate[expert_id].to(torch.float32).t()) # [num_tokens, intermediate_size] + up_output = torch.mm(local_input_expert, up[expert_id].to(torch.float32).t()) # [num_tokens, intermediate_size] + + # 应用激活函数 + gate_output_activated = silu_fwd(gate_output) * up_output + + local_gate_output[expert_id, :num_tokens] = gate_output + local_up_output[expert_id, :num_tokens] = up_output + + for expert_id in range(expert_num): + num_tokens = expert_token_counts[expert_id] + if num_tokens == 0: + continue + # print(f"local_down_output_grad_E_{expert_id}: {local_down_output_grad[expert_id, :num_tokens]}") + # print(f"shape:{local_down_output_grad[expert_id, :num_tokens].shape}") + # torch.save(local_down_output_grad[expert_id, :num_tokens], f"debug/py_layer0_E_End{expert_id}_down_output_grad_.pt") + # torch.save(local_gate_output[expert_id, :num_tokens], f"debug/py_layer0_E_End{expert_id}_gate_output_.pt") + # torch.save(local_up_output[expert_id, :num_tokens], f"debug/py_layer0_E_End{expert_id}_up_output_.pt") + + # 6. 计算down_input_grad（模拟C++中的down_t_bc_计算） + for expert_id in range(expert_num): + num_tokens = expert_token_counts[expert_id] + if num_tokens == 0: + continue + # down_input_grad = down_proj_t @ output_grad + down_input_grad = torch.mm(local_down_output_grad[expert_id, :num_tokens], down[expert_id].to(torch.float32)) # [num_tokens, intermediate_size] + local_down_input_grad[expert_id, :num_tokens] = down_input_grad + + for expert_id in range(expert_num): + num_tokens = expert_token_counts[expert_id] + if num_tokens == 0: + continue + # torch.save(local_gate_output_grad[expert_id, :num_tokens], f"debug/py_layer0_E_End{expert_id}_gate_output_grad_.pt") + # torch.save(local_up_output_grad[expert_id, :num_tokens], f"debug/py_layer0_E_End{expert_id}_up_output_grad_.pt") + torch.save(local_down_output_grad[expert_id, :num_tokens], f"debug/py_layer0_E_End{expert_id}_down_output_grad_.pt") + # torch.save(down[expert_id].to(torch.float32), f"debug/py_layer0_E_End{expert_id}_down_weight_.pt") + torch.save(local_gate_output[expert_id, :num_tokens], f"debug/py_layer0_E_End{expert_id}_gate_output_.pt") + + # 7. 计算gate_output_grad和up_output_grad（模拟C++中的核心计算） + for expert_id in range(expert_num): + num_tokens = expert_token_counts[expert_id] + if num_tokens == 0: + continue + + for local_idx in range(num_tokens): + token_idx = expert_token_indices[expert_id][local_idx] + expert_pos = expert_token_positions[expert_id][local_idx] + weight = w[token_idx, expert_pos].item() + + # 只为第一个token的第一个expert输出调试信息 + should_print = (token_idx == 0 and expert_pos == 0) + + # 获取当前token的中间结果 + gate_u = local_gate_output[expert_id, local_idx] # [intermediate_size] + up_v = local_up_output[expert_id, local_idx] # [intermediate_size] + down_input_grad_token = local_down_input_grad[expert_id, local_idx] # [intermediate_size] + + # 应用weight + down_input_grad_token = down_input_grad_token * weight + + if should_print: + print(f" down_input_grad前5个值: {down_input_grad_token[:5].tolist()}") + + # gate_output_grad = down_input_grad * up_v * silu_grad(gate_u) + gate_output_grad = down_input_grad_token * up_v * silu_grad(gate_u) + + # up_output_grad = down_input_grad * silu_fwd(gate_u) + up_output_grad = down_input_grad_token * silu_fwd(gate_u) + + if should_print: + print(f" gate_output_grad前5个值: {gate_output_grad[:5].tolist()}") + print(f" up_output_grad前5个值: {up_output_grad[:5].tolist()}") + + local_gate_output_grad[expert_id, local_idx] = gate_output_grad + local_up_output_grad[expert_id, local_idx] = up_output_grad + + # 8. 计算gate_input_grad和up_input_grad（模拟C++中的矩阵乘法） + for expert_id in range(expert_num): + num_tokens = expert_token_counts[expert_id] + if num_tokens == 0: + continue + + # gate_input_grad = gate_proj_t @ gate_output_grad + gate_input_grad = torch.mm(local_gate_output_grad[expert_id, :num_tokens], + gate[expert_id].to(torch.float32)) # [num_tokens, hidden_size] + + # up_input_grad = up_proj_t @ up_output_grad + up_input_grad = torch.mm(local_up_output_grad[expert_id, :num_tokens], + up[expert_id].to(torch.float32)) # [num_tokens, hidden_size] + + local_gate_input_grad[expert_id, :num_tokens] = gate_input_grad + local_up_input_grad[expert_id, :num_tokens] = up_input_grad + + # 输出第一个token的调试信息 + if expert_id == 0 and num_tokens > 0: + token_idx = expert_token_indices[expert_id][0] + expert_pos = expert_token_positions[expert_id][0] + if token_idx == 0 and expert_pos == 0: + print(f" gate_input_grad前5个值: {gate_input_grad[0, :5].tolist()}") + print(f" up_input_grad前5个值: {up_input_grad[0, :5].tolist()}") + + # 9. 累加所有expert的梯度到最终输出（模拟C++中的最终累加） + for token_idx in range(T): + token_grad = torch.zeros(hidden_size, dtype=torch.float32) + + for expert_pos in range(k): + expert_id = int(eid[token_idx, expert_pos].item()) + + # 找到这个token在expert_id中的本地索引 + local_idx = expert_token_indices[expert_id].index(token_idx) + + # 累加gate和up的输入梯度 + token_grad += local_gate_input_grad[expert_id, local_idx] + token_grad += local_up_input_grad[expert_id, local_idx] + + grad_input[token_idx] = token_grad + + # 输出第一个token的最终结果 + if token_idx == 0: + print(f" Token 0 最终input_grad前5个值: {token_grad[:5].tolist()}") + + # print(f"gate_input_grad:{gate_input_grad}") + # print(f"up_input_grad:{up_input_grad}") + # print(f"grad_input:{grad_input}") + + return grad_input + +# --------------------------- 主测试 --------------------------- +def test_amx_moe_two_round(): + # ------------ 构造权重 ------------ + gate_proj = torch.randn(expert_num, intermediate_size, hidden_size, + dtype=torch.bfloat16, requires_grad=True).contiguous() + up_proj = torch.randn_like(gate_proj) + down_proj = torch.randn(expert_num, hidden_size, intermediate_size, + dtype=torch.bfloat16, requires_grad=True).contiguous() + + # gate_proj_t = gate_proj.transpose(1, 2).contiguous() # 形状: (E, H, I) + # up_proj_t = up_proj.transpose(1, 2).contiguous() + # down_proj_t = down_proj.transpose(1, 2).contiguous() + + # ------------ SFT-AMX 对象 ------------ + cfg = cpuinfer_ext.sft_moe.SFT_AMX_MOEConfig( + expert_num, n_routed_experts, + hidden_size, intermediate_size, + max_len, + gate_proj.data_ptr(), up_proj.data_ptr(), down_proj.data_ptr() + ) + moe_cpp = cpuinfer_ext.sft_moe.SFT_AMXInt8_MOE(cfg) + + + cpu_infer = cpuinfer_ext.CPUInfer(num_threads) + + cpu_infer.submit(moe_cpp.load_weights()) + cpu_infer.sync() # ATTENTION: DO NOT FORGET sync after load weights + + expert_ids = torch.stack( + [torch.randperm(expert_num)[:n_routed_experts] for _ in range(qlen)]).contiguous() + + weights = torch.rand(qlen, n_routed_experts, dtype=torch.float32).contiguous() + + input_pt = (torch.randn((qlen, hidden_size), dtype=dtype) / 100)\ + .detach().requires_grad_(True).contiguous() + input_cpp = input_pt.detach().clone().requires_grad_(True).contiguous() + + # ------------- forward ------------- + # Torch reference + out_ref = moe_torch(input_pt, expert_ids, weights, + gate_proj, up_proj, down_proj, True) + out_ref.retain_grad() + + # 缓存forward中间结果用于python backward + gate_u_cache = [] + up_v_cache = [] + + # 模拟forward过程并缓存中间结果 + for token_idx in range(qlen): + token_gate_u = [] + token_up_v = [] + for expert_pos in range(n_routed_experts): + expert_id = int(expert_ids[token_idx, expert_pos].item()) + # 计算gate和up的输出 + gate_u = torch.mm(input_pt[token_idx:token_idx+1].to(torch.float32), gate_proj[expert_id].to(torch.float32).t()).squeeze() + up_v = torch.mm(input_pt[token_idx:token_idx+1].to(torch.float32), up_proj[expert_id].to(torch.float32).t()).squeeze() + token_gate_u.append(gate_u) + token_up_v.append(up_v) + gate_u_cache.append(token_gate_u) + up_v_cache.append(token_up_v) + + flop_fwd = 6 * qlen * n_routed_experts * hidden_size * intermediate_size + flop_bwd = 18 * qlen * n_routed_experts * hidden_size * intermediate_size + + # C++ AMX forward + out_cpp = torch.empty_like(out_ref, dtype=dtype).contiguous() + t0 = time.time() + cpu_infer.submit(moe_cpp.forward( + qlen, n_routed_experts, + expert_ids.data_ptr(), weights.data_ptr(), + input_cpp.data_ptr(), out_cpp.data_ptr())) + cpu_infer.sync() + t1 = time.time() + diff_fwd = (out_cpp.to(torch.float32) - out_ref.to(torch.float32)).abs() + print(f"out_cpp.to(torch.float32):{out_cpp.to(torch.float32)}, out_ref.to(torch.float32):{out_ref.to(torch.float32)}") + rel_fwd = diff_fwd.mean() / out_ref.abs().mean() + print(f"Forward diff: {rel_fwd.item():.3e} | time {t1-t0:.4f}s | " + f"TFLOPS {flop_fwd/(t1-t0)/1e12:.2f}") + + + # ------------- backward ------------- + grad_out = torch.randn_like(out_ref, dtype=gradtype).contiguous() + grad_out_cpp = grad_out.clone().contiguous() + grad_in_cpp = torch.zeros_like(input_cpp, dtype=gradtype).contiguous() + + # # Torch backward + for p in (gate_proj, up_proj, down_proj, input_pt): + if p.grad is not None: + p.grad.zero_() + t2 = time.time() + out_ref.backward(grad_out, retain_graph=True) + t3 = time.time() + print(f"PyTorch backward time {t3-t2:.4f}s | " + f"TFLOPS {flop_bwd/(t3-t2)/1e12:.2f}") + + # Python backward（模拟C++逻辑）- 详细版本 + t4_py = time.time() + grad_in_python = moe_backward_python( + input_pt, expert_ids, weights, + gate_proj, up_proj, down_proj, + grad_out.to(torch.float32), gate_u_cache, up_v_cache) + t5_py = time.time() + print(f"Python backward time {t5_py-t4_py:.4f}s | " + f"TFLOPS {flop_bwd/(t5_py-t4_py)/1e12:.2f}") + + # C++ backward + t4 = time.time() + print("Before backward") + cpu_infer.submit(moe_cpp.backward( + qlen, n_routed_experts, + expert_ids.data_ptr(), weights.data_ptr(), input_cpp.data_ptr(), + grad_out_cpp.data_ptr(), + grad_in_cpp.data_ptr())) + cpu_infer.sync() + t5 = time.time() + print("After backward") + print(f"C++ backward time {t5-t4:.4f}s | " + f"TFLOPS {flop_bwd/(t5-t4)/1e12:.2f}") + + # 三种backward结果对比 + gcpp = grad_in_cpp.to(torch.float32) + gref = input_pt.grad.to(torch.float32) if input_pt.grad is not None else torch.zeros_like(input_pt, dtype=torch.float32) + gpy = grad_in_python.to(torch.float32) + + print(f"C++ AMX backward:{gcpp}", '\n', '\n', f"python backward:{gpy}") + + # 对比结果 + rel_bwd_cpp = (gcpp - gref).abs().mean() / gref.abs().mean() + rel_bwd_py = (gpy - gref).abs().mean() / gref.abs().mean() + rel_bwd_cpp_py = (gcpp - gpy).abs().mean() / gpy.abs().mean() + + print(f"Torch vs C++: {rel_bwd_cpp.item():.3e}") + print(f"Torch vs Python: {rel_bwd_py.item():.3e}") + print(f"C++ vs Python: {rel_bwd_cpp_py.item():.3e}") + + # 检查是否对拍成功 + if rel_bwd_cpp_py.item() < 5e-2: + print("✅ C++和Python backward对拍成功!") + else: + print("❌ C++和Python backward对拍失败，存在显著差异") + + + # manual_check(expert_ids) + +def load_bf16(stub, shape): + with open(stub + ".bf16", "rb") as f: + return torch.frombuffer(f.read(), dtype=torch.bfloat16).view(shape).float() +def load_f16(stub, shape): + with open(stub+".f16",'rb') as f: + return torch.frombuffer(f.read(), dtype=torch.float16).view(shape).float() +def load_f32(stub, shape): + with open(stub+".f32",'rb') as f: + return torch.frombuffer(f.read(), dtype=torch.float32).view(shape) +def load_uint8(stub, shape): + with open(stub+".uint8",'rb') as f: + return torch.frombuffer(f.read(), dtype=torch.uint8).view(shape) +def load_int8(stub, shape): + with open(stub+".int8",'rb') as f: + return torch.frombuffer(f.read(), dtype=torch.int8).view(shape) + +# 通用加载函数 +def load_dump_tensor(experts_idx: int, name: str, shape: tuple, Ename: str = "E_Before"): + """ + 根据 experts_idx / name / shape 读取 dump 文件，并返回 torch.Tensor + """ + stub = DUMP_DIR / f"layer{LAYER_IDX}_{Ename}{experts_idx}_{name}" + if stub.with_suffix(".bf16").exists(): + return load_bf16(str(stub), shape) + elif stub.with_suffix(".f16").exists(): + return load_f16(str(stub), shape) + elif stub.with_suffix(".f32").exists(): + return load_f32(str(stub), shape) + elif stub.with_suffix(".uint8").exists(): + return load_uint8(str(stub), shape) + elif stub.with_suffix(".int8").exists(): + return load_int8(str(stub), shape) + else: + raise FileNotFoundError(f"{stub}（bf16/f16/f32/u8/i8 均不存在）") + +def load_bin(path, n, k): + # 从文件读出 n*k 个 float32 + data = np.fromfile(path, dtype=np.float32) + assert data.size == n * k + data = data.reshape(n, k) + return torch.from_numpy(data).to(torch.bfloat16) + +def check_nan(name, shape): + stub1 = DUMP_DIR / f"{name}" + if stub1.with_suffix(".bf16").exists(): + cpp_bef = load_bf16(str(stub1), shape) + elif stub1.with_suffix(".f16").exists(): + cpp_bef = load_f16(str(stub1), shape) + elif stub1.with_suffix(".f32").exists(): + cpp_bef = load_f32(str(stub1), shape) + elif stub1.with_suffix(".int8").exists(): + return load_int8(str(stub1), shape) + else: + print("dump 缺失/未知类型"); return + + print(f"{name}:{cpp_bef}") + print(f" shape : {cpp_bef.shape}") + print(f" dtype : {cpp_bef.dtype}") + + finite_mask = torch.isfinite(cpp_bef) + if finite_mask.any(): + t_finite = cpp_bef[finite_mask] + t_max = t_finite.max().item() + t_min = t_finite.min().item() + print(f" max : {t_max:.6e}") + print(f" min : {t_min:.6e}") + else: + print(" max/min: 所有元素均为 NaN / Inf") + + for nan_name, t in [(f"{name}", cpp_bef)]: + nan_cnt = torch.isnan(t).sum().item() + inf_cnt = torch.isinf(t).sum().item() + if nan_cnt or inf_cnt: + print(f"{name} 含 NaN={nan_cnt}、Inf={inf_cnt}") + else: + print("NO NaN or Inf exist") + +def get_tensor(name, shape) -> torch.Tensor: + stub1 = DUMP_DIR / f"{name}" + if stub1.with_suffix(".bf16").exists(): + cpp_bef = load_bf16(str(stub1), shape) + elif stub1.with_suffix(".f16").exists(): + cpp_bef = load_f16(str(stub1), shape) + elif stub1.with_suffix(".f32").exists(): + cpp_bef = load_f32(str(stub1), shape) + elif stub1.with_suffix(".int8").exists(): + return load_int8(str(stub1), shape) + else: + print("dump 缺失/未知类型"); return + + return cpp_bef + +def check_py_cpp(name1, name2, shape): + print(f"compare {name1} with {name2}, at shape{shape}") + stub1 = DUMP_DIR / f"{name1}" + py_bef = torch.load(f"{stub1}") + if not isinstance(py_bef, torch.Tensor): + print(f"⚠️ {name1} 不是 Tensor，而是 {type(py_bef)}") + return + stub2 = DUMP_DIR / f"{name2}" + if stub2.with_suffix(".bf16").exists(): + cpp_bef = load_bf16(str(stub2), shape) + elif stub2.with_suffix(".f16").exists(): + cpp_bef = load_f16(str(stub2), shape) + elif stub2.with_suffix(".f32").exists(): + cpp_bef = load_f32(str(stub2), shape) + elif stub2.with_suffix(".int8").exists(): + return load_int8(str(stub2), shape) + else: + print(f"dump 缺失/未知类型: {stub2}"); return + + for t in [py_bef]: + nan_cnt = torch.isnan(t).sum().item() + inf_cnt = torch.isinf(t).sum().item() + if nan_cnt or inf_cnt: + print(f"{name1} 含 NaN={nan_cnt}、Inf={inf_cnt}") + else: + print("NO NaN or Inf exist") + for t in [cpp_bef]: + nan_cnt = torch.isnan(t).sum().item() + inf_cnt = torch.isinf(t).sum().item() + if nan_cnt or inf_cnt: + print(f"{name2} 含 NaN={nan_cnt}、Inf={inf_cnt}") + else: + print("NO NaN or Inf exist") + + if py_bef.shape != cpp_bef.shape: + print(f"shape 不一致: py_bef {py_bef.shape}, cpp_bef {cpp_bef.shape}") + else: + # 计算绝对差、相对差、最大差值 + eps = 1e-6 # 防止除以 0 + denominator = torch.abs(py_bef) + eps + rel_diff = torch.abs(py_bef - cpp_bef) / denominator + + # 找出大于 2% 的项 + mask = rel_diff > 0.02 + num_large_diff = mask.sum().item() + total = rel_diff.numel() + + if num_large_diff == 0: + print("✅ 所有元素相对误差都在 2% 范围内") + flat_rel_diff = rel_diff.view(-1) + max_idx = torch.argmax(flat_rel_diff) + max_val = flat_rel_diff[max_idx].item() + + # 还原成多维索引 + max_pos = tuple(torch.unravel_index(max_idx, py_bef.shape)) + + # 获取原始值 + py_val = py_bef[max_pos].item() + cpp_val = cpp_bef[max_pos].item() + + print(f" 最大相对误差 = {max_val:.2%}") + print(f" 最大相对误差位置: {max_pos}, py = {py_val:.6f}, cpp = {cpp_val:.6f}") + else: + print(f"❗ 相对误差 > 2% 的元素数量: {num_large_diff} / {total}") + print(f"{name1}: {py_bef}") + print(f"{name2}: {cpp_bef}") + +# 汇总要查哪些内容 +def manual_check(experts_ids): + expert_token_counts = torch.zeros(expert_num, dtype=torch.int64) + T, k = experts_ids.shape + for i in range(T): + for j in range(k): + expert_token_counts[experts_ids[i, j]] += 1 + for experts_idx in range(expert_num): + # input1 = get_tensor(f"cpp_layer0_E_End{experts_idx}_down_t_ba_", (expert_token_counts[experts_idx], hidden_size)) + # # down_ba_new = get_tensor(f"cpp_layer0_E_End{experts_idx}_down_ba_new_", (expert_token_counts[experts_idx], intermediate_size)) + # weight1 = get_tensor(f"cpp_layer0_E_End{experts_idx}_down_t_bb_", (hidden_size, intermediate_size)) + # output1 = torch.matmul(input1, weight1) + # print(f"input1:{input1}, shape:{input1.shape}") + # # print(f"down_ba_new:{down_ba_new}, shape:{down_ba_new.shape}") + # print(f"weight1:{weight1}, shape:{weight1.shape}") + # print(f"output1:{output1}, shape:{output1.shape}") + + # shape=(expert_token_counts[experts_idx], intermediate_size) + # stub_bc = DUMP_DIR / f"cpp_layer0_E_End{experts_idx}_down_t_bc_" + # if stub_bc.with_suffix(".bf16").exists(): + # output1_5 = load_bf16(str(stub_bc), shape) + # elif stub_bc.with_suffix(".f16").exists(): + # output1_5 = load_f16(str(stub_bc), shape) + # elif stub_bc.with_suffix(".f32").exists(): + # output1_5 = load_f32(str(stub_bc), shape) + # elif stub_bc.with_suffix(".int8").exists(): + # return load_int8(str(stub_bc), shape) + # else: + # print(f"dump 缺失/未知类型: {stub_bc}"); return + # print(f"output1_5:{output1_5}, shape:{output1_5.shape}") + + # torch.set_printoptions(profile="full") + + down_ba_ori = get_tensor(f"cpp_layer0_E_End{experts_idx}_down_ba_ori_", (expert_token_counts[experts_idx], intermediate_size)) + + # with open(f"/home/lpl/kt-sft/debug/cpp_{experts_idx}_down_ba_ori_view.txt", "w") as f: + # f.write(str(down_ba_ori)) + + + down_output_grad = get_tensor(f"cpp_layer0_E_End{experts_idx}_down_output_grad_", (expert_token_counts[experts_idx], hidden_size)) + + # with open(f"/home/lpl/kt-sft/debug/cpp_{experts_idx}_down_t_ba_ori_view.txt", "w") as f: + # f.write(str(down_output_grad)) + + + # input2 = torch.load(f"debug/py_layer0_E_End{experts_idx}_down_output_grad_.pt") + # weight2 = torch.load(f"debug/py_layer0_E_End{experts_idx}_down_weight_.pt") + # output2 = torch.load(f"debug/py_layer0_E_End{experts_idx}_down_input_grad_.pt") + # print(f"input2:{input2}, shape:{input2.shape}") + # print(f"weight2:{weight2}, shape:{weight2.shape}") + # print(f"output2:{output2}, shape:{output2.shape}") + + # down_t_ba_new = load_bin(f'debug/{experts_idx}_down_ba_t_debug3.bin', expert_token_counts[experts_idx], hidden_size) + + # print(f"input3: {down_t_ba_new}, shape: {down_t_ba_new.shape}") + + py_down_t_ba = torch.load(f"debug/py_layer0_E_End{experts_idx}_down_output_grad_.pt") + py_down_ba = torch.load(f"debug/py_layer0_E_End{experts_idx}_gate_output_.pt") + + # with open(f"/home/lpl/kt-sft/debug/py_{experts_idx}_down_t_ba_ori_view.txt", "w") as f: + # f.write(str(py_down_t_ba)) + + # with open(f"/home/lpl/kt-sft/debug/py_{experts_idx}_down_ba_ori_view.txt", "w") as f: + # f.write(str(py_down_ba)) + + print(f"cpp_{experts_idx}_down_ba_ori_:{down_ba_ori}") + print(f"py_{experts_idx}_down_ba_ori_view: {py_down_ba}") + print(f"cpp_{experts_idx}_down_t_ba_ori_view:{down_output_grad}") + print(f"py_{experts_idx}_down_t_ba_ori_view: {py_down_t_ba}") + + # torch.set_printoptions(profile="default") + + +if __name__ == "__main__": + torch.manual_seed(42) + test_amx_moe_two_round() \ No newline at end of file diff --git a/archive/kt-sft/csrc/ktransformers_ext/examples/test_sft_moe.py b/archive/kt-sft/csrc/ktransformers_ext/examples/test_sft_moe.py new file mode 100644 index 0000000..3041e22 --- /dev/null +++ b/archive/kt-sft/csrc/ktransformers_ext/examples/test_sft_moe.py @@ -0,0 +1,756 @@ +#!/usr/bin/env python +# coding=utf-8 +''' +Description : +Author : chenht2022 +Date : 2024-07-25 10:32:05 +Version : 1.0.0 +LastEditors : chenht2022 +LastEditTime : 2024-08-06 10:38:05 +Copyright (c) 2024 by KVCache.AI, All Rights Reserved. +''' +import os, sys +import time +sys.path.append(os.path.dirname(__file__) + '/../build') +import cpuinfer_ext +import torch + +expert_num = 10 +hidden_size = 5120 +intermediate_size = 1536 +stride = 32 +group_min_len = 10 +group_max_len = 1024 +gate_type = 1 # ggml_type::GGML_TYPE_F16 +up_type = 1 # ggml_type::GGML_TYPE_F16 +down_type = 1 # ggml_type::GGML_TYPE_F16 +hidden_type = 1 # ggml_type::GGML_TYPE_F16 +n_routed_experts = 2 +qlen = 30 +layer_num = 10 +CPUInfer = cpuinfer_ext.CPUInfer(48) +validation_iter = 100 + +dtype = torch.float16 +gradtype = torch.bfloat16 + +def act_fn(x): + return x / (1.0 + torch.exp(-x)) + +# 定义SiLU激活函数的可微版本（带梯度） +class SiLU(torch.autograd.Function): + @staticmethod + def forward(ctx, input): + ctx.save_for_backward(input) + return input / (1.0 + torch.exp(-input)) + + @staticmethod + def backward(ctx, grad_output): + input, = ctx.saved_tensors + sigmoid = 1.0 / (1.0 + torch.exp(-input)) + return grad_output * (sigmoid + input * sigmoid * (1 - sigmoid)) + +silu = SiLU.apply + +def mlp_torch(input, gate_proj, up_proj, down_proj, requires_grad=False): + gate_buf = torch.mm(input, gate_proj.t()) + up_buf = torch.mm(input, up_proj.t()) + + # 使用可微的SiLU或者原来的函数，取决于是否需要梯度 + if requires_grad: + intermediate = silu(gate_buf) * up_buf + else: + intermediate = act_fn(gate_buf) * up_buf + + ret = torch.mm(intermediate, down_proj.t()) + return ret + +def moe_torch(input, expert_ids, weights, gate_proj, up_proj, down_proj, requires_grad=False): + cnts = expert_ids.new_zeros((expert_ids.shape[0], expert_num)) + cnts.scatter_(1, expert_ids, 1) + tokens_per_expert = cnts.sum(dim=0) + idxs = expert_ids.view(-1).argsort() + sorted_tokens = input[idxs // expert_ids.shape[1]] + + outputs = [] + start_idx = 0 + for i, num_tokens in enumerate(tokens_per_expert): + end_idx = start_idx + num_tokens + if num_tokens == 0: + continue + tokens_for_this_expert = sorted_tokens[start_idx:end_idx] + expert_out = mlp_torch(tokens_for_this_expert, gate_proj[i], up_proj[i], down_proj[i], requires_grad) + outputs.append(expert_out) + start_idx = end_idx + + outs = torch.cat(outputs, dim=0) if len(outputs) else sorted_tokens.new_empty(0) + + new_x = torch.empty_like(outs) + new_x[idxs] = outs + t_output = ( + new_x.view(*expert_ids.shape, -1) + .type(weights.dtype) + .mul_(weights.unsqueeze(dim=-1)) + .sum(dim=1) + .type(new_x.dtype) + ) + return t_output + +# 前向传播验证 +def test_forward(): + with torch.inference_mode(mode=True): + moes = [] + gate_projs = [] + up_projs = [] + down_projs = [] + for _ in range(layer_num): + gate_proj = torch.randn((expert_num, intermediate_size, hidden_size), dtype=dtype, device = "cuda").to("cpu").contiguous() + up_proj = torch.randn((expert_num, intermediate_size, hidden_size), dtype=dtype, device = "cuda").to("cpu").contiguous() + down_proj = torch.randn((expert_num, hidden_size, intermediate_size), dtype=dtype, device = "cuda").to("cpu").contiguous() + config = cpuinfer_ext.sft_moe.SFT_MOEConfig(expert_num, n_routed_experts, hidden_size, intermediate_size, stride, group_min_len, group_max_len, gate_proj.data_ptr(), up_proj.data_ptr(), down_proj.data_ptr(), gate_type, up_type, down_type, hidden_type, 0) + moe = cpuinfer_ext.sft_moe.SFT_MOE(config) + gate_projs.append(gate_proj) + up_projs.append(up_proj) + down_projs.append(down_proj) + moes.append(moe) + + # validation + for i in range(validation_iter): + expert_ids = torch.stack([torch.randperm(expert_num)[:n_routed_experts] for _ in range(qlen)]).contiguous() + weights = torch.rand((qlen, n_routed_experts), dtype=torch.float32).contiguous() + input = torch.randn((qlen, hidden_size), dtype=dtype).contiguous() + output = torch.empty((qlen, hidden_size), dtype=dtype).contiguous() + input = input / 100 + + moe = moes[i % layer_num] + CPUInfer.submit( + moe.forward( + qlen, + n_routed_experts, + expert_ids.data_ptr(), + weights.data_ptr(), + input.data_ptr(), + output.data_ptr() + ) + ) + CPUInfer.sync() + # print('cpuinfer output', output) + + gate_proj = gate_projs[i%layer_num] + up_proj = up_projs[i%layer_num] + down_proj = down_projs[i%layer_num] + t_output = moe_torch(input, expert_ids, weights, gate_proj, up_proj, down_proj) + # print('torch output', t_output) + + diff = torch.mean(torch.abs(output - t_output)) / torch.mean(torch.abs(t_output)) + print('diff = ', diff) + assert(diff < 0.001) + +# 反向传播验证 +def test_backward(): + # 先测试backward是否能正常调用 + print("\n===== Testing Backward Pass =====") + # 创建一个单层MOE用于测试 + gate_proj = torch.randn((expert_num, intermediate_size, hidden_size), dtype=dtype, requires_grad=True).contiguous() + up_proj = torch.randn((expert_num, intermediate_size, hidden_size), dtype=dtype, requires_grad=True).contiguous() + down_proj = torch.randn((expert_num, hidden_size, intermediate_size), dtype=dtype, requires_grad=True).contiguous() + # 创建MOE实例 + config = cpuinfer_ext.sft_moe.SFT_MOEConfig(expert_num, n_routed_experts, hidden_size, intermediate_size, + stride, group_min_len, group_max_len, + gate_proj.data_ptr(), up_proj.data_ptr(), down_proj.data_ptr(), + gate_type, up_type, down_type, hidden_type) # 使用float16类型(0=GGML_TYPE_F16) + moe = cpuinfer_ext.sft_moe.SFT_MOE(config) + + # 创建输入数据 + expert_ids = torch.stack([torch.randperm(expert_num)[:n_routed_experts] for _ in range(qlen)]).contiguous() + weights = torch.rand((qlen, n_routed_experts), dtype=torch.float32).contiguous() + + # 使用相同的输入进行torch和C++算子的计算 + input = torch.randn((qlen, hidden_size), dtype=dtype, requires_grad=True).contiguous() + input = (input / 100).detach().requires_grad_(True) + input_cpp = input.clone().detach().requires_grad_(True).contiguous() + + # 计算PyTorch参考输出 + t_output = moe_torch(input, expert_ids, weights, gate_proj, up_proj, down_proj, requires_grad=True) + # 确保非叶子张量保留梯度 + t_output.retain_grad() + + # 计算C++算子输出 + output_cpp = torch.empty((qlen, hidden_size), dtype=dtype).contiguous() + + # 前向传播 + forward_start_time = time.time() + CPUInfer.submit( + moe.forward( + qlen, + n_routed_experts, + expert_ids.data_ptr(), + weights.data_ptr(), + input_cpp.data_ptr(), + output_cpp.data_ptr() + ) + ) + CPUInfer.sync() + forward_end_time = time.time() + print(f"C++ forward 耗时: {forward_end_time - forward_start_time:.4f} 秒") + + FLOPs_fwd = 6 * qlen * n_routed_experts * hidden_size * intermediate_size + KT_TFLOPS_fwd = FLOPs_fwd / (forward_end_time - forward_start_time) / 1e12 + + # 验证前向传播结果 + forward_diff = torch.mean(torch.abs(output_cpp - t_output)) / torch.mean(torch.abs(t_output)) + print(f"Forward diff: {forward_diff.item()}") + assert forward_diff < 0.001, f"Forward diff too large: {forward_diff.item()}" + print("✅ Forward test passed!") + + grad_input_cpp = torch.empty_like(input_cpp, dtype=gradtype).contiguous() + grad_output = torch.randn_like(t_output, dtype=gradtype).contiguous() + grad_output_cpp = grad_output.clone() + + print("-- pytorch backward --") + # PyTorch反向传播性能测试 + pytorch_start_time = time.time() + + t_output.backward(grad_output, retain_graph=True) + + pytorch_end_time = time.time() + pytorch_time = (pytorch_end_time - pytorch_start_time) + + print("-- c++ backward --") + # C++反向传播性能测试 + CPUInfer.submit( + moe.backward( + qlen, + n_routed_experts, + expert_ids.data_ptr(), + weights.data_ptr(), + input_cpp.data_ptr(), + grad_output_cpp.data_ptr(), + grad_input_cpp.data_ptr() + ) + ) + CPUInfer.sync() + + cpp_start_time = time.time() + CPUInfer.submit( + moe.backward( + qlen, + n_routed_experts, + expert_ids.data_ptr(), + weights.data_ptr(), + input_cpp.data_ptr(), + grad_output_cpp.data_ptr(), + grad_input_cpp.data_ptr() + ) + ) + CPUInfer.sync() + + cpp_end_time = time.time() + cpp_time = (cpp_end_time - cpp_start_time) + print(f"PyTorch backward 耗时: {pytorch_time:.4f} 秒") + print(f"C++ backward 耗时: {cpp_time:.4f} 秒") + print(f"性能比较: PyTorch/C++ = {pytorch_time/cpp_time:.2f}x") + + + print(f"qlen:{qlen}, n_exp:{n_routed_experts}, hidden:{hidden_size}, inter:{intermediate_size}") + FLOPs_bwd = 18 * qlen * n_routed_experts * hidden_size * intermediate_size + torch_TFLOPS_bwd = FLOPs_bwd / pytorch_time / 1e12 + KT_TFLOPS_bwd = FLOPs_bwd / cpp_time / 1e12 + + print(f"PyTorch backward TFLOPS: {torch_TFLOPS_bwd}") + print(f"KT forward TFLOPS: {KT_TFLOPS_fwd}") + print(f"KT backward TFLOPS: {KT_TFLOPS_bwd}") + + # ================== TFLOPS 统计 ================== + total_flops_fwd = 6 * qlen * n_routed_experts * hidden_size * intermediate_size + total_flops_bwd = 18 * qlen * n_routed_experts * hidden_size * intermediate_size + + tflops_fwd_cpp = total_flops_fwd / (forward_end_time - forward_start_time) / 1e12 + tflops_bwd_cpp = total_flops_bwd / cpp_time / 1e12 + tflops_bwd_torch = total_flops_bwd / pytorch_time / 1e12 + + print(f"\n=== TFLOPS ===") + print(f"CPUInfer forward : {tflops_fwd_cpp:.2f} TFLOPS") + print(f"CPUInfer backward : {tflops_bwd_cpp:.2f} TFLOPS") + print(f"Torch backward : {tflops_bwd_torch:.2f} TFLOPS") + + + # 验证梯度结果 + backward_diff = torch.mean(torch.abs(grad_input_cpp - input.grad)) / torch.mean(torch.abs(input.grad)) + print(f"Backward diff: {backward_diff.item()}") + assert backward_diff < 0.005, f"Backward diff too large: {backward_diff.item()}" # FIXME: 0.005 是不是太大了？ + print("✅ Backward pass test passed!") + +def test_backward_2round_with_tflops(): + """ + 跑两轮 forward+backward，对比 PyTorch 与 C++ 实现的正确性和性能， + 并输出每轮及总体的 TFLOPS 与耗时信息。 + 依赖：已在全局定义 expert_num、n_routed_experts、hidden_size、intermediate_size、 + stride、group_min_len、group_max_len、gate_type、up_type、down_type、 + hidden_type、qlen、dtype、gradtype 以及 moe_torch、cpuinfer_ext、CPUInfer。 + """ + # ------------- 初始化可训练参数（保持与单轮测试一致）------------- + gate_proj = torch.randn((expert_num, intermediate_size, hidden_size), + dtype=dtype, requires_grad=True).contiguous() + up_proj = torch.randn((expert_num, intermediate_size, hidden_size), + dtype=dtype, requires_grad=True).contiguous() + down_proj = torch.randn((expert_num, hidden_size, intermediate_size), + dtype=dtype, requires_grad=True).contiguous() + + config = cpuinfer_ext.sft_moe.SFT_MOEConfig( + expert_num, n_routed_experts, hidden_size, intermediate_size, + stride, group_min_len, group_max_len, + gate_proj.data_ptr(), up_proj.data_ptr(), down_proj.data_ptr(), + gate_type, up_type, down_type, hidden_type + ) + moe = cpuinfer_ext.sft_moe.SFT_MOE(config) + + # ----------- 预先计算 FLOPs（与 KT 公式保持一致）----------- + FLOPs_fwd = 6 * qlen * n_routed_experts * hidden_size * intermediate_size + FLOPs_bwd = 18 * qlen * n_routed_experts * hidden_size * intermediate_size + + # ----------- 统计两轮测试的信息 ----------- + summary = [] # 每轮: dict(round, fwd_time, bwd_torch_time, bwd_cpp_time, diffs, TFLOPS...) + + for round_idx in range(2): + print(f"\n================ Round {round_idx+1}/2 ================") + + # ---------- 随机构造输入 ---------- + expert_ids = torch.stack( + [torch.randperm(expert_num)[:n_routed_experts] for _ in range(qlen)] + ).contiguous() + weights = torch.rand((qlen, n_routed_experts), dtype=torch.float32).contiguous() + + input_pt = (torch.randn((qlen, hidden_size), dtype=dtype) / 100)\ + .detach().requires_grad_(True).contiguous() + input_cpp = input_pt.clone().detach().requires_grad_(True).contiguous() + + # ================= 前向传播 ================= + # Torch 参考实现 + t_output = moe_torch( + input_pt, expert_ids, weights, + gate_proj, up_proj, down_proj, requires_grad=True + ) + t_output.retain_grad() + + # C++ 实现 + output_cpp = torch.empty((qlen, hidden_size), dtype=dtype).contiguous() + fwd_start = time.time() + CPUInfer.submit( + moe.forward( + qlen, n_routed_experts, + expert_ids.data_ptr(), weights.data_ptr(), + input_cpp.data_ptr(), output_cpp.data_ptr() + ) + ) + CPUInfer.sync() + fwd_end = time.time() + fwd_time = fwd_end - fwd_start + print(f"C++ forward 耗时: {fwd_time:.4f} s") + + # 结果比对 + fwd_diff = torch.mean(torch.abs(output_cpp - t_output)) \ + / torch.mean(torch.abs(t_output)) + print(f"Forward diff: {fwd_diff.item():.4e}") + + # ================= 反向传播 ================= + grad_output = torch.randn_like(t_output, dtype=gradtype).contiguous() + grad_output_cpp = grad_output.clone().contiguous() + grad_input_cpp = torch.zeros_like(input_cpp, dtype=gradtype).contiguous() + + # -- PyTorch backward -- + for p in (gate_proj, up_proj, down_proj, input_pt): + if p.grad is not None: + p.grad.zero_() + pyt_start = time.time() + t_output.backward(grad_output, retain_graph=True) + pyt_end = time.time() + pyt_time = pyt_end - pyt_start + print(f"PyTorch backward 耗时: {pyt_time:.4f} s") + + # # -- C++ backward（保持两次调用顺序） -- + # CPUInfer.submit( + # moe.backward( + # round_idx, + # qlen, n_routed_experts, + # expert_ids.data_ptr(), weights.data_ptr(), + # input_cpp.data_ptr(), + # grad_output_cpp.data_ptr(), + # grad_input_cpp.data_ptr() + # ) + # ) + # CPUInfer.sync() + + cpp_start = time.time() + CPUInfer.submit( + moe.backward( + round_idx, + qlen, n_routed_experts, + expert_ids.data_ptr(), weights.data_ptr(), + input_cpp.data_ptr(), + grad_output_cpp.data_ptr(), + grad_input_cpp.data_ptr() + ) + ) + CPUInfer.sync() + cpp_end = time.time() + cpp_time = cpp_end - cpp_start + print(f"C++ backward(第2次) 耗时: {cpp_time:.4f} s") + + # 反向结果比对 - 修复类型不匹配问题 + # grad_input_cpp是BF16，input_pt.grad是FP16，需要转换为相同类型 + if input_pt.grad is None: + print("错误：input_pt.grad为None，PyTorch反向传播可能失败") + bwd_diff = float('nan') + else: + # 添加详细调试信息 + print(f"[DEBUG] PyTorch grad shape: {input_pt.grad.shape}, dtype: {input_pt.grad.dtype}") + print(f"[DEBUG] C++ grad shape: {grad_input_cpp.shape}, dtype: {grad_input_cpp.dtype}") + + # 检查PyTorch梯度是否包含NaN + pt_grad_has_nan = torch.isnan(input_pt.grad).any() + print(f"[DEBUG] PyTorch grad contains NaN: {pt_grad_has_nan}") + if pt_grad_has_nan: + print(f"[DEBUG] PyTorch grad NaN count: {torch.isnan(input_pt.grad).sum().item()}") + + # 检查C++梯度是否包含NaN + cpp_grad_has_nan = torch.isnan(grad_input_cpp).any() + print(f"[DEBUG] C++ grad contains NaN: {cpp_grad_has_nan}") + if cpp_grad_has_nan: + print(f"[DEBUG] C++ grad NaN count: {torch.isnan(grad_input_cpp).sum().item()}") + + # 转换为FP32进行比较 + grad_input_cpp_fp32 = grad_input_cpp.to(torch.float32) + input_pt_grad_fp32 = input_pt.grad.to(torch.float32) + + # 再次检查转换后是否有NaN + cpp_fp32_has_nan = torch.isnan(grad_input_cpp_fp32).any() + pt_fp32_has_nan = torch.isnan(input_pt_grad_fp32).any() + print(f"[DEBUG] After FP32 conversion - PyTorch NaN: {pt_fp32_has_nan}, C++ NaN: {cpp_fp32_has_nan}") + + if pt_fp32_has_nan or cpp_fp32_has_nan: + bwd_diff = float('nan') + print(f"[DEBUG] 检测到NaN，跳过diff计算") + else: + diff_tensor = torch.abs(grad_input_cpp_fp32 - input_pt_grad_fp32) + denominator = torch.mean(torch.abs(input_pt_grad_fp32)) + + print(f"[DEBUG] Diff stats - max: {diff_tensor.max().item():.6f}, mean: {diff_tensor.mean().item():.6f}") + print(f"[DEBUG] Denominator: {denominator.item():.6f}") + + bwd_diff = torch.mean(diff_tensor) / denominator + if isinstance(bwd_diff, torch.Tensor): + print(f"Backward diff: {bwd_diff.item():.4e}") + elif isinstance(bwd_diff, float): + print(f"Backward diff: {bwd_diff:.4e}") + else: + print(f"Backward diff: {bwd_diff}") + + # ================= TFLOPS 统计 ================= + tflops_fwd_cpp = FLOPs_fwd / fwd_time / 1e12 + tflops_bwd_cpp = FLOPs_bwd / cpp_time / 1e12 + tflops_bwd_torch = FLOPs_bwd / pyt_time / 1e12 + + print(f"\n--- Round {round_idx+1} TFLOPS ---") + print(f"CPUInfer forward : {tflops_fwd_cpp:.2f} TFLOPS") + print(f"CPUInfer backward : {tflops_bwd_cpp:.2f} TFLOPS") + print(f"Torch backward : {tflops_bwd_torch:.2f} TFLOPS") + + # 保存本轮结果 + summary.append(dict( + round = round_idx+1, + fwd_time = fwd_time, + pyt_bwd_time = pyt_time, + cpp_bwd_time = cpp_time, + fwd_diff = fwd_diff.item(), + bwd_diff = bwd_diff.item() if isinstance(bwd_diff, torch.Tensor) else bwd_diff, + tflops_fwd_cpp = tflops_fwd_cpp, + tflops_bwd_cpp = tflops_bwd_cpp, + tflops_bwd_torch = tflops_bwd_torch, + )) + + # ================= 汇总输出 ================= + print("\n================= Two-Round Summary =================") + for item in summary: + print(f"Round {item['round']}: " + f"fwd {item['fwd_time']:.4f}s | " + f"bwd_torch {item['pyt_bwd_time']:.4f}s | " + f"bwd_cpp {item['cpp_bwd_time']:.4f}s | " + f"diff(fwd/bwd) {item['fwd_diff']:.2e}/{item['bwd_diff']:.2e} | " + f"TFLOPS(cpp fwd/bwd) {item['tflops_fwd_cpp']:.2f}/{item['tflops_bwd_cpp']:.2f}") +def test_backward_10round_5layer(): + """ + 创建 5 个独立 SFT-MOE 层，连续跑 10 轮 forward+backward。 + 第 n 轮使用第 n % 5 层，逐轮验证 C++ 与 PyTorch 的数值一致性， + 同时统计 TFLOPS / 耗时。全程不修改任何全局变量。 + """ + num_layers = 5 + num_rounds = 10 + + # ---------- 1. 为 5 层分别初始化权重 ---------- + gate_projs, up_projs, down_projs, moes = [], [], [], [] + for _ in range(num_layers): + gp = torch.randn((expert_num, intermediate_size, hidden_size), + dtype=dtype, requires_grad=True).contiguous() + up = torch.randn_like(gp, requires_grad=True) # 同形状 + dp = torch.randn((expert_num, hidden_size, intermediate_size), + dtype=dtype, requires_grad=True).contiguous() + + cfg = cpuinfer_ext.sft_moe.SFT_MOEConfig( + expert_num, n_routed_experts, + hidden_size, intermediate_size, + stride, group_min_len, group_max_len, + gp.data_ptr(), up.data_ptr(), dp.data_ptr(), + gate_type, up_type, down_type, hidden_type + ) + moes.append(cpuinfer_ext.sft_moe.SFT_MOE(cfg)) + gate_projs.append(gp); up_projs.append(up); down_projs.append(dp) + + # ---------- 2. FLOPs 常数 ---------- + FLOPs_fwd = 6 * qlen * n_routed_experts * hidden_size * intermediate_size + FLOPs_bwd = 18 * qlen * n_routed_experts * hidden_size * intermediate_size + + summary = [] + + for r in range(num_rounds): + layer_id = r % num_layers + moe = moes[layer_id] + gp, up, dp = gate_projs[layer_id], up_projs[layer_id], down_projs[layer_id] + + print(f"\n================ Round {r+1}/{num_rounds} " + f"(use layer {layer_id}) ================") + + # ---------- 3. 构造输入 ---------- + expert_ids = torch.stack( + [torch.randperm(expert_num)[:n_routed_experts] for _ in range(qlen)] + ).contiguous() + weights = torch.rand((qlen, n_routed_experts), dtype=torch.float32).contiguous() + + inp_pt = (torch.randn((qlen, hidden_size), dtype=dtype) / 100 + ).detach().requires_grad_(True).contiguous() + inp_cpp = inp_pt.clone().detach().requires_grad_(True).contiguous() + + # ================= 前向 ================= + t_out = moe_torch(inp_pt, expert_ids, weights, gp, up, dp, requires_grad=True) + t_out.retain_grad() + + out_cpp = torch.empty_like(t_out).contiguous() + t0 = time.time() + CPUInfer.submit( + moe.forward(qlen, n_routed_experts, + expert_ids.data_ptr(), weights.data_ptr(), + inp_cpp.data_ptr(), out_cpp.data_ptr()) + ) + CPUInfer.sync() + fwd_time = time.time() - t0 + + fwd_diff = (out_cpp - t_out).abs().mean() / t_out.abs().mean() + print(f"Forward diff = {fwd_diff.item():.3e} | " + f"C++ fwd {fwd_time:.3f}s") + + # ================= 反向 ================= + grad_out = torch.randn_like(t_out, dtype=gradtype).contiguous() + grad_out_cpp = grad_out.clone().contiguous() + grad_inp_cpp = torch.empty_like(inp_cpp, dtype=gradtype).contiguous() + + # PyTorch backward + for p in (gp, up, dp, inp_pt): + if p.grad is not None: + p.grad.zero_() + t1 = time.time() + t_out.backward(grad_out, retain_graph=True) + pyt_time = time.time() - t1 + + # C++ backward + t2 = time.time() + CPUInfer.submit( + moe.backward(r, qlen, n_routed_experts, + expert_ids.data_ptr(), weights.data_ptr(), + inp_cpp.data_ptr(), + grad_out_cpp.data_ptr(), grad_inp_cpp.data_ptr()) + ) + CPUInfer.sync() + cpp_time = time.time() - t2 + + bwd_diff = (grad_inp_cpp - inp_pt.grad).abs().mean() / inp_pt.grad.abs().mean() + print(f"Backward diff = {bwd_diff.item():.3e} | " + f"PyTorch bwd {pyt_time:.3f}s | C++ bwd {cpp_time:.3f}s") + + # ================= TFLOPS ================= + tflops_fwd_cpp = FLOPs_fwd / fwd_time / 1e12 + tflops_bwd_cpp = FLOPs_bwd / cpp_time / 1e12 + tflops_bwd_torch = FLOPs_bwd / pyt_time / 1e12 + + summary.append(dict( + rd=r+1, layer=layer_id, + fwd_time=fwd_time, pyt_time=pyt_time, cpp_time=cpp_time, + fwd_diff=fwd_diff.item(), bwd_diff=bwd_diff.item(), + tf_fwd=tflops_fwd_cpp, tf_bwd_cpp=tflops_bwd_cpp, + tf_bwd_torch=tflops_bwd_torch + )) + + # ---------- 4. 汇总 ---------- + print("\n================ 10-Round Summary ================") + for s in summary: + print(f"R{s['rd']:02d}(L{s['layer']}) | " + f"Δf {s['fwd_diff']:.2e} / {s['bwd_diff']:.2e} | " + f"t fwd {s['fwd_time']:.3f}s " + f"bwd Torch {s['pyt_time']:.3f}s / C++ {s['cpp_time']:.3f}s | " + f"TFLOPS C++ f/b {s['tf_fwd']:.2f}/{s['tf_bwd_cpp']:.2f} " + f"Torch bwd {s['tf_bwd_torch']:.2f}") + + print("\n✅ 10 轮 5 层测试完成，全部差异在可接受范围内！") + +def test_backward_one_vs_many_comparison(): + """ + 专门对比 backward_one 和 backward_many 的结果差异 + """ + print("\n=== Backward One vs Many Comparison ===") + + # 初始化权重（固定随机种子确保一致性） + torch.manual_seed(42) + gate_proj = torch.randn((expert_num, intermediate_size, hidden_size), + dtype=dtype, requires_grad=True).contiguous() + up_proj = torch.randn((expert_num, intermediate_size, hidden_size), + dtype=dtype, requires_grad=True).contiguous() + down_proj = torch.randn((expert_num, hidden_size, intermediate_size), + dtype=dtype, requires_grad=True).contiguous() + + # 创建两个不同的配置：一个强制使用backward_one，一个使用backward_many + config_one = cpuinfer_ext.sft_moe.SFT_MOEConfig( + expert_num, n_routed_experts, hidden_size, intermediate_size, + stride, 10000000, group_max_len, # 设置超大的group_min_len强制使用backward_one + gate_proj.data_ptr(), up_proj.data_ptr(), down_proj.data_ptr(), + gate_type, up_type, down_type, hidden_type + ) + config_many = cpuinfer_ext.sft_moe.SFT_MOEConfig( + expert_num, n_routed_experts, hidden_size, intermediate_size, + stride, group_min_len, group_max_len, # 正常配置使用backward_many + gate_proj.data_ptr(), up_proj.data_ptr(), down_proj.data_ptr(), + gate_type, up_type, down_type, hidden_type + ) + moe_one = cpuinfer_ext.sft_moe.SFT_MOE(config_one) + moe_many = cpuinfer_ext.sft_moe.SFT_MOE(config_many) + + # 固定输入数据 + torch.manual_seed(123) + expert_ids = torch.stack( + [torch.randperm(expert_num)[:n_routed_experts] for _ in range(qlen)] + ).contiguous() + weights = torch.rand((qlen, n_routed_experts), dtype=torch.float32).contiguous() + + input_one = (torch.randn((qlen, hidden_size), dtype=dtype) / 100)\ + .detach().requires_grad_(True).contiguous() + input_many = input_one.clone().detach().requires_grad_(True).contiguous() + + # Forward passes (应该是一样的) + output_one = torch.empty((qlen, hidden_size), dtype=dtype).contiguous() + output_many = torch.empty((qlen, hidden_size), dtype=dtype).contiguous() + + CPUInfer.submit( + moe_one.forward( + qlen, n_routed_experts, + expert_ids.data_ptr(), weights.data_ptr(), + input_one.data_ptr(), output_one.data_ptr() + ) + ) + CPUInfer.sync() + + CPUInfer.submit( + moe_many.forward( + qlen, n_routed_experts, + expert_ids.data_ptr(), weights.data_ptr(), + input_many.data_ptr(), output_many.data_ptr() + ) + ) + CPUInfer.sync() + + print(f"Forward outputs identical: {torch.allclose(output_one, output_many, atol=1e-6)}") + if not torch.allclose(output_one, output_many, atol=1e-6): + print(f"Forward diff: {torch.mean(torch.abs(output_one - output_many))}") + + # Backward passes + grad_output = torch.randn_like(output_one, dtype=gradtype).contiguous() + grad_output_one = grad_output.clone().contiguous() + grad_output_many = grad_output.clone().contiguous() + + grad_input_one = torch.zeros_like(input_one, dtype=gradtype).contiguous() + grad_input_many = torch.zeros_like(input_many, dtype=gradtype).contiguous() + + print("\n--- Testing backward_one (force group_min_len = 10000000) ---") + + CPUInfer.submit( + moe_one.backward( + 0, # layer_idx + qlen, n_routed_experts, + expert_ids.data_ptr(), weights.data_ptr(), + input_one.data_ptr(), + grad_output_one.data_ptr(), + grad_input_one.data_ptr() + ) + ) + CPUInfer.sync() + + # 检查backward_one结果 + one_has_nan = torch.isnan(grad_input_one).any() + print(f"backward_one result has NaN: {one_has_nan}") + if one_has_nan: + print(f"backward_one NaN count: {torch.isnan(grad_input_one).sum().item()}/{grad_input_one.numel()}") + else: + print(f"backward_one grad_input stats: min={grad_input_one.min():.6f}, max={grad_input_one.max():.6f}, mean={grad_input_one.mean():.6f}") + + print("\n--- Testing backward_many (normal group_min_len) ---") + + CPUInfer.submit( + moe_many.backward( + 0, # layer_idx + qlen, n_routed_experts, + expert_ids.data_ptr(), weights.data_ptr(), + input_many.data_ptr(), + grad_output_many.data_ptr(), + grad_input_many.data_ptr() + ) + ) + CPUInfer.sync() + + # 检查backward_many结果 + many_has_nan = torch.isnan(grad_input_many).any() + print(f"backward_many result has NaN: {many_has_nan}") + if many_has_nan: + print(f"backward_many NaN count: {torch.isnan(grad_input_many).sum().item()}/{grad_input_many.numel()}") + else: + print(f"backward_many grad_input stats: min={grad_input_many.min():.6f}, max={grad_input_many.max():.6f}, mean={grad_input_many.mean():.6f}") + + # 对比结果 + if not one_has_nan and not many_has_nan: + print(f"\n--- Comparison ---") + grad_one_fp32 = grad_input_one.to(torch.float32) + grad_many_fp32 = grad_input_many.to(torch.float32) + print(f"Results identical: {torch.allclose(grad_one_fp32, grad_many_fp32, atol=1e-6)}") + diff = torch.abs(grad_one_fp32 - grad_many_fp32) + print(f"Max absolute difference: {diff.max():.6f}") + print(f"Mean absolute difference: {diff.mean():.6f}") + + # 找到最大差异的位置 + max_diff_idx = torch.argmax(diff.flatten()) + token_idx = max_diff_idx // hidden_size + feature_idx = max_diff_idx % hidden_size + print(f"Max diff at token {token_idx}, feature {feature_idx}: " + f"one={grad_one_fp32.flatten()[max_diff_idx]:.6f}, " + f"many={grad_many_fp32.flatten()[max_diff_idx]:.6f}") + elif not one_has_nan and many_has_nan: + print(f"\n--- backward_one正常，backward_many有NaN ---") + print("这确认了问题出在backward_many实现上") + elif one_has_nan and not many_has_nan: + print(f"\n--- backward_one有NaN，backward_many正常 ---") + print("这很奇怪，需要进一步调查") + else: + print(f"\n--- 两者都有NaN ---") + print("问题可能在更基础的地方") + + +if __name__ == "__main__": + # test_backward_2round_with_tflops() + # test_backward_10round_5layer() + test_backward_one_vs_many_comparison() + \ No newline at end of file diff --git a/archive/kt-sft/csrc/ktransformers_ext/ext_bindings.cpp b/archive/kt-sft/csrc/ktransformers_ext/ext_bindings.cpp new file mode 100644 index 0000000..a5c75cf --- /dev/null +++ b/archive/kt-sft/csrc/ktransformers_ext/ext_bindings.cpp @@ -0,0 +1,1110 @@ +/** + * @Description : + * @Author : chenht2022, Jianwei Dong + * @Date : 2024-07-22 02:03:22 + * @Version : 1.0.0 + * @LastEditors : Jianwei Dong + * @LastEditTime : 2024-08-26 22:47:06 + * @Copyright (c) 2024 by KVCache.AI, All Rights Reserved. + **/ +// Python bindings +#include "cpu_backend/cpuinfer.h" +#if !defined(KTRANSFORMERS_USE_ROCM) && !defined(KTRANSFORMERS_USE_XPU) +#include "device_launch_parameters.h" +#endif +#include "llamafile/flags.h" +#include "operators/kvcache/kvcache.h" +#include "operators/llamafile/linear.h" +#include "operators/llamafile/mlp.h" +#include "operators/llamafile/moe.h" +#include "operators/llamafile/sft_moe.h" + +#if defined(__x86_64__) && defined(__HAS_AVX512F__) && defined(__HAS_AMX__) +#include "operators/amx/moe.hpp" +#include "operators/amx/sft_moe.hpp" +#endif + +#include "pybind11/functional.h" +#include "pybind11/operators.h" +#include "pybind11/pybind11.h" +#include "pybind11/stl.h" +#include +#include +#include + +namespace py = pybind11; +using namespace pybind11::literals; + +// Binding functions for the KVCache class +class KVCacheBindings { + public: + class AttnBindings { + public: + struct Args { + CPUInfer *cpuinfer; + KVCache *kv_cache; + const ggml_fp16_t *q_in; + ggml_fp16_t *output; + float *attn_lse; + int layer_idx; + int generate_token_idx; + int q_len; + int batch_size; + int max_block_num; + int *block_table; + int *cache_seqlens; + int pick_block_num; + int init_block_num; + int local_block_num; + }; + static void inner(void *args) { + Args *args_ = (Args *)args; + args_->cpuinfer->enqueue( + &KVCache::attn, args_->kv_cache, args_->q_in, args_->output, + args_->attn_lse, args_->layer_idx, args_->generate_token_idx, + args_->q_len, args_->batch_size, args_->max_block_num, + args_->block_table, args_->cache_seqlens, args_->pick_block_num, + args_->init_block_num, args_->local_block_num); + } + static std::pair + cpuinfer_interface(KVCache &kv_cache, intptr_t q_in, intptr_t output, + intptr_t attn_lse, int layer_idx, + int generate_token_idx, int q_len, int batch_size, + int max_block_num, intptr_t block_table, + intptr_t cache_seqlens, int pick_block_num, + int init_block_num, int local_block_num) { + Args *args = new Args{nullptr, + &kv_cache, + (const ggml_fp16_t *)q_in, + (ggml_fp16_t *)output, + (float *)attn_lse, + layer_idx, + generate_token_idx, + q_len, + batch_size, + max_block_num, + (int *)block_table, + (int *)cache_seqlens, + pick_block_num, + init_block_num, + local_block_num}; + return std::make_pair((intptr_t)&inner, (intptr_t)args); + } + }; + + class GetAllKVCacheOneLayerBindings { + public: + struct Args { + CPUInfer *cpuinfer; + KVCache *kv_cache; + int layer_id; + ggml_fp16_t *k_in; + ggml_fp16_t *v_in; + }; + static void inner(void *args) { + Args *args_ = (Args *)args; + args_->cpuinfer->enqueue(&KVCache::get_all_kvcache_one_layer, + args_->kv_cache, args_->layer_id, + args_->k_in, args_->v_in); + } + static std::pair + cpuinfer_interface(KVCache &kv_cache, intptr_t k_in, intptr_t v_in, + int layer_id) { + Args *args = new Args{nullptr, &kv_cache, layer_id, + (ggml_fp16_t *)k_in, (ggml_fp16_t *)v_in}; + return std::make_pair((intptr_t)&inner, (intptr_t)args); + } + }; + + class GetAndUpdateKVCacheFp16Bindings { + public: + struct Args { + CPUInfer *cpuinfer; + KVCache *kv_cache; + ggml_fp16_t *k_in; + ggml_fp16_t *v_in; + int layer_id; + int *block_table; + int batch_size; + int max_block_num; + int *cache_seqlens; + int q_len; + }; + static void inner(void *args) { + Args *args_ = (Args *)args; + args_->cpuinfer->enqueue(&KVCache::get_and_update_kvcache_fp16, + args_->kv_cache, args_->k_in, args_->v_in, + args_->layer_id, args_->block_table, + args_->batch_size, args_->max_block_num, + args_->cache_seqlens, args_->q_len); + } + static std::pair + cpuinfer_interface(KVCache &kv_cache, intptr_t k_in, intptr_t v_in, + int layer_id, intptr_t block_table, int batch_size, + int max_block_num, intptr_t cache_seqlens, + int q_len) { + Args *args = new Args{nullptr, + &kv_cache, + (ggml_fp16_t *)k_in, + (ggml_fp16_t *)v_in, + layer_id, + (int *)block_table, + batch_size, + max_block_num, + (int *)cache_seqlens, + q_len}; + return std::make_pair((intptr_t)&inner, (intptr_t)args); + } + }; + class GetKVCacheFp16Bindings { + public: + struct Args { + CPUInfer *cpuinfer; + KVCache *kv_cache; + ggml_fp16_t *k_in; + ggml_fp16_t *v_in; + int layer_id; + int *block_table; + int batch_size; + int max_block_num; + int *cache_seqlens; + }; + static void inner(void *args) { + Args *args_ = (Args *)args; + args_->cpuinfer->enqueue( + &KVCache::get_kvcache_fp16, args_->kv_cache, args_->k_in, + args_->v_in, args_->layer_id, args_->block_table, + args_->batch_size, args_->max_block_num, args_->cache_seqlens); + } + static std::pair + cpuinfer_interface(KVCache &kv_cache, intptr_t k_in, intptr_t v_in, + int layer_id, intptr_t block_table, int batch_size, + int max_block_num, intptr_t cache_seqlens) { + Args *args = new Args{nullptr, + &kv_cache, + (ggml_fp16_t *)k_in, + (ggml_fp16_t *)v_in, + layer_id, + (int *)block_table, + batch_size, + max_block_num, + (int *)cache_seqlens}; + return std::make_pair((intptr_t)&inner, (intptr_t)args); + } + }; + + class UpdateKVCacheFp16Bindings { + public: + struct Args { + CPUInfer *cpuinfer; + KVCache *kv_cache; + ggml_fp16_t *k_in; + ggml_fp16_t *v_in; + int layer_id; + int *block_table; + int batch_size; + int max_block_num; + int *cache_seqlens; + int q_len; + }; + static void inner(void *args) { + Args *args_ = (Args *)args; + args_->cpuinfer->enqueue(&KVCache::update_kvcache_fp16, + args_->kv_cache, args_->k_in, args_->v_in, + args_->layer_id, args_->block_table, + args_->batch_size, args_->max_block_num, + args_->cache_seqlens, args_->q_len); + } + static std::pair + cpuinfer_interface(KVCache &kv_cache, intptr_t k_in, intptr_t v_in, + int layer_id, intptr_t block_table, int batch_size, + int max_block_num, intptr_t cache_seqlens, + int q_len) { + Args *args = new Args{nullptr, + &kv_cache, + (ggml_fp16_t *)k_in, + (ggml_fp16_t *)v_in, + layer_id, + (int *)block_table, + batch_size, + max_block_num, + (int *)cache_seqlens, + q_len}; + return std::make_pair((intptr_t)&inner, (intptr_t)args); + } + }; + + class UpdateImportanceBindings { + public: + struct Args { + CPUInfer *cpuinfer; + KVCache *kv_cache; + const ggml_fp16_t *importance; + int layer_id; + int *block_table; + int batch_size; + int max_block_num; + int *offset; + int width; + }; + static void inner(void *args) { + Args *args_ = (Args *)args; + args_->cpuinfer->enqueue( + &KVCache::update_importance, args_->kv_cache, args_->importance, + args_->layer_id, args_->block_table, args_->batch_size, + args_->max_block_num, args_->offset, args_->width); + } + static std::pair + cpuinfer_interface(KVCache &kv_cache, intptr_t importance, int layer_id, + intptr_t block_table, int batch_size, + int max_block_num, intptr_t offset, int width) { + Args *args = new Args{nullptr, + &kv_cache, + (const ggml_fp16_t *)importance, + layer_id, + (int *)block_table, + batch_size, + max_block_num, + (int *)offset, + width}; + return std::make_pair((intptr_t)&inner, (intptr_t)args); + } + }; + + class AttnWithKVCacheBindings { + public: + struct Args { + CPUInfer *cpuinfer; + KVCache *kv_cache; + const ggml_fp16_t *q_in; + const ggml_fp16_t *k_in; + const ggml_fp16_t *v_in; + ggml_fp16_t *output; + float *attn_lse; + int layer_idx; + int generate_token_idx; + int q_len; + int batch_size; + int max_block_num; + int *block_table; + int *cache_seqlens; + int topk; + int local; + }; + static void inner(void *args) { + Args *args_ = (Args *)args; + args_->cpuinfer->enqueue( + &KVCache::attn_with_kvcache, args_->kv_cache, args_->q_in, + args_->k_in, args_->v_in, args_->output, args_->attn_lse, + args_->layer_idx, args_->generate_token_idx, args_->q_len, + args_->batch_size, args_->max_block_num, args_->block_table, + args_->cache_seqlens, args_->topk, args_->local); + } + static std::pair + cpuinfer_interface(KVCache &kv_cache, intptr_t q_in, intptr_t k_in, + intptr_t v_in, intptr_t output, intptr_t attn_lse, + int layer_idx, int generate_token_idx, int q_len, + int batch_size, int max_block_num, + intptr_t block_table, intptr_t cache_seqlens, + int topk, int local) { + Args *args = new Args{nullptr, + &kv_cache, + (const ggml_fp16_t *)q_in, + (const ggml_fp16_t *)k_in, + (const ggml_fp16_t *)v_in, + (ggml_fp16_t *)output, + (float *)attn_lse, + layer_idx, + generate_token_idx, + q_len, + batch_size, + max_block_num, + (int *)block_table, + (int *)cache_seqlens, + topk, + local}; + return std::make_pair((intptr_t)&inner, (intptr_t)args); + } + }; + + class ClearImportanceAllLayersBindings { + public: + struct Args { + CPUInfer *cpuinfer; + KVCache *kv_cache; + int *block_table; + int *cache_seqlens; + int batch_size; + int max_block_num; + }; + static void inner(void *args) { + Args *args_ = (Args *)args; + args_->cpuinfer->enqueue(&KVCache::clear_importance_all_layers, + args_->kv_cache, args_->block_table, + args_->cache_seqlens, args_->batch_size, + args_->max_block_num); + } + static std::pair + cpuinfer_interface(KVCache &kv_cache, intptr_t block_table, + intptr_t cache_seqlens, int batch_size, + int max_block_num) { + Args *args = new Args{nullptr, + &kv_cache, + (int *)block_table, + (int *)cache_seqlens, + batch_size, + max_block_num}; + return std::make_pair((intptr_t)&inner, (intptr_t)args); + } + }; + + class CalcAnchorAllLayersBindinds { + public: + struct Args { + CPUInfer *cpuinfer; + KVCache *kv_cache; + int *block_table; + int *cache_seqlens; + int batch_size; + int max_block_num; + }; + static void inner(void *args) { + Args *args_ = (Args *)args; + args_->cpuinfer->enqueue(&KVCache::calc_anchor_all_layers, + args_->kv_cache, args_->block_table, + args_->cache_seqlens, args_->batch_size, + args_->max_block_num); + } + static std::pair + cpuinfer_interface(KVCache &kv_cache, intptr_t block_table, + intptr_t cache_seqlens, int batch_size, + int max_block_num) { + Args *args = new Args{nullptr, + &kv_cache, + (int *)block_table, + (int *)cache_seqlens, + batch_size, + max_block_num}; + return std::make_pair((intptr_t)&inner, (intptr_t)args); + } + }; + + class LoadKVCacheBindings { + public: + struct Args { + CPUInfer *cpuinfer; + KVCache *kv_cache; + std::string tensor_file_path; + }; + static void inner(void *args) { + Args *args_ = (Args *)args; + args_->cpuinfer->enqueue(&KVCache::load_kvcache, args_->kv_cache, + args_->tensor_file_path); + } + static std::pair + cpuinfer_interface(KVCache &kv_cache, std::string tensor_file_path) { + Args *args = + new Args{nullptr, &kv_cache, (std::string)tensor_file_path}; + return std::make_pair((intptr_t)&inner, (intptr_t)args); + } + }; + class DumpKVCacheBindings { + public: + struct Args { + CPUInfer *cpuinfer; + KVCache *kv_cache; + int *block_table; + int cache_total_len; + std::string tensor_file_path; + }; + static void inner(void *args) { + Args *args_ = (Args *)args; + args_->cpuinfer->enqueue(&KVCache::dump_kvcache, args_->kv_cache, + args_->block_table, args_->cache_total_len, + args_->tensor_file_path); + } + static std::pair + cpuinfer_interface(KVCache &kv_cache, intptr_t block_table, + int cache_total_len, std::string tensor_file_path) { + Args *args = + new Args{nullptr, &kv_cache, (int *)block_table, + cache_total_len, (std::string)tensor_file_path}; + return std::make_pair((intptr_t)&inner, (intptr_t)args); + } + }; +}; + +class LinearBindings { + public: + class WarmUpBindinds { + public: + struct Args { + CPUInfer *cpuinfer; + Linear *linear; + }; + static void inner(void *args) { + Args *args_ = (Args *)args; + args_->cpuinfer->enqueue(&Linear::warm_up, args_->linear); + } + static std::pair + cpuinfer_interface(Linear &linear) { + Args *args = new Args{nullptr, &linear}; + return std::make_pair((intptr_t)&inner, (intptr_t)args); + } + }; + class ForwardBindings { + public: + struct Args { + CPUInfer *cpuinfer; + Linear *linear; + int qlen; + const void *input; + void *output; + }; + static void inner(void *args) { + Args *args_ = (Args *)args; + args_->cpuinfer->enqueue(&Linear::forward, args_->linear, + args_->qlen, args_->input, args_->output); + } + static std::pair + cpuinfer_interface(Linear &linear, int qlen, intptr_t input, + intptr_t output) { + Args *args = new Args{nullptr, &linear, qlen, (const void *)input, + (void *)output}; + return std::make_pair((intptr_t)&inner, (intptr_t)args); + } + }; +}; + +class MLPBindings { + public: + class WarmUpBindinds { + public: + struct Args { + CPUInfer *cpuinfer; + MLP *mlp; + }; + static void inner(void *args) { + Args *args_ = (Args *)args; + args_->cpuinfer->enqueue(&MLP::warm_up, args_->mlp); + } + static std::pair cpuinfer_interface(MLP &mlp) { + Args *args = new Args{nullptr, &mlp}; + return std::make_pair((intptr_t)&inner, (intptr_t)args); + } + }; + class ForwardBindings { + public: + struct Args { + CPUInfer *cpuinfer; + MLP *mlp; + int qlen; + const void *input; + void *output; + }; + static void inner(void *args) { + Args *args_ = (Args *)args; + args_->cpuinfer->enqueue(&MLP::forward, args_->mlp, args_->qlen, + args_->input, args_->output); + } + static std::pair + cpuinfer_interface(MLP &mlp, int qlen, intptr_t input, + intptr_t output) { + Args *args = new Args{nullptr, &mlp, qlen, (const void *)input, + (void *)output}; + return std::make_pair((intptr_t)&inner, (intptr_t)args); + } + }; +}; + +class MOEBindings { + public: + class WarmUpBindinds { + public: + struct Args { + CPUInfer *cpuinfer; + MOE *moe; + }; + static void inner(void *args) { + Args *args_ = (Args *)args; + args_->cpuinfer->enqueue(&MOE::warm_up, args_->moe); + } + static std::pair cpuinfer_interface(MOE &moe) { + Args *args = new Args{nullptr, &moe}; + return std::make_pair((intptr_t)&inner, (intptr_t)args); + } + }; + class ForwardBindings { + public: + struct Args { + CPUInfer *cpuinfer; + MOE *moe; + int qlen; + int k; + const uint64_t *expert_ids; + const float *weights; + const void *input; + void *output; + int *batch_size_tensor; + }; + static void inner(void *args) { + Args *args_ = (Args *)args; + args_->cpuinfer->enqueue( + &MOE::forward, args_->moe, args_->qlen, args_->k, + args_->expert_ids, args_->weights, args_->input, args_->output, args_->batch_size_tensor); + } + static std::pair + cpuinfer_interface(MOE &moe, int qlen, int k, intptr_t expert_ids, + intptr_t weights, intptr_t input, intptr_t output, intptr_t batch_size_tensor) { + Args *args = new Args{nullptr, + &moe, + qlen, + k, + (const uint64_t *)expert_ids, + (const float *)weights, + (const void *)input, + (void *)output, + (int *)batch_size_tensor}; + return std::make_pair((intptr_t)&inner, (intptr_t)args); + } + }; +}; + +namespace { + inline void sft_moe_forward_wrapper( + SFT_MOE& self, + int qlen, int k, + const uint64_t* expert_ids, + const float* weights, + const void* input, + void* output, + Backend* backend) + { + self.ensure_fwd_cache(qlen, k); + self.forward(qlen, k, expert_ids, weights, + input, output, + backend, + self.fwd_cache_ptr()); + } + + inline void sft_moe_backward_wrapper( + SFT_MOE& self, + int layer_idx, + int qlen, int k, + const uint64_t* expert_ids, + const float* weights, + const void* input, + const void* grad_output, + void* grad_input, + Backend* backend) + { + self.backward(layer_idx, qlen, k, expert_ids, weights, + input, grad_output, grad_input, + backend, + self.fwd_cache_ptr()); + } +} + +class SFT_MOEBindings { + public: + class WarmUpBindinds { + public: + struct Args { + CPUInfer *cpuinfer; + SFT_MOE *moe; + }; + static void inner(void *args) { + Args *args_ = (Args *)args; + args_->cpuinfer->enqueue(&SFT_MOE::warm_up, args_->moe); + } + static std::pair cpuinfer_interface(SFT_MOE &moe) { + Args *args = new Args{nullptr, &moe}; + return std::make_pair((intptr_t)&inner, (intptr_t)args); + } + }; + class ForwardBindings { + public: + struct Args { + CPUInfer *cpuinfer; + SFT_MOE *moe; + int qlen; + int k; + const uint64_t *expert_ids; + const float *weights; + const void *input; + void *output; + }; + // static void inner(void *args) { + // Args *args_ = (Args *)args; + // args_->cpuinfer->enqueue( + // &SFT_MOE::forward, args_->moe, args_->qlen, args_->k, + // args_->expert_ids, args_->weights, args_->input, args_->output); + // } + static void inner(void *args) { + Args *args_ = static_cast(args); + args_->cpuinfer->enqueue( + &sft_moe_forward_wrapper, // 使用包装函数 + args_->moe, + args_->qlen, args_->k, + args_->expert_ids, + args_->weights, + args_->input, + args_->output); + } + static std::pair + cpuinfer_interface(SFT_MOE &moe, int qlen, int k, intptr_t expert_ids, + intptr_t weights, intptr_t input, intptr_t output) { + Args *args = new Args{nullptr, + &moe, + qlen, + k, + (const uint64_t *)expert_ids, + (const float *)weights, + (const void *)input, + (void *)output}; + return std::make_pair((intptr_t)&inner, (intptr_t)args); + } + }; + // FIXME: need fit the args setting with the backward of MoE + class BackwardBindings { + public: + struct Args { + CPUInfer* cpuinfer; + SFT_MOE* moe; + int layer_idx; + int qlen; + int k; + const uint64_t* expert_ids; + const float* weights; + const void* input; + const void* grad_output; + void* grad_input; + }; + + // static void inner(void* args) { + // Args* args_ = static_cast(args); + // args_->cpuinfer->enqueue(&SFT_MOE::backward, args_->moe, + // args_->qlen, args_->k, + // args_->expert_ids, args_->weights, + // args_->input, + // args_->grad_output, + // args_->grad_input); + // } + + static void inner(void *args) { + Args *args_ = static_cast(args); + args_->cpuinfer->enqueue( + &sft_moe_backward_wrapper, // 使用包装函数 + args_->moe, + args_->layer_idx, + args_->qlen, args_->k, + args_->expert_ids, + args_->weights, + args_->input, + args_->grad_output, + args_->grad_input); + } + + static std::pair cpuinfer_interface( + SFT_MOE& moe, int layer_idx, int qlen, int k, + intptr_t expert_ids, intptr_t weights, + intptr_t input, + intptr_t grad_output, intptr_t grad_input) { + + Args* args = new Args{ + nullptr, &moe, layer_idx, qlen, k, + reinterpret_cast(expert_ids), + reinterpret_cast(weights), + reinterpret_cast(input), + reinterpret_cast(grad_output), + reinterpret_cast(grad_input) + }; + return std::make_pair( + reinterpret_cast(&inner), + reinterpret_cast(args)); + } + }; +}; + +#if defined(__x86_64__) && defined(__HAS_AVX512F__) && defined(__HAS_AMX__) +template +class AMX_MOEBindings { + public: + class WarmUpBindings { + public: + struct Args { + CPUInfer *cpuinfer; + AMX_MOE *moe; + }; + static void inner(void *args) { + Args *args_ = (Args *)args; + args_->cpuinfer->enqueue(&AMX_MOE::warm_up, args_->moe); + } + static std::pair cpuinfer_interface(AMX_MOE &moe) { + Args *args = new Args{nullptr, &moe}; + return std::make_pair((intptr_t)&inner, (intptr_t)args); + } + }; + class LoadWeightsBindings { + public: + struct Args { + CPUInfer *cpuinfer; + AMX_MOE *moe; + }; + static void inner(void *args) { + Args *args_ = (Args *)args; + args_->cpuinfer->enqueue(&AMX_MOE::load_weights, args_->moe); + } + static std::pair cpuinfer_interface(AMX_MOE &moe) { + Args *args = new Args{nullptr, &moe}; + return std::make_pair((intptr_t)&inner, (intptr_t)args); + } + }; + class ForwardBindings { + public: + struct Args { + CPUInfer *cpuinfer; + AMX_MOE *moe; + int qlen; + int k; + const uint64_t *expert_ids; + const float *weights; + const void *input; + void *output; + int *batch_size_tensor; + }; + static void inner(void *args) { + Args *args_ = (Args *)args; + args_->cpuinfer->enqueue( + &AMX_MOE::forward, args_->moe, args_->qlen, args_->k, + args_->expert_ids, args_->weights, args_->input, args_->output, args_->batch_size_tensor); + } + static std::pair + cpuinfer_interface(AMX_MOE &moe, int qlen, int k, intptr_t expert_ids, + intptr_t weights, intptr_t input, intptr_t output, intptr_t batch_size_tensor) { + Args *args = new Args{nullptr, + &moe, + qlen, + k, + (const uint64_t *)expert_ids, + (const float *)weights, + (const void *)input, + (void *)output, + (int *)batch_size_tensor}; + return std::make_pair((intptr_t)&inner, (intptr_t)args); + } + }; +}; +#endif + +#if defined(__x86_64__) && defined(__HAS_AVX512F__) && defined(__HAS_AMX__) +template +class SFT_AMX_MOEBindings { + public: + class WarmUpBindings { + public: + struct Args { + CPUInfer *cpuinfer; + SFT_AMX_MOE *moe; + }; + static void inner(void *args) { + Args *args_ = (Args *)args; + args_->cpuinfer->enqueue(&SFT_AMX_MOE::warm_up, args_->moe); + } + static std::pair cpuinfer_interface(SFT_AMX_MOE &moe) { + Args *args = new Args{nullptr, &moe}; + return std::make_pair((intptr_t)&inner, (intptr_t)args); + } + }; + class LoadWeightsBindings { + public: + struct Args { + CPUInfer *cpuinfer; + SFT_AMX_MOE *moe; + }; + static void inner(void *args) { + Args *args_ = (Args *)args; + args_->cpuinfer->enqueue(&SFT_AMX_MOE::load_weights, args_->moe); + } + static std::pair cpuinfer_interface(SFT_AMX_MOE &moe) { + Args *args = new Args{nullptr, &moe}; + return std::make_pair((intptr_t)&inner, (intptr_t)args); + } + }; + class ForwardBindings { + public: + struct Args { + CPUInfer *cpuinfer; + SFT_AMX_MOE *moe; + int qlen; + int k; + const uint64_t *expert_ids; + const float *weights; + const void *input; + void *output; + }; + static void inner(void *args) { + Args *args_ = (Args *)args; + args_->cpuinfer->enqueue( + &SFT_AMX_MOE::forward, args_->moe, args_->qlen, args_->k, + args_->expert_ids, args_->weights, args_->input, args_->output); + } + static std::pair + cpuinfer_interface(SFT_AMX_MOE &moe, int qlen, int k, intptr_t expert_ids, + intptr_t weights, intptr_t input, intptr_t output) { + Args *args = new Args{nullptr, + &moe, + qlen, + k, + (const uint64_t *)expert_ids, + (const float *)weights, + (const void *)input, + (void *)output}; + return std::make_pair((intptr_t)&inner, (intptr_t)args); + } + }; + + class BackwardBindings { + public: + struct Args { + CPUInfer* cpuinfer; + SFT_AMX_MOE *moe; + int qlen; + int k; + const uint64_t* expert_ids; + const float* weights; + const void* input; + const void* output_grad; + void* input_grad; + }; + + static void inner(void *args) { + Args *args_ = static_cast(args); + args_->cpuinfer->enqueue( + &SFT_AMX_MOE::backward, + args_->moe, + args_->qlen, args_->k, + args_->expert_ids, + args_->weights, + args_->input, + args_->output_grad, + args_->input_grad); + } + + static std::pair cpuinfer_interface( + SFT_AMX_MOE &moe, int qlen, int k, + intptr_t expert_ids, intptr_t weights, + intptr_t input, + intptr_t output_grad, intptr_t input_grad) { + + Args* args = new Args{ + nullptr, &moe, qlen, k, + (const uint64_t*)expert_ids, + (const float*)weights, + (const void*)input, + (const void*)output_grad, + (void*)input_grad + }; + return std::make_pair( + (intptr_t)&inner, + (intptr_t)args); + } + }; +}; +#endif + +PYBIND11_MODULE(cpuinfer_ext, m) { + py::class_(m, "CPUInfer") + .def(py::init()) + .def("submit", &CPUInfer::submit) + .def("submit_with_cuda_stream", &CPUInfer::submit_with_cuda_stream) + .def("sync", &CPUInfer::sync) + .def("sync_with_cuda_stream", &CPUInfer::sync_with_cuda_stream); + + auto linear_module = m.def_submodule("linear"); + py::class_(linear_module, "LinearConfig") + .def(py::init([](int hidden_size, int intermediate_size, int stride, + int group_max_len, intptr_t proj, int proj_type, + int hidden_type) { + return LinearConfig(hidden_size, intermediate_size, stride, + group_max_len, (void *)proj, + (ggml_type)proj_type, (ggml_type)hidden_type); + })); + py::class_(linear_module, "Linear") + .def(py::init()) + .def("warm_up", &LinearBindings::WarmUpBindinds::cpuinfer_interface) + .def("forward", &LinearBindings::ForwardBindings::cpuinfer_interface); + + auto mlp_module = m.def_submodule("mlp"); + py::class_(mlp_module, "MLPConfig") + .def(py::init([](int hidden_size, int intermediate_size, int stride, + int group_max_len, intptr_t gate_proj, + intptr_t up_proj, intptr_t down_proj, int gate_type, + int up_type, int down_type, int hidden_type) { + return MLPConfig(hidden_size, intermediate_size, stride, + group_max_len, (void *)gate_proj, (void *)up_proj, + (void *)down_proj, (ggml_type)gate_type, + (ggml_type)up_type, (ggml_type)down_type, + (ggml_type)hidden_type); + })); + py::class_(mlp_module, "MLP") + .def(py::init()) + .def("warm_up", &MLPBindings::WarmUpBindinds::cpuinfer_interface) + .def("forward", &MLPBindings::ForwardBindings::cpuinfer_interface); + + auto moe_module = m.def_submodule("moe"); + py::class_(moe_module, "MOEConfig") + .def(py::init([](int expert_num, int routed_expert_num, int hidden_size, + int intermediate_size, int stride, int group_min_len, + int group_max_len, intptr_t gate_proj, + intptr_t up_proj, intptr_t down_proj, int gate_type, + int up_type, int down_type, int hidden_type) { + return MOEConfig(expert_num, routed_expert_num, hidden_size, + intermediate_size, stride, group_min_len, + group_max_len, (void *)gate_proj, (void *)up_proj, + (void *)down_proj, (ggml_type)gate_type, + (ggml_type)up_type, (ggml_type)down_type, + (ggml_type)hidden_type); + })); + py::class_(moe_module, "MOE") + .def(py::init()) + .def("warm_up", &MOEBindings::WarmUpBindinds::cpuinfer_interface) + .def("forward", &MOEBindings::ForwardBindings::cpuinfer_interface); + + auto sft_moe_module = m.def_submodule("sft_moe"); + py::class_(sft_moe_module, "SFT_MOEConfig") + .def(py::init([](int expert_num, int routed_expert_num, int hidden_size, + int intermediate_size, int stride, int group_min_len, + int group_max_len, intptr_t gate_proj, + intptr_t up_proj, intptr_t down_proj, int gate_type, + int up_type, int down_type, int hidden_type) { + return SFT_MOEConfig(expert_num, routed_expert_num, hidden_size, + intermediate_size, stride, group_min_len, + group_max_len, (void *)gate_proj, (void *)up_proj, + (void *)down_proj, (ggml_type)gate_type, + (ggml_type)up_type, (ggml_type)down_type, + (ggml_type)hidden_type); + })); + py::class_(sft_moe_module, "SFT_MOE") + .def(py::init()) + .def("warm_up", &SFT_MOEBindings::WarmUpBindinds::cpuinfer_interface) + .def("forward", &SFT_MOEBindings::ForwardBindings::cpuinfer_interface) + .def("backward", &SFT_MOEBindings::BackwardBindings::cpuinfer_interface); + + #if defined(__x86_64__) && defined(__HAS_AVX512F__) && defined(__HAS_AMX__) + py::class_(moe_module, "AMX_MOEConfig") + .def(py::init([](int expert_num, int routed_expert_num, int hidden_size, + int intermediate_size, + int max_len, intptr_t gate_proj, + intptr_t up_proj, intptr_t down_proj) { + return AMX_MOEConfig(expert_num, routed_expert_num, hidden_size, + intermediate_size, + max_len, (void *)gate_proj, + (void *)up_proj, (void *)down_proj); + })); + + py::class_>(moe_module, "AMXBF16_MOE") + .def(py::init()) + .def("warm_up", &AMX_MOEBindings::WarmUpBindings::cpuinfer_interface) + .def("load_weights", &AMX_MOEBindings::LoadWeightsBindings::cpuinfer_interface) + .def("forward", &AMX_MOEBindings::ForwardBindings::cpuinfer_interface); + py::class_>(moe_module, "AMXInt8_MOE") + .def(py::init()) + .def("warm_up", &AMX_MOEBindings::WarmUpBindings::cpuinfer_interface) + .def("load_weights", &AMX_MOEBindings::LoadWeightsBindings::cpuinfer_interface) + .def("forward", &AMX_MOEBindings::ForwardBindings::cpuinfer_interface); + + #endif + + #if defined(__x86_64__) && defined(__HAS_AVX512F__) && defined(__HAS_AMX__) + py::class_(sft_moe_module, "SFT_AMX_MOEConfig") + .def(py::init([](int expert_num, int routed_expert_num, int hidden_size, + int intermediate_size, + int max_len, intptr_t gate_proj, + intptr_t up_proj, intptr_t down_proj) { + return SFT_AMX_MOEConfig(expert_num, routed_expert_num, hidden_size, + intermediate_size, + max_len, (void *)gate_proj, + (void *)up_proj, (void *)down_proj); + })); + + py::class_>(sft_moe_module, "SFT_AMXBF16_MOE") + .def(py::init()) + .def("warm_up", &SFT_AMX_MOEBindings::WarmUpBindings::cpuinfer_interface) + .def("load_weights", &SFT_AMX_MOEBindings::LoadWeightsBindings::cpuinfer_interface) + .def("forward", &SFT_AMX_MOEBindings::ForwardBindings::cpuinfer_interface) + .def("backward", &SFT_AMX_MOEBindings::BackwardBindings::cpuinfer_interface); + + py::class_>(sft_moe_module, "SFT_AMXInt8_MOE") + .def(py::init()) + .def("warm_up", &SFT_AMX_MOEBindings::WarmUpBindings::cpuinfer_interface) + .def("load_weights", &SFT_AMX_MOEBindings::LoadWeightsBindings::cpuinfer_interface) + .def("forward", &SFT_AMX_MOEBindings::ForwardBindings::cpuinfer_interface) + .def("backward", &SFT_AMX_MOEBindings::BackwardBindings::cpuinfer_interface); + + #endif + + auto kvcache_module = m.def_submodule("kvcache"); + + py::enum_(kvcache_module, "AnchorType") + .value("FIXED", AnchorType::FIXED_ANCHOR) + .value("DYNAMIC", AnchorType::DYNAMIC) + .value("QUEST", AnchorType::QUEST) + .value("BLOCK_MAX", AnchorType::BLOCK_MAX) + .value("BLOCK_MEAN", AnchorType::BLOCK_MEAN); + py::enum_(kvcache_module, "ggml_type") + .value("FP16", ggml_type::GGML_TYPE_F16) + .value("FP32", ggml_type::GGML_TYPE_F32) + .value("Q4_0", ggml_type::GGML_TYPE_Q4_0) + .value("Q8_0", ggml_type::GGML_TYPE_Q8_0); + py::enum_(kvcache_module, "RetrievalType") + .value("LAYER", RetrievalType::LAYER) + .value("KVHEAD", RetrievalType::KVHEAD) + .value("QHEAD", RetrievalType::QHEAD); + + py::class_(kvcache_module, "KVCacheConfig") + .def(py::init()) + .def_readwrite("layer_num", &KVCacheConfig::layer_num) + .def_readwrite("kv_head_num", &KVCacheConfig::kv_head_num) + .def_readwrite("q_head_num", &KVCacheConfig::q_head_num) + .def_readwrite("head_dim", &KVCacheConfig::head_dim) + .def_readwrite("block_len", &KVCacheConfig::block_len) + .def_readwrite("anchor_num", &KVCacheConfig::anchor_num) + .def_readwrite("anchor_type", &KVCacheConfig::anchor_type) + .def_readwrite("kv_type", &KVCacheConfig::kv_type) + .def_readwrite("retrieval_type", &KVCacheConfig::retrieval_type) + .def_readwrite("layer_step", &KVCacheConfig::layer_step) + .def_readwrite("token_step", &KVCacheConfig::token_step) + .def_readwrite("layer_offset", &KVCacheConfig::layer_offset) + .def_readwrite("max_block_num", &KVCacheConfig::max_block_num) + .def_readwrite("max_batch_size", &KVCacheConfig::max_batch_size) + .def_readwrite("max_thread_num", &KVCacheConfig::max_thread_num); + py::class_(kvcache_module, "KVCache") + .def(py::init()) + .def("get_cache_total_len", &KVCache::get_cache_total_len) + .def("update_cache_total_len", + [](KVCache &kvcache, int cache_total_len) { + kvcache.update_cache_total_len(cache_total_len); + }) + .def("attn", &KVCacheBindings::AttnBindings::cpuinfer_interface) + .def( + "get_all_kvcache_one_layer", + &KVCacheBindings::GetAllKVCacheOneLayerBindings::cpuinfer_interface) + .def("get_and_update_kvcache_fp16", + &KVCacheBindings::GetAndUpdateKVCacheFp16Bindings:: + cpuinfer_interface) + .def("get_kvcache_fp16", + &KVCacheBindings::GetKVCacheFp16Bindings::cpuinfer_interface) + .def("update_kvcache_fp16", + &KVCacheBindings::UpdateKVCacheFp16Bindings::cpuinfer_interface) + .def("update_importance", + &KVCacheBindings::UpdateImportanceBindings::cpuinfer_interface) + .def("attn_with_kvcache", + &KVCacheBindings::AttnWithKVCacheBindings::cpuinfer_interface) + .def("clear_importance_all_layers", + &KVCacheBindings::ClearImportanceAllLayersBindings:: + cpuinfer_interface) + .def("calc_anchor_all_layers", + &KVCacheBindings::CalcAnchorAllLayersBindinds::cpuinfer_interface); +} \ No newline at end of file diff --git a/archive/kt-sft/csrc/ktransformers_ext/operators/amx/debug_sft_moe.hpp b/archive/kt-sft/csrc/ktransformers_ext/operators/amx/debug_sft_moe.hpp new file mode 100644 index 0000000..9f64a41 --- /dev/null +++ b/archive/kt-sft/csrc/ktransformers_ext/operators/amx/debug_sft_moe.hpp @@ -0,0 +1,927 @@ +/** + * @Description : Mainly used for dev debug, with no numa version + * @Author : chenht2022 + * @Date : 2025-04-25 18:28:12 + * @Version : 1.0.0 + * @LastEditors : chenht2022 + * @LastEditTime : 2025-04-25 18:28:12 + * @Copyright (c) 2024 by KVCache.AI, All Rights Reserved. + **/ +#ifndef CPUINFER_OPERATOR_SFT_AMX_MOE_H +#define CPUINFER_OPERATOR_SFT_AMX_MOE_H + +#include +#include +#include +#include +#include +#include +#include + +#include "debug_sft_moe.hpp" + +#include "../../cpu_backend/backend.h" +#include "../../cpu_backend/shared_mem_buffer.h" +#include "llama.cpp/ggml-impl.h" +#include "llama.cpp/ggml-quants.h" +#include "llama.cpp/ggml.h" +#include "llamafile/sgemm.h" + +#include "la/amx.hpp" + +#ifdef USE_NUMA +#include +#include +void *numa_alloc_aligned(size_t size, int node, size_t alignment) { + void *ptr = numa_alloc_onnode(size, node); + assert(reinterpret_cast(ptr) % 64 == 0); + return ptr; +} +#endif + +// static inline __m512 exp_avx512(__m512 x) { +// const __m512 log2e = _mm512_set1_ps(1.44269504089f); +// const __m512 c1 = _mm512_set1_ps(0.69314718056f); + +// __m512 y = _mm512_mul_ps(x, log2e); +// __m512i int_part = _mm512_cvtps_epi32(y); +// __m512 frac_part = _mm512_sub_ps(y, _mm512_cvtepi32_ps(int_part)); + +// const __m512 poly_1 = _mm512_set1_ps(0.9999999995f); +// const __m512 poly_2 = _mm512_set1_ps(0.6931471805f); +// const __m512 poly_3 = _mm512_set1_ps(0.2402265069f); +// const __m512 poly_4 = _mm512_set1_ps(0.0555041087f); +// const __m512 poly_5 = _mm512_set1_ps(0.0096181291f); +// const __m512 poly_6 = _mm512_set1_ps(0.0013333558f); + +// __m512 frac_exp = _mm512_fmadd_ps( +// frac_part, poly_6, +// _mm512_fmadd_ps(frac_part, poly_5, +// _mm512_fmadd_ps(frac_part, poly_4, +// _mm512_fmadd_ps(frac_part, poly_3, _mm512_fmadd_ps(frac_part, poly_2, poly_1))))); + +// __m512 two_pow_i = _mm512_scalef_ps(_mm512_set1_ps(1.0f), _mm512_cvtepi32_ps(int_part)); +// return _mm512_mul_ps(two_pow_i, frac_exp); +// } + +// static inline __m512 act_fn(__m512 gate_val, __m512 up_val) { +// __m512 neg_gate_val = _mm512_sub_ps(_mm512_setzero_ps(), gate_val); +// __m512 exp_neg_gate = exp_avx512(neg_gate_val); +// __m512 denom = _mm512_add_ps(_mm512_set1_ps(1.0f), exp_neg_gate); +// __m512 act_val = _mm512_div_ps(gate_val, denom); + +// return _mm512_mul_ps(act_val, up_val); +// } + +static inline __m512 sigmoid(__m512 x) { + __m512 neg = _mm512_sub_ps(_mm512_setzero_ps(), x); + __m512 e = exp_avx512(neg); + __m512 denom = _mm512_add_ps(_mm512_set1_ps(1.0f), e); + return _mm512_div_ps(_mm512_set1_ps(1.0f), denom); +} + +static inline __m512 act_fn_1(__m512 x) { + __m512 sigmoid_val = sigmoid(x); + return _mm512_mul_ps(sigmoid_val, x); +} + +static inline __m512 act_fn_grad(__m512 x) { + // sigmoid(x) * (1 + x * (1 - sigmoid(x))) + __m512 sigmoid_val = sigmoid(x); + __m512 one_minus_sigmoid = _mm512_sub_ps(_mm512_set1_ps(1.0f), sigmoid_val); + __m512 x_term = _mm512_mul_ps(x, one_minus_sigmoid); + __m512 one_plus_x_term = _mm512_add_ps(_mm512_set1_ps(1.0f), x_term); + return _mm512_mul_ps(sigmoid_val, one_plus_x_term); +} + +// static inline float bf16_to_fp32(ggml_bf16_t v) { +// uint16_t lo16; +// std::memcpy(&lo16, &v, sizeof(lo16)); // 取出 16 bit 数据 +// uint32_t tmp = uint32_t(lo16) << 16; // 放到高 16 位 +// float out; +// std::memcpy(&out, &tmp, sizeof(float)); +// return out; +// } + +// 把 ggml_bf16_t 数组转换成可读字符串（逗号分隔） +std::string int8_row_to_string(const int8_t* row, int len) { + std::string s; + for (int i = 0; i < len; ++i) { + if (i) s += ", "; + s += std::to_string(row[i]); + } + return s; +} + +struct SFT_AMX_MOEConfig { + int expert_num; + int routed_expert_num; + int hidden_size; + int intermediate_size; + int max_len; + void *gate_proj; + void *up_proj; + void *down_proj; + + SFT_AMX_MOEConfig() {} + + SFT_AMX_MOEConfig(int expert_num, int routed_expert_num, int hidden_size, int intermediate_size, int max_len, + void *gate_proj, void *up_proj, void *down_proj) + : expert_num(expert_num), routed_expert_num(routed_expert_num), hidden_size(hidden_size), + intermediate_size(intermediate_size), max_len(max_len), gate_proj(gate_proj), up_proj(up_proj), + down_proj(down_proj) {} +}; + +template class SFT_AMX_MOE { +private: + SFT_AMX_MOEConfig config_; + void *gate_proj_; // [expert_num * intermediate_size * hidden_size ( /32 if quantized)] + void *up_proj_; // [expert_num * intermediate_size * hidden_size ( /32 if quantized)] + void *down_proj_; // [expert_num * hidden_size * intermediate_size ( /32 if quantized)] + + void *gate_proj_t_; // [expert_num * intermediate_size * hidden_size] + void *up_proj_t_; // [expert_num * intermediate_size * hidden_size] + void *down_proj_t_; // [expert_num * hidden_size * intermediate_size] + + ggml_bf16_t *m_local_input_; // [routed_expert_num * max_len * hidden_size] + ggml_bf16_t *m_local_gate_output_; // [routed_expert_num * max_len * intermediate_size] + ggml_bf16_t *m_local_up_output_; // [routed_expert_num * max_len * intermediate_size] + ggml_bf16_t *m_local_down_output_; // [routed_expert_num * max_len * hidden_size] + + std::vector> m_local_pos_; // [max_len, routed_expert_num] + std::vector m_local_num_; // [expert_num] + std::vector m_expert_id_map_; // [expert_num] + std::vector m_local_input_ptr_; // [expert_num] + std::vector m_local_gate_output_ptr_; // [expert_num] + std::vector m_local_up_output_ptr_; // [expert_num] + std::vector m_local_down_output_ptr_; // [expert_num] + + std::vector> gate_up_ba_; + std::vector> gate_bc_; + std::vector> up_bc_; + std::vector> down_ba_; + std::vector> down_bc_; + +#ifdef USE_NUMA + std::vector>> gate_bb_numa_; + std::vector>> up_bb_numa_; + std::vector>> down_bb_numa_; +#else + std::vector> gate_bb_; + std::vector> up_bb_; + std::vector> down_bb_; +#endif + + ggml_bf16_t *m_local_down_output_grad_; // [routed_expert_num * max_len * hidden_size] + ggml_bf16_t *m_local_down_input_grad_; // [routed_expert_num * max_len * intermediate_size] + ggml_bf16_t *m_local_gate_output_grad_; // [routed_expert_num * max_len * intermediate_size] + ggml_bf16_t *m_local_up_output_grad_; // [routed_expert_num * max_len * intermediate_size] + ggml_bf16_t *m_local_gate_input_grad_; // [routed_expert_num * max_len * hidden_size] + ggml_bf16_t *m_local_up_input_grad_; // [routed_expert_num * max_len * hidden_size] + + std::vector m_local_down_output_grad_ptr_; // [expert_num] + std::vector m_local_down_input_grad_ptr_; // [expert_num] + std::vector m_local_gate_output_grad_ptr_; // [expert_num] + std::vector m_local_up_output_grad_ptr_; // [expert_num] + std::vector m_local_gate_input_grad_ptr_; // [expert_num] + std::vector m_local_up_input_grad_ptr_; // [expert_num] + + std::vector> gate_t_ba_; + std::vector> gate_t_bc_; + std::vector> up_t_ba_; + std::vector> up_t_bc_; + std::vector> down_t_ba_; + std::vector> down_t_bc_; + + // TODO: NUMA + std::vector> gate_t_bb_; + std::vector> up_t_bb_; + std::vector> down_t_bb_; + + int* m_local_token_indices_; // [routed_expert_num * max_len] + int* m_local_expert_positions_; // [routed_expert_num * max_len] + std::vector m_local_token_indices_ptr_; // [expert_num] + std::vector m_local_expert_positions_ptr_; // [expert_num] + +public: + SFT_AMX_MOE(SFT_AMX_MOEConfig config) { + config_ = config; + gate_proj_ = config_.gate_proj; + up_proj_ = config_.up_proj; + down_proj_ = config_.down_proj; + + std::vector> m_mem_requests; + m_mem_requests.push_back({(void **)&m_local_input_, + sizeof(ggml_bf16_t) * config_.routed_expert_num * config_.max_len * config_.hidden_size}); + m_mem_requests.push_back({(void **)&m_local_gate_output_, sizeof(ggml_bf16_t) * config_.routed_expert_num * + config_.max_len * config_.intermediate_size}); + m_mem_requests.push_back({(void **)&m_local_up_output_, sizeof(ggml_bf16_t) * config_.routed_expert_num * + config_.max_len * config_.intermediate_size}); + m_mem_requests.push_back({(void **)&m_local_down_output_, + sizeof(ggml_bf16_t) * config_.routed_expert_num * config_.max_len * config_.hidden_size}); + std::vector gate_up_ba_ptr(config_.expert_num); + std::vector gate_bc_ptr(config_.expert_num); + std::vector up_bc_ptr(config_.expert_num); + std::vector down_ba_ptr(config_.expert_num); + std::vector down_bc_ptr(config_.expert_num); + for (int i = 0; i < config_.expert_num; i++) { + m_mem_requests.push_back( + {(void **)&gate_up_ba_ptr[i], T::BufferA::required_size(config_.max_len, config_.hidden_size)}); + m_mem_requests.push_back( + {(void **)&gate_bc_ptr[i], T::BufferC::required_size(config_.max_len, config_.intermediate_size)}); + m_mem_requests.push_back( + {(void **)&up_bc_ptr[i], T::BufferC::required_size(config_.max_len, config_.intermediate_size)}); + m_mem_requests.push_back( + {(void **)&down_ba_ptr[i], T::BufferA::required_size(config_.max_len, config_.intermediate_size)}); + m_mem_requests.push_back( + {(void **)&down_bc_ptr[i], T::BufferC::required_size(config_.max_len, config_.hidden_size)}); + } + + m_mem_requests.push_back({(void **)&gate_proj_t_, + sizeof(ggml_bf16_t) * config_.expert_num * config_.intermediate_size * config_.hidden_size}); + m_mem_requests.push_back({(void **)&up_proj_t_, + sizeof(ggml_bf16_t) * config_.expert_num * config_.intermediate_size * config_.hidden_size}); + m_mem_requests.push_back({(void **)&down_proj_t_, + sizeof(ggml_bf16_t) * config_.expert_num * config_.hidden_size * config_.intermediate_size}); + + m_mem_requests.push_back({(void **)&m_local_down_output_grad_, + sizeof(ggml_bf16_t) * config_.routed_expert_num * config_.max_len * config_.hidden_size}); + m_mem_requests.push_back({(void **)&m_local_down_input_grad_, + sizeof(ggml_bf16_t) * config_.routed_expert_num * config_.max_len * config_.intermediate_size}); + m_mem_requests.push_back({(void **)&m_local_gate_output_grad_, + sizeof(ggml_bf16_t) * config_.routed_expert_num * config_.max_len * config_.intermediate_size}); + m_mem_requests.push_back({(void **)&m_local_up_output_grad_, + sizeof(ggml_bf16_t) * config_.routed_expert_num * config_.max_len * config_.intermediate_size}); + m_mem_requests.push_back({(void **)&m_local_gate_input_grad_, + sizeof(ggml_bf16_t) * config_.routed_expert_num * config_.max_len * config_.hidden_size}); + m_mem_requests.push_back({(void **)&m_local_up_input_grad_, + sizeof(ggml_bf16_t) * config_.routed_expert_num * config_.max_len * config_.hidden_size}); + m_mem_requests.push_back({(void **)&m_local_token_indices_, + sizeof(int) * config_.routed_expert_num * config_.max_len}); + m_mem_requests.push_back({(void **)&m_local_expert_positions_, + sizeof(int) * config_.routed_expert_num * config_.max_len}); + std::vector gate_t_ba_ptr(config_.expert_num); + std::vector gate_t_bc_ptr(config_.expert_num); + std::vector up_t_ba_ptr(config_.expert_num); + std::vector up_t_bc_ptr(config_.expert_num); + std::vector down_t_ba_ptr(config_.expert_num); + std::vector down_t_bc_ptr(config_.expert_num); + for (int i = 0; i < config_.expert_num; i++) { + m_mem_requests.push_back( + {(void **)&gate_t_ba_ptr[i], T::BufferA::required_size(config_.max_len, config_.intermediate_size)}); + m_mem_requests.push_back( + {(void **)&gate_t_bc_ptr[i], T::BufferC::required_size(config_.max_len, config_.hidden_size)}); + m_mem_requests.push_back( + {(void **)&up_t_ba_ptr[i], T::BufferA::required_size(config_.max_len, config_.intermediate_size)}); + m_mem_requests.push_back( + {(void **)&up_t_bc_ptr[i], T::BufferC::required_size(config_.max_len, config_.hidden_size)}); + m_mem_requests.push_back( + {(void **)&down_t_ba_ptr[i], T::BufferA::required_size(config_.max_len, config_.hidden_size)}); + m_mem_requests.push_back( + {(void **)&down_t_bc_ptr[i], T::BufferC::required_size(config_.max_len, config_.intermediate_size)}); + } + + shared_mem_buffer.alloc(this, m_mem_requests); + + m_local_pos_.resize(config_.max_len); + for (int i = 0; i < config_.max_len; i++) { + m_local_pos_[i].resize(config_.routed_expert_num); + } + m_expert_id_map_.resize(config_.expert_num); + m_local_num_.resize(config_.expert_num); + m_local_input_ptr_.resize(config_.expert_num); + m_local_gate_output_ptr_.resize(config_.expert_num); + m_local_up_output_ptr_.resize(config_.expert_num); + m_local_down_output_ptr_.resize(config_.expert_num); + m_local_down_output_grad_ptr_.resize(config_.expert_num); + m_local_down_input_grad_ptr_.resize(config_.expert_num); + m_local_gate_output_grad_ptr_.resize(config_.expert_num); + m_local_up_output_grad_ptr_.resize(config_.expert_num); + m_local_gate_input_grad_ptr_.resize(config_.expert_num); + m_local_up_input_grad_ptr_.resize(config_.expert_num); + + for (uint64_t i = 0; i < config_.expert_num; i++) { + gate_up_ba_.push_back( + std::make_shared(config_.max_len, config_.hidden_size, gate_up_ba_ptr[i])); + gate_bc_.push_back( + std::make_shared(config_.max_len, config_.intermediate_size, gate_bc_ptr[i])); + up_bc_.push_back(std::make_shared(config_.max_len, config_.intermediate_size, up_bc_ptr[i])); + down_ba_.push_back( + std::make_shared(config_.max_len, config_.intermediate_size, down_ba_ptr[i])); + down_bc_.push_back(std::make_shared(config_.max_len, config_.hidden_size, down_bc_ptr[i])); + +#ifdef USE_NUMA + int numa_nodes = numa_num_configured_nodes(); + gate_bb_numa_.resize(numa_nodes); + up_bb_numa_.resize(numa_nodes); + down_bb_numa_.resize(numa_nodes); + for (int j = 0; j < numa_nodes; j++) { + void *gate_bb_ptr = + numa_alloc_aligned(T::BufferB::required_size(config_.intermediate_size, config_.hidden_size), j, 64); + gate_bb_numa_[j].push_back( + std::make_shared(config_.intermediate_size, config_.hidden_size, gate_bb_ptr)); + void *up_bb_ptr = + numa_alloc_aligned(T::BufferB::required_size(config_.intermediate_size, config_.hidden_size), j, 64); + up_bb_numa_[j].push_back( + std::make_shared(config_.intermediate_size, config_.hidden_size, up_bb_ptr)); + void *down_bb_ptr = + numa_alloc_aligned(T::BufferB::required_size(config_.hidden_size, config_.intermediate_size), j, 64); + down_bb_numa_[j].push_back( + std::make_shared(config_.hidden_size, config_.intermediate_size, down_bb_ptr)); + } +#else + void *gate_bb_ptr = + std::aligned_alloc(64, T::BufferB::required_size(config_.intermediate_size, config_.hidden_size)); + gate_bb_.push_back( + std::make_shared(config_.intermediate_size, config_.hidden_size, gate_bb_ptr)); + + void *up_bb_ptr = + std::aligned_alloc(64, T::BufferB::required_size(config_.intermediate_size, config_.hidden_size)); + up_bb_.push_back( + std::make_shared(config_.intermediate_size, config_.hidden_size, up_bb_ptr)); + + void *down_bb_ptr = + std::aligned_alloc(64, T::BufferB::required_size(config_.hidden_size, config_.intermediate_size)); + down_bb_.push_back( + std::make_shared(config_.hidden_size, config_.intermediate_size, down_bb_ptr)); +#endif + } + + for (uint64_t i = 0; i < config_.expert_num; i++) { + gate_t_ba_.push_back( + std::make_shared(config_.max_len, config_.intermediate_size, gate_t_ba_ptr[i])); + gate_t_bc_.push_back( + std::make_shared(config_.max_len, config_.hidden_size, gate_t_bc_ptr[i])); + up_t_ba_.push_back(std::make_shared(config_.max_len, config_.intermediate_size, up_t_ba_ptr[i])); + up_t_bc_.push_back(std::make_shared(config_.max_len, config_.hidden_size, up_t_bc_ptr[i])); + down_t_ba_.push_back( + std::make_shared(config_.max_len, config_.hidden_size, down_t_ba_ptr[i])); + down_t_bc_.push_back(std::make_shared(config_.max_len, config_.intermediate_size, down_t_bc_ptr[i])); + + // TODO: NUMA + void *gate_t_bb_ptr = + std::aligned_alloc(64, T::BufferB::required_size(config_.hidden_size, config_.intermediate_size)); + gate_t_bb_.push_back( + std::make_shared(config_.hidden_size, config_.intermediate_size, gate_t_bb_ptr)); + + void *up_t_bb_ptr = + std::aligned_alloc(64, T::BufferB::required_size(config_.hidden_size, config_.intermediate_size)); + up_t_bb_.push_back( + std::make_shared(config_.hidden_size, config_.intermediate_size, up_t_bb_ptr)); + + void *down_t_bb_ptr = + std::aligned_alloc(64, T::BufferB::required_size(config_.intermediate_size, config_.hidden_size)); + down_t_bb_.push_back( + std::make_shared(config_.intermediate_size, config_.hidden_size, down_t_bb_ptr)); + } + + m_local_token_indices_ptr_.resize(config_.expert_num); + m_local_expert_positions_ptr_.resize(config_.expert_num); + } + + ~SFT_AMX_MOE() { shared_mem_buffer.dealloc(this); } + + void transpose_expert(const void* src, void* dst, int R, int C, Backend* backend) { + backend->do_work_stealing_job( + config_.expert_num, nullptr, + [&](uint64_t expert_idx) { + for (int r = 0; r < R; ++r) { + for (int c = 0; c < C; ++c) { + memcpy( + (uint8_t*)dst + (expert_idx * R * C + (c * R + r)) * sizeof(ggml_bf16_t), + (uint8_t*)src + (expert_idx * R * C + (r * C + c)) * sizeof(ggml_bf16_t), + sizeof(ggml_bf16_t)); + } + } + }, + nullptr); + } + + void load_weights(Backend *backend) { + transpose_expert(config_.gate_proj, gate_proj_t_, config_.intermediate_size, config_.hidden_size, backend); + transpose_expert(config_.up_proj, up_proj_t_, config_.intermediate_size, config_.hidden_size, backend); + transpose_expert(config_.down_proj, down_proj_t_, config_.hidden_size, config_.intermediate_size, backend); + + int nth = T::recommended_nth(config_.intermediate_size); + backend->do_work_stealing_job( + nth * config_.expert_num, nullptr, + [&](int task_id) { + uint64_t expert_idx = task_id / nth; + int ith = task_id % nth; +#ifdef USE_NUMA + int numa_nodes = numa_num_configured_nodes(); + for (int j = 0; j < numa_nodes; j++) { + gate_bb_numa_[j][expert_idx]->from_mat((ggml_bf16_t *)config_.gate_proj + + expert_idx * config_.intermediate_size * config_.hidden_size, + ith, nth); + up_bb_numa_[j][expert_idx]->from_mat((ggml_bf16_t *)config_.up_proj + + expert_idx * config_.intermediate_size * config_.hidden_size, + ith, nth); + } +#else + gate_bb_[expert_idx]->from_mat((ggml_bf16_t *)config_.gate_proj + + expert_idx * config_.intermediate_size * config_.hidden_size, + ith, nth); + up_bb_[expert_idx]->from_mat( + (ggml_bf16_t *)config_.up_proj + expert_idx * config_.intermediate_size * config_.hidden_size, ith, nth); +#endif + }, + nullptr); + nth = T::recommended_nth(config_.intermediate_size); + backend->do_work_stealing_job( + nth * config_.expert_num, nullptr, + [&](int task_id) { + uint64_t expert_idx = task_id / nth; + int ith = task_id % nth; + down_t_bb_[expert_idx]->from_mat((ggml_bf16_t *)down_proj_t_ + + expert_idx * config_.intermediate_size * config_.hidden_size, + ith, nth); + }, + nullptr); + + // if constexpr (std::is_same_v) { + // // 确保 debug/ 目录存在 + // std::filesystem::create_directories("debug"); + + // int tail_cols = 1024; + // for (int expert_idx = 0; expert_idx < config_.expert_num; ++expert_idx) { + // auto buf = down_t_bb_[expert_idx].get(); + // std::cout << "k: " << buf->k << "; n: " << buf->n << std::endl; + // // 打开对应 expert 的文件 + // std::string path = "debug/" + std::to_string(expert_idx) + "_down_bb_t_debug.txt"; + // std::ofstream ofs(path, std::ios::out); + // if (!ofs) { + // std::cerr << "Failed to open " << path << " for writing\n"; + // continue; + // } + + // ofs << "==== Expert " << expert_idx << " ====\n"; + // for (int n_idx = 0; n_idx < buf->k; ++n_idx) { + // // 明确当作 int8 读 + // const int8_t* row = reinterpret_cast(buf->b) + n_idx * buf->n; + + // // 写整行 + // ofs << "row[" << n_idx << "] = { " + // << int8_row_to_string(row, buf->n) + // << " }\n"; + // } + + // ofs.close(); + // } + // } + + // if constexpr (std::is_same_v) { + // for (int expert_idx = 0; expert_idx < config_.expert_num; ++expert_idx) { + // auto buf = down_t_bb_[expert_idx].get(); + + // // 打开对应 expert 的文件 + // std::string path = "debug/" + std::to_string(expert_idx) + "_down_bb_t_debug3.bin"; + // std::ofstream ofs(path, std::ios::binary); + // for (int n_idx = 0; n_idx < buf->k; ++n_idx) { + // const int8_t* row = reinterpret_cast(buf->b) + n_idx * buf->n; + // for (int j = 0; j < buf->n; ++j) { + // float v = row[j]; + // ofs.write(reinterpret_cast(&v), sizeof(v)); + // } + // } + // ofs.close(); + // } + // } + + + // for (uint64_t expert_idx = 0; expert_idx < (uint64_t)config_.expert_num; ++expert_idx) { + // // 打开对应 expert 的文件 + // std::string path = "debug/" + std::to_string(expert_idx) + "_up_proj_t.bin"; + // std::ofstream ofs(path, std::ios::binary); + // std::cout << "config_.hidden_size: " << config_.hidden_size << std::endl; + // std::cout << "config_.intermediate_size: " << config_.intermediate_size << std::endl; + // for (int n_idx = 0; n_idx < config_.intermediate_size; ++n_idx) { + // const int8_t* row = reinterpret_cast(config_.down_proj + expert_idx * n_idx * config_.hidden_size); + // for (int j = 0; j < config_.hidden_size; ++j) { + // float v = row[j]; + // ofs.write(reinterpret_cast(&v), sizeof(v)); + // } + // } + // ofs.close(); + // } + + nth = T::recommended_nth(config_.hidden_size); + backend->do_work_stealing_job( + nth * config_.expert_num, nullptr, + [&](int task_id) { + uint64_t expert_idx = task_id / nth; + int ith = task_id % nth; +#ifdef USE_NUMA + int numa_nodes = numa_num_configured_nodes(); + for (int j = 0; j < numa_nodes; j++) { + down_bb_numa_[j][expert_idx]->from_mat((ggml_bf16_t *)config_.down_proj + + expert_idx * config_.hidden_size * config_.intermediate_size, + ith, nth); + } +#else + down_bb_[expert_idx]->from_mat((ggml_bf16_t *)config_.down_proj + + expert_idx * config_.hidden_size * config_.intermediate_size, + ith, nth); +#endif + up_t_bb_[expert_idx]->from_mat((ggml_bf16_t *)up_proj_t_ + + expert_idx * config_.hidden_size * config_.intermediate_size, + ith, nth); + gate_t_bb_[expert_idx]->from_mat((ggml_bf16_t *)gate_proj_t_ + + expert_idx * config_.hidden_size * config_.intermediate_size, + ith, nth); + }, + nullptr); + } + + void warm_up(Backend *backend) {} + + void forward(int qlen, int k, const uint64_t *expert_ids, const float *weights, const void *input, void *output, Backend *backend) { + bool use_amx = (qlen > 4 * config_.expert_num / config_.routed_expert_num); + int activated_expert = 0; + for (int i = 0; i < config_.expert_num; i++) { + m_local_num_[i] = 0; + } + for (int i = 0; i < qlen; i++) { + for (int j = 0; j < k; j++) { + m_local_pos_[i][j] = m_local_num_[expert_ids[i * k + j]]++; + } + } + for (int i = 0; i < config_.expert_num; i++) { + if (m_local_num_[i] > 0) { + m_expert_id_map_[activated_expert] = i; + activated_expert++; + } + } + uint64_t offset = 0; + for (int i = 0; i < config_.expert_num; i++) { + m_local_input_ptr_[i] = m_local_input_ + offset * config_.hidden_size; + m_local_gate_output_ptr_[i] = m_local_gate_output_ + offset * config_.intermediate_size; + m_local_up_output_ptr_[i] = m_local_up_output_ + offset * config_.intermediate_size; + m_local_down_output_ptr_[i] = m_local_down_output_ + offset * config_.hidden_size; + offset += m_local_num_[i]; + } + backend->do_work_stealing_job( + qlen, nullptr, + [&](int i) { + for (int j = 0; j < k; j++) { + memcpy(m_local_input_ptr_[expert_ids[i * k + j]] + m_local_pos_[i][j] * config_.hidden_size, + (ggml_bf16_t *)input + i * config_.hidden_size, sizeof(ggml_bf16_t) * config_.hidden_size); + } + }, + nullptr); + backend->do_work_stealing_job( + activated_expert, nullptr, + [&](int task_id) { + int expert_idx = m_expert_id_map_[task_id]; + + gate_up_ba_[expert_idx]->from_mat(m_local_num_[expert_idx], m_local_input_ptr_[expert_idx], 0, 1); + }, + nullptr); + int nth = T::recommended_nth(config_.intermediate_size); + backend->do_work_stealing_job( + nth * activated_expert, [&](int _) { T::config(); }, + [&](int task_id) { + int expert_idx = m_expert_id_map_[task_id / nth]; + int ith = task_id % nth; +#ifdef USE_NUMA + amx::mat_mul(m_local_num_[expert_idx], config_.intermediate_size, config_.hidden_size, + gate_up_ba_[expert_idx], gate_bb_numa_[Backend::numa_node][expert_idx], gate_bc_[expert_idx], + ith, nth, use_amx); + amx::mat_mul(m_local_num_[expert_idx], config_.intermediate_size, config_.hidden_size, + gate_up_ba_[expert_idx], up_bb_numa_[Backend::numa_node][expert_idx], up_bc_[expert_idx], ith, + nth, use_amx); +#else + amx::mat_mul(m_local_num_[expert_idx], config_.intermediate_size, config_.hidden_size, + gate_up_ba_[expert_idx], gate_bb_[expert_idx], gate_bc_[expert_idx], ith, nth, use_amx); + amx::mat_mul(m_local_num_[expert_idx], config_.intermediate_size, config_.hidden_size, + gate_up_ba_[expert_idx], up_bb_[expert_idx], up_bc_[expert_idx], ith, nth, use_amx); +#endif + gate_bc_[expert_idx]->to_mat(m_local_num_[expert_idx], m_local_gate_output_ptr_[expert_idx], ith, nth); + up_bc_[expert_idx]->to_mat(m_local_num_[expert_idx], m_local_up_output_ptr_[expert_idx], ith, nth); + auto [n_start, n_end] = T::split_range_n(config_.intermediate_size, ith, nth); + for (int i = 0; i < m_local_num_[expert_idx]; i++) { + ggml_bf16_t *gate_output_ptr = &m_local_gate_output_ptr_[expert_idx][i * config_.intermediate_size]; + ggml_bf16_t *up_output_ptr = &m_local_up_output_ptr_[expert_idx][i * config_.intermediate_size]; + for (int j = n_start; j < n_end; j += 32) { + __m512 gate_val0, gate_val1, up_val0, up_val1; + avx512_32xbf16_to_32xfp32((__m512i *)(gate_output_ptr + j), &gate_val0, &gate_val1); + avx512_32xbf16_to_32xfp32((__m512i *)(up_output_ptr + j), &up_val0, &up_val1); + __m512 result0 = act_fn(gate_val0, up_val0); + __m512 result1 = act_fn(gate_val1, up_val1); + avx512_32xfp32_to_32xbf16(&result0, &result1, (__m512i *)(gate_output_ptr + j)); + } + } + }, + nullptr); + backend->do_work_stealing_job( + activated_expert, nullptr, + [&](int task_id) { + int expert_idx = m_expert_id_map_[task_id]; + down_ba_[expert_idx]->from_mat(m_local_num_[expert_idx], m_local_gate_output_ptr_[expert_idx], 0, 1); + }, + nullptr); + + // for (uint64_t expert_idx = 0; expert_idx < (uint64_t)config_.expert_num; ++expert_idx) { + // dump_grad_bin("cpp_layer0_E_End"+std::to_string(expert_idx)+"_down_ba_ori_", (ggml_bf16_t*)m_local_gate_output_ptr_[expert_idx], config_.intermediate_size * m_local_num_[expert_idx], GGML_TYPE_BF16); + // } + // if constexpr (std::is_same_v) { + // std::cout << "GO INTO forward output" << std::endl; + // // 确保 debug/ 目录存在 + // std::filesystem::create_directories("debug"); + + // for (int expert_idx = 0; expert_idx < config_.expert_num; ++expert_idx) { + // auto buf = down_ba_[expert_idx].get(); + // // std::cout << "k: " << buf->k << "; n: " << buf->n << std::endl; + // // 打开对应 expert 的文件 + // std::string path = "debug/" + std::to_string(expert_idx) + "_down_ba_debug.txt"; + // std::ofstream ofs(path, std::ios::out); + // if (!ofs) { + // std::cerr << "Failed to open " << path << " for writing\n"; + // continue; + // } + + // ofs << "==== Expert " << expert_idx << " ====\n"; + // ofs << "buf_k: " << buf->k << "\n"; + // for (int n_idx = 0; n_idx < m_local_num_[expert_idx]; ++n_idx) { + // // 明确当作 bfloat16 读 + // const int8_t* row = reinterpret_cast(buf->a) + n_idx * buf->k; + + // // 写整行 + // ofs << "row[" << n_idx << "] = { " + // << int8_row_to_string(row, buf->k) + // << " }\n"; + // } + + // ofs.close(); + // } + // std::cout << "OUT INTO forward output" << std::endl; + // } + nth = T::recommended_nth(config_.hidden_size); + backend->do_work_stealing_job( + nth * activated_expert, [&](int _) { T::config(); }, + [&](int task_id) { + int expert_idx = m_expert_id_map_[task_id / nth]; + int ith = task_id % nth; +#ifdef USE_NUMA + amx::mat_mul(m_local_num_[expert_idx], config_.hidden_size, config_.intermediate_size, down_ba_[expert_idx], + down_bb_numa_[Backend::numa_node][expert_idx], down_bc_[expert_idx], ith, nth, use_amx); +#else + amx::mat_mul(m_local_num_[expert_idx], config_.hidden_size, config_.intermediate_size, down_ba_[expert_idx], + down_bb_[expert_idx], down_bc_[expert_idx], ith, nth, use_amx); +#endif + down_bc_[expert_idx]->to_mat(m_local_num_[expert_idx], m_local_down_output_ptr_[expert_idx], ith, nth); + }, + nullptr); + backend->do_work_stealing_job( + qlen, nullptr, + [&](int i) { + for (int e = 0; e < config_.hidden_size; e += 32) { + __m512 x0 = _mm512_setzero_ps(); + __m512 x1 = _mm512_setzero_ps(); + for (int j = 0; j < k; j++) { + __m512 weight = _mm512_set1_ps(weights[i * k + j]); + __m512 down_output0, down_output1; + avx512_32xbf16_to_32xfp32((__m512i *)(m_local_down_output_ptr_[expert_ids[i * k + j]] + + m_local_pos_[i][j] * config_.hidden_size + e), + &down_output0, &down_output1); + x0 = _mm512_fmadd_ps(down_output0, weight, x0); + x1 = _mm512_fmadd_ps(down_output1, weight, x1); + } + avx512_32xfp32_to_32xbf16(&x0, &x1, (__m512i *)((ggml_bf16_t *)output + i * config_.hidden_size + e)); + } + }, + nullptr); + } + + void backward(int qlen, int k, const uint64_t *expert_ids, const float *weights, const void* input, const void *output_grad, void *input_grad, Backend *backend) { + bool use_amx = (qlen > 4 * config_.expert_num / config_.routed_expert_num); + int activated_expert = 0; + for (int i = 0; i < config_.expert_num; i++) { + m_local_num_[i] = 0; + } + for (int i = 0; i < qlen; i++) { + for (int j = 0; j < k; j++) { + m_local_pos_[i][j] = m_local_num_[expert_ids[i * k + j]]++; + } + } + for (int i = 0; i < config_.expert_num; i++) { + if (m_local_num_[i] > 0) { + m_expert_id_map_[activated_expert] = i; + activated_expert++; + } + } + uint64_t offset = 0; + for (int i = 0; i < config_.expert_num; i++) { + m_local_input_ptr_[i] = m_local_input_ + offset * config_.hidden_size; + m_local_gate_output_ptr_[i] = m_local_gate_output_ + offset * config_.intermediate_size; + m_local_up_output_ptr_[i] = m_local_up_output_ + offset * config_.intermediate_size; + + m_local_down_output_grad_ptr_[i] = m_local_down_output_grad_ + offset * config_.hidden_size; + m_local_down_input_grad_ptr_[i] = m_local_down_input_grad_ + offset * config_.intermediate_size; + m_local_gate_output_grad_ptr_[i] = m_local_gate_output_grad_ + offset * config_.intermediate_size; + m_local_up_output_grad_ptr_[i] = m_local_up_output_grad_ + offset * config_.intermediate_size; + m_local_gate_input_grad_ptr_[i] = m_local_gate_input_grad_ + offset * config_.hidden_size; + m_local_up_input_grad_ptr_[i] = m_local_up_input_grad_ + offset * config_.hidden_size; + m_local_token_indices_ptr_[i] = m_local_token_indices_ + offset; + m_local_expert_positions_ptr_[i] = m_local_expert_positions_ + offset; + offset += m_local_num_[i]; + } + + // TODO: cache + backend->do_work_stealing_job( + qlen, nullptr, + [&](int i) { + for (int j = 0; j < k; j++) { + uint64_t expert_id = expert_ids[i * k + j]; + int local_row = m_local_pos_[i][j]; + memcpy(m_local_input_ptr_[expert_id] + local_row * config_.hidden_size, + (ggml_bf16_t *)input + i * config_.hidden_size, sizeof(ggml_bf16_t) * config_.hidden_size); // TODO: cache + memcpy(m_local_down_output_grad_ptr_[expert_id] + local_row * config_.hidden_size, + (ggml_bf16_t *)output_grad + i * config_.hidden_size, sizeof(ggml_bf16_t) * config_.hidden_size); + m_local_token_indices_ptr_[expert_id][local_row] = i; + m_local_expert_positions_ptr_[expert_id][local_row] = j; + } + }, + nullptr); + + backend->do_work_stealing_job( + activated_expert, nullptr, + [&](int task_id) { + int expert_idx = m_expert_id_map_[task_id]; + gate_up_ba_[expert_idx]->from_mat(m_local_num_[expert_idx], m_local_input_ptr_[expert_idx], 0, 1); // TODO: cache + down_t_ba_[expert_idx]->from_mat(m_local_num_[expert_idx], m_local_down_output_grad_ptr_[expert_idx], 0, 1); + }, + nullptr); + + // // for debug + // for (uint64_t expert_idx = 0; expert_idx < (uint64_t)config_.expert_num; ++expert_idx) { + // dump_grad_bin("cpp_layer0_E_End"+std::to_string(expert_idx)+"_down_output_grad_", (ggml_bf16_t*)m_local_down_output_grad_ptr_[expert_idx], config_.hidden_size * m_local_num_[expert_idx], GGML_TYPE_BF16); + // } + + // if constexpr (std::is_same_v) { + // // 确保 debug/ 目录存在 + // std::filesystem::create_directories("debug"); + + // for (int expert_idx = 0; expert_idx < config_.expert_num; ++expert_idx) { + // auto buf = down_t_ba_[expert_idx].get(); + // // std::cout << "k: " << buf->k << "; n: " << buf->n << std::endl; + // // 打开对应 expert 的文件 + // std::string path = "debug/" + std::to_string(expert_idx) + "_down_ba_t_debug.txt"; + // std::ofstream ofs(path, std::ios::out); + // if (!ofs) { + // std::cerr << "Failed to open " << path << " for writing\n"; + // continue; + // } + + // ofs << "==== Expert " << expert_idx << " ====\n"; + // for (int n_idx = 0; n_idx < m_local_num_[expert_idx]; ++n_idx) { + // // 明确当作 bfloat16 读 + // const int8_t* row = reinterpret_cast(buf->a) + n_idx * buf->k; + + // // 写整行 + // ofs << "row[" << n_idx << "] = { " + // << int8_row_to_string(row, buf->k) + // << " }\n"; + // } + + // ofs.close(); + // } + // } + + int nth = T::recommended_nth(config_.intermediate_size); + backend->do_work_stealing_job( + nth * activated_expert, [&](int _) { T::config(); }, + [&](int task_id) { + int expert_idx = m_expert_id_map_[task_id / nth]; + int ith = task_id % nth; + + // // TODO: cache + amx::mat_mul(m_local_num_[expert_idx], config_.intermediate_size, config_.hidden_size, + gate_up_ba_[expert_idx], gate_bb_[expert_idx], gate_bc_[expert_idx], ith, nth, use_amx); + amx::mat_mul(m_local_num_[expert_idx], config_.intermediate_size, config_.hidden_size, + gate_up_ba_[expert_idx], up_bb_[expert_idx], up_bc_[expert_idx], ith, nth, use_amx); + gate_bc_[expert_idx]->to_mat(m_local_num_[expert_idx], m_local_gate_output_ptr_[expert_idx], ith, nth); + up_bc_[expert_idx]->to_mat(m_local_num_[expert_idx], m_local_up_output_ptr_[expert_idx], ith, nth); + + amx::mat_mul(m_local_num_[expert_idx], config_.intermediate_size, config_.hidden_size, + down_t_ba_[expert_idx], down_t_bb_[expert_idx], down_t_bc_[expert_idx], ith, nth, use_amx); + down_t_bc_[expert_idx]->to_mat(m_local_num_[expert_idx], m_local_down_input_grad_ptr_[expert_idx], ith, nth); + + + auto [n_start, n_end] = T::split_range_n(config_.intermediate_size, ith, nth); + for (int i = 0; i < m_local_num_[expert_idx]; i++) { + ggml_bf16_t *gate_output_ptr = &m_local_gate_output_ptr_[expert_idx][i * config_.intermediate_size]; + ggml_bf16_t *up_output_ptr = &m_local_up_output_ptr_[expert_idx][i * config_.intermediate_size]; + ggml_bf16_t *down_input_grad_ptr = &m_local_down_input_grad_ptr_[expert_idx][i * config_.intermediate_size]; + ggml_bf16_t *gate_output_grad_ptr = &m_local_gate_output_grad_ptr_[expert_idx][i * config_.intermediate_size]; + ggml_bf16_t *up_output_grad_ptr = &m_local_up_output_grad_ptr_[expert_idx][i * config_.intermediate_size]; + + int token_idx = m_local_token_indices_ptr_[expert_idx][i]; + int expert_pos = m_local_expert_positions_ptr_[expert_idx][i]; + __m512 weight = _mm512_set1_ps(weights[token_idx * k + expert_pos]); + + for (int j = n_start; j < n_end; j += 32) { + __m512 gate_val0, gate_val1, up_val0, up_val1, down_input_grad0, down_input_grad1; + avx512_32xbf16_to_32xfp32((__m512i *)(gate_output_ptr + j), &gate_val0, &gate_val1); + avx512_32xbf16_to_32xfp32((__m512i *)(up_output_ptr + j), &up_val0, &up_val1); + avx512_32xbf16_to_32xfp32((__m512i *)(down_input_grad_ptr + j), &down_input_grad0, &down_input_grad1); + + down_input_grad0 = _mm512_mul_ps(down_input_grad0, weight); + down_input_grad1 = _mm512_mul_ps(down_input_grad1, weight); + + // gate_output_grad = δ_zji ⊙ v_ji ⊙ σ'(u_ji) + __m512 gate_grad0 = _mm512_mul_ps(down_input_grad0, + _mm512_mul_ps(up_val0, act_fn_grad(gate_val0))); + __m512 gate_grad1 = _mm512_mul_ps(down_input_grad1, + _mm512_mul_ps(up_val1, act_fn_grad(gate_val1))); + + // up_output_grad = δ_zji ⊙ σ(u_ji) + __m512 up_grad0 = _mm512_mul_ps(down_input_grad0, act_fn_1(gate_val0)); + __m512 up_grad1 = _mm512_mul_ps(down_input_grad1, act_fn_1(gate_val1)); + + avx512_32xfp32_to_32xbf16(&gate_grad0, &gate_grad1, (__m512i *)(gate_output_grad_ptr + j)); + avx512_32xfp32_to_32xbf16(&up_grad0, &up_grad1, (__m512i *)(up_output_grad_ptr + j)); + } + } + }, + nullptr); + + // for debug + // if constexpr (std::is_same_v) { + // for (int expert_idx = 0; expert_idx < config_.expert_num; ++expert_idx) { + // auto buf = down_t_ba_[expert_idx].get(); + + // // 打开对应 expert 的文件 + // std::string path = "debug/" + std::to_string(expert_idx) + "_down_ba_t_debug3.bin"; + // std::ofstream ofs(path, std::ios::binary); + // for (int n_idx = 0; n_idx < m_local_num_[expert_idx]; ++n_idx) { + // const ggml_bf16_t* row = reinterpret_cast(buf->a) + n_idx * buf->k; + // for (int j = 0; j < buf->k; ++j) { + // float v = row[j]; + // ofs.write(reinterpret_cast(&v), sizeof(v)); + // } + // } + // ofs.close(); + // } + // } + + // for (uint64_t expert_idx = 0; expert_idx < (uint64_t)config_.expert_num; ++expert_idx) { + // dump_grad_bin("cpp_layer0_E_End"+std::to_string(expert_idx)+"_down_t_ba_", (ggml_bf16_t*)m_local_down_output_grad_ptr_[expert_idx], config_.hidden_size * m_local_num_[expert_idx], GGML_TYPE_BF16); + // } + + // for (uint64_t expert_idx = 0; expert_idx < (uint64_t)config_.expert_num; ++expert_idx) { + // dump_grad_bin("cpp_layer0_E_End"+std::to_string(expert_idx)+"_down_t_bb_", (ggml_bf16_t *)down_proj_t_ + expert_idx * config_.intermediate_size * config_.hidden_size, config_.hidden_size * config_.intermediate_size, GGML_TYPE_BF16); + // } + + // for (uint64_t expert_idx = 0; expert_idx < (uint64_t)config_.expert_num; ++expert_idx) { + // dump_grad_bin("cpp_layer0_E_End"+std::to_string(expert_idx)+"_down_t_bc_", (ggml_bf16_t*)m_local_down_input_grad_ptr_[expert_idx], config_.intermediate_size * m_local_num_[expert_idx], GGML_TYPE_BF16); + // } + + backend->do_work_stealing_job( + activated_expert, nullptr, + [&](int task_id) { + int expert_idx = m_expert_id_map_[task_id]; + gate_t_ba_[expert_idx]->from_mat(m_local_num_[expert_idx], m_local_gate_output_grad_ptr_[expert_idx], 0, 1); + up_t_ba_[expert_idx]->from_mat(m_local_num_[expert_idx], m_local_up_output_grad_ptr_[expert_idx], 0, 1); + }, + nullptr); + nth = T::recommended_nth(config_.hidden_size); + backend->do_work_stealing_job( + nth * activated_expert, [&](int _) { T::config(); }, + [&](int task_id) { + int expert_idx = m_expert_id_map_[task_id / nth]; + int ith = task_id % nth; + amx::mat_mul(m_local_num_[expert_idx], config_.hidden_size, config_.intermediate_size, + gate_t_ba_[expert_idx], gate_t_bb_[expert_idx], gate_t_bc_[expert_idx], ith, nth, use_amx); + amx::mat_mul(m_local_num_[expert_idx], config_.hidden_size, config_.intermediate_size, + up_t_ba_[expert_idx], up_t_bb_[expert_idx], up_t_bc_[expert_idx], ith, nth, use_amx); + gate_t_bc_[expert_idx]->to_mat(m_local_num_[expert_idx], m_local_gate_input_grad_ptr_[expert_idx], ith, nth); + up_t_bc_[expert_idx]->to_mat(m_local_num_[expert_idx], m_local_up_input_grad_ptr_[expert_idx], ith, nth); + }, + nullptr); + backend->do_work_stealing_job( + qlen, nullptr, + [&](int i) { + for (int e = 0; e < config_.hidden_size; e += 32) { + __m512 x0 = _mm512_setzero_ps(); + __m512 x1 = _mm512_setzero_ps(); + for (int j = 0; j < k; j++) { + __m512 gate_input_grad0, gate_input_grad1, up_input_grad0, up_input_grad1; + avx512_32xbf16_to_32xfp32((__m512i *)(m_local_gate_input_grad_ptr_[expert_ids[i * k + j]] + + m_local_pos_[i][j] * config_.hidden_size + e), + &gate_input_grad0, &gate_input_grad1); + avx512_32xbf16_to_32xfp32((__m512i *)(m_local_up_input_grad_ptr_[expert_ids[i * k + j]] + + m_local_pos_[i][j] * config_.hidden_size + e), + &up_input_grad0, &up_input_grad1); + x0 = _mm512_add_ps(gate_input_grad0, x0); + x1 = _mm512_add_ps(gate_input_grad1, x1); + x0 = _mm512_add_ps(up_input_grad0, x0); + x1 = _mm512_add_ps(up_input_grad1, x1); + } + avx512_32xfp32_to_32xbf16(&x0, &x1, (__m512i *)((ggml_bf16_t *)input_grad + i * config_.hidden_size + e)); + } + }, + nullptr); + } +}; +#endif \ No newline at end of file diff --git a/archive/kt-sft/csrc/ktransformers_ext/operators/amx/debug_tools_sft_moe.hpp b/archive/kt-sft/csrc/ktransformers_ext/operators/amx/debug_tools_sft_moe.hpp new file mode 100644 index 0000000..cb11fc1 --- /dev/null +++ b/archive/kt-sft/csrc/ktransformers_ext/operators/amx/debug_tools_sft_moe.hpp @@ -0,0 +1,77 @@ +#ifndef SFT_DEBUG_HPP +#define SFT_DEBUG_HPP + +#include +#include +#include +#include +#include + +inline std::string get_env_or_default(const char *var_name, const std::string &default_value) { + const char *value = std::getenv(var_name); + return (value != nullptr) ? std::string(value) : default_value; +} + +/* use example: + for (uint64_t expert_idx = 0; expert_idx < (uint64_t)config_.expert_num; ++expert_idx) { + dump_grad_bin("layer0_E_End"+std::to_string(expert_idx)+"_gate_proj_out_trans_", (uint8_t*)gate_proj_t_ + expert_idx * config_.hidden_size * config_.intermediate_size * ggml_type_size(config_.grad_type), config_.hidden_size * config_.intermediate_size, config_.grad_type); + std::cout << "gate_proj_t_:" << static_cast((uint8_t*)gate_proj_t_ + expert_idx * config_.hidden_size * config_.intermediate_size) << ", grad_type: " << config_.grad_type << std::endl; + } +*/ +inline void dump_grad_bin(const std::string &file_name, + const void *data, + size_t elem_cnt, + ggml_type dtype, + std::streamoff offset_bytes = 0) +{ + std::string path = get_env_or_default("SFT_DEBUG_PATH","debug") + "/" + file_name; + switch (dtype) { + case GGML_TYPE_F32: path += ".f32"; break; + case GGML_TYPE_F16: path += ".f16"; break; + case GGML_TYPE_BF16: path += ".bf16"; break; + case GGML_TYPE_I8: path += ".int8"; break; + default: path += ".raw"; break; + } + std::fstream f(path, std::ios::in | std::ios::out | std::ios::binary); + if (!f.is_open()) { + std::ofstream tmp(path, std::ios::out | std::ios::binary); + tmp.close(); + f.open(path, std::ios::in | std::ios::out | std::ios::binary); + } + + f.seekp(offset_bytes * ggml_type_size(dtype)); + // std::cout << "seekp: " << offset_bytes * ggml_type_size(dtype) << std::endl; + + f.write(reinterpret_cast(data), static_cast(elem_cnt * ggml_type_size(dtype))); + f.close(); +} + +// inline void dump_bin(std::string file_name, float16_t *data, size_t count) { +// file_name = get_env_or_default("SFT_DEBUG_PATH", "debug") + "/" + file_name + ".f16"; +// std::ofstream f(file_name, std::ios::binary); +// f.write(reinterpret_cast(data), count * sizeof(*data)); +// f.close(); +// } +inline void dump_bin(std::string file_name, float *data, size_t count) { + file_name = get_env_or_default("SFT_DEBUG_PATH", "debug") + "/" + file_name + ".f32"; + std::cout << file_name << std::endl; + std::ofstream f(file_name, std::ios::binary); + f.write(reinterpret_cast(data), count * sizeof(*data)); + f.close(); +} +inline void dump_bin(std::string file_name, int64_t *data, size_t count) { + file_name = get_env_or_default("SFT_DEBUG_PATH", "debug") + "/" + file_name + ".int64"; + std::cout << file_name << std::endl; + std::ofstream f(file_name, std::ios::binary); + f.write(reinterpret_cast(data), count * sizeof(*data)); + f.close(); +} +inline void dump_bin(std::string file_name, uint8_t *data, size_t count) { + file_name = get_env_or_default("SFT_DEBUG_PATH", "debug") + "/" + file_name + ".uint8"; + std::cout << file_name << std::endl; + std::ofstream f(file_name, std::ios::binary); + f.write(reinterpret_cast(data), count * sizeof(*data)); + f.close(); +} + +#endif \ No newline at end of file diff --git a/archive/kt-sft/csrc/ktransformers_ext/operators/amx/la/amx.hpp b/archive/kt-sft/csrc/ktransformers_ext/operators/amx/la/amx.hpp new file mode 100644 index 0000000..9745817 --- /dev/null +++ b/archive/kt-sft/csrc/ktransformers_ext/operators/amx/la/amx.hpp @@ -0,0 +1,976 @@ +/** + * @Description : + * @Author : chenht2022 + * @Date : 2025-04-25 18:28:12 + * @Version : 1.0.0 + * @LastEditors : chenht2022 + * @LastEditTime : 2025-04-25 18:28:12 + * @Copyright (c) 2024 by KVCache.AI, All Rights Reserved. + **/ + +#pragma once +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "utils.hpp" +#include + +#if (defined(_WIN32) || defined(_WIN64)) +#define RESTRICT __restrict +#else +#define RESTRICT __restrict__ +#endif + +#if (defined(_WIN32) || defined(_WIN64)) +#define ALWAYS_INLINE __forceinline +#elif __has_attribute(always_inline) || defined(__GNUC__) +#define ALWAYS_INLINE __attribute__((__always_inline__)) inline +#else +#define ALWAYS_INLINE inline +#endif + +namespace amx { + +#define ARCH_GET_XCOMP_PERM 0x1022 +#define ARCH_REQ_XCOMP_PERM 0x1023 +#define XFEATURE_XTILECFG 17 +#define XFEATURE_XTILEDATA 18 + +const int TMMCount = 8; +const int MaxTileHeight = 16; +const int MaxTileWidth = 64; + +const int AMX_BLK_SIZE = 32; + +#define TMM0 0 +#define TMM1 1 +#define TMM2 2 +#define TMM3 3 +#define TMM4 4 +#define TMM5 5 +#define TMM6 6 +#define TMM7 7 + +inline bool enable_amx() { + static thread_local bool initialized = false; + if (initialized) { + return true; + } + initialized = true; + + if (syscall(SYS_arch_prctl, ARCH_REQ_XCOMP_PERM, XFEATURE_XTILEDATA)) { + printf("\n Fail to do XFEATURE_XTILEDATA \n\n"); + return false; + } else { + // printf("\n TILE DATA USE SET - OK \n\n"); + return true; + } + return true; +} + +struct alignas(64) TileConfig { + uint8_t palette; + uint8_t start_row; + std::array __0 = {}; + std::array colsb; + std::array __1 = {}; + std::array rows; + std::array __2 = {}; + + TileConfig() { + palette = 1; + start_row = 0; + for (int i = 0; i < 8; i++) { + set_row_col(i, 0, 0); + } + } + + void set_row_col(int i, uint8_t row, uint16_t col) { + colsb[i] = col; + rows[i] = row; + } + + void set_config() { _tile_loadconfig(this); } + + static void load_data(int to, void *from, size_t stride) { + switch (to) { + case 0: + _tile_loadd(0, from, stride); + break; + case 1: + _tile_loadd(1, from, stride); + break; + case 2: + _tile_loadd(2, from, stride); + break; + case 3: + _tile_loadd(3, from, stride); + break; + case 4: + _tile_loadd(4, from, stride); + break; + case 5: + _tile_loadd(5, from, stride); + break; + case 6: + _tile_loadd(6, from, stride); + break; + case 7: + _tile_loadd(7, from, stride); + break; + default: + throw std::runtime_error("no such tile"); + } + } + + static void store_data(int from, void *to, size_t stride) { + switch (from) { + case 0: + _tile_stored(0, to, stride); + break; + case 1: + _tile_stored(1, to, stride); + break; + case 2: + _tile_stored(2, to, stride); + break; + case 3: + _tile_stored(3, to, stride); + break; + case 4: + _tile_stored(4, to, stride); + break; + case 5: + _tile_stored(5, to, stride); + break; + case 6: + _tile_stored(6, to, stride); + break; + case 7: + _tile_stored(7, to, stride); + break; + default: + throw std::runtime_error("no such tile"); + } + } +}; + +static_assert(sizeof(TileConfig) == 64); + +inline void debug_tile(int t) { + printf("Tile %d\n", t); + uint8_t data[16][64] = {}; + TileConfig::store_data(t, data, 64); + for (int i = 0; i < 16; i++) { + for (int j = 0; j < 64; j++) { + printf("%3d ", data[i][j]); + } + printf("\n"); + } + printf("\n"); +} + +inline void debug_tiles(int to = 8) { + for (int i = 0; i < to; i++) { + debug_tile(i); + } +} + +inline void debug_m512(__m512 x) { + float data[16]; + _mm512_storeu_ps(data, x); + for (int i = 0; i < 16; i++) { + printf("%f ", data[i]); + } + printf("\n"); +} + +// transpose utils +inline void transpose_16x16_32bit(__m512i *v) { + __m512i v1[16]; + v1[0] = _mm512_unpacklo_epi32(v[0], v[1]); + v1[1] = _mm512_unpackhi_epi32(v[0], v[1]); + v1[2] = _mm512_unpacklo_epi32(v[2], v[3]); + v1[3] = _mm512_unpackhi_epi32(v[2], v[3]); + v1[4] = _mm512_unpacklo_epi32(v[4], v[5]); + v1[5] = _mm512_unpackhi_epi32(v[4], v[5]); + v1[6] = _mm512_unpacklo_epi32(v[6], v[7]); + v1[7] = _mm512_unpackhi_epi32(v[6], v[7]); + v1[8] = _mm512_unpacklo_epi32(v[8], v[9]); + v1[9] = _mm512_unpackhi_epi32(v[8], v[9]); + v1[10] = _mm512_unpacklo_epi32(v[10], v[11]); + v1[11] = _mm512_unpackhi_epi32(v[10], v[11]); + v1[12] = _mm512_unpacklo_epi32(v[12], v[13]); + v1[13] = _mm512_unpackhi_epi32(v[12], v[13]); + v1[14] = _mm512_unpacklo_epi32(v[14], v[15]); + v1[15] = _mm512_unpackhi_epi32(v[14], v[15]); + + v[0] = _mm512_unpacklo_epi64(v1[0], v1[2]); + v[1] = _mm512_unpackhi_epi64(v1[0], v1[2]); + v[2] = _mm512_unpacklo_epi64(v1[1], v1[3]); + v[3] = _mm512_unpackhi_epi64(v1[1], v1[3]); + v[4] = _mm512_unpacklo_epi64(v1[4], v1[6]); + v[5] = _mm512_unpackhi_epi64(v1[4], v1[6]); + v[6] = _mm512_unpacklo_epi64(v1[5], v1[7]); + v[7] = _mm512_unpackhi_epi64(v1[5], v1[7]); + v[8] = _mm512_unpacklo_epi64(v1[8], v1[10]); + v[9] = _mm512_unpackhi_epi64(v1[8], v1[10]); + v[10] = _mm512_unpacklo_epi64(v1[9], v1[11]); + v[11] = _mm512_unpackhi_epi64(v1[9], v1[11]); + v[12] = _mm512_unpacklo_epi64(v1[12], v1[14]); + v[13] = _mm512_unpackhi_epi64(v1[12], v1[14]); + v[14] = _mm512_unpacklo_epi64(v1[13], v1[15]); + v[15] = _mm512_unpackhi_epi64(v1[13], v1[15]); + + v1[0] = _mm512_shuffle_i32x4(v[0], v[4], 0x88); + v1[1] = _mm512_shuffle_i32x4(v[1], v[5], 0x88); + v1[2] = _mm512_shuffle_i32x4(v[2], v[6], 0x88); + v1[3] = _mm512_shuffle_i32x4(v[3], v[7], 0x88); + v1[4] = _mm512_shuffle_i32x4(v[0], v[4], 0xdd); + v1[5] = _mm512_shuffle_i32x4(v[1], v[5], 0xdd); + v1[6] = _mm512_shuffle_i32x4(v[2], v[6], 0xdd); + v1[7] = _mm512_shuffle_i32x4(v[3], v[7], 0xdd); + v1[8] = _mm512_shuffle_i32x4(v[8], v[12], 0x88); + v1[9] = _mm512_shuffle_i32x4(v[9], v[13], 0x88); + v1[10] = _mm512_shuffle_i32x4(v[10], v[14], 0x88); + v1[11] = _mm512_shuffle_i32x4(v[11], v[15], 0x88); + v1[12] = _mm512_shuffle_i32x4(v[8], v[12], 0xdd); + v1[13] = _mm512_shuffle_i32x4(v[9], v[13], 0xdd); + v1[14] = _mm512_shuffle_i32x4(v[10], v[14], 0xdd); + v1[15] = _mm512_shuffle_i32x4(v[11], v[15], 0xdd); + + v[0] = _mm512_shuffle_i32x4(v1[0], v1[8], 0x88); + v[1] = _mm512_shuffle_i32x4(v1[1], v1[9], 0x88); + v[2] = _mm512_shuffle_i32x4(v1[2], v1[10], 0x88); + v[3] = _mm512_shuffle_i32x4(v1[3], v1[11], 0x88); + v[4] = _mm512_shuffle_i32x4(v1[4], v1[12], 0x88); + v[5] = _mm512_shuffle_i32x4(v1[5], v1[13], 0x88); + v[6] = _mm512_shuffle_i32x4(v1[6], v1[14], 0x88); + v[7] = _mm512_shuffle_i32x4(v1[7], v1[15], 0x88); + v[8] = _mm512_shuffle_i32x4(v1[0], v1[8], 0xdd); + v[9] = _mm512_shuffle_i32x4(v1[1], v1[9], 0xdd); + v[10] = _mm512_shuffle_i32x4(v1[2], v1[10], 0xdd); + v[11] = _mm512_shuffle_i32x4(v1[3], v1[11], 0xdd); + v[12] = _mm512_shuffle_i32x4(v1[4], v1[12], 0xdd); + v[13] = _mm512_shuffle_i32x4(v1[5], v1[13], 0xdd); + v[14] = _mm512_shuffle_i32x4(v1[6], v1[14], 0xdd); + v[15] = _mm512_shuffle_i32x4(v1[7], v1[15], 0xdd); +} + +/* + Transpose 16x16 32-bit elements + Note that v must be 64 byte aligned +*/ +inline void transpose_16x16_32bit(__m512i *v, size_t stride) { + assert(reinterpret_cast(v) % 64 == 0 && "v must be 64 aligned"); + + auto stride_v = [=](int i) { return offset_pointer(v, i * stride); }; + __m512i v1[16]; + + v1[0] = _mm512_unpacklo_epi32(*stride_v(0), *stride_v(1)); + v1[1] = _mm512_unpackhi_epi32(*stride_v(0), *stride_v(1)); + v1[2] = _mm512_unpacklo_epi32(*stride_v(2), *stride_v(3)); + v1[3] = _mm512_unpackhi_epi32(*stride_v(2), *stride_v(3)); + v1[4] = _mm512_unpacklo_epi32(*stride_v(4), *stride_v(5)); + v1[5] = _mm512_unpackhi_epi32(*stride_v(4), *stride_v(5)); + v1[6] = _mm512_unpacklo_epi32(*stride_v(6), *stride_v(7)); + v1[7] = _mm512_unpackhi_epi32(*stride_v(6), *stride_v(7)); + v1[8] = _mm512_unpacklo_epi32(*stride_v(8), *stride_v(9)); + v1[9] = _mm512_unpackhi_epi32(*stride_v(8), *stride_v(9)); + v1[10] = _mm512_unpacklo_epi32(*stride_v(10), *stride_v(11)); + v1[11] = _mm512_unpackhi_epi32(*stride_v(10), *stride_v(11)); + v1[12] = _mm512_unpacklo_epi32(*stride_v(12), *stride_v(13)); + v1[13] = _mm512_unpackhi_epi32(*stride_v(12), *stride_v(13)); + v1[14] = _mm512_unpacklo_epi32(*stride_v(14), *stride_v(15)); + v1[15] = _mm512_unpackhi_epi32(*stride_v(14), *stride_v(15)); + + *stride_v(0) = _mm512_unpacklo_epi64(v1[0], v1[2]); + *stride_v(1) = _mm512_unpackhi_epi64(v1[0], v1[2]); + *stride_v(2) = _mm512_unpacklo_epi64(v1[1], v1[3]); + *stride_v(3) = _mm512_unpackhi_epi64(v1[1], v1[3]); + *stride_v(4) = _mm512_unpacklo_epi64(v1[4], v1[6]); + *stride_v(5) = _mm512_unpackhi_epi64(v1[4], v1[6]); + *stride_v(6) = _mm512_unpacklo_epi64(v1[5], v1[7]); + *stride_v(7) = _mm512_unpackhi_epi64(v1[5], v1[7]); + *stride_v(8) = _mm512_unpacklo_epi64(v1[8], v1[10]); + *stride_v(9) = _mm512_unpackhi_epi64(v1[8], v1[10]); + *stride_v(10) = _mm512_unpacklo_epi64(v1[9], v1[11]); + *stride_v(11) = _mm512_unpackhi_epi64(v1[9], v1[11]); + *stride_v(12) = _mm512_unpacklo_epi64(v1[12], v1[14]); + *stride_v(13) = _mm512_unpackhi_epi64(v1[12], v1[14]); + *stride_v(14) = _mm512_unpacklo_epi64(v1[13], v1[15]); + *stride_v(15) = _mm512_unpackhi_epi64(v1[13], v1[15]); + + v1[0] = _mm512_shuffle_i32x4(*stride_v(0), *stride_v(4), 0x88); + v1[1] = _mm512_shuffle_i32x4(*stride_v(1), *stride_v(5), 0x88); + v1[2] = _mm512_shuffle_i32x4(*stride_v(2), *stride_v(6), 0x88); + v1[3] = _mm512_shuffle_i32x4(*stride_v(3), *stride_v(7), 0x88); + v1[4] = _mm512_shuffle_i32x4(*stride_v(0), *stride_v(4), 0xdd); + v1[5] = _mm512_shuffle_i32x4(*stride_v(1), *stride_v(5), 0xdd); + v1[6] = _mm512_shuffle_i32x4(*stride_v(2), *stride_v(6), 0xdd); + v1[7] = _mm512_shuffle_i32x4(*stride_v(3), *stride_v(7), 0xdd); + v1[8] = _mm512_shuffle_i32x4(*stride_v(8), *stride_v(12), 0x88); + v1[9] = _mm512_shuffle_i32x4(*stride_v(9), *stride_v(13), 0x88); + v1[10] = _mm512_shuffle_i32x4(*stride_v(10), *stride_v(14), 0x88); + v1[11] = _mm512_shuffle_i32x4(*stride_v(11), *stride_v(15), 0x88); + v1[12] = _mm512_shuffle_i32x4(*stride_v(8), *stride_v(12), 0xdd); + v1[13] = _mm512_shuffle_i32x4(*stride_v(9), *stride_v(13), 0xdd); + v1[14] = _mm512_shuffle_i32x4(*stride_v(10), *stride_v(14), 0xdd); + v1[15] = _mm512_shuffle_i32x4(*stride_v(11), *stride_v(15), 0xdd); + + *stride_v(0) = _mm512_shuffle_i32x4(v1[0], v1[8], 0x88); + *stride_v(1) = _mm512_shuffle_i32x4(v1[1], v1[9], 0x88); + *stride_v(2) = _mm512_shuffle_i32x4(v1[2], v1[10], 0x88); + *stride_v(3) = _mm512_shuffle_i32x4(v1[3], v1[11], 0x88); + *stride_v(4) = _mm512_shuffle_i32x4(v1[4], v1[12], 0x88); + *stride_v(5) = _mm512_shuffle_i32x4(v1[5], v1[13], 0x88); + *stride_v(6) = _mm512_shuffle_i32x4(v1[6], v1[14], 0x88); + *stride_v(7) = _mm512_shuffle_i32x4(v1[7], v1[15], 0x88); + *stride_v(8) = _mm512_shuffle_i32x4(v1[0], v1[8], 0xdd); + *stride_v(9) = _mm512_shuffle_i32x4(v1[1], v1[9], 0xdd); + *stride_v(10) = _mm512_shuffle_i32x4(v1[2], v1[10], 0xdd); + *stride_v(11) = _mm512_shuffle_i32x4(v1[3], v1[11], 0xdd); + *stride_v(12) = _mm512_shuffle_i32x4(v1[4], v1[12], 0xdd); + *stride_v(13) = _mm512_shuffle_i32x4(v1[5], v1[13], 0xdd); + *stride_v(14) = _mm512_shuffle_i32x4(v1[6], v1[14], 0xdd); + *stride_v(15) = _mm512_shuffle_i32x4(v1[7], v1[15], 0xdd); +} + +struct GemmKernel224BF { + using dt = ggml_bf16_t; + using output_t = float; + static const int TILE_M = 16; + static const int TILE_K = 32; + static const int TILE_N = 16; + static const int VNNI_BLK = 2; + + static inline constexpr int M_STEP = TILE_M * 2; + static inline constexpr int N_STEP = TILE_N * 2; + static inline constexpr int K_STEP = TILE_K; + + static inline const int N_BLOCK = 256; + static inline const int K_BLOCK = 1792; + + static int recommended_nth(int n) { return (n + N_BLOCK - 1) / N_BLOCK; } + + static std::pair split_range_n(int n, int ith, int nth) { + int n_start = N_BLOCK * ith; + int n_end = std::min(n, N_BLOCK * (ith + 1)); + return {n_start, n_end}; + } + + static void config() { + enable_amx(); + TileConfig tile_config; + + // size is 16 x 32 + for (int i = 0; i < 2; i++) + tile_config.set_row_col(i, TILE_M, TILE_K * sizeof(dt)); + + // size is 16 x 32 + for (int i = 2; i < 4; i++) + tile_config.set_row_col(i, TILE_K / VNNI_BLK, TILE_N * VNNI_BLK * sizeof(dt)); + + // size is 16 x 16 + for (int i = 4; i < 8; i++) + tile_config.set_row_col(i, TILE_M, TILE_N * sizeof(output_t)); + + tile_config.set_config(); + } + + static void load_a(dt *a, size_t lda) { + _tile_loadd(0, a, lda); + _tile_loadd(1, offset_pointer(a, lda * TILE_M), lda); + } + + static void load_b(dt *b, size_t ldb) { + _tile_loadd(2, b, ldb); + _tile_loadd(3, offset_pointer(b, ldb * TILE_N), ldb); + } + + static void clean_c() { + _tile_zero(4); + _tile_zero(5); + _tile_zero(6); + _tile_zero(7); + } + + static void load_c(output_t *c, size_t ldc) { + _tile_loadd(4, c, ldc); + _tile_loadd(5, offset_pointer(c, TILE_N * sizeof(output_t)), ldc); + _tile_loadd(6, offset_pointer(c, ldc * TILE_M), ldc); + _tile_loadd(7, offset_pointer(c, ldc * TILE_M + TILE_N * sizeof(output_t)), ldc); + } + + static void store_c(output_t *c, size_t ldc) { + _tile_stored(4, c, ldc); + _tile_stored(5, offset_pointer(c, TILE_N * sizeof(output_t)), ldc); + _tile_stored(6, offset_pointer(c, ldc * TILE_M), ldc); + _tile_stored(7, offset_pointer(c, ldc * TILE_M + TILE_N * sizeof(output_t)), ldc); + } + + static void run_tile() { + _tile_dpbf16ps(4, 0, 2); + _tile_dpbf16ps(5, 0, 3); + _tile_dpbf16ps(6, 1, 2); + _tile_dpbf16ps(7, 1, 3); + } + + struct BufferA { + ggml_bf16_t *a; + int max_m, k; + + static size_t required_size(int max_m, int k) { return max_m * k * sizeof(ggml_bf16_t); } + + BufferA(int max_m, int k, void *ptr) : max_m(max_m), k(k) { + assert(reinterpret_cast(ptr) % 64 == 0); + assert(max_m % M_STEP == 0); + assert(k % K_STEP == 0); + a = reinterpret_cast(ptr); + } + + void from_mat(int m, ggml_bf16_t *src, int ith, int nth) { + assert(m <= max_m); + assert(ith == 0 && nth == 1); + int m_block_size = (m + M_STEP - 1) / M_STEP * M_STEP; + for (int m_begin = 0; m_begin < m; m_begin += M_STEP) { + for (int k_block_begin = 0; k_block_begin < k; k_block_begin += K_BLOCK) { + int k_block_size = std::min(K_BLOCK, k - k_block_begin); + for (int k_begin = 0; k_begin < k_block_size; k_begin += K_STEP) { + for (int i = 0; i < M_STEP && m_begin + i < m; i++) { + __m512i *s = (__m512i *)(src + (m_begin + i) * k + k_block_begin + k_begin); + __m512i *d = (__m512i *)(a + k_block_begin * m_block_size + m_begin * k_block_size + k_begin * M_STEP + + i * K_STEP); + avx512_copy_32xbf16(s, d); + } + } + } + } + } + + ggml_bf16_t *get_submat(int m, int k, int m_begin, int k_begin) { + int m_block_size = (m + M_STEP - 1) / M_STEP * M_STEP; + int k_block_begin = k_begin / K_BLOCK * K_BLOCK; + k_begin -= k_block_begin; + int k_block_size = std::min(K_BLOCK, k - k_block_begin); + return a + k_block_begin * m_block_size + m_begin * k_block_size + k_begin * M_STEP; + } + }; + + struct BufferB { + ggml_bf16_t *b; + int n, k; + + static size_t required_size(int n, int k) { return n * k * sizeof(ggml_bf16_t); } + + BufferB(int n, int k, void *ptr) : n(n), k(k) { + assert(reinterpret_cast(ptr) % 64 == 0); + assert(n % N_STEP == 0); + assert(k % K_STEP == 0); + b = reinterpret_cast(ptr); + } + + void from_mat(ggml_bf16_t *src, int ith, int nth) { + auto [n_start, n_end] = split_range_n(n, ith, nth); + int n_block_begin = n_start; + int n_block_size = n_end - n_block_begin; + for (int n_begin = 0; n_begin < n_block_size; n_begin += N_STEP) { + for (int k_block_begin = 0; k_block_begin < k; k_block_begin += K_BLOCK) { + int k_block_size = std::min(K_BLOCK, k - k_block_begin); + for (int k_begin = 0; k_begin < k_block_size; k_begin += K_STEP) { + for (int i = 0; i < N_STEP; i++) { + __m512i *s = (__m512i *)(src + (n_block_begin + n_begin + i) * k + k_block_begin + k_begin); + __m512i *d = (__m512i *)(b + n_block_begin * k + k_block_begin * n_block_size + n_begin * k_block_size + + k_begin * N_STEP + i * K_STEP); + avx512_copy_32xbf16(s, d); + } + transpose_16x16_32bit((__m512i *)(b + n_block_begin * k + k_block_begin * n_block_size + + n_begin * k_block_size + k_begin * N_STEP)); + transpose_16x16_32bit((__m512i *)(b + n_block_begin * k + k_block_begin * n_block_size + + n_begin * k_block_size + k_begin * N_STEP + TILE_N * K_STEP)); + } + } + } + } + + ggml_bf16_t *get_submat(int n, int k, int n_begin, int k_begin) { + int n_block_begin = n_begin / N_BLOCK * N_BLOCK; + n_begin -= n_block_begin; + int n_block_size = std::min(N_BLOCK, n - n_block_begin); + int k_block_begin = k_begin / K_BLOCK * K_BLOCK; + k_begin -= k_block_begin; + int k_block_size = std::min(K_BLOCK, k - k_block_begin); + return b + n_block_begin * k + k_block_begin * n_block_size + n_begin * k_block_size + k_begin * N_STEP; + } + }; + + struct BufferC { + float *c; + int max_m, n; + + static size_t required_size(int max_m, int n) { return max_m * n * sizeof(float); } + + BufferC(int max_m, int n, void *ptr) : max_m(max_m), n(n) { + assert(reinterpret_cast(ptr) % 64 == 0); + assert(max_m % M_STEP == 0); + assert(n % N_STEP == 0); + c = reinterpret_cast(ptr); + } + + void to_mat(int m, ggml_bf16_t *dst, int ith, int nth) { + assert(m <= max_m); + auto [n_start, n_end] = split_range_n(n, ith, nth); + int m_block_size = (m + M_STEP - 1) / M_STEP * M_STEP; + int n_block_begin = n_start; + int n_block_size = n_end - n_block_begin; + for (int m_begin = 0; m_begin < m; m_begin += M_STEP) { + for (int n_begin = 0; n_begin < n_block_size; n_begin += N_STEP) { + for (int i = 0; i < M_STEP && m_begin + i < m; i++) { + __m512 *x0 = + (__m512 *)(c + m_block_size * n_block_begin + m_begin * n_block_size + n_begin * M_STEP + i * N_STEP); + __m512 *x1 = (__m512 *)(c + m_block_size * n_block_begin + m_begin * n_block_size + n_begin * M_STEP + + i * N_STEP + 16); + avx512_32xfp32_to_32xbf16(x0, x1, (__m512i *)(dst + (m_begin + i) * n + n_block_begin + n_begin)); + } + } + } + } + + float *get_submat(int m, int n, int m_begin, int n_begin) { + int m_block_size = (m + M_STEP - 1) / M_STEP * M_STEP; + int n_block_begin = n_begin / N_BLOCK * N_BLOCK; + int n_block_size = std::min(N_BLOCK, n - n_block_begin); + n_begin -= n_block_begin; + return c + m_block_size * n_block_begin + m_begin * n_block_size + n_begin * M_STEP; + } + }; +}; + +struct GemmKernel224Int8 { + using dt = int8_t; + using output_t = int32_t; + static const int TILE_M = 16; + static const int TILE_K = 64; + static const int TILE_N = 16; + static const int VNNI_BLK = 4; + + static inline constexpr int M_STEP = TILE_M * 2; + static inline constexpr int N_STEP = TILE_N * 2; + static inline constexpr int K_STEP = TILE_K; + + static inline const int N_BLOCK = 256; + static inline const int K_BLOCK = 3584; + + static int recommended_nth(int n) { return (n + N_BLOCK - 1) / N_BLOCK; } + + static std::pair split_range_n(int n, int ith, int nth) { + int n_start = N_BLOCK * ith; + int n_end = std::min(n, N_BLOCK * (ith + 1)); + return {n_start, n_end}; + } + + static void config() { + enable_amx(); + TileConfig tile_config; + + // size is 16 x 64 + for (int i = 0; i < 2; i++) + tile_config.set_row_col(i, TILE_M, TILE_K * sizeof(dt)); + + // size is 16 x 64 + for (int i = 2; i < 4; i++) + tile_config.set_row_col(i, TILE_K / VNNI_BLK, TILE_N * VNNI_BLK * sizeof(dt)); + + // size is 16 x 16 + for (int i = 4; i < 8; i++) + tile_config.set_row_col(i, TILE_M, TILE_N * sizeof(output_t)); + + tile_config.set_config(); + } + + static void load_a(dt *a, size_t lda) { + _tile_loadd(0, a, lda); + _tile_loadd(1, offset_pointer(a, lda * TILE_M), lda); + } + + static void load_b(dt *b, size_t ldb) { + _tile_loadd(2, b, ldb); + _tile_loadd(3, offset_pointer(b, ldb * TILE_N), ldb); + } + + static void clean_c() { + _tile_zero(4); + _tile_zero(5); + _tile_zero(6); + _tile_zero(7); + } + + static void load_c(output_t *c, size_t ldc) { + _tile_loadd(4, c, ldc); + _tile_loadd(5, offset_pointer(c, TILE_N * sizeof(output_t)), ldc); + _tile_loadd(6, offset_pointer(c, ldc * TILE_M), ldc); + _tile_loadd(7, offset_pointer(c, ldc * TILE_M + TILE_N * sizeof(output_t)), ldc); + } + + static void store_c(output_t *c, size_t ldc) { + _tile_stored(4, c, ldc); + _tile_stored(5, offset_pointer(c, TILE_N * sizeof(output_t)), ldc); + _tile_stored(6, offset_pointer(c, ldc * TILE_M), ldc); + _tile_stored(7, offset_pointer(c, ldc * TILE_M + TILE_N * sizeof(output_t)), ldc); + } + + static void run_tile() { + _tile_dpbssd(4, 0, 2); + _tile_dpbssd(5, 0, 3); + _tile_dpbssd(6, 1, 2); + _tile_dpbssd(7, 1, 3); + } + + struct BufferA { + int8_t *a; + float *d; + int max_m, k; + + static size_t required_size(int max_m, int k) { return max_m * k * sizeof(int8_t) + max_m * sizeof(float); } + + BufferA(int max_m, int k, void *ptr) : max_m(max_m), k(k) { + assert(reinterpret_cast(ptr) % 64 == 0); + assert(max_m % M_STEP == 0); + assert(k % K_STEP == 0); + a = reinterpret_cast(ptr); + d = reinterpret_cast(a + max_m * k); + } + + void from_mat(int m, ggml_bf16_t *src, int ith, int nth) { + assert(m <= max_m); + assert(ith == 0 && nth == 1); + for (int m_begin = 0; m_begin < m; m_begin += M_STEP) { + for (int i = 0; i < M_STEP && m_begin + i < m; i++) { + float amax = 0.0f; + for (int j = 0; j < k; j += 32) { + __m512 f0, f1; + avx512_32xbf16_to_32xfp32((__m512i *)(src + (m_begin + i) * k + j), &f0, &f1); + amax = MAX(amax, _mm512_reduce_max_ps(_mm512_abs_ps(f0))); + amax = MAX(amax, _mm512_reduce_max_ps(_mm512_abs_ps(f1))); + } + d[m_begin + i] = amax / ((1 << 7) - 1); + } + } + int m_block_size = (m + M_STEP - 1) / M_STEP * M_STEP; + for (int m_begin = 0; m_begin < m; m_begin += M_STEP) { + for (int k_block_begin = 0; k_block_begin < k; k_block_begin += K_BLOCK) { + int k_block_size = std::min(K_BLOCK, k - k_block_begin); + for (int k_begin = 0; k_begin < k_block_size; k_begin += K_STEP) { + for (int i = 0; i < M_STEP && m_begin + i < m; i++) { + __m512 id = _mm512_set1_ps(d[m_begin + i] ? 1.0f / d[m_begin + i] : 0.0f); + int8_t *dst = a + k_block_begin * m_block_size + m_begin * k_block_size + k_begin * M_STEP + i * K_STEP; + __m512 f0, f1, f2, f3; + avx512_32xbf16_to_32xfp32((__m512i *)(src + (m_begin + i) * k + k_block_begin + k_begin), &f0, &f1); + avx512_32xbf16_to_32xfp32((__m512i *)(src + (m_begin + i) * k + k_block_begin + k_begin) + 1, &f2, &f3); + __m512i i0 = _mm512_cvtps_epi32(_mm512_mul_ps(f0, id)); + __m512i i1 = _mm512_cvtps_epi32(_mm512_mul_ps(f1, id)); + __m512i i2 = _mm512_cvtps_epi32(_mm512_mul_ps(f2, id)); + __m512i i3 = _mm512_cvtps_epi32(_mm512_mul_ps(f3, id)); + __m128i s0 = _mm512_cvtsepi32_epi8(i0); + __m128i s1 = _mm512_cvtsepi32_epi8(i1); + __m128i s2 = _mm512_cvtsepi32_epi8(i2); + __m128i s3 = _mm512_cvtsepi32_epi8(i3); + _mm_storeu_si128((__m128i *)dst, s0); + _mm_storeu_si128((__m128i *)(dst + 16), s1); + _mm_storeu_si128((__m128i *)(dst + 32), s2); + _mm_storeu_si128((__m128i *)(dst + 48), s3); + } + } + } + } + } + + int8_t *get_submat(int m, int k, int m_begin, int k_begin) { + int m_block_size = (m + M_STEP - 1) / M_STEP * M_STEP; + int k_block_begin = k_begin / K_BLOCK * K_BLOCK; + k_begin -= k_block_begin; + int k_block_size = std::min(K_BLOCK, k - k_block_begin); + return a + k_block_begin * m_block_size + m_begin * k_block_size + k_begin * M_STEP; + } + + float *get_scale(int m, int m_begin) { return d + m_begin; } + }; + + struct BufferB { + int8_t *b; + float *d; + int n, k; + + static size_t required_size(int n, int k) { return n * k * sizeof(int8_t) + n * sizeof(float); } + + BufferB(int n, int k, void *ptr) : n(n), k(k) { + assert(reinterpret_cast(ptr) % 64 == 0); + assert(n % N_STEP == 0); + assert(k % K_STEP == 0); + b = reinterpret_cast(ptr); + d = reinterpret_cast(b + n * k); + } + + void from_mat(ggml_bf16_t *src, int ith, int nth) { + auto [n_start, n_end] = split_range_n(n, ith, nth); + int n_block_begin = n_start; + int n_block_size = n_end - n_block_begin; + for (int n_begin = 0; n_begin < n_block_size; n_begin += N_STEP) { + for (int i = 0; i < N_STEP; i++) { + float amax = 0.0f; + for (int j = 0; j < k; j += 32) { + __m512 f0, f1; + avx512_32xbf16_to_32xfp32((__m512i *)(src + (n_block_begin + n_begin + i) * k + j), &f0, &f1); + amax = MAX(amax, _mm512_reduce_max_ps(_mm512_abs_ps(f0))); + amax = MAX(amax, _mm512_reduce_max_ps(_mm512_abs_ps(f1))); + } + d[n_block_begin + n_begin + i] = amax / ((1 << 7) - 1); + } + } + for (int n_begin = 0; n_begin < n_block_size; n_begin += N_STEP) { + for (int k_block_begin = 0; k_block_begin < k; k_block_begin += K_BLOCK) { + int k_block_size = std::min(K_BLOCK, k - k_block_begin); + for (int k_begin = 0; k_begin < k_block_size; k_begin += K_STEP) { + for (int i = 0; i < N_STEP; i++) { + __m512 id = _mm512_set1_ps(d[n_block_begin + n_begin + i] ? 1.0f / d[n_block_begin + n_begin + i] : 0.0f); + int8_t *dst = b + n_block_begin * k + k_block_begin * n_block_size + n_begin * k_block_size + + k_begin * N_STEP + i * K_STEP; + __m512 f0, f1, f2, f3; + avx512_32xbf16_to_32xfp32((__m512i *)(src + (n_block_begin + n_begin + i) * k + k_block_begin + k_begin), + &f0, &f1); + avx512_32xbf16_to_32xfp32( + (__m512i *)(src + (n_block_begin + n_begin + i) * k + k_block_begin + k_begin) + 1, &f2, &f3); + __m512i i0 = _mm512_cvtps_epi32(_mm512_mul_ps(f0, id)); + __m512i i1 = _mm512_cvtps_epi32(_mm512_mul_ps(f1, id)); + __m512i i2 = _mm512_cvtps_epi32(_mm512_mul_ps(f2, id)); + __m512i i3 = _mm512_cvtps_epi32(_mm512_mul_ps(f3, id)); + __m128i s0 = _mm512_cvtsepi32_epi8(i0); + __m128i s1 = _mm512_cvtsepi32_epi8(i1); + __m128i s2 = _mm512_cvtsepi32_epi8(i2); + __m128i s3 = _mm512_cvtsepi32_epi8(i3); + _mm_storeu_si128((__m128i *)dst, s0); + _mm_storeu_si128((__m128i *)(dst + 16), s1); + _mm_storeu_si128((__m128i *)(dst + 32), s2); + _mm_storeu_si128((__m128i *)(dst + 48), s3); + } + transpose_16x16_32bit((__m512i *)(b + n_block_begin * k + k_block_begin * n_block_size + + n_begin * k_block_size + k_begin * N_STEP)); + transpose_16x16_32bit((__m512i *)(b + n_block_begin * k + k_block_begin * n_block_size + + n_begin * k_block_size + k_begin * N_STEP + TILE_N * K_STEP)); + } + } + } + } + + int8_t *get_submat(int n, int k, int n_begin, int k_begin) { + int n_block_begin = n_begin / N_BLOCK * N_BLOCK; + n_begin -= n_block_begin; + int n_block_size = std::min(N_BLOCK, n - n_block_begin); + int k_block_begin = k_begin / K_BLOCK * K_BLOCK; + k_begin -= k_block_begin; + int k_block_size = std::min(K_BLOCK, k - k_block_begin); + return b + n_block_begin * k + k_block_begin * n_block_size + n_begin * k_block_size + k_begin * N_STEP; + } + + float *get_scale(int n, int n_begin) { return d + n_begin; } + }; + + struct BufferC { + float *c; + int max_m, n; + + static size_t required_size(int max_m, int n) { return max_m * n * sizeof(float); } + + BufferC(int max_m, int n, void *ptr) : max_m(max_m), n(n) { + assert(reinterpret_cast(ptr) % 64 == 0); + assert(max_m % M_STEP == 0); + assert(n % N_STEP == 0); + c = reinterpret_cast(ptr); + } + + void to_mat(int m, ggml_bf16_t *dst, int ith, int nth) { + assert(m <= max_m); + auto [n_start, n_end] = split_range_n(n, ith, nth); + int m_block_size = (m + M_STEP - 1) / M_STEP * M_STEP; + int n_block_begin = n_start; + int n_block_size = n_end - n_block_begin; + for (int m_begin = 0; m_begin < m; m_begin += M_STEP) { + for (int n_begin = 0; n_begin < n_block_size; n_begin += N_STEP) { + for (int i = 0; i < M_STEP && m_begin + i < m; i++) { + __m512 *x0 = + (__m512 *)(c + m_block_size * n_block_begin + m_begin * n_block_size + n_begin * M_STEP + i * N_STEP); + __m512 *x1 = (__m512 *)(c + m_block_size * n_block_begin + m_begin * n_block_size + n_begin * M_STEP + + i * N_STEP + 16); + avx512_32xfp32_to_32xbf16(x0, x1, (__m512i *)(dst + (m_begin + i) * n + n_block_begin + n_begin)); + } + } + } + } + + float *get_submat(int m, int n, int m_begin, int n_begin) { + int m_block_size = (m + M_STEP - 1) / M_STEP * M_STEP; + int n_block_begin = n_begin / N_BLOCK * N_BLOCK; + int n_block_size = std::min(N_BLOCK, n - n_block_begin); + n_begin -= n_block_begin; + return c + m_block_size * n_block_begin + m_begin * n_block_size + n_begin * M_STEP; + } + }; +}; + +inline void mat_mul(int m, int n, int k, std::shared_ptr ba, + std::shared_ptr bb, std::shared_ptr bc, int ith, + int nth, bool use_amx) { +// std::cout << "mat_mul in BF16!!!!" << std::endl; + using K = GemmKernel224BF; + assert(n % K::N_STEP == 0); + assert(k % K::K_STEP == 0); + + auto [n_start, n_end] = K::split_range_n(n, ith, nth); + + for (int k_block_begin = 0; k_block_begin < k; k_block_begin += K::K_BLOCK) { + for (int m_begin = 0; m_begin < m; m_begin += K::M_STEP) { + for (int n_begin = n_start; n_begin < n_end; n_begin += K::N_STEP) { + + float *c = bc->get_submat(m, n, m_begin, n_begin); + if (!use_amx) { + __m512 *c512 = (__m512 *)c; + if (k_block_begin == 0) { + for (int m_i = 0; m_i < m && m_i < K::M_STEP; m_i++) { + c512[m_i * 2] = _mm512_setzero_ps(); + c512[m_i * 2 + 1] = _mm512_setzero_ps(); + } + } + + for (int k_begin = 0; k_begin < K::K_BLOCK && k_block_begin + k_begin < k; k_begin += K::K_STEP) { + int32_t *a32 = (int32_t *)ba->get_submat(m, k, m_begin, k_block_begin + k_begin); + __m512bh *b512 = (__m512bh *)bb->get_submat(n, k, n_begin, k_block_begin + k_begin); + for (int m_i = 0; m_i < m && m_i < K::M_STEP; m_i++) { + for (int k_i = 0; k_i < 16; k_i++) { + __m512bh ma = (__m512bh)_mm512_set1_epi32(a32[m_i * 16 + k_i]); + for (int n_i = 0; n_i < 2; n_i++) { + c512[m_i * 2 + n_i] = _mm512_dpbf16_ps(c512[m_i * 2 + n_i], ma, b512[n_i * 16 + k_i]); + } + } + } + } + + } else { + if (k_block_begin == 0) { + K::clean_c(); + } else { + K::load_c(c, K::N_STEP * sizeof(float)); + } + for (int k_begin = 0; k_begin < K::K_BLOCK && k_block_begin + k_begin < k; k_begin += K::K_STEP) { + K::load_a(ba->get_submat(m, k, m_begin, k_block_begin + k_begin), K::K_STEP * sizeof(ggml_bf16_t)); + K::load_b(bb->get_submat(n, k, n_begin, k_block_begin + k_begin), K::K_STEP * sizeof(ggml_bf16_t)); + K::run_tile(); + } + K::store_c(c, K::N_STEP * sizeof(float)); + } + } + } + } +} + +inline __m512i _mm512_dpbssd_epi32(__m512i src, __m512i a, __m512i b) { + __m256i a_lo = _mm512_extracti64x4_epi64(a, 0); + __m256i a_hi = _mm512_extracti64x4_epi64(a, 1); + __m256i b_lo = _mm512_extracti64x4_epi64(b, 0); + __m256i b_hi = _mm512_extracti64x4_epi64(b, 1); + + b_lo = _mm256_sign_epi8(b_lo, a_lo); + b_hi = _mm256_sign_epi8(b_hi, a_hi); + + b = _mm512_inserti64x4(b, b_lo, 0); + b = _mm512_inserti64x4(b, b_hi, 1); + + a = _mm512_abs_epi8(a); + + return _mm512_dpbusd_epi32(src, a, b); +} + +inline void mat_mul(int m, int n, int k, std::shared_ptr ba, + std::shared_ptr bb, std::shared_ptr bc, + int ith, int nth, bool use_amx) { +// std::cout << "mat_mul in INT8!!!!" << std::endl; + using K = GemmKernel224Int8; + assert(n % K::N_STEP == 0); + assert(k % K::K_STEP == 0); + + auto [n_start, n_end] = K::split_range_n(n, ith, nth); + + for (int k_block_begin = 0; k_block_begin < k; k_block_begin += K::K_BLOCK) { + for (int m_begin = 0; m_begin < m; m_begin += K::M_STEP) { + for (int n_begin = n_start; n_begin < n_end; n_begin += K::N_STEP) { + float *c = bc->get_submat(m, n, m_begin, n_begin); + + if (!use_amx) { + __m512i *c512 = (__m512i *)c; + if (k_block_begin == 0) { + for (int m_i = 0; m_i < m && m_i < K::M_STEP; m_i++) { + c512[m_i * 2] = _mm512_setzero_si512(); + c512[m_i * 2 + 1] = _mm512_setzero_si512(); + } + } + + for (int k_begin = 0; k_begin < K::K_BLOCK && k_block_begin + k_begin < k; k_begin += K::K_STEP) { + static_assert(K::K_STEP * sizeof(int8_t) == sizeof(__m512i)); + static_assert(K::N_STEP / K::TILE_N == 2, "Must be lke this"); + + int32_t *a32 = (int32_t *)ba->get_submat(m, k, m_begin, k_block_begin + k_begin); + __m512i *b512 = (__m512i *)bb->get_submat(n, k, n_begin, k_block_begin + k_begin); + for (int m_i = 0; m_i < m && m_i < K::M_STEP; m_i++) { + for (int k_i = 0; k_i < 16; k_i++) { + __m512i ma = _mm512_set1_epi32(a32[m_i * 16 + k_i]); + for (int n_i = 0; n_i < 2; n_i++) { + c512[m_i * 2 + n_i] = _mm512_dpbssd_epi32(c512[m_i * 2 + n_i], ma, b512[n_i * 16 + k_i]); + } + } + } + } + } else { + if (k_block_begin == 0) { + K::clean_c(); + } else { + K::load_c((int32_t *)c, K::N_STEP * sizeof(int32_t)); + } + for (int k_begin = 0; k_begin < K::K_BLOCK && k_block_begin + k_begin < k; k_begin += K::K_STEP) { + K::load_a(ba->get_submat(m, k, m_begin, k_block_begin + k_begin), K::K_STEP * sizeof(int8_t)); + K::load_b(bb->get_submat(n, k, n_begin, k_block_begin + k_begin), K::K_STEP * sizeof(int8_t)); + K::run_tile(); + } + K::store_c((int32_t *)c, K::N_STEP * sizeof(int32_t)); + } + + if (k_block_begin + K::K_BLOCK >= k) { + int to = m - m_begin; + if (m - m_begin > K::M_STEP) { + to = K::M_STEP; + } + for (int i = 0; i < to; i++) { + __m512 as = _mm512_set1_ps(*ba->get_scale(m, m_begin + i)); + __m512 bs = _mm512_load_ps(bb->get_scale(n, n_begin)); + __m512i now = _mm512_load_si512((__m512i *)(c + i * K::N_STEP)); + __m512 result = _mm512_mul_ps(_mm512_mul_ps(as, bs), _mm512_cvtepi32_ps(now)); + _mm512_store_ps((__m512 *)(c + i * K::N_STEP), result); + bs = _mm512_load_ps(bb->get_scale(n, n_begin) + K::TILE_N); + now = _mm512_load_si512((__m512i *)(c + i * K::N_STEP + K::TILE_N)); + result = _mm512_mul_ps(_mm512_mul_ps(as, bs), _mm512_cvtepi32_ps(now)); + _mm512_store_ps((__m512 *)(c + i * K::N_STEP + K::TILE_N), result); + } + } + } + } + } +} + +} // namespace amx \ No newline at end of file diff --git a/archive/kt-sft/csrc/ktransformers_ext/operators/amx/la/utils.hpp b/archive/kt-sft/csrc/ktransformers_ext/operators/amx/la/utils.hpp new file mode 100644 index 0000000..a95a2e7 --- /dev/null +++ b/archive/kt-sft/csrc/ktransformers_ext/operators/amx/la/utils.hpp @@ -0,0 +1,46 @@ +/** + * @Description : + * @Author : chenht2022 + * @Date : 2025-04-25 18:28:12 + * @Version : 1.0.0 + * @LastEditors : chenht2022 + * @LastEditTime : 2025-04-25 18:28:12 + * @Copyright (c) 2024 by KVCache.AI, All Rights Reserved. + **/ + +#pragma once +#include + + +template +T* offset_pointer(T* ptr, std::size_t byte_offset) { + return reinterpret_cast(reinterpret_cast(ptr) + byte_offset); +} + +template +const T* offset_pointer(const T* ptr, std::size_t byte_offset) { + return reinterpret_cast(reinterpret_cast(ptr) + byte_offset); +} + +template +T* offset_pointer_row_major(T* t, int row, int col, std::size_t ld) { + return offset_pointer(t, row * ld) + col; +} + +template +T* offset_pointer_col_major(T* t, int row, int col, std::size_t ld) { + return offset_pointer(t, col * ld) + row; +} + +static inline void avx512_copy_32xbf16(__m512i* src, __m512i* dst) { + _mm512_storeu_si512(dst, _mm512_loadu_si512(src)); +} + +static inline void avx512_32xfp32_to_32xbf16(__m512* src0, __m512* src1, __m512i* dst) { + _mm512_storeu_si512(dst, __m512i(_mm512_cvtne2ps_pbh(*src1, *src0))); +} + +static inline void avx512_32xbf16_to_32xfp32(__m512i* src, __m512* dst0, __m512* dst1) { + _mm512_storeu_ps(dst0, _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu16_epi32(_mm256_loadu_si256((const __m256i *)(src))), 16))); + _mm512_storeu_ps(dst1, _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu16_epi32(_mm256_loadu_si256((const __m256i *)(src) + 1)), 16))); +} \ No newline at end of file diff --git a/archive/kt-sft/csrc/ktransformers_ext/operators/amx/moe.hpp b/archive/kt-sft/csrc/ktransformers_ext/operators/amx/moe.hpp new file mode 100644 index 0000000..cf3b2f9 --- /dev/null +++ b/archive/kt-sft/csrc/ktransformers_ext/operators/amx/moe.hpp @@ -0,0 +1,398 @@ +/** + * @Description : + * @Author : chenht2022 + * @Date : 2025-04-25 18:28:12 + * @Version : 1.0.0 + * @LastEditors : chenht2022 + * @LastEditTime : 2025-04-25 18:28:12 + * @Copyright (c) 2024 by KVCache.AI, All Rights Reserved. + **/ +#ifndef CPUINFER_OPERATOR_AMX_MOE_H +#define CPUINFER_OPERATOR_AMX_MOE_H + +#include +#include +#include +#include +#include + +#include "../../cpu_backend/backend.h" +#include "../../cpu_backend/shared_mem_buffer.h" +#include "llama.cpp/ggml-impl.h" +#include "llama.cpp/ggml-quants.h" +#include "llama.cpp/ggml.h" +#include "llamafile/sgemm.h" + +#include "la/amx.hpp" + +#ifdef USE_NUMA +#include +#include +void *numa_alloc_aligned(size_t size, int node, size_t alignment) { + void *ptr = numa_alloc_onnode(size, node); + assert(reinterpret_cast(ptr) % 64 == 0); + return ptr; +} +#endif + +static inline __m512 exp_avx512(__m512 x) { + const __m512 log2e = _mm512_set1_ps(1.44269504089f); + const __m512 c1 = _mm512_set1_ps(0.69314718056f); + + __m512 y = _mm512_mul_ps(x, log2e); + __m512i int_part = _mm512_cvtps_epi32(y); + __m512 frac_part = _mm512_sub_ps(y, _mm512_cvtepi32_ps(int_part)); + + const __m512 poly_1 = _mm512_set1_ps(0.9999999995f); + const __m512 poly_2 = _mm512_set1_ps(0.6931471805f); + const __m512 poly_3 = _mm512_set1_ps(0.2402265069f); + const __m512 poly_4 = _mm512_set1_ps(0.0555041087f); + const __m512 poly_5 = _mm512_set1_ps(0.0096181291f); + const __m512 poly_6 = _mm512_set1_ps(0.0013333558f); + + __m512 frac_exp = _mm512_fmadd_ps( + _mm512_fmadd_ps(_mm512_fmadd_ps(_mm512_fmadd_ps(_mm512_fmadd_ps(poly_6, frac_part, poly_5), frac_part, poly_4), + frac_part, poly_3), + frac_part, poly_2), + frac_part, poly_1); + + __m512 two_pow_i = _mm512_scalef_ps(_mm512_set1_ps(1.0f), _mm512_cvtepi32_ps(int_part)); + return _mm512_mul_ps(two_pow_i, frac_exp); +} + +static inline __m512 act_fn(__m512 gate_val, __m512 up_val) { + __m512 neg_gate_val = _mm512_sub_ps(_mm512_setzero_ps(), gate_val); + __m512 exp_neg_gate = exp_avx512(neg_gate_val); + __m512 denom = _mm512_add_ps(_mm512_set1_ps(1.0f), exp_neg_gate); + __m512 act_val = _mm512_div_ps(gate_val, denom); + + return _mm512_mul_ps(act_val, up_val); +} + +struct AMX_MOEConfig { + int expert_num; + int routed_expert_num; + int hidden_size; + int intermediate_size; + int max_len; + void *gate_proj; + void *up_proj; + void *down_proj; + + AMX_MOEConfig() {} + + AMX_MOEConfig(int expert_num, int routed_expert_num, int hidden_size, int intermediate_size, int max_len, + void *gate_proj, void *up_proj, void *down_proj) + : expert_num(expert_num), routed_expert_num(routed_expert_num), hidden_size(hidden_size), + intermediate_size(intermediate_size), max_len(max_len), gate_proj(gate_proj), up_proj(up_proj), + down_proj(down_proj) {} +}; + +template class AMX_MOE { +private: + AMX_MOEConfig config_; + void *gate_proj_; // [expert_num * intermediate_size * hidden_size ( /32 if quantized)] + void *up_proj_; // [expert_num * intermediate_size * hidden_size ( /32 if quantized)] + void *down_proj_; // [expert_num * hidden_size * intermediate_size ( /32 if quantized)] + + ggml_bf16_t *m_local_input_; // [routed_expert_num * max_len * hidden_size] + ggml_bf16_t *m_local_gate_output_; // [routed_expert_num * max_len * intermediate_size] + ggml_bf16_t *m_local_up_output_; // [routed_expert_num * max_len * intermediate_size] + ggml_bf16_t *m_local_down_output_; // [routed_expert_num * max_len * hidden_size] + + std::vector> m_local_pos_; // [max_len, routed_expert_num] + std::vector m_local_num_; // [expert_num] + std::vector m_expert_id_map_; // [expert_num] + std::vector m_local_input_ptr_; // [expert_num] + std::vector m_local_gate_output_ptr_; // [expert_num] + std::vector m_local_up_output_ptr_; // [expert_num] + std::vector m_local_down_output_ptr_; // [expert_num] + + std::vector> gate_up_ba_; + std::vector> gate_bc_; + std::vector> up_bc_; + std::vector> down_ba_; + std::vector> down_bc_; + +#ifdef USE_NUMA + std::vector>> gate_bb_numa_; + std::vector>> up_bb_numa_; + std::vector>> down_bb_numa_; +#else + std::vector> gate_bb_; + std::vector> up_bb_; + std::vector> down_bb_; +#endif + +public: + AMX_MOE(AMX_MOEConfig config) { + config_ = config; + gate_proj_ = config_.gate_proj; + up_proj_ = config_.up_proj; + down_proj_ = config_.down_proj; + + std::vector> m_mem_requests; + m_mem_requests.push_back({(void **)&m_local_input_, + sizeof(ggml_bf16_t) * config_.routed_expert_num * config_.max_len * config_.hidden_size}); + m_mem_requests.push_back({(void **)&m_local_gate_output_, sizeof(ggml_bf16_t) * config_.routed_expert_num * + config_.max_len * config_.intermediate_size}); + m_mem_requests.push_back({(void **)&m_local_up_output_, sizeof(ggml_bf16_t) * config_.routed_expert_num * + config_.max_len * config_.intermediate_size}); + m_mem_requests.push_back({(void **)&m_local_down_output_, + sizeof(ggml_bf16_t) * config_.routed_expert_num * config_.max_len * config_.hidden_size}); + std::vector gate_up_ba_ptr(config_.expert_num); + std::vector gate_bc_ptr(config_.expert_num); + std::vector up_bc_ptr(config_.expert_num); + std::vector down_ba_ptr(config_.expert_num); + std::vector down_bc_ptr(config_.expert_num); + for (int i = 0; i < config_.expert_num; i++) { + m_mem_requests.push_back( + {(void **)&gate_up_ba_ptr[i], T::BufferA::required_size(config_.max_len, config_.hidden_size)}); + m_mem_requests.push_back( + {(void **)&gate_bc_ptr[i], T::BufferC::required_size(config_.max_len, config_.intermediate_size)}); + m_mem_requests.push_back( + {(void **)&up_bc_ptr[i], T::BufferC::required_size(config_.max_len, config_.intermediate_size)}); + m_mem_requests.push_back( + {(void **)&down_ba_ptr[i], T::BufferA::required_size(config_.max_len, config_.intermediate_size)}); + m_mem_requests.push_back( + {(void **)&down_bc_ptr[i], T::BufferC::required_size(config_.max_len, config_.hidden_size)}); + } + shared_mem_buffer.alloc(this, m_mem_requests); + + m_local_pos_.resize(config_.max_len); + for (int i = 0; i < config_.max_len; i++) { + m_local_pos_[i].resize(config_.routed_expert_num); + } + m_expert_id_map_.resize(config_.expert_num); + m_local_num_.resize(config_.expert_num); + m_local_input_ptr_.resize(config_.expert_num); + m_local_gate_output_ptr_.resize(config_.expert_num); + m_local_up_output_ptr_.resize(config_.expert_num); + m_local_down_output_ptr_.resize(config_.expert_num); + + for (uint64_t i = 0; i < config_.expert_num; i++) { + gate_up_ba_.push_back( + std::make_shared(config_.max_len, config_.hidden_size, gate_up_ba_ptr[i])); + gate_bc_.push_back( + std::make_shared(config_.max_len, config_.intermediate_size, gate_bc_ptr[i])); + up_bc_.push_back(std::make_shared(config_.max_len, config_.intermediate_size, up_bc_ptr[i])); + down_ba_.push_back( + std::make_shared(config_.max_len, config_.intermediate_size, down_ba_ptr[i])); + down_bc_.push_back(std::make_shared(config_.max_len, config_.hidden_size, down_bc_ptr[i])); + +#ifdef USE_NUMA + int numa_nodes = numa_num_configured_nodes(); + gate_bb_numa_.resize(numa_nodes); + up_bb_numa_.resize(numa_nodes); + down_bb_numa_.resize(numa_nodes); + for (int j = 0; j < numa_nodes; j++) { + void *gate_bb_ptr = + numa_alloc_aligned(T::BufferB::required_size(config_.intermediate_size, config_.hidden_size), j, 64); + gate_bb_numa_[j].push_back( + std::make_shared(config_.intermediate_size, config_.hidden_size, gate_bb_ptr)); + void *up_bb_ptr = + numa_alloc_aligned(T::BufferB::required_size(config_.intermediate_size, config_.hidden_size), j, 64); + up_bb_numa_[j].push_back( + std::make_shared(config_.intermediate_size, config_.hidden_size, up_bb_ptr)); + void *down_bb_ptr = + numa_alloc_aligned(T::BufferB::required_size(config_.hidden_size, config_.intermediate_size), j, 64); + down_bb_numa_[j].push_back( + std::make_shared(config_.hidden_size, config_.intermediate_size, down_bb_ptr)); + } +#else + void *gate_bb_ptr = + std::aligned_alloc(64, T::BufferB::required_size(config_.intermediate_size, config_.hidden_size)); + gate_bb_.push_back( + std::make_shared(config_.intermediate_size, config_.hidden_size, gate_bb_ptr)); + + void *up_bb_ptr = + std::aligned_alloc(64, T::BufferB::required_size(config_.intermediate_size, config_.hidden_size)); + up_bb_.push_back( + std::make_shared(config_.intermediate_size, config_.hidden_size, up_bb_ptr)); + + void *down_bb_ptr = + std::aligned_alloc(64, T::BufferB::required_size(config_.hidden_size, config_.intermediate_size)); + down_bb_.push_back( + std::make_shared(config_.hidden_size, config_.intermediate_size, down_bb_ptr)); +#endif + } + } + + ~AMX_MOE() { shared_mem_buffer.dealloc(this); } + + void load_weights(Backend *backend) { + int nth = T::recommended_nth(config_.intermediate_size); + backend->do_work_stealing_job( + nth * config_.expert_num, nullptr, + [&](int task_id) { + uint64_t expert_idx = task_id / nth; + int ith = task_id % nth; +#ifdef USE_NUMA + int numa_nodes = numa_num_configured_nodes(); + for (int j = 0; j < numa_nodes; j++) { + gate_bb_numa_[j][expert_idx]->from_mat((ggml_bf16_t *)config_.gate_proj + + expert_idx * config_.intermediate_size * config_.hidden_size, + ith, nth); + up_bb_numa_[j][expert_idx]->from_mat((ggml_bf16_t *)config_.up_proj + + expert_idx * config_.intermediate_size * config_.hidden_size, + ith, nth); + } +#else + gate_bb_[expert_idx]->from_mat((ggml_bf16_t *)config_.gate_proj + + expert_idx * config_.intermediate_size * config_.hidden_size, + ith, nth); + up_bb_[expert_idx]->from_mat( + (ggml_bf16_t *)config_.up_proj + expert_idx * config_.intermediate_size * config_.hidden_size, ith, nth); +#endif + }, + nullptr); + nth = T::recommended_nth(config_.hidden_size); + backend->do_work_stealing_job( + nth * config_.expert_num, nullptr, + [&](int task_id) { + uint64_t expert_idx = task_id / nth; + int ith = task_id % nth; +#ifdef USE_NUMA + int numa_nodes = numa_num_configured_nodes(); + for (int j = 0; j < numa_nodes; j++) { + down_bb_numa_[j][expert_idx]->from_mat((ggml_bf16_t *)config_.down_proj + + expert_idx * config_.hidden_size * config_.intermediate_size, + ith, nth); + } +#else + down_bb_[expert_idx]->from_mat((ggml_bf16_t *)config_.down_proj + + expert_idx * config_.hidden_size * config_.intermediate_size, + ith, nth); +#endif + }, + nullptr); + } + + void warm_up(Backend *backend) {} + + void forward(int qlen, int k, const uint64_t *expert_ids, const float *weights, const void *input, void *output, + int *batch_size_tensor, Backend *backend) { + qlen = batch_size_tensor[0]; + bool use_amx = (qlen > 4 * config_.expert_num / config_.routed_expert_num); + int activated_expert = 0; + for (int i = 0; i < config_.expert_num; i++) { + m_local_num_[i] = 0; + } + for (int i = 0; i < qlen; i++) { + for (int j = 0; j < k; j++) { + m_local_pos_[i][j] = m_local_num_[expert_ids[i * k + j]]++; + } + } + for (int i = 0; i < config_.expert_num; i++) { + if (m_local_num_[i] > 0) { + m_expert_id_map_[activated_expert] = i; + activated_expert++; + } + } + uint64_t offset = 0; + for (int i = 0; i < config_.expert_num; i++) { + m_local_input_ptr_[i] = m_local_input_ + offset * config_.hidden_size; + m_local_gate_output_ptr_[i] = m_local_gate_output_ + offset * config_.intermediate_size; + m_local_up_output_ptr_[i] = m_local_up_output_ + offset * config_.intermediate_size; + m_local_down_output_ptr_[i] = m_local_down_output_ + offset * config_.hidden_size; + offset += m_local_num_[i]; + } + backend->do_work_stealing_job( + qlen, nullptr, + [&](int i) { + for (int j = 0; j < k; j++) { + memcpy(m_local_input_ptr_[expert_ids[i * k + j]] + m_local_pos_[i][j] * config_.hidden_size, + (ggml_bf16_t *)input + i * config_.hidden_size, sizeof(ggml_bf16_t) * config_.hidden_size); + } + }, + nullptr); + backend->do_work_stealing_job( + activated_expert, nullptr, + [&](int task_id) { + int expert_idx = m_expert_id_map_[task_id]; + + gate_up_ba_[expert_idx]->from_mat(m_local_num_[expert_idx], m_local_input_ptr_[expert_idx], 0, 1); + }, + nullptr); + int nth = T::recommended_nth(config_.intermediate_size); + backend->do_work_stealing_job( + nth * activated_expert, [&](int _) { T::config(); }, + [&](int task_id) { + int expert_idx = m_expert_id_map_[task_id / nth]; + int ith = task_id % nth; +#ifdef USE_NUMA + amx::mat_mul(m_local_num_[expert_idx], config_.intermediate_size, config_.hidden_size, + gate_up_ba_[expert_idx], gate_bb_numa_[Backend::numa_node][expert_idx], gate_bc_[expert_idx], + ith, nth, use_amx); + amx::mat_mul(m_local_num_[expert_idx], config_.intermediate_size, config_.hidden_size, + gate_up_ba_[expert_idx], up_bb_numa_[Backend::numa_node][expert_idx], up_bc_[expert_idx], ith, + nth, use_amx); +#else + amx::mat_mul(m_local_num_[expert_idx], config_.intermediate_size, config_.hidden_size, + gate_up_ba_[expert_idx], gate_bb_[expert_idx], gate_bc_[expert_idx], ith, nth, use_amx); + amx::mat_mul(m_local_num_[expert_idx], config_.intermediate_size, config_.hidden_size, + gate_up_ba_[expert_idx], up_bb_[expert_idx], up_bc_[expert_idx], ith, nth, use_amx); +#endif + gate_bc_[expert_idx]->to_mat(m_local_num_[expert_idx], m_local_gate_output_ptr_[expert_idx], ith, nth); + up_bc_[expert_idx]->to_mat(m_local_num_[expert_idx], m_local_up_output_ptr_[expert_idx], ith, nth); + auto [n_start, n_end] = T::split_range_n(config_.intermediate_size, ith, nth); + for (int i = 0; i < m_local_num_[expert_idx]; i++) { + ggml_bf16_t *gate_output_ptr = &m_local_gate_output_ptr_[expert_idx][i * config_.intermediate_size]; + ggml_bf16_t *up_output_ptr = &m_local_up_output_ptr_[expert_idx][i * config_.intermediate_size]; + for (int j = n_start; j < n_end; j += 32) { + __m512 gate_val0, gate_val1, up_val0, up_val1; + avx512_32xbf16_to_32xfp32((__m512i *)(gate_output_ptr + j), &gate_val0, &gate_val1); + avx512_32xbf16_to_32xfp32((__m512i *)(up_output_ptr + j), &up_val0, &up_val1); + __m512 result0 = act_fn(gate_val0, up_val0); + __m512 result1 = act_fn(gate_val1, up_val1); + avx512_32xfp32_to_32xbf16(&result0, &result1, (__m512i *)(gate_output_ptr + j)); + } + } + }, + nullptr); + backend->do_work_stealing_job( + activated_expert, nullptr, + [&](int task_id) { + int expert_idx = m_expert_id_map_[task_id]; + down_ba_[expert_idx]->from_mat(m_local_num_[expert_idx], m_local_gate_output_ptr_[expert_idx], 0, 1); + }, + nullptr); + nth = T::recommended_nth(config_.hidden_size); + backend->do_work_stealing_job( + nth * activated_expert, [&](int _) { T::config(); }, + [&](int task_id) { + int expert_idx = m_expert_id_map_[task_id / nth]; + int ith = task_id % nth; +#ifdef USE_NUMA + amx::mat_mul(m_local_num_[expert_idx], config_.hidden_size, config_.intermediate_size, down_ba_[expert_idx], + down_bb_numa_[Backend::numa_node][expert_idx], down_bc_[expert_idx], ith, nth, use_amx); +#else + amx::mat_mul(m_local_num_[expert_idx], config_.hidden_size, config_.intermediate_size, down_ba_[expert_idx], + down_bb_[expert_idx], down_bc_[expert_idx], ith, nth, use_amx); +#endif + down_bc_[expert_idx]->to_mat(m_local_num_[expert_idx], m_local_down_output_ptr_[expert_idx], ith, nth); + }, + nullptr); + backend->do_work_stealing_job( + qlen, nullptr, + [&](int i) { + for (int e = 0; e < config_.hidden_size; e += 32) { + __m512 x0 = _mm512_setzero_ps(); + __m512 x1 = _mm512_setzero_ps(); + for (int j = 0; j < k; j++) { + __m512 weight = _mm512_set1_ps(weights[i * k + j]); + __m512 down_output0, down_output1; + avx512_32xbf16_to_32xfp32((__m512i *)(m_local_down_output_ptr_[expert_ids[i * k + j]] + + m_local_pos_[i][j] * config_.hidden_size + e), + &down_output0, &down_output1); + x0 = _mm512_fmadd_ps(down_output0, weight, x0); + x1 = _mm512_fmadd_ps(down_output1, weight, x1); + } + avx512_32xfp32_to_32xbf16(&x0, &x1, (__m512i *)((ggml_bf16_t *)output + i * config_.hidden_size + e)); + } + }, + nullptr); + } +}; + +#endif diff --git a/archive/kt-sft/csrc/ktransformers_ext/operators/amx/sft_moe.hpp b/archive/kt-sft/csrc/ktransformers_ext/operators/amx/sft_moe.hpp new file mode 100644 index 0000000..5bac021 --- /dev/null +++ b/archive/kt-sft/csrc/ktransformers_ext/operators/amx/sft_moe.hpp @@ -0,0 +1,805 @@ +/** + * @Description : + * @Author : chenht2022 + * @Date : 2025-04-25 18:28:12 + * @Version : 1.0.0 + * @LastEditors : chenht2022 + * @LastEditTime : 2025-04-25 18:28:12 + * @Copyright (c) 2024 by KVCache.AI, All Rights Reserved. + **/ +#ifndef CPUINFER_OPERATOR_SFT_AMX_MOE_H +#define CPUINFER_OPERATOR_SFT_AMX_MOE_H + +#include +#include +#include +#include +#include +#include +#include + +#include "debug_sft_moe.hpp" + +#include "../../cpu_backend/backend.h" +#include "../../cpu_backend/shared_mem_buffer.h" +#include "llama.cpp/ggml-impl.h" +#include "llama.cpp/ggml-quants.h" +#include "llama.cpp/ggml.h" +#include "llamafile/sgemm.h" + +#include "la/amx.hpp" + +#ifdef USE_NUMA +#include +#include +// void *numa_alloc_aligned(size_t size, int node, size_t alignment) { +// void *ptr = numa_alloc_onnode(size, node); +// assert(reinterpret_cast(ptr) % 64 == 0); +// return ptr; +// } +#endif + +static inline __m512 sigmoid(__m512 x) { + __m512 neg = _mm512_sub_ps(_mm512_setzero_ps(), x); + __m512 e = exp_avx512(neg); + __m512 denom = _mm512_add_ps(_mm512_set1_ps(1.0f), e); + return _mm512_div_ps(_mm512_set1_ps(1.0f), denom); +} + +static inline __m512 act_fn_1(__m512 x) { + __m512 sigmoid_val = sigmoid(x); + return _mm512_mul_ps(sigmoid_val, x); +} + +static inline __m512 act_fn_grad(__m512 x) { + // sigmoid(x) * (1 + x * (1 - sigmoid(x))) + __m512 sigmoid_val = sigmoid(x); + __m512 one_minus_sigmoid = _mm512_sub_ps(_mm512_set1_ps(1.0f), sigmoid_val); + __m512 x_term = _mm512_mul_ps(x, one_minus_sigmoid); + __m512 one_plus_x_term = _mm512_add_ps(_mm512_set1_ps(1.0f), x_term); + return _mm512_mul_ps(sigmoid_val, one_plus_x_term); +} + +struct SFT_AMX_MOEConfig { + int expert_num; + int routed_expert_num; + int hidden_size; + int intermediate_size; + int max_len; + void *gate_proj; + void *up_proj; + void *down_proj; + + SFT_AMX_MOEConfig() {} + + SFT_AMX_MOEConfig(int expert_num, int routed_expert_num, int hidden_size, int intermediate_size, int max_len, + void *gate_proj, void *up_proj, void *down_proj) + : expert_num(expert_num), routed_expert_num(routed_expert_num), hidden_size(hidden_size), + intermediate_size(intermediate_size), max_len(max_len), gate_proj(gate_proj), up_proj(up_proj), + down_proj(down_proj) {} +}; + +template class SFT_AMX_MOE { +private: + SFT_AMX_MOEConfig config_; + void *gate_proj_; // [expert_num * intermediate_size * hidden_size ( /32 if quantized)] + void *up_proj_; // [expert_num * intermediate_size * hidden_size ( /32 if quantized)] + void *down_proj_; // [expert_num * hidden_size * intermediate_size ( /32 if quantized)] + + void *gate_proj_t_; // [expert_num * hidden_size * intermediate_size] + void *up_proj_t_; // [expert_num * hidden_size * intermediate_size] + void *down_proj_t_; // [expert_num * intermediate_size * hidden_size] + + ggml_bf16_t *m_local_input_; // [routed_expert_num * max_len * hidden_size] + ggml_bf16_t *m_local_gate_output_; // [routed_expert_num * max_len * intermediate_size] + ggml_bf16_t *m_local_up_output_; // [routed_expert_num * max_len * intermediate_size] + ggml_bf16_t *m_local_down_output_; // [routed_expert_num * max_len * hidden_size] + + std::vector> m_local_pos_; // [max_len, routed_expert_num] + std::vector m_local_num_; // [expert_num] + std::vector m_expert_id_map_; // [expert_num] + std::vector m_local_input_ptr_; // [expert_num] + std::vector m_local_gate_output_ptr_; // [expert_num] + std::vector m_local_up_output_ptr_; // [expert_num] + std::vector m_local_down_output_ptr_; // [expert_num] + + std::vector> gate_up_ba_; + std::vector> gate_bc_; + std::vector> up_bc_; + std::vector> down_ba_; + std::vector> down_bc_; + +#ifdef USE_NUMA + std::vector>> gate_bb_numa_; + std::vector>> up_bb_numa_; + std::vector>> down_bb_numa_; +#else + std::vector> gate_bb_; + std::vector> up_bb_; + std::vector> down_bb_; +#endif + + ggml_bf16_t *m_local_down_output_grad_; // [routed_expert_num * max_len * hidden_size] + ggml_bf16_t *m_local_down_input_grad_; // [routed_expert_num * max_len * intermediate_size] + ggml_bf16_t *m_local_gate_output_grad_; // [routed_expert_num * max_len * intermediate_size] + ggml_bf16_t *m_local_up_output_grad_; // [routed_expert_num * max_len * intermediate_size] + ggml_bf16_t *m_local_gate_input_grad_; // [routed_expert_num * max_len * hidden_size] + ggml_bf16_t *m_local_up_input_grad_; // [routed_expert_num * max_len * hidden_size] + + std::vector m_local_down_output_grad_ptr_; // [expert_num] + std::vector m_local_down_input_grad_ptr_; // [expert_num] + std::vector m_local_gate_output_grad_ptr_; // [expert_num] + std::vector m_local_up_output_grad_ptr_; // [expert_num] + std::vector m_local_gate_input_grad_ptr_; // [expert_num] + std::vector m_local_up_input_grad_ptr_; // [expert_num] + + std::vector> gate_t_ba_; + std::vector> gate_t_bc_; + std::vector> up_t_ba_; + std::vector> up_t_bc_; + std::vector> down_t_ba_; + std::vector> down_t_bc_; + + // TODO: NUMA +#ifdef USE_NUMA + std::vector>> gate_t_bb_numa_; + std::vector>> up_t_bb_numa_; + std::vector>> down_t_bb_numa_; +#else + std::vector> gate_t_bb_; + std::vector> up_t_bb_; + std::vector> down_t_bb_; +#endif + + int* m_local_token_indices_; // [routed_expert_num * max_len] + int* m_local_expert_positions_; // [routed_expert_num * max_len] + std::vector m_local_token_indices_ptr_; // [expert_num] + std::vector m_local_expert_positions_ptr_; // [expert_num] + +public: + SFT_AMX_MOE(SFT_AMX_MOEConfig config) { + config_ = config; + gate_proj_ = config_.gate_proj; + up_proj_ = config_.up_proj; + down_proj_ = config_.down_proj; + + std::vector> m_mem_requests; + m_mem_requests.push_back({(void **)&m_local_input_, + sizeof(ggml_bf16_t) * config_.routed_expert_num * config_.max_len * config_.hidden_size}); + m_mem_requests.push_back({(void **)&m_local_gate_output_, sizeof(ggml_bf16_t) * config_.routed_expert_num * + config_.max_len * config_.intermediate_size}); + m_mem_requests.push_back({(void **)&m_local_up_output_, sizeof(ggml_bf16_t) * config_.routed_expert_num * + config_.max_len * config_.intermediate_size}); + m_mem_requests.push_back({(void **)&m_local_down_output_, + sizeof(ggml_bf16_t) * config_.routed_expert_num * config_.max_len * config_.hidden_size}); + std::vector gate_up_ba_ptr(config_.expert_num); + std::vector gate_bc_ptr(config_.expert_num); + std::vector up_bc_ptr(config_.expert_num); + std::vector down_ba_ptr(config_.expert_num); + std::vector down_bc_ptr(config_.expert_num); + for (int i = 0; i < config_.expert_num; i++) { + m_mem_requests.push_back( + {(void **)&gate_up_ba_ptr[i], T::BufferA::required_size(config_.max_len, config_.hidden_size)}); + m_mem_requests.push_back( + {(void **)&gate_bc_ptr[i], T::BufferC::required_size(config_.max_len, config_.intermediate_size)}); + m_mem_requests.push_back( + {(void **)&up_bc_ptr[i], T::BufferC::required_size(config_.max_len, config_.intermediate_size)}); + m_mem_requests.push_back( + {(void **)&down_ba_ptr[i], T::BufferA::required_size(config_.max_len, config_.intermediate_size)}); + m_mem_requests.push_back( + {(void **)&down_bc_ptr[i], T::BufferC::required_size(config_.max_len, config_.hidden_size)}); + } + + m_mem_requests.push_back({(void **)&gate_proj_t_, + sizeof(ggml_bf16_t) * config_.expert_num * config_.intermediate_size * config_.hidden_size}); + m_mem_requests.push_back({(void **)&up_proj_t_, + sizeof(ggml_bf16_t) * config_.expert_num * config_.intermediate_size * config_.hidden_size}); + m_mem_requests.push_back({(void **)&down_proj_t_, + sizeof(ggml_bf16_t) * config_.expert_num * config_.hidden_size * config_.intermediate_size}); + + m_mem_requests.push_back({(void **)&m_local_down_output_grad_, + sizeof(ggml_bf16_t) * config_.routed_expert_num * config_.max_len * config_.hidden_size}); + m_mem_requests.push_back({(void **)&m_local_down_input_grad_, + sizeof(ggml_bf16_t) * config_.routed_expert_num * config_.max_len * config_.intermediate_size}); + m_mem_requests.push_back({(void **)&m_local_gate_output_grad_, + sizeof(ggml_bf16_t) * config_.routed_expert_num * config_.max_len * config_.intermediate_size}); + m_mem_requests.push_back({(void **)&m_local_up_output_grad_, + sizeof(ggml_bf16_t) * config_.routed_expert_num * config_.max_len * config_.intermediate_size}); + m_mem_requests.push_back({(void **)&m_local_gate_input_grad_, + sizeof(ggml_bf16_t) * config_.routed_expert_num * config_.max_len * config_.hidden_size}); + m_mem_requests.push_back({(void **)&m_local_up_input_grad_, + sizeof(ggml_bf16_t) * config_.routed_expert_num * config_.max_len * config_.hidden_size}); + m_mem_requests.push_back({(void **)&m_local_token_indices_, + sizeof(int) * config_.routed_expert_num * config_.max_len}); + m_mem_requests.push_back({(void **)&m_local_expert_positions_, + sizeof(int) * config_.routed_expert_num * config_.max_len}); + std::vector gate_t_ba_ptr(config_.expert_num); + std::vector gate_t_bc_ptr(config_.expert_num); + std::vector up_t_ba_ptr(config_.expert_num); + std::vector up_t_bc_ptr(config_.expert_num); + std::vector down_t_ba_ptr(config_.expert_num); + std::vector down_t_bc_ptr(config_.expert_num); + for (int i = 0; i < config_.expert_num; i++) { + m_mem_requests.push_back( + {(void **)&gate_t_ba_ptr[i], T::BufferA::required_size(config_.max_len, config_.intermediate_size)}); + m_mem_requests.push_back( + {(void **)&gate_t_bc_ptr[i], T::BufferC::required_size(config_.max_len, config_.hidden_size)}); + m_mem_requests.push_back( + {(void **)&up_t_ba_ptr[i], T::BufferA::required_size(config_.max_len, config_.intermediate_size)}); + m_mem_requests.push_back( + {(void **)&up_t_bc_ptr[i], T::BufferC::required_size(config_.max_len, config_.hidden_size)}); + m_mem_requests.push_back( + {(void **)&down_t_ba_ptr[i], T::BufferA::required_size(config_.max_len, config_.hidden_size)}); + m_mem_requests.push_back( + {(void **)&down_t_bc_ptr[i], T::BufferC::required_size(config_.max_len, config_.intermediate_size)}); + } + + shared_mem_buffer.alloc(this, m_mem_requests); + + m_local_pos_.resize(config_.max_len); + for (int i = 0; i < config_.max_len; i++) { + m_local_pos_[i].resize(config_.routed_expert_num); + } + m_expert_id_map_.resize(config_.expert_num); + m_local_num_.resize(config_.expert_num); + m_local_input_ptr_.resize(config_.expert_num); + m_local_gate_output_ptr_.resize(config_.expert_num); + m_local_up_output_ptr_.resize(config_.expert_num); + m_local_down_output_ptr_.resize(config_.expert_num); + m_local_down_output_grad_ptr_.resize(config_.expert_num); + m_local_down_input_grad_ptr_.resize(config_.expert_num); + m_local_gate_output_grad_ptr_.resize(config_.expert_num); + m_local_up_output_grad_ptr_.resize(config_.expert_num); + m_local_gate_input_grad_ptr_.resize(config_.expert_num); + m_local_up_input_grad_ptr_.resize(config_.expert_num); + + for (uint64_t i = 0; i < config_.expert_num; i++) { + gate_up_ba_.push_back( + std::make_shared(config_.max_len, config_.hidden_size, gate_up_ba_ptr[i])); + gate_bc_.push_back( + std::make_shared(config_.max_len, config_.intermediate_size, gate_bc_ptr[i])); + up_bc_.push_back(std::make_shared(config_.max_len, config_.intermediate_size, up_bc_ptr[i])); + down_ba_.push_back( + std::make_shared(config_.max_len, config_.intermediate_size, down_ba_ptr[i])); + down_bc_.push_back(std::make_shared(config_.max_len, config_.hidden_size, down_bc_ptr[i])); + +#ifdef USE_NUMA + int numa_nodes = numa_num_configured_nodes(); + gate_bb_numa_.resize(numa_nodes); + up_bb_numa_.resize(numa_nodes); + down_bb_numa_.resize(numa_nodes); + for (int j = 0; j < numa_nodes; j++) { + void *gate_bb_ptr = + numa_alloc_aligned(T::BufferB::required_size(config_.intermediate_size, config_.hidden_size), j, 64); + gate_bb_numa_[j].push_back( + std::make_shared(config_.intermediate_size, config_.hidden_size, gate_bb_ptr)); + void *up_bb_ptr = + numa_alloc_aligned(T::BufferB::required_size(config_.intermediate_size, config_.hidden_size), j, 64); + up_bb_numa_[j].push_back( + std::make_shared(config_.intermediate_size, config_.hidden_size, up_bb_ptr)); + void *down_bb_ptr = + numa_alloc_aligned(T::BufferB::required_size(config_.hidden_size, config_.intermediate_size), j, 64); + down_bb_numa_[j].push_back( + std::make_shared(config_.hidden_size, config_.intermediate_size, down_bb_ptr)); + } +#else + void *gate_bb_ptr = + std::aligned_alloc(64, T::BufferB::required_size(config_.intermediate_size, config_.hidden_size)); + gate_bb_.push_back( + std::make_shared(config_.intermediate_size, config_.hidden_size, gate_bb_ptr)); + + void *up_bb_ptr = + std::aligned_alloc(64, T::BufferB::required_size(config_.intermediate_size, config_.hidden_size)); + up_bb_.push_back( + std::make_shared(config_.intermediate_size, config_.hidden_size, up_bb_ptr)); + + void *down_bb_ptr = + std::aligned_alloc(64, T::BufferB::required_size(config_.hidden_size, config_.intermediate_size)); + down_bb_.push_back( + std::make_shared(config_.hidden_size, config_.intermediate_size, down_bb_ptr)); +#endif + } + + for (uint64_t i = 0; i < config_.expert_num; i++) { + gate_t_ba_.push_back( + std::make_shared(config_.max_len, config_.intermediate_size, gate_t_ba_ptr[i])); + gate_t_bc_.push_back( + std::make_shared(config_.max_len, config_.hidden_size, gate_t_bc_ptr[i])); + up_t_ba_.push_back(std::make_shared(config_.max_len, config_.intermediate_size, up_t_ba_ptr[i])); + up_t_bc_.push_back(std::make_shared(config_.max_len, config_.hidden_size, up_t_bc_ptr[i])); + down_t_ba_.push_back( + std::make_shared(config_.max_len, config_.hidden_size, down_t_ba_ptr[i])); + down_t_bc_.push_back(std::make_shared(config_.max_len, config_.intermediate_size, down_t_bc_ptr[i])); + +#ifdef USE_NUMA + int numa_nodes = numa_num_configured_nodes(); + gate_t_bb_numa_.resize(numa_nodes); + up_t_bb_numa_.resize(numa_nodes); + down_t_bb_numa_.resize(numa_nodes); + for (int j = 0; j < numa_nodes; j++) { + void *gate_t_bb_ptr = + numa_alloc_aligned(T::BufferB::required_size(config_.hidden_size, config_.intermediate_size), j, 64); + gate_t_bb_numa_[j].push_back( + std::make_shared(config_.hidden_size, config_.intermediate_size, gate_t_bb_ptr)); + void *up_t_bb_ptr = + numa_alloc_aligned(T::BufferB::required_size(config_.hidden_size, config_.intermediate_size), j, 64); + up_t_bb_numa_[j].push_back( + std::make_shared(config_.hidden_size, config_.intermediate_size, up_t_bb_ptr)); + void *down_t_bb_ptr = + numa_alloc_aligned(T::BufferB::required_size(config_.intermediate_size, config_.hidden_size), j, 64); + down_t_bb_numa_[j].push_back( + std::make_shared(config_.intermediate_size, config_.hidden_size, down_t_bb_ptr)); + } +#else + void *gate_t_bb_ptr = + std::aligned_alloc(64, T::BufferB::required_size(config_.hidden_size, config_.intermediate_size)); + gate_t_bb_.push_back( + std::make_shared(config_.hidden_size, config_.intermediate_size, gate_t_bb_ptr)); + + void *up_t_bb_ptr = + std::aligned_alloc(64, T::BufferB::required_size(config_.hidden_size, config_.intermediate_size)); + up_t_bb_.push_back( + std::make_shared(config_.hidden_size, config_.intermediate_size, up_t_bb_ptr)); + + void *down_t_bb_ptr = + std::aligned_alloc(64, T::BufferB::required_size(config_.intermediate_size, config_.hidden_size)); + down_t_bb_.push_back( + std::make_shared(config_.intermediate_size, config_.hidden_size, down_t_bb_ptr)); +#endif + } + + m_local_token_indices_ptr_.resize(config_.expert_num); + m_local_expert_positions_ptr_.resize(config_.expert_num); + } + + ~SFT_AMX_MOE() { shared_mem_buffer.dealloc(this); } + + void transpose_expert(const void* src, void* dst, int R, int C, Backend* backend) { + backend->do_work_stealing_job( + config_.expert_num, nullptr, + [&](uint64_t expert_idx) { + for (int r = 0; r < R; ++r) { + for (int c = 0; c < C; ++c) { + memcpy( + (uint8_t*)dst + (expert_idx * R * C + (c * R + r)) * sizeof(ggml_bf16_t), + (uint8_t*)src + (expert_idx * R * C + (r * C + c)) * sizeof(ggml_bf16_t), + sizeof(ggml_bf16_t)); + } + } + }, + nullptr); + } + + void load_weights(Backend *backend) { + transpose_expert(config_.gate_proj, gate_proj_t_, config_.intermediate_size, config_.hidden_size, backend); + transpose_expert(config_.up_proj, up_proj_t_, config_.intermediate_size, config_.hidden_size, backend); + transpose_expert(config_.down_proj, down_proj_t_, config_.hidden_size, config_.intermediate_size, backend); + + int nth = T::recommended_nth(config_.intermediate_size); + backend->do_work_stealing_job( + nth * config_.expert_num, nullptr, + [&](int task_id) { + uint64_t expert_idx = task_id / nth; + int ith = task_id % nth; +#ifdef USE_NUMA + int numa_nodes = numa_num_configured_nodes(); + for (int j = 0; j < numa_nodes; j++) { + gate_bb_numa_[j][expert_idx]->from_mat((ggml_bf16_t *)config_.gate_proj + + expert_idx * config_.intermediate_size * config_.hidden_size, + ith, nth); + up_bb_numa_[j][expert_idx]->from_mat((ggml_bf16_t *)config_.up_proj + + expert_idx * config_.intermediate_size * config_.hidden_size, + ith, nth); + } +#else + gate_bb_[expert_idx]->from_mat((ggml_bf16_t *)config_.gate_proj + + expert_idx * config_.intermediate_size * config_.hidden_size, + ith, nth); + up_bb_[expert_idx]->from_mat( + (ggml_bf16_t *)config_.up_proj + expert_idx * config_.intermediate_size * config_.hidden_size, ith, nth); +#endif + }, + nullptr); + nth = T::recommended_nth(config_.intermediate_size); + backend->do_work_stealing_job( + nth * config_.expert_num, nullptr, + [&](int task_id) { + uint64_t expert_idx = task_id / nth; + int ith = task_id % nth; +#ifdef USE_NUMA + int numa_nodes = numa_num_configured_nodes(); + for (int j = 0; j < numa_nodes; j++) { + down_t_bb_numa_[j][expert_idx]->from_mat((ggml_bf16_t *)down_proj_t_ + + expert_idx * config_.intermediate_size * config_.hidden_size, + ith, nth); + } +#else + down_t_bb_[expert_idx]->from_mat((ggml_bf16_t *)down_proj_t_ + + expert_idx * config_.intermediate_size * config_.hidden_size, + ith, nth); +#endif + }, + nullptr); + + nth = T::recommended_nth(config_.hidden_size); + backend->do_work_stealing_job( + nth * config_.expert_num, nullptr, + [&](int task_id) { + uint64_t expert_idx = task_id / nth; + int ith = task_id % nth; +#ifdef USE_NUMA + int numa_nodes = numa_num_configured_nodes(); + for (int j = 0; j < numa_nodes; j++) { + down_bb_numa_[j][expert_idx]->from_mat((ggml_bf16_t *)config_.down_proj + + expert_idx * config_.hidden_size * config_.intermediate_size, + ith, nth); + } +#else + down_bb_[expert_idx]->from_mat((ggml_bf16_t *)config_.down_proj + + expert_idx * config_.hidden_size * config_.intermediate_size, + ith, nth); +#endif +#ifdef USE_NUMA + for (int j = 0; j < numa_nodes; j++) { + gate_t_bb_numa_[j][expert_idx]->from_mat((ggml_bf16_t *)gate_proj_t_ + + expert_idx * config_.hidden_size * config_.intermediate_size, + ith, nth); + up_t_bb_numa_[j][expert_idx]->from_mat((ggml_bf16_t *)up_proj_t_ + + expert_idx * config_.hidden_size * config_.intermediate_size, + ith, nth); + } +#else + gate_t_bb_[expert_idx]->from_mat((ggml_bf16_t *)gate_proj_t_ + + expert_idx * config_.hidden_size * config_.intermediate_size, + ith, nth); + up_t_bb_[expert_idx]->from_mat((ggml_bf16_t *)up_proj_t_ + + expert_idx * config_.hidden_size * config_.intermediate_size, + ith, nth); +#endif + }, + nullptr); + } + + void warm_up(Backend *backend) {} + + void forward(int qlen, int k, const uint64_t *expert_ids, const float *weights, const void *input, void *output, Backend *backend) { + bool use_amx = (qlen > 4 * config_.expert_num / config_.routed_expert_num); + int activated_expert = 0; + for (int i = 0; i < config_.expert_num; i++) { + m_local_num_[i] = 0; + } + for (int i = 0; i < qlen; i++) { + for (int j = 0; j < k; j++) { + m_local_pos_[i][j] = m_local_num_[expert_ids[i * k + j]]++; + } + } + for (int i = 0; i < config_.expert_num; i++) { + if (m_local_num_[i] > 0) { + m_expert_id_map_[activated_expert] = i; + activated_expert++; + } + } + uint64_t offset = 0; + for (int i = 0; i < config_.expert_num; i++) { + m_local_input_ptr_[i] = m_local_input_ + offset * config_.hidden_size; + m_local_gate_output_ptr_[i] = m_local_gate_output_ + offset * config_.intermediate_size; + m_local_up_output_ptr_[i] = m_local_up_output_ + offset * config_.intermediate_size; + m_local_down_output_ptr_[i] = m_local_down_output_ + offset * config_.hidden_size; + offset += m_local_num_[i]; + } + backend->do_work_stealing_job( + qlen, nullptr, + [&](int i) { + for (int j = 0; j < k; j++) { + memcpy(m_local_input_ptr_[expert_ids[i * k + j]] + m_local_pos_[i][j] * config_.hidden_size, + (ggml_bf16_t *)input + i * config_.hidden_size, sizeof(ggml_bf16_t) * config_.hidden_size); + } + }, + nullptr); + backend->do_work_stealing_job( + activated_expert, nullptr, + [&](int task_id) { + int expert_idx = m_expert_id_map_[task_id]; + + gate_up_ba_[expert_idx]->from_mat(m_local_num_[expert_idx], m_local_input_ptr_[expert_idx], 0, 1); + }, + nullptr); + int nth = T::recommended_nth(config_.intermediate_size); + backend->do_work_stealing_job( + nth * activated_expert, [&](int _) { T::config(); }, + [&](int task_id) { + int expert_idx = m_expert_id_map_[task_id / nth]; + int ith = task_id % nth; +#ifdef USE_NUMA + amx::mat_mul(m_local_num_[expert_idx], config_.intermediate_size, config_.hidden_size, + gate_up_ba_[expert_idx], gate_bb_numa_[Backend::numa_node][expert_idx], gate_bc_[expert_idx], + ith, nth, use_amx); + amx::mat_mul(m_local_num_[expert_idx], config_.intermediate_size, config_.hidden_size, + gate_up_ba_[expert_idx], up_bb_numa_[Backend::numa_node][expert_idx], up_bc_[expert_idx], ith, + nth, use_amx); +#else + amx::mat_mul(m_local_num_[expert_idx], config_.intermediate_size, config_.hidden_size, + gate_up_ba_[expert_idx], gate_bb_[expert_idx], gate_bc_[expert_idx], ith, nth, use_amx); + amx::mat_mul(m_local_num_[expert_idx], config_.intermediate_size, config_.hidden_size, + gate_up_ba_[expert_idx], up_bb_[expert_idx], up_bc_[expert_idx], ith, nth, use_amx); +#endif + gate_bc_[expert_idx]->to_mat(m_local_num_[expert_idx], m_local_gate_output_ptr_[expert_idx], ith, nth); + up_bc_[expert_idx]->to_mat(m_local_num_[expert_idx], m_local_up_output_ptr_[expert_idx], ith, nth); + auto [n_start, n_end] = T::split_range_n(config_.intermediate_size, ith, nth); + for (int i = 0; i < m_local_num_[expert_idx]; i++) { + ggml_bf16_t *gate_output_ptr = &m_local_gate_output_ptr_[expert_idx][i * config_.intermediate_size]; + ggml_bf16_t *up_output_ptr = &m_local_up_output_ptr_[expert_idx][i * config_.intermediate_size]; + for (int j = n_start; j < n_end; j += 32) { + __m512 gate_val0, gate_val1, up_val0, up_val1; + avx512_32xbf16_to_32xfp32((__m512i *)(gate_output_ptr + j), &gate_val0, &gate_val1); + avx512_32xbf16_to_32xfp32((__m512i *)(up_output_ptr + j), &up_val0, &up_val1); + __m512 result0 = act_fn(gate_val0, up_val0); + __m512 result1 = act_fn(gate_val1, up_val1); + avx512_32xfp32_to_32xbf16(&result0, &result1, (__m512i *)(gate_output_ptr + j)); + } + } + }, + nullptr); + backend->do_work_stealing_job( + activated_expert, nullptr, + [&](int task_id) { + int expert_idx = m_expert_id_map_[task_id]; + down_ba_[expert_idx]->from_mat(m_local_num_[expert_idx], m_local_gate_output_ptr_[expert_idx], 0, 1); + }, + nullptr); + + nth = T::recommended_nth(config_.hidden_size); + backend->do_work_stealing_job( + nth * activated_expert, [&](int _) { T::config(); }, + [&](int task_id) { + int expert_idx = m_expert_id_map_[task_id / nth]; + int ith = task_id % nth; +#ifdef USE_NUMA + amx::mat_mul(m_local_num_[expert_idx], config_.hidden_size, config_.intermediate_size, down_ba_[expert_idx], + down_bb_numa_[Backend::numa_node][expert_idx], down_bc_[expert_idx], ith, nth, use_amx); +#else + amx::mat_mul(m_local_num_[expert_idx], config_.hidden_size, config_.intermediate_size, down_ba_[expert_idx], + down_bb_[expert_idx], down_bc_[expert_idx], ith, nth, use_amx); +#endif + down_bc_[expert_idx]->to_mat(m_local_num_[expert_idx], m_local_down_output_ptr_[expert_idx], ith, nth); + }, + nullptr); + backend->do_work_stealing_job( + qlen, nullptr, + [&](int i) { + for (int e = 0; e < config_.hidden_size; e += 32) { + __m512 x0 = _mm512_setzero_ps(); + __m512 x1 = _mm512_setzero_ps(); + for (int j = 0; j < k; j++) { + __m512 weight = _mm512_set1_ps(weights[i * k + j]); + __m512 down_output0, down_output1; + avx512_32xbf16_to_32xfp32((__m512i *)(m_local_down_output_ptr_[expert_ids[i * k + j]] + + m_local_pos_[i][j] * config_.hidden_size + e), + &down_output0, &down_output1); + x0 = _mm512_fmadd_ps(down_output0, weight, x0); + x1 = _mm512_fmadd_ps(down_output1, weight, x1); + } + avx512_32xfp32_to_32xbf16(&x0, &x1, (__m512i *)((ggml_bf16_t *)output + i * config_.hidden_size + e)); + } + }, + nullptr); + } + + void backward(int qlen, int k, const uint64_t *expert_ids, const float *weights, const void* input, const void *output_grad, void *input_grad, Backend *backend) { + bool use_amx = (qlen > 4 * config_.expert_num / config_.routed_expert_num); + int activated_expert = 0; + for (int i = 0; i < config_.expert_num; i++) { + m_local_num_[i] = 0; + } + for (int i = 0; i < qlen; i++) { + for (int j = 0; j < k; j++) { + m_local_pos_[i][j] = m_local_num_[expert_ids[i * k + j]]++; + } + } + for (int i = 0; i < config_.expert_num; i++) { + if (m_local_num_[i] > 0) { + m_expert_id_map_[activated_expert] = i; + activated_expert++; + } + } + uint64_t offset = 0; + for (int i = 0; i < config_.expert_num; i++) { + m_local_input_ptr_[i] = m_local_input_ + offset * config_.hidden_size; + m_local_gate_output_ptr_[i] = m_local_gate_output_ + offset * config_.intermediate_size; + m_local_up_output_ptr_[i] = m_local_up_output_ + offset * config_.intermediate_size; + + m_local_down_output_grad_ptr_[i] = m_local_down_output_grad_ + offset * config_.hidden_size; + m_local_down_input_grad_ptr_[i] = m_local_down_input_grad_ + offset * config_.intermediate_size; + m_local_gate_output_grad_ptr_[i] = m_local_gate_output_grad_ + offset * config_.intermediate_size; + m_local_up_output_grad_ptr_[i] = m_local_up_output_grad_ + offset * config_.intermediate_size; + m_local_gate_input_grad_ptr_[i] = m_local_gate_input_grad_ + offset * config_.hidden_size; + m_local_up_input_grad_ptr_[i] = m_local_up_input_grad_ + offset * config_.hidden_size; + m_local_token_indices_ptr_[i] = m_local_token_indices_ + offset; + m_local_expert_positions_ptr_[i] = m_local_expert_positions_ + offset; + offset += m_local_num_[i]; + } + + // TODO: cache + backend->do_work_stealing_job( + qlen, nullptr, + [&](int i) { + for (int j = 0; j < k; j++) { + uint64_t expert_id = expert_ids[i * k + j]; + int local_row = m_local_pos_[i][j]; + memcpy(m_local_input_ptr_[expert_id] + local_row * config_.hidden_size, + (ggml_bf16_t *)input + i * config_.hidden_size, sizeof(ggml_bf16_t) * config_.hidden_size); // TODO: cache + memcpy(m_local_down_output_grad_ptr_[expert_id] + local_row * config_.hidden_size, + (ggml_bf16_t *)output_grad + i * config_.hidden_size, sizeof(ggml_bf16_t) * config_.hidden_size); + m_local_token_indices_ptr_[expert_id][local_row] = i; + m_local_expert_positions_ptr_[expert_id][local_row] = j; + } + }, + nullptr); + + backend->do_work_stealing_job( + activated_expert, nullptr, + [&](int task_id) { + int expert_idx = m_expert_id_map_[task_id]; + gate_up_ba_[expert_idx]->from_mat(m_local_num_[expert_idx], m_local_input_ptr_[expert_idx], 0, 1); // TODO: cache + down_t_ba_[expert_idx]->from_mat(m_local_num_[expert_idx], m_local_down_output_grad_ptr_[expert_idx], 0, 1); + }, + nullptr); + + int nth = T::recommended_nth(config_.intermediate_size); + backend->do_work_stealing_job( + nth * activated_expert, [&](int _) { T::config(); }, + [&](int task_id) { + int expert_idx = m_expert_id_map_[task_id / nth]; + int ith = task_id % nth; + + // TODO: cache +#ifdef USE_NUMA + amx::mat_mul(m_local_num_[expert_idx], config_.intermediate_size, config_.hidden_size, + gate_up_ba_[expert_idx], gate_bb_numa_[Backend::numa_node][expert_idx], gate_bc_[expert_idx], + ith, nth, use_amx); + amx::mat_mul(m_local_num_[expert_idx], config_.intermediate_size, config_.hidden_size, + gate_up_ba_[expert_idx], up_bb_numa_[Backend::numa_node][expert_idx], up_bc_[expert_idx], ith, + nth, use_amx); +#else + amx::mat_mul(m_local_num_[expert_idx], config_.intermediate_size, config_.hidden_size, + gate_up_ba_[expert_idx], gate_bb_[expert_idx], gate_bc_[expert_idx], ith, nth, use_amx); + amx::mat_mul(m_local_num_[expert_idx], config_.intermediate_size, config_.hidden_size, + gate_up_ba_[expert_idx], up_bb_[expert_idx], up_bc_[expert_idx], ith, nth, use_amx); +#endif + gate_bc_[expert_idx]->to_mat(m_local_num_[expert_idx], m_local_gate_output_ptr_[expert_idx], ith, nth); + up_bc_[expert_idx]->to_mat(m_local_num_[expert_idx], m_local_up_output_ptr_[expert_idx], ith, nth); + +#ifdef USE_NUMA + amx::mat_mul(m_local_num_[expert_idx], config_.intermediate_size, config_.hidden_size, + down_t_ba_[expert_idx], down_t_bb_numa_[Backend::numa_node][expert_idx], down_t_bc_[expert_idx], ith, nth, use_amx); +#else + amx::mat_mul(m_local_num_[expert_idx], config_.intermediate_size, config_.hidden_size, + down_t_ba_[expert_idx], down_t_bb_[expert_idx], down_t_bc_[expert_idx], ith, nth, use_amx); +#endif + down_t_bc_[expert_idx]->to_mat(m_local_num_[expert_idx], m_local_down_input_grad_ptr_[expert_idx], ith, nth); + + + auto [n_start, n_end] = T::split_range_n(config_.intermediate_size, ith, nth); + for (int i = 0; i < m_local_num_[expert_idx]; i++) { + ggml_bf16_t *gate_output_ptr = &m_local_gate_output_ptr_[expert_idx][i * config_.intermediate_size]; + ggml_bf16_t *up_output_ptr = &m_local_up_output_ptr_[expert_idx][i * config_.intermediate_size]; + ggml_bf16_t *down_input_grad_ptr = &m_local_down_input_grad_ptr_[expert_idx][i * config_.intermediate_size]; + ggml_bf16_t *gate_output_grad_ptr = &m_local_gate_output_grad_ptr_[expert_idx][i * config_.intermediate_size]; + ggml_bf16_t *up_output_grad_ptr = &m_local_up_output_grad_ptr_[expert_idx][i * config_.intermediate_size]; + + int token_idx = m_local_token_indices_ptr_[expert_idx][i]; + int expert_pos = m_local_expert_positions_ptr_[expert_idx][i]; + __m512 weight = _mm512_set1_ps(weights[token_idx * k + expert_pos]); + + for (int j = n_start; j < n_end; j += 32) { + __m512 gate_val0, gate_val1, up_val0, up_val1, down_input_grad0, down_input_grad1; + avx512_32xbf16_to_32xfp32((__m512i *)(gate_output_ptr + j), &gate_val0, &gate_val1); + avx512_32xbf16_to_32xfp32((__m512i *)(up_output_ptr + j), &up_val0, &up_val1); + avx512_32xbf16_to_32xfp32((__m512i *)(down_input_grad_ptr + j), &down_input_grad0, &down_input_grad1); + + down_input_grad0 = _mm512_mul_ps(down_input_grad0, weight); + down_input_grad1 = _mm512_mul_ps(down_input_grad1, weight); + + // gate_output_grad = δ_zji ⊙ v_ji ⊙ σ'(u_ji) + __m512 gate_grad0 = _mm512_mul_ps(down_input_grad0, + _mm512_mul_ps(up_val0, act_fn_grad(gate_val0))); + __m512 gate_grad1 = _mm512_mul_ps(down_input_grad1, + _mm512_mul_ps(up_val1, act_fn_grad(gate_val1))); + + // up_output_grad = δ_zji ⊙ σ(u_ji) + __m512 up_grad0 = _mm512_mul_ps(down_input_grad0, act_fn_1(gate_val0)); + __m512 up_grad1 = _mm512_mul_ps(down_input_grad1, act_fn_1(gate_val1)); + + avx512_32xfp32_to_32xbf16(&gate_grad0, &gate_grad1, (__m512i *)(gate_output_grad_ptr + j)); + avx512_32xfp32_to_32xbf16(&up_grad0, &up_grad1, (__m512i *)(up_output_grad_ptr + j)); + } + } + }, + nullptr); + + + backend->do_work_stealing_job( + activated_expert, nullptr, + [&](int task_id) { + int expert_idx = m_expert_id_map_[task_id]; + gate_t_ba_[expert_idx]->from_mat(m_local_num_[expert_idx], m_local_gate_output_grad_ptr_[expert_idx], 0, 1); + up_t_ba_[expert_idx]->from_mat(m_local_num_[expert_idx], m_local_up_output_grad_ptr_[expert_idx], 0, 1); + }, + nullptr); + nth = T::recommended_nth(config_.hidden_size); + backend->do_work_stealing_job( + nth * activated_expert, [&](int _) { T::config(); }, + [&](int task_id) { + int expert_idx = m_expert_id_map_[task_id / nth]; + int ith = task_id % nth; +#ifdef USE_NUMA + amx::mat_mul(m_local_num_[expert_idx], config_.hidden_size, config_.intermediate_size, + gate_t_ba_[expert_idx], gate_t_bb_numa_[Backend::numa_node][expert_idx], gate_t_bc_[expert_idx], ith, nth, use_amx); + amx::mat_mul(m_local_num_[expert_idx], config_.hidden_size, config_.intermediate_size, + up_t_ba_[expert_idx], up_t_bb_numa_[Backend::numa_node][expert_idx], up_t_bc_[expert_idx], ith, nth, use_amx); +#else + amx::mat_mul(m_local_num_[expert_idx], config_.hidden_size, config_.intermediate_size, + gate_t_ba_[expert_idx], gate_t_bb_[expert_idx], gate_t_bc_[expert_idx], ith, nth, use_amx); + amx::mat_mul(m_local_num_[expert_idx], config_.hidden_size, config_.intermediate_size, + up_t_ba_[expert_idx], up_t_bb_[expert_idx], up_t_bc_[expert_idx], ith, nth, use_amx); +#endif + gate_t_bc_[expert_idx]->to_mat(m_local_num_[expert_idx], m_local_gate_input_grad_ptr_[expert_idx], ith, nth); + up_t_bc_[expert_idx]->to_mat(m_local_num_[expert_idx], m_local_up_input_grad_ptr_[expert_idx], ith, nth); + }, + nullptr); + backend->do_work_stealing_job( + qlen, nullptr, + [&](int i) { + for (int e = 0; e < config_.hidden_size; e += 32) { + __m512 x0 = _mm512_setzero_ps(); + __m512 x1 = _mm512_setzero_ps(); + for (int j = 0; j < k; j++) { + __m512 gate_input_grad0, gate_input_grad1, up_input_grad0, up_input_grad1; + avx512_32xbf16_to_32xfp32((__m512i *)(m_local_gate_input_grad_ptr_[expert_ids[i * k + j]] + + m_local_pos_[i][j] * config_.hidden_size + e), + &gate_input_grad0, &gate_input_grad1); + avx512_32xbf16_to_32xfp32((__m512i *)(m_local_up_input_grad_ptr_[expert_ids[i * k + j]] + + m_local_pos_[i][j] * config_.hidden_size + e), + &up_input_grad0, &up_input_grad1); + x0 = _mm512_add_ps(gate_input_grad0, x0); + x1 = _mm512_add_ps(gate_input_grad1, x1); + x0 = _mm512_add_ps(up_input_grad0, x0); + x1 = _mm512_add_ps(up_input_grad1, x1); + } + avx512_32xfp32_to_32xbf16(&x0, &x1, (__m512i *)((ggml_bf16_t *)input_grad + i * config_.hidden_size + e)); + } + }, + nullptr); + } +}; +#endif + +// for debug +// if constexpr (std::is_same_v) { +// for (int expert_idx = 0; expert_idx < config_.expert_num; ++expert_idx) { +// auto buf = down_t_ba_[expert_idx].get(); + +// std::string path = "debug/" + std::to_string(expert_idx) + "_down_ba_t_debug3.bin"; +// std::ofstream ofs(path, std::ios::binary); +// for (int n_idx = 0; n_idx < m_local_num_[expert_idx]; ++n_idx) { +// const ggml_bf16_t* row = reinterpret_cast(buf->a) + n_idx * buf->k; +// for (int j = 0; j < buf->k; ++j) { +// float v = row[j]; +// ofs.write(reinterpret_cast(&v), sizeof(v)); +// } +// } +// ofs.close(); +// } +// } + +// for (uint64_t expert_idx = 0; expert_idx < (uint64_t)config_.expert_num; ++expert_idx) { +// dump_grad_bin("cpp_layer0_E_End"+std::to_string(expert_idx)+"_down_t_ba_", (ggml_bf16_t*)m_local_down_output_grad_ptr_[expert_idx], config_.hidden_size * m_local_num_[expert_idx], GGML_TYPE_BF16); +// } + +// for (uint64_t expert_idx = 0; expert_idx < (uint64_t)config_.expert_num; ++expert_idx) { +// dump_grad_bin("cpp_layer0_E_End"+std::to_string(expert_idx)+"_down_t_bb_", (ggml_bf16_t *)down_proj_t_ + expert_idx * config_.intermediate_size * config_.hidden_size, config_.hidden_size * config_.intermediate_size, GGML_TYPE_BF16); +// } + +// for (uint64_t expert_idx = 0; expert_idx < (uint64_t)config_.expert_num; ++expert_idx) { +// dump_grad_bin("cpp_layer0_E_End"+std::to_string(expert_idx)+"_down_t_bc_", (ggml_bf16_t*)m_local_down_input_grad_ptr_[expert_idx], config_.intermediate_size * m_local_num_[expert_idx], GGML_TYPE_BF16); +// } \ No newline at end of file diff --git a/archive/kt-sft/csrc/ktransformers_ext/operators/kvcache/kvcache.h b/archive/kt-sft/csrc/ktransformers_ext/operators/kvcache/kvcache.h new file mode 100644 index 0000000..ac91778 --- /dev/null +++ b/archive/kt-sft/csrc/ktransformers_ext/operators/kvcache/kvcache.h @@ -0,0 +1,727 @@ +/** + * @Description : + * @Author : Jianwei Dong + * @Date : 2024-08-26 22:47:06 + * @Version : 1.0.0 + * @LastEditors : Jianwei Dong + * @LastEditTime : 2024-08-26 22:47:06 + * @Copyright (c) 2024 by KVCache.AI, All Rights Reserved. + **/ + +#ifndef CPUINFER_OPERATOR_KVCACHE_H +#define CPUINFER_OPERATOR_KVCACHE_H + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "../../cpu_backend/backend.h" +#include "llama.cpp/ggml-common.h" +#include "llama.cpp/ggml-impl.h" +#include "llama.cpp/ggml-quants.h" +#include "llama.cpp/ggml.h" +#include "llamafile/sgemm.h" + +#define CHUNK_SIZE 32 + +/** + * @brief Converts a ggml_type enum value to its corresponding string + * representation. + * + * This function provides a human-readable string representation for a given + * ggml_type enum value. The string can be used for logging, debugging, or + * displaying information in a user interface. + * + * @param type The ggml_type enum value to convert. + * @return A string representation of the enum value. + */ +std::string ggml_type_to_string(ggml_type type); + +/** + * @enum AnchorType + * @brief Defines the types of anchors used in attention mechanisms. + * + * This enum specifies different types of anchors that can be used in attention + * mechanisms, such as fixed anchors, dynamic anchors, or special anchors like + * QUEST, BLOCK_MEAN, or BLOCK_MAX. + */ +enum AnchorType { + FIXED_ANCHOR, /**< A fixed anchor that does not change. */ + DYNAMIC, /**< A dynamic anchor that can change over time. */ + QUEST, /**< A special anchor type used for QUEST (Query and Embedding Space + Transformation). */ + BLOCK_MEAN, /**< An anchor based on the mean of a block of data. */ + BLOCK_MAX /**< An anchor based on the maximum value within a block of data. + */ +}; + +/** + * @brief Converts an AnchorType enum value to its corresponding string + * representation. + * + * This function provides a human-readable string representation for a given + * AnchorType enum value. The string can be used for logging, debugging, or + * displaying information in a user interface. + * + * @param anchor_type The AnchorType enum value to convert. + * @return A string representation of the enum value. + */ +std::string AnchorTypeToString(AnchorType anchor_type); + +/** + * @enum RetrievalType + * @brief Defines the types of retrieval strategies in attention mechanisms. + * + * This enum specifies different retrieval strategies that can be used in + * attention mechanisms, such as layer-level retrieval, key-value head-level + * retrieval, or query head-level retrieval. + */ +enum RetrievalType { + LAYER, /**< Retrieval at the layer level. */ + KVHEAD, /**< Retrieval at the key-value head level. */ + QHEAD /**< Retrieval at the query head level. */ +}; + +/** + * @brief Converts a RetrievalType enum value to its corresponding string + * representation. + * + * This function provides a human-readable string representation for a given + * RetrievalType enum value. The string can be used for logging, debugging, or + * displaying information in a user interface. + * + * @param retrieval_type The RetrievalType enum value to convert. + * @return A string representation of the enum value. + */ +std::string RetrievalTypeToString(RetrievalType retrieval_type); + +/** + * @struct KVCacheConfig + * @brief Configuration structure for Key-Value (KV) Cache. + * + * This structure holds configuration parameters for setting up and managing + * a Key-Value (KV) Cache used in various attention mechanisms. It includes + * parameters such as the number of layers, the number of heads, the dimension + * of each head, block length, anchor information, and memory-related settings. + */ +struct KVCacheConfig { + int layer_num; /**< Number of layers in the model. */ + int kv_head_num; /**< Number of heads in the KV Cache. */ + int q_head_num; /**< Number of heads in the query. */ + int head_dim; /**< Dimension of each head. */ + int block_len; /**< Length of each block in the cache. */ + int anchor_num; /**< Number of anchors used in attention. */ + + ggml_type kv_type; /**< Data type of the KV Cache (e.g., fp16, q8_0). */ + + // Controls the pre-allocated memory size + int max_block_num; /**< Maximum number of blocks that can be allocated. */ + int max_batch_size; /**< Maximum batch size that can be processed. */ + int max_thread_num; /**< Maximum number of threads that can be used. */ + + AnchorType + anchor_type; /**< Type of anchors used in the attention mechanism. */ + RetrievalType + retrieval_type; /**< Type of retrieval strategy used in the cache. */ + + int layer_step; /**< Step size between layers. */ + int token_step; /**< Step size between tokens. */ + int layer_offset; /**< Offset value for layers. */ + + /** + * @brief Default constructor for KVCacheConfig. + * + * Initializes the configuration with default values. This constructor + * does not initialize any member variables explicitly. + */ + KVCacheConfig() = default; + + /** + * @brief Parameterized constructor for KVCacheConfig. + * + * This constructor initializes the configuration with specific values + * for all member variables. + * + * @param layer_num The number of layers in the model. + * @param kv_head_num The number of heads in the KV Cache. + * @param q_head_num The number of heads in the query. + * @param head_dim The dimension of each head. + * @param block_len The length of each block in the cache. + * @param anchor_num The number of anchors used in attention. + * @param anchor_type The type of anchors used in the attention mechanism. + * @param kv_type The data type of the KV Cache (e.g., fp16, q8_0). + * @param retrieval_type The type of retrieval strategy used in the cache. + * @param layer_step The step size between layers. + * @param token_step The step size between tokens. + * @param layer_offset The offset value for layers. + * @param max_block_num The maximum number of blocks that can be allocated. + * @param max_batch_size The maximum batch size that can be processed. + * @param max_thread_num The maximum number of threads that can be used. + */ + KVCacheConfig(int layer_num, int kv_head_num, int q_head_num, int head_dim, + int block_len, int anchor_num, AnchorType anchor_type, + ggml_type kv_type, RetrievalType retrieval_type, + int layer_step, int token_step, int layer_offset, + int max_block_num, int max_batch_size, int max_thread_num); +}; + +/** + * @class KVCache + * @brief Manages the Key-Value (KV) Cache used in attention mechanisms. + * + * The KVCache class provides functionality for managing the Key-Value Cache, + * including resizing the cache, retrieving configuration parameters, and + * updating internal states. This class is typically used in transformer models + * to store and manage past key and value states for efficient attention + * computations. + */ +class KVCache { + public: + /** + * @brief Constructs a KVCache object with the given configuration. + * + * Initializes the KVCache with the specified configuration parameters, + * such as the number of layers, heads, head dimensions, and other + * relevant settings. + * + * @param config The configuration object containing initialization + * parameters. + */ + KVCache(KVCacheConfig config); + + /** + * @brief Resizes the number of threads used by the cache. + * + * This function adjusts the number of threads that the cache can utilize. + * It allows dynamic reconfiguration of the parallel processing capabilities + * based on the current workload or system resources. + * + * @param thread_num The new number of threads to use. + */ + void ThreadResize(int thread_num); + + /** + * @brief Resizes the batch size managed by the cache. + * + * This function adjusts the batch size that the cache can handle. It + * is useful when the input batch size changes dynamically, allowing + * the cache to be reconfigured accordingly. + * + * @param batch_size The new batch size. + */ + void BatchResize(int batch_size); + + /** + * @brief Resizes the number of blocks managed by the cache. + * + * This function adjusts the number of blocks that the cache can manage. + * It allows dynamic reconfiguration of the block structure based on the + * current sequence length or other factors. + * + * @param block_num The new number of blocks. + */ + void BlockResize(int block_num); + + /** + * @brief Gets the number of layers in the cache. + * + * @return The number of layers configured in the cache. + */ + int get_layer_num() { return config_.layer_num; } + + /** + * @brief Gets the number of KV heads in the cache. + * + * @return The number of KV heads configured in the cache. + */ + int get_kv_head_num() { return config_.kv_head_num; } + + /** + * @brief Gets the number of query heads in the cache. + * + * @return The number of query heads configured in the cache. + */ + int get_q_head_num() { return config_.q_head_num; } + + /** + * @brief Gets the dimension of each head in the cache. + * + * @return The dimension of each head. + */ + int get_head_dim() { return config_.head_dim; } + + /** + * @brief Gets the length of each block in the cache. + * + * @return The length of each block. + */ + int get_block_len() { return config_.block_len; } + + /** + * @brief Gets the number of blocks for a specific layer. + * + * @param layer_id The ID of the layer for which to retrieve the block + * number. + * @return The number of blocks in the specified layer. + */ + int get_block_num(int layer_id) { return past_block_num_[layer_id]; } + + /** + * @brief Gets the number of anchors in the cache. + * + * @return The number of anchors configured in the cache. + */ + int get_anchor_num() { return config_.anchor_num; } + + /** + * @brief Gets the total length of the cache. + * + * @return The total length of the cache. + */ + int get_cache_total_len() { return cache_total_len_; } + + /** + * @brief Gets the total number of blocks in the cache. + * + * This function computes and returns the total number of blocks in the + * cache based on the total cache length and the block length configuration. + * + * @return The total number of blocks in the cache. + */ + int get_cache_total_block_num() { + return (cache_total_len_ + config_.block_len - 1) / config_.block_len; + } + + /** + * @brief Updates the total length of the cache. + * + * This function sets a new total length for the cache, allowing dynamic + * adjustment of the cache size during runtime. + * + * @param cache_total_len The new total length of the cache. + */ + void update_cache_total_len(int cache_total_len) { + cache_total_len_ = cache_total_len; + } + void attn(const ggml_fp16_t *q_in, ggml_fp16_t *output, float *attn_lse, + int layer_idx, int generate_token_idx, int q_len, int batch_size, + int max_block_num, int *block_table, int *cache_seqlens, + int pick_block_num, int init_block_num, int local_block_num, + Backend *backend); + + void update_kvcache_one_block_fp16(const ggml_fp16_t *k_in, + const ggml_fp16_t *v_in, int layer_id, + int block_idx, Backend *backend); + + void get_kvcache_one_block_fp16(ggml_fp16_t *k_in, ggml_fp16_t *v_in, + int layer_id, int block_idx, + Backend *backend); + + void update_importance_one_block(const ggml_fp16_t *importance, + int layer_id, int block_idx, + Backend *backend); + void get_importance_one_block(ggml_fp16_t *importance, int layer_id, + int block_idx, Backend *backend); + + void get_anchor_one_block(ggml_fp16_t *anchor, int layer_id, int block_idx, + Backend *backend); + + void update_anchor_one_block(const ggml_fp16_t *anchor, int layer_id, + int block_idx, Backend *backend); + + void calc_anchor_all_layers(int *block_table, int *cache_seqlens, + int batch_size, int max_block_num, + Backend *backend); + + void load_kvcache(std::string tensor_file_path, Backend *backend); + void dump_kvcache(int *block_table, int cache_total_len, + std::string tensor_file_path, Backend *backend); + + void get_and_update_kvcache_fp16(ggml_fp16_t *k_in, ggml_fp16_t *v_in, + int layer_id, int *block_table, + int batch_size, int max_block_num, + int *cache_seqlens, int q_len, + Backend *backend); + + void get_kvcache_fp16(ggml_fp16_t *k_in, ggml_fp16_t *v_in, int layer_id, + int *block_table, int batch_size, int max_block_num, + int *cache_seqlens, Backend *backend); + + void update_kvcache_fp16(const ggml_fp16_t *k_in, const ggml_fp16_t *v_in, + int layer_id, int *block_table, int batch_size, + int max_block_num, int *cache_seqlens, int q_len, + Backend *backend); + + void update_importance(const ggml_fp16_t *importance, int layer_id, + int *block_table, int batch_size, int max_block_num, + int *offset, int width, Backend *backend); + + void attn_with_kvcache(const ggml_fp16_t *q_in, const ggml_fp16_t *k_in, + const ggml_fp16_t *v_in, ggml_fp16_t *output, + float *attn_lse, int layer_idx, + int generate_token_idx, int q_len, int batch_size, + int max_block_num, int *block_table, + int *cache_seqlens, int topk, int local, + Backend *backend); + + void clear_importance_all_layers(int *block_table, int *cache_seqlens, + int batch_size, int max_block_num, + Backend *backend); + + void clear_kvcache_all_layers(int *block_table, int *cache_seqlens, + int batch_size, int max_block_num, + Backend *backend); + + void get_sincos(ggml_fp16_t *sin, ggml_fp16_t *cos, int seqlen); + + void get_attn_sparsity(const ggml_fp16_t *q_in, float *attn_sparsity, + int layer_idx, int generate_token_idx, int q_len, + int batch_size, int max_block_num, int *block_table, + int *cache_seqlens, int *block_table_origin, + int *cache_seqlens_origin, int max_block_num_origin, + int topk, int local, Backend *backend); + + void get_all_kvcache_one_layer(int layer_id, ggml_fp16_t *k_in, + ggml_fp16_t *v_in, Backend *backend); + + private: + // Persistent data + KVCacheConfig config_; + int n_gqa_; // q_head_num / kv_head_num + int cache_total_len_; // Number of tokens in cache + std::vector past_block_num_; // [layer_num] + std::vector>>> + k_cache_q4; // [layer_num, kv_head_num, past_block_num, block_len * + // (head_dim / QK_4)] + std::vector>>> + v_cache_q4; // [layer_num, kv_head_num, past_block_num, head_dim * + // (block_len / QK_4)] + std::vector>>> + k_cache_q8; // [layer_num, kv_head_num, past_block_num, block_len * + // (head_dim / QK_8)] + std::vector>>> + v_cache_q8; // [layer_num, kv_head_num, past_block_num, head_dim * + // (block_len / QK_8)] + + std::vector>>> + k_cache_fp16_; // [layer_num, kv_head_num, past_block_num, block_len * + // head_dim] + std::vector>>> + v_cache_fp16_; // [layer_num, kv_head_num, past_block_num, head_dim * + // block_len] + + std::vector>>> + importance_; // [layer_num, past_block_num, block_len, + // attention_head_num] + + std::vector + anchor_; // [layer_num * past_block_num * anchor_num * + // attention_head_num * head_dim] + + // Runtime data + int64_t layer_id_; + int64_t block_idx_; + int *block_table_; + uint64_t block_num_; + int max_block_num_after_retrieval_; + + // Rotary positional embeddings + std::vector> sin_; // [seq_len, head_dim] + std::vector> cos_; // [seq_len, head_dim] + + // update/get + int seq_len_; + uint16_t *k_scales_; // q4_0 + uint8_t *k_in_; // q4_0 + uint16_t *v_scales_; // q4_0 + uint8_t *v_in_; // q4_0 + uint16_t *k_data_; // fp16 + uint16_t *v_data_; // fp16 + uint16_t *importance_data_; // fp16 + uint16_t *anchor_data_; // fp16 + + // sparsity = (sigma(block lse / lse)) + std::vector>> + block_lse_; // [batch_size, max_block_num, q_head_num] + std::vector> attn_sparsity_; // [batch_size, q_head_num] + + // attn + std::vector> + avg_q; // [batch_size, q_head_num * head_dim] + + std::vector> + avg_q_fp16; // [batch_size, q_head_num * head_dim] + std::vector< + std::priority_queue, + std::vector>, std::greater<>>> + top_similar_block_; + + std::vector> block_similar_; + std::vector>> block_similar_kv_head_; + std::vector>> block_similar_q_head_; + + std::vector cache_seqlens_; // [batch_size] + std::vector selected_blocks_num_history_; // [layer_num // layer_step] + + std::vector>> selected_blocks_history_; + // [layer_num // layer_step, batch_size, max_block_num] + + std::vector>>> + selected_blocks_history_kvhead_; // [layer_num // layer_step, + // batch_size, max_block_num, + // kv_head_num] + + std::vector> + block_table_before_retrieval_; // [batch_size, max_block_num] + std::vector> + block_table_after_retrieval_; // [batch_size, pick_block_num] + + std::vector>> + block_table_before_retrieval_qhead_; // [batch_size, max_block_num, + // q_head_num] + std::vector>> + block_table_after_retrieval_qhead_; // [batch_size, pick_block_num, + // q_head_num] + + std::vector>> + block_table_before_retrieval_kvhead_; // [batch_size, max_block_num, + // kv_head_num] + std::vector>> + block_table_after_retrieval_kvhead_; // [batch_size, pick_block_num, + // kv_head_num] + + std::vector>> + mutex_; // [batch_size, kv_head_num] + std::vector>> + q_q8_0_; // [batch_size, kv_head_num, n_gqa * head_dim / QK8_0] + std::vector>> + q_fp32_; // [batch_size, kv_head_num, n_gqa * head_dim] + + std::vector>> + output_fp32_; // [batch_size, kv_head_num, n_gqa * head_dim] + std::vector>> + attn_lse_; // [batch_size, kv_head_num, n_gqa] + + std::vector> thread_cur_head_idx_; // [thread_num] + + std::vector> + thread_local_output_q8_0_; // [thread_num, n_gqa * head_dim / QK8_0] + std::vector> + thread_local_attn_score_; // [thread_num, n_gqa * block_len] + std::vector> + thread_local_output_fp32_; // [thread_num, n_gqa * head_dim] + std::vector> + thread_local_attn_lse_; // [thread_num, n_gqa] + std::vector> + thread_local_cur_output_fp32_; // [thread_num, n_gqa * head_dim] + std::vector> + thread_local_cur_attn_lse_; // [thread_num, n_gqa] + std::vector> + thread_local_attn_mask_; // [thread_num, block_len // 8] + std::vector> + thread_local_draft_; // [thread_num, 2 * n_gqa * block_len + 6 * n_gqa * + // head_dim + 2 * block_len * head_dim] + + // tmp space + std::vector q_fp32; // [n_gqa * head_dim] + + void quantize_q_(const uint16_t *q_in_data, int batch_size); + void attn_initialize_layer_(int batch_size, int layer_idx, int *block_table, + int &max_block_num, int *cache_seqlens); + void attn_initialize_kvhead_(int batch_size, int layer_idx, + int *block_table, int &max_block_num, + int *cache_seqlens); + void retrieval_kvcache_layer_(const uint16_t *q_in_data, int init_block_num, + int local_block_num, int pick_block_num, + int q_len, int generate_token_idx, + int batch_size, int layer_idx, + int *cache_seqlens, int &max_block_num, + Backend *backend); + void retrieval_kvcache_kvhead_(const uint16_t *q_in_data, + int init_block_num, int local_block_num, + int pick_block_num, int q_len, + int generate_token_idx, int batch_size, + int layer_idx, int *cache_seqlens, + int &max_block_num, Backend *backend); + + void calculate_block_similarity_layer_( + const uint16_t *q_in_data, int batch_size, int layer_idx, int q_len, + int max_block_num, int *cache_seqlens, int init_block_num, + int local_block_num, int pick_block_num, Backend *backend); + void calculate_block_similarity_kvhead_( + const uint16_t *q_in_data, int batch_size, int layer_idx, int q_len, + int max_block_num, int *cache_seqlens, int init_block_num, + int local_block_num, int pick_block_num, Backend *backend); + + void select_block_layer_(int batch_size, int layer_idx, int max_block_num, + int init_block_num, int local_block_num, + int pick_block_num); + void select_block_kvhead_(int batch_size, int layer_idx, int max_block_num, + int init_block_num, int local_block_num, + int pick_block_num); + + void calculate_sparsity_layer_(const uint16_t *q_in_data, + float *attn_sparsity, int batch_size, + int max_block_num, int *block_table, + int *cache_seqlens, Backend *backend); + void calculate_sparsity_kvhead_(const uint16_t *q_in_data, + float *attn_sparsity, int batch_size, + int max_block_num, int *block_table, + int *cache_seqlens, Backend *backend); + + void attention_kvhead_(const uint16_t *q_in_data, ggml_fp16_t *output, + float *attn_lse, int batch_size, Backend *backend); + void attention_layer_(const uint16_t *q_in_data, ggml_fp16_t *output, + float *attn_lse, int batch_size, Backend *backend); + + /** + * @brief Computes attention with KV cache for one block. + * + * This function performs attention computation for one block using KV + * cache. The function supports different data types for Q, K, and V caches, + * and provides options for quantization. The function does not perform any + * dynamic memory allocation internally, so all necessary buffers must be + * pre-allocated externally. + * + * @param head_dim The dimension of the head. + * @param bsz The batch size. + * @param q_type The data type of Q (GGML data type). Only supports fp16 and + * q8_0. + * @param q Pointer to the Q tensor [bsz, head_dim]. The quantization is + * always applied along the head_dim dimension. The size must be + * bsz * head_dim/32 * qtype_size. If head_dim % 32 != 0, an error + * will be raised. + * @param past_kv_len The length of the past KV cache. + * @param past_kv_offset The offset in the past KV cache. + * @param is_full_attn Boolean flag indicating whether to use full attention + * (true for full 1 mask). + * @param attn_mask Pointer to the attention mask [bsz, past_kv_len]. If + * is_full_attn = false, a bit matrix is passed to + * represent the mask. + * @param k_type The data type of K cache (GGML data type). Only supports + * fp16, q4_0, and q8_0. + * @param k_quant_type Quantization type for K cache. 0 for per_token, 1 for + * per_channel. Other values will raise an error. + * @param k_cache Pointer to the K cache tensor [seq_len, head_dim]. If + * quant_type == 0, head_dim % 32 must be 0. If quant_type == + * 1, seq_len % 32 must be 0. + * @param num_k_anchor The number of K anchors. If num_k_anchor == 0, it + * means no anchor is present. + * @param k_cache_anchors Pointer to the K cache anchors [num_k_anchor, + * head_dim]. The k_anchor_type must be fp16. + * @param k_cache_anchor_pos Pointer to the K cache anchor positions. Each + * token is associated with the nearest previous anchor position. + * @param v_type The data type of V cache (GGML data type). + * @param v_quant_type Quantization type for V cache. + * @param v_cache Pointer to the V cache tensor [head_dim, seq_len]. + * @param num_v_anchor The number of V anchors. + * @param v_cache_anchors Pointer to the V cache anchors. + * @param v_cache_anchor_pos Pointer to the V cache anchor positions. + * @param attn_score Pre-allocated buffer for attention scores [bsz, + * past_kv_len]. + * @param output Output tensor [bsz, head_dim] with the same type as q_type. + * @param lse Pre-allocated buffer [bsz] for the log-sum-exp of the + * attention scores. + * @param draft Pre-allocated temporary buffer. The buffer size should be + * enough to hold (2 * bsz * past_kv_len + 6 * bsz * head_dim + 2 * + * past_kv_len * head_dim + past_kv_len * head_dim / 32) bytes. + * @param rotary_angle Pointer to the rotary angle tensor. + * @param rotary_cos Pointer to the cosine values for rotary embedding. + * @param rotary_sin Pointer to the sine values for rotary embedding. + */ + void attn_with_kvcache_one_block_( + int head_dim, int bsz, + ggml_type q_type, // GGML data type of `Q`, only supports fp16 and q8_0 + // [bsz, head_dim] + // Quantization is always on the head_dim dimension (per_token). If + // head_dim % 32 != 0, an error will be raised. The size must be bsz * + // head_dim/32 * qtype_size. + const void *q, + + int past_kv_len, int past_kv_offset, + bool is_full_attn, // true indicates a full 1 mask + // If is_full_attn = false, a bit matrix representing the mask is + // passed. [bsz, past_kv_len] + const uint8_t *attn_mask, + + ggml_type k_type, // GGML data type of `K Cache`, only supports fp16, + // q4_0, q8_0 + int k_quant_type, // 0 for per_token, 1 for per_channel, others raise an + // error + // [seq_len, head_dim] + // If quant_type == 0, head_dim % 32 must be 0. + // If quant_type == 1, seq_len % 32 must be 0. + const void *k_cache, + + // k_anchor_type must be fp16 + int num_k_anchor, // num_k_anchor == 0 indicates no anchor + // [num_k_anchor, head_dim] + const void *k_cache_anchors, + // Each token is associated with the nearest previous position's anchor, + // with the same distance. + const int *k_cache_anchor_pos, + + // v_cache similar to k_cache + ggml_type v_type, int v_quant_type, + // [head_dim, seq_len] + const void *v_cache, int num_v_anchor, const void *v_cache_anchors, + const int *v_cache_anchor_pos, + + // Pre-allocated buffer for intermediate calculations [bsz, + // past_kv_len]. No malloc is performed inside this function. + float *attn_score, + + // Output: [bsz, head_dim], with the same type as q_type + void *output, + // [bsz] + float *lse, + + // Pre-allocated temporary buffer with sufficient size: + // (2 * bsz * past_kv_len + 6 * bsz * head_dim + 2 * past_kv_len * + // head_dim + past_kv_len * head_dim / 32) bytes. + void *draft, + + // Apply rotary embedding online + const int *rotary_angle, const void *rotary_cos, const void *rotary_sin + // rotary_cos=None, + // rotary_sin=None, + // cache_seqlens: Optional[Union[(int, torch.Tensor)]] = None, + // cache_batch_idx: Optional[torch.Tensor] = None, + // rotary_interleaved=True, + + // // Not supported for now + // window_size=(-1, -1), # -1 means infinite context window + // alibi_slopes=None, + ); +}; + +/** + * @brief Scales a float32 vector by a given scalar value. + * + * This function multiplies each element of the input vector `y` by a scalar + * `v`. It uses platform-specific optimizations if available, such as Apple's + * Accelerate framework or SIMD instructions. If no specific optimization is + * available, the function falls back to a simple scalar multiplication loop. + * + * @param n The number of elements in the vector `y`. + * @param y The input vector to be scaled. The result will be stored in the same + * vector. + * @param v The scalar value by which to scale the vector. + */ +void ggml_vec_scale_f32(const int n, float *y, const float v); +#endif \ No newline at end of file diff --git a/archive/kt-sft/csrc/ktransformers_ext/operators/kvcache/kvcache_attn.cpp b/archive/kt-sft/csrc/ktransformers_ext/operators/kvcache/kvcache_attn.cpp new file mode 100644 index 0000000..4190c03 --- /dev/null +++ b/archive/kt-sft/csrc/ktransformers_ext/operators/kvcache/kvcache_attn.cpp @@ -0,0 +1,2535 @@ +/** + * @Description : + * @Author : Jianwei Dong + * @Date : 2024-08-26 22:47:06 + * @Version : 1.0.0 + * @LastEditors : Jianwei Dong + * @LastEditTime : 2024-08-26 22:47:06 + * @Copyright (c) 2024 by KVCache.AI, All Rights Reserved. + **/ + +#include "kvcache.h" + +#include + +void KVCache::attention_kvhead_(const uint16_t *q_in_data, ggml_fp16_t *output, + float *attn_lse, int batch_size, + Backend *backend) { + // Timer start + auto start = std::chrono::high_resolution_clock::now(); + seq_len_ = config_.block_len; + + backend->do_work_stealing_job( + batch_size * config_.kv_head_num * max_block_num_after_retrieval_, + [&](int thread_id) { + thread_cur_head_idx_[thread_id].first = -1; + thread_cur_head_idx_[thread_id].second = -1; + }, + [&](int task_id) { + int batch_id = task_id / (config_.kv_head_num * + max_block_num_after_retrieval_); + int head_id = (task_id % (config_.kv_head_num * + max_block_num_after_retrieval_)) / + max_block_num_after_retrieval_; + int block_id = task_id % max_block_num_after_retrieval_; + int thread_id = Backend::thread_local_id; + + // If the block is out of the sequence length, skip it. + if (cache_seqlens_[batch_id] / config_.block_len < block_id) { + return; + } + int block_idx = + block_table_after_retrieval_kvhead_[batch_id][block_id] + [head_id]; + if (cache_seqlens_[batch_id] / config_.block_len == block_id) { + int seq_len = cache_seqlens_[batch_id] % config_.block_len; + if (seq_len == 0) + return; + + // Prepare the attention mask for the last block. + int full_blocks = seq_len / 8; + int remaining_bits = seq_len % 8; + // Fill full blocks with 1s + for (int i = 0; i < full_blocks; ++i) { + thread_local_attn_mask_[thread_id][i] = 0xFF; + } + // Fill the remaining bits in the next block + if (remaining_bits > 0 && full_blocks < seq_len_ / 8) { + thread_local_attn_mask_[thread_id][full_blocks] = + (1 << remaining_bits) - 1; + } else { + thread_local_attn_mask_[thread_id][full_blocks] = 0; + } + + for (int i = full_blocks + 1; i < seq_len_ / 8; ++i) { + thread_local_attn_mask_[thread_id][i] = 0; + } + if (config_.kv_type == ggml_type::GGML_TYPE_F16) { + attn_with_kvcache_one_block_( + config_.head_dim, + config_.q_head_num / config_.kv_head_num, GGML_TYPE_F16, + (void *)&q_in_data[batch_id * config_.kv_head_num * + n_gqa_ * config_.head_dim + + head_id * n_gqa_ * config_.head_dim], + seq_len_, 0, false, + thread_local_attn_mask_[thread_id].data(), + GGML_TYPE_F16, 0, + k_cache_fp16_[layer_id_][head_id][block_idx].data(), 0, + nullptr, nullptr, GGML_TYPE_F16, 1, + v_cache_fp16_[layer_id_][head_id][block_idx].data(), 0, + nullptr, nullptr, + thread_local_attn_score_[thread_id].data(), + thread_local_output_fp32_[thread_id].data(), + thread_local_attn_lse_[thread_id].data(), + thread_local_draft_[thread_id].data(), nullptr, + cos_.data(), sin_.data()); + } else if (config_.kv_type == ggml_type::GGML_TYPE_Q4_0) { + attn_with_kvcache_one_block_( + config_.head_dim, + config_.q_head_num / config_.kv_head_num, + GGML_TYPE_Q8_0, q_q8_0_[batch_id][head_id].data(), + seq_len_, 0, false, + thread_local_attn_mask_[thread_id].data(), + GGML_TYPE_Q4_0, 0, + k_cache_q4[layer_id_][head_id][block_idx].data(), 0, + nullptr, nullptr, GGML_TYPE_Q4_0, 1, + v_cache_q4[layer_id_][head_id][block_idx].data(), 0, + nullptr, nullptr, + thread_local_attn_score_[thread_id].data(), + thread_local_output_q8_0_[thread_id].data(), + thread_local_attn_lse_[thread_id].data(), + thread_local_draft_[thread_id].data(), nullptr, + cos_.data(), sin_.data()); + dequantize_row_q8_0( + thread_local_output_q8_0_[thread_id].data(), + thread_local_output_fp32_[thread_id].data(), + n_gqa_ * config_.head_dim); + } else if (config_.kv_type == ggml_type::GGML_TYPE_Q8_0) { + attn_with_kvcache_one_block_( + config_.head_dim, + config_.q_head_num / config_.kv_head_num, + GGML_TYPE_Q8_0, q_q8_0_[batch_id][head_id].data(), + seq_len_, 0, false, + thread_local_attn_mask_[thread_id].data(), + GGML_TYPE_Q8_0, 0, + k_cache_q8[layer_id_][head_id][block_idx].data(), 0, + nullptr, nullptr, GGML_TYPE_Q8_0, 1, + v_cache_q8[layer_id_][head_id][block_idx].data(), 0, + nullptr, nullptr, + thread_local_attn_score_[thread_id].data(), + thread_local_output_q8_0_[thread_id].data(), + thread_local_attn_lse_[thread_id].data(), + thread_local_draft_[thread_id].data(), nullptr, + cos_.data(), sin_.data()); + dequantize_row_q8_0( + thread_local_output_q8_0_[thread_id].data(), + thread_local_output_fp32_[thread_id].data(), + n_gqa_ * config_.head_dim); + } + } else { + if (config_.kv_type == ggml_type::GGML_TYPE_F16) { + attn_with_kvcache_one_block_( + config_.head_dim, + config_.q_head_num / config_.kv_head_num, GGML_TYPE_F16, + (void *)&q_in_data[batch_id * config_.kv_head_num * + n_gqa_ * config_.head_dim + + head_id * n_gqa_ * config_.head_dim], + seq_len_, 0, true, nullptr, GGML_TYPE_F16, 0, + k_cache_fp16_[layer_id_][head_id][block_idx].data(), 0, + nullptr, nullptr, GGML_TYPE_F16, 1, + v_cache_fp16_[layer_id_][head_id][block_idx].data(), 0, + nullptr, nullptr, + thread_local_attn_score_[thread_id].data(), + thread_local_output_fp32_[thread_id].data(), + thread_local_attn_lse_[thread_id].data(), + thread_local_draft_[thread_id].data(), nullptr, + cos_.data(), sin_.data()); + + } else if (config_.kv_type == ggml_type::GGML_TYPE_Q4_0) { + attn_with_kvcache_one_block_( + config_.head_dim, + config_.q_head_num / config_.kv_head_num, + GGML_TYPE_Q8_0, q_q8_0_[batch_id][head_id].data(), + seq_len_, 0, true, nullptr, GGML_TYPE_Q4_0, 0, + k_cache_q4[layer_id_][head_id][block_idx].data(), 0, + nullptr, nullptr, GGML_TYPE_Q4_0, 1, + v_cache_q4[layer_id_][head_id][block_idx].data(), 0, + nullptr, nullptr, + thread_local_attn_score_[thread_id].data(), + thread_local_output_q8_0_[thread_id].data(), + thread_local_attn_lse_[thread_id].data(), + thread_local_draft_[thread_id].data(), nullptr, + cos_.data(), sin_.data()); + dequantize_row_q8_0( + thread_local_output_q8_0_[thread_id].data(), + thread_local_output_fp32_[thread_id].data(), + n_gqa_ * config_.head_dim); + } else if (config_.kv_type == ggml_type::GGML_TYPE_Q8_0) { + attn_with_kvcache_one_block_( + config_.head_dim, + config_.q_head_num / config_.kv_head_num, + GGML_TYPE_Q8_0, q_q8_0_[batch_id][head_id].data(), + seq_len_, 0, true, nullptr, GGML_TYPE_Q8_0, 0, + k_cache_q8[layer_id_][head_id][block_idx].data(), 0, + nullptr, nullptr, GGML_TYPE_Q8_0, 1, + v_cache_q8[layer_id_][head_id][block_idx].data(), 0, + nullptr, nullptr, + thread_local_attn_score_[thread_id].data(), + thread_local_output_q8_0_[thread_id].data(), + thread_local_attn_lse_[thread_id].data(), + thread_local_draft_[thread_id].data(), nullptr, + cos_.data(), sin_.data()); + dequantize_row_q8_0( + thread_local_output_q8_0_[thread_id].data(), + thread_local_output_fp32_[thread_id].data(), + n_gqa_ * config_.head_dim); + } + } + int cur_batch_idx = thread_cur_head_idx_[thread_id].first; + int cur_head_id = thread_cur_head_idx_[thread_id].second; + if (batch_id == cur_batch_idx && head_id == cur_head_id) { + for (int i = 0; i < n_gqa_; i++) { + float new_attn_lse = + thread_local_cur_attn_lse_[thread_id][i] + + std::log( + 1.0 + + std::exp(thread_local_attn_lse_[thread_id][i] - + thread_local_cur_attn_lse_[thread_id][i])); + ggml_vec_scale_f32( + config_.head_dim, + thread_local_cur_output_fp32_[thread_id].data() + + i * config_.head_dim, + std::exp(thread_local_cur_attn_lse_[thread_id][i] - + new_attn_lse)); + ggml_vec_scale_f32( + config_.head_dim, + thread_local_output_fp32_[thread_id].data() + + i * config_.head_dim, + std::exp(thread_local_attn_lse_[thread_id][i] - + new_attn_lse)); + for (int j = 0; j < config_.head_dim; j++) { + thread_local_cur_output_fp32_[thread_id] + [i * config_.head_dim + + j] += + thread_local_output_fp32_[thread_id] + [i * config_.head_dim + j]; + } + thread_local_cur_attn_lse_[thread_id][i] = new_attn_lse; + } + } else { + if (cur_batch_idx != -1) { + mutex_[cur_batch_idx][cur_head_id]->lock(); + for (int i = 0; i < n_gqa_; i++) { + if (std::abs(attn_lse_[cur_batch_idx][cur_head_id][i]) < + 1e-6) { + attn_lse_[cur_batch_idx][cur_head_id][i] = + thread_local_cur_attn_lse_[thread_id][i]; + for (int j = 0; j < config_.head_dim; j++) { + output_fp32_[cur_batch_idx][cur_head_id] + [i * config_.head_dim + j] = + thread_local_cur_output_fp32_ + [thread_id] + [i * config_.head_dim + j]; + } + continue; + } + float new_attn_lse = + attn_lse_[cur_batch_idx][cur_head_id][i] + + std::log( + 1.0 + + std::exp( + thread_local_cur_attn_lse_[thread_id][i] - + attn_lse_[cur_batch_idx][cur_head_id][i])); + ggml_vec_scale_f32( + config_.head_dim, + output_fp32_[cur_batch_idx][cur_head_id].data() + + i * config_.head_dim, + std::exp(attn_lse_[cur_batch_idx][cur_head_id][i] - + new_attn_lse)); + ggml_vec_scale_f32( + config_.head_dim, + thread_local_cur_output_fp32_[thread_id].data() + + i * config_.head_dim, + std::exp(thread_local_cur_attn_lse_[thread_id][i] - + new_attn_lse)); + for (int j = 0; j < config_.head_dim; j++) { + output_fp32_[cur_batch_idx][cur_head_id] + [i * config_.head_dim + j] += + thread_local_cur_output_fp32_ + [thread_id][i * config_.head_dim + j]; + } + attn_lse_[cur_batch_idx][cur_head_id][i] = new_attn_lse; + } + mutex_[cur_batch_idx][cur_head_id]->unlock(); + } + thread_cur_head_idx_[thread_id].first = batch_id; + thread_cur_head_idx_[thread_id].second = head_id; + for (int i = 0; i < n_gqa_; i++) { + thread_local_cur_attn_lse_[thread_id][i] = + thread_local_attn_lse_[thread_id][i]; + for (int j = 0; j < config_.head_dim; j++) { + thread_local_cur_output_fp32_ + [thread_id][i * config_.head_dim + j] = + thread_local_output_fp32_[thread_id] + [i * config_.head_dim + + j]; + } + } + } + }, + // Merge the results of the remaining blocks. + [&](int thread_id) { + int cur_batch_idx = thread_cur_head_idx_[thread_id].first; + int cur_head_id = thread_cur_head_idx_[thread_id].second; + if (cur_head_id != -1) { + mutex_[cur_batch_idx][cur_head_id]->lock(); + for (int i = 0; i < n_gqa_; i++) { + float new_attn_lse; + if (std::abs(attn_lse_[cur_batch_idx][cur_head_id][i]) < + 1e-6) { + attn_lse_[cur_batch_idx][cur_head_id][i] = + thread_local_cur_attn_lse_[thread_id][i]; + for (int j = 0; j < config_.head_dim; j++) { + output_fp32_[cur_batch_idx][cur_head_id] + [i * config_.head_dim + j] = + thread_local_cur_output_fp32_ + [thread_id] + [i * config_.head_dim + j]; + } + continue; + } + new_attn_lse = + attn_lse_[cur_batch_idx][cur_head_id][i] + + std::log( + 1.0 + + std::exp(thread_local_cur_attn_lse_[thread_id][i] - + attn_lse_[cur_batch_idx][cur_head_id][i])); + ggml_vec_scale_f32( + config_.head_dim, + output_fp32_[cur_batch_idx][cur_head_id].data() + + i * config_.head_dim, + std::exp(attn_lse_[cur_batch_idx][cur_head_id][i] - + new_attn_lse)); + ggml_vec_scale_f32( + config_.head_dim, + thread_local_cur_output_fp32_[thread_id].data() + + i * config_.head_dim, + std::exp(thread_local_cur_attn_lse_[thread_id][i] - + new_attn_lse)); + for (int j = 0; j < config_.head_dim; j++) { + output_fp32_[cur_batch_idx][cur_head_id] + [i * config_.head_dim + j] += + thread_local_cur_output_fp32_[thread_id] + [i * config_.head_dim + + j]; + } + attn_lse_[cur_batch_idx][cur_head_id][i] = new_attn_lse; + } + mutex_[cur_batch_idx][cur_head_id]->unlock(); + } + }); + // move the results to output and attn_lse + uint16_t *output_data = reinterpret_cast(output); + float *attn_lse_data = attn_lse; + for (int batch_idx = 0; batch_idx < batch_size; batch_idx++) { + for (int i = 0; i < config_.kv_head_num; i++) { + for (int j = 0; j < n_gqa_ * config_.head_dim; j++) { + output_data[batch_idx * config_.kv_head_num * n_gqa_ * + config_.head_dim + + i * n_gqa_ * config_.head_dim + j] = + GGML_FP32_TO_FP16(output_fp32_[batch_idx][i][j]); + } + for (int j = 0; j < n_gqa_; j++) { + attn_lse_data[batch_idx * config_.kv_head_num * n_gqa_ + + i * n_gqa_ + j] = attn_lse_[batch_idx][i][j]; + } + } + } + + // Timer end + auto end = std::chrono::high_resolution_clock::now(); + std::chrono::duration diff = end - start; + // printf("layer %d time of computing attention: %f s\n", layer_idx, + // diff.count()); +} + +void KVCache::attention_layer_(const uint16_t *q_in_data, ggml_fp16_t *output, + float *attn_lse, int batch_size, + Backend *backend) { + // Timer start + auto start = std::chrono::high_resolution_clock::now(); + seq_len_ = config_.block_len; + backend->do_work_stealing_job( + batch_size * config_.kv_head_num * max_block_num_after_retrieval_, + [&](int thread_id) { + thread_cur_head_idx_[thread_id].first = -1; + thread_cur_head_idx_[thread_id].second = -1; + }, + [&](int task_id) { + int batch_id = task_id / (config_.kv_head_num * + max_block_num_after_retrieval_); + int head_id = (task_id % (config_.kv_head_num * + max_block_num_after_retrieval_)) / + max_block_num_after_retrieval_; + int block_id = task_id % max_block_num_after_retrieval_; + int thread_id = Backend::thread_local_id; + // If the block is out of the sequence length, skip it. + if (cache_seqlens_[batch_id] / config_.block_len < block_id) { + return; + } + int block_idx = block_table_after_retrieval_[batch_id][block_id]; + if (cache_seqlens_[batch_id] / config_.block_len == block_id) { + int seq_len = cache_seqlens_[batch_id] % config_.block_len; + if (seq_len == 0) + return; + + // Prepare the attention mask for the last block. + int full_blocks = seq_len / 8; + int remaining_bits = seq_len % 8; + + // Fill full blocks with 1s + for (int i = 0; i < full_blocks; ++i) { + thread_local_attn_mask_[thread_id][i] = 0xFF; + } + // Fill the remaining bits in the next block + if (remaining_bits > 0 && full_blocks < seq_len_ / 8) { + thread_local_attn_mask_[thread_id][full_blocks] = + (1 << remaining_bits) - 1; + } else { + thread_local_attn_mask_[thread_id][full_blocks] = 0; + } + + for (int i = full_blocks + 1; i < seq_len_ / 8; ++i) { + thread_local_attn_mask_[thread_id][i] = 0; + } + if (config_.kv_type == ggml_type::GGML_TYPE_F16) { + attn_with_kvcache_one_block_( + config_.head_dim, + config_.q_head_num / config_.kv_head_num, GGML_TYPE_F16, + (void *)&q_in_data[batch_id * config_.kv_head_num * + n_gqa_ * config_.head_dim + + head_id * n_gqa_ * config_.head_dim], + seq_len_, 0, false, + thread_local_attn_mask_[thread_id].data(), + GGML_TYPE_F16, 0, + k_cache_fp16_[layer_id_][head_id][block_idx].data(), 0, + nullptr, nullptr, GGML_TYPE_F16, 1, + v_cache_fp16_[layer_id_][head_id][block_idx].data(), 0, + nullptr, nullptr, + thread_local_attn_score_[thread_id].data(), + thread_local_output_fp32_[thread_id].data(), + thread_local_attn_lse_[thread_id].data(), + thread_local_draft_[thread_id].data(), nullptr, + cos_.data(), sin_.data()); + } else if (config_.kv_type == ggml_type::GGML_TYPE_Q4_0) { + attn_with_kvcache_one_block_( + config_.head_dim, + config_.q_head_num / config_.kv_head_num, + GGML_TYPE_Q8_0, q_q8_0_[batch_id][head_id].data(), + seq_len_, 0, false, + thread_local_attn_mask_[thread_id].data(), + GGML_TYPE_Q4_0, 0, + k_cache_q4[layer_id_][head_id][block_idx].data(), 0, + nullptr, nullptr, GGML_TYPE_Q4_0, 1, + v_cache_q4[layer_id_][head_id][block_idx].data(), 0, + nullptr, nullptr, + thread_local_attn_score_[thread_id].data(), + thread_local_output_q8_0_[thread_id].data(), + thread_local_attn_lse_[thread_id].data(), + thread_local_draft_[thread_id].data(), nullptr, + cos_.data(), sin_.data()); + dequantize_row_q8_0( + thread_local_output_q8_0_[thread_id].data(), + thread_local_output_fp32_[thread_id].data(), + n_gqa_ * config_.head_dim); + } else if (config_.kv_type == ggml_type::GGML_TYPE_Q8_0) { + attn_with_kvcache_one_block_( + config_.head_dim, + config_.q_head_num / config_.kv_head_num, + GGML_TYPE_Q8_0, q_q8_0_[batch_id][head_id].data(), + seq_len_, 0, false, + thread_local_attn_mask_[thread_id].data(), + GGML_TYPE_Q8_0, 0, + k_cache_q8[layer_id_][head_id][block_idx].data(), 0, + nullptr, nullptr, GGML_TYPE_Q8_0, 1, + v_cache_q8[layer_id_][head_id][block_idx].data(), 0, + nullptr, nullptr, + thread_local_attn_score_[thread_id].data(), + thread_local_output_q8_0_[thread_id].data(), + thread_local_attn_lse_[thread_id].data(), + thread_local_draft_[thread_id].data(), nullptr, + cos_.data(), sin_.data()); + dequantize_row_q8_0( + thread_local_output_q8_0_[thread_id].data(), + thread_local_output_fp32_[thread_id].data(), + n_gqa_ * config_.head_dim); + } + } else { + if (config_.kv_type == ggml_type::GGML_TYPE_F16) { + attn_with_kvcache_one_block_( + config_.head_dim, + config_.q_head_num / config_.kv_head_num, GGML_TYPE_F16, + (void *)&q_in_data[batch_id * config_.kv_head_num * + n_gqa_ * config_.head_dim + + head_id * n_gqa_ * config_.head_dim], + seq_len_, 0, true, nullptr, GGML_TYPE_F16, 0, + k_cache_fp16_[layer_id_][head_id][block_idx].data(), 0, + nullptr, nullptr, GGML_TYPE_F16, 1, + v_cache_fp16_[layer_id_][head_id][block_idx].data(), 0, + nullptr, nullptr, + thread_local_attn_score_[thread_id].data(), + thread_local_output_fp32_[thread_id].data(), + thread_local_attn_lse_[thread_id].data(), + thread_local_draft_[thread_id].data(), nullptr, + cos_.data(), sin_.data()); + + } else if (config_.kv_type == ggml_type::GGML_TYPE_Q4_0) { + attn_with_kvcache_one_block_( + config_.head_dim, + config_.q_head_num / config_.kv_head_num, + GGML_TYPE_Q8_0, q_q8_0_[batch_id][head_id].data(), + seq_len_, 0, true, nullptr, GGML_TYPE_Q4_0, 0, + k_cache_q4[layer_id_][head_id][block_idx].data(), 0, + nullptr, nullptr, GGML_TYPE_Q4_0, 1, + v_cache_q4[layer_id_][head_id][block_idx].data(), 0, + nullptr, nullptr, + thread_local_attn_score_[thread_id].data(), + thread_local_output_q8_0_[thread_id].data(), + thread_local_attn_lse_[thread_id].data(), + thread_local_draft_[thread_id].data(), nullptr, + cos_.data(), sin_.data()); + dequantize_row_q8_0( + thread_local_output_q8_0_[thread_id].data(), + thread_local_output_fp32_[thread_id].data(), + n_gqa_ * config_.head_dim); + } else if (config_.kv_type == ggml_type::GGML_TYPE_Q8_0) { + attn_with_kvcache_one_block_( + config_.head_dim, + config_.q_head_num / config_.kv_head_num, + GGML_TYPE_Q8_0, q_q8_0_[batch_id][head_id].data(), + seq_len_, 0, true, nullptr, GGML_TYPE_Q8_0, 0, + k_cache_q8[layer_id_][head_id][block_idx].data(), 0, + nullptr, nullptr, GGML_TYPE_Q8_0, 1, + v_cache_q8[layer_id_][head_id][block_idx].data(), 0, + nullptr, nullptr, + thread_local_attn_score_[thread_id].data(), + thread_local_output_q8_0_[thread_id].data(), + thread_local_attn_lse_[thread_id].data(), + thread_local_draft_[thread_id].data(), nullptr, + cos_.data(), sin_.data()); + dequantize_row_q8_0( + thread_local_output_q8_0_[thread_id].data(), + thread_local_output_fp32_[thread_id].data(), + n_gqa_ * config_.head_dim); + } + } + int cur_batch_idx = thread_cur_head_idx_[thread_id].first; + int cur_head_id = thread_cur_head_idx_[thread_id].second; + if (batch_id == cur_batch_idx && head_id == cur_head_id) { + for (int i = 0; i < n_gqa_; i++) { + float new_attn_lse = + thread_local_cur_attn_lse_[thread_id][i] + + std::log( + 1.0 + + std::exp(thread_local_attn_lse_[thread_id][i] - + thread_local_cur_attn_lse_[thread_id][i])); + ggml_vec_scale_f32( + config_.head_dim, + thread_local_cur_output_fp32_[thread_id].data() + + i * config_.head_dim, + std::exp(thread_local_cur_attn_lse_[thread_id][i] - + new_attn_lse)); + ggml_vec_scale_f32( + config_.head_dim, + thread_local_output_fp32_[thread_id].data() + + i * config_.head_dim, + std::exp(thread_local_attn_lse_[thread_id][i] - + new_attn_lse)); + for (int j = 0; j < config_.head_dim; j++) { + thread_local_cur_output_fp32_[thread_id] + [i * config_.head_dim + + j] += + thread_local_output_fp32_[thread_id] + [i * config_.head_dim + j]; + } + thread_local_cur_attn_lse_[thread_id][i] = new_attn_lse; + } + } else { + if (cur_batch_idx != -1) { + mutex_[cur_batch_idx][cur_head_id]->lock(); + for (int i = 0; i < n_gqa_; i++) { + if (std::abs(attn_lse_[cur_batch_idx][cur_head_id][i]) < + 1e-6) { + attn_lse_[cur_batch_idx][cur_head_id][i] = + thread_local_cur_attn_lse_[thread_id][i]; + for (int j = 0; j < config_.head_dim; j++) { + output_fp32_[cur_batch_idx][cur_head_id] + [i * config_.head_dim + j] = + thread_local_cur_output_fp32_ + [thread_id] + [i * config_.head_dim + j]; + } + continue; + } + float new_attn_lse = + attn_lse_[cur_batch_idx][cur_head_id][i] + + std::log( + 1.0 + + std::exp( + thread_local_cur_attn_lse_[thread_id][i] - + attn_lse_[cur_batch_idx][cur_head_id][i])); + ggml_vec_scale_f32( + config_.head_dim, + output_fp32_[cur_batch_idx][cur_head_id].data() + + i * config_.head_dim, + std::exp(attn_lse_[cur_batch_idx][cur_head_id][i] - + new_attn_lse)); + ggml_vec_scale_f32( + config_.head_dim, + thread_local_cur_output_fp32_[thread_id].data() + + i * config_.head_dim, + std::exp(thread_local_cur_attn_lse_[thread_id][i] - + new_attn_lse)); + for (int j = 0; j < config_.head_dim; j++) { + output_fp32_[cur_batch_idx][cur_head_id] + [i * config_.head_dim + j] += + thread_local_cur_output_fp32_ + [thread_id][i * config_.head_dim + j]; + } + attn_lse_[cur_batch_idx][cur_head_id][i] = new_attn_lse; + } + mutex_[cur_batch_idx][cur_head_id]->unlock(); + } + thread_cur_head_idx_[thread_id].first = batch_id; + thread_cur_head_idx_[thread_id].second = head_id; + for (int i = 0; i < n_gqa_; i++) { + thread_local_cur_attn_lse_[thread_id][i] = + thread_local_attn_lse_[thread_id][i]; + for (int j = 0; j < config_.head_dim; j++) { + thread_local_cur_output_fp32_ + [thread_id][i * config_.head_dim + j] = + thread_local_output_fp32_[thread_id] + [i * config_.head_dim + + j]; + } + } + } + }, + // Merge the results of the remaining blocks. + [&](int thread_id) { + int cur_batch_idx = thread_cur_head_idx_[thread_id].first; + int cur_head_id = thread_cur_head_idx_[thread_id].second; + if (cur_head_id != -1) { + mutex_[cur_batch_idx][cur_head_id]->lock(); + for (int i = 0; i < n_gqa_; i++) { + float new_attn_lse; + if (std::abs(attn_lse_[cur_batch_idx][cur_head_id][i]) < + 1e-6) { + attn_lse_[cur_batch_idx][cur_head_id][i] = + thread_local_cur_attn_lse_[thread_id][i]; + for (int j = 0; j < config_.head_dim; j++) { + output_fp32_[cur_batch_idx][cur_head_id] + [i * config_.head_dim + j] = + thread_local_cur_output_fp32_ + [thread_id] + [i * config_.head_dim + j]; + } + continue; + } + new_attn_lse = + attn_lse_[cur_batch_idx][cur_head_id][i] + + std::log( + 1.0 + + std::exp(thread_local_cur_attn_lse_[thread_id][i] - + attn_lse_[cur_batch_idx][cur_head_id][i])); + ggml_vec_scale_f32( + config_.head_dim, + output_fp32_[cur_batch_idx][cur_head_id].data() + + i * config_.head_dim, + std::exp(attn_lse_[cur_batch_idx][cur_head_id][i] - + new_attn_lse)); + ggml_vec_scale_f32( + config_.head_dim, + thread_local_cur_output_fp32_[thread_id].data() + + i * config_.head_dim, + std::exp(thread_local_cur_attn_lse_[thread_id][i] - + new_attn_lse)); + for (int j = 0; j < config_.head_dim; j++) { + output_fp32_[cur_batch_idx][cur_head_id] + [i * config_.head_dim + j] += + thread_local_cur_output_fp32_[thread_id] + [i * config_.head_dim + + j]; + } + attn_lse_[cur_batch_idx][cur_head_id][i] = new_attn_lse; + } + mutex_[cur_batch_idx][cur_head_id]->unlock(); + } + }); + + // move the results to output and attn_lse + uint16_t *output_data = reinterpret_cast(output); + float *attn_lse_data = attn_lse; + for (int batch_idx = 0; batch_idx < batch_size; batch_idx++) { + for (int i = 0; i < config_.kv_head_num; i++) { + for (int j = 0; j < n_gqa_ * config_.head_dim; j++) { + output_data[batch_idx * config_.kv_head_num * n_gqa_ * + config_.head_dim + + i * n_gqa_ * config_.head_dim + j] = + GGML_FP32_TO_FP16(output_fp32_[batch_idx][i][j]); + } + for (int j = 0; j < n_gqa_; j++) { + attn_lse_data[batch_idx * config_.kv_head_num * n_gqa_ + + i * n_gqa_ + j] = attn_lse_[batch_idx][i][j]; + } + } + } + // Timer end + auto end = std::chrono::high_resolution_clock::now(); + std::chrono::duration diff = end - start; + // printf("layer %d time of computing attention: %f s\n", layer_id_, + // diff.count()); +} + +void KVCache::attn(const ggml_fp16_t *q_in, ggml_fp16_t *output, + float *attn_lse, int layer_idx, int generate_token_idx, + int q_len, int batch_size, int max_block_num, + int *block_table, int *cache_seqlens, int pick_block_num, + int init_block_num, int local_block_num, Backend *backend) { + + // Timer start + auto start = std::chrono::high_resolution_clock::now(); + layer_id_ = layer_idx; + batch_size = batch_size * q_len; + + const uint16_t *q_in_data = const_cast(q_in); + + quantize_q_(q_in_data, batch_size); + if (config_.retrieval_type == RetrievalType::LAYER) { + attn_initialize_layer_(batch_size, layer_idx, block_table, + max_block_num, cache_seqlens); + retrieval_kvcache_layer_(q_in_data, init_block_num, local_block_num, + pick_block_num, q_len, generate_token_idx, + batch_size, layer_idx, cache_seqlens, + max_block_num, backend); + attention_layer_(q_in_data, output, attn_lse, batch_size, backend); + } else if (config_.retrieval_type == RetrievalType::KVHEAD) { + attn_initialize_kvhead_(batch_size, layer_idx, block_table, + max_block_num, cache_seqlens); + retrieval_kvcache_kvhead_(q_in_data, init_block_num, local_block_num, + pick_block_num, q_len, generate_token_idx, + batch_size, layer_idx, cache_seqlens, + max_block_num, backend); + attention_kvhead_(q_in_data, output, attn_lse, batch_size, backend); + } + + // Timer end + auto end = std::chrono::high_resolution_clock::now(); + std::chrono::duration diff = end - start; + // printf("layer %d time of computing attention: %f s\n", layer_idx, + // diff.count()); +} + +void KVCache::attn_with_kvcache( + const ggml_fp16_t *q_in, const ggml_fp16_t *k_in, const ggml_fp16_t *v_in, + ggml_fp16_t *output, float *attn_lse, int layer_idx, int generate_token_idx, + int q_len, int batch_size, int max_block_num, int *block_table, + int *cache_seqlens, int topk, int local, Backend *backend) { + // printf("attn_with_kvcache start\n"); + assert(q_len == 1); + // Timer start + auto start = std::chrono::high_resolution_clock::now(); + + layer_id_ = layer_idx; + + update_kvcache_fp16(k_in, v_in, layer_idx, block_table, batch_size, + max_block_num, cache_seqlens, q_len, backend); + // printf("update finished.\n"); + + // cache_seqlens memory is modified. + for (int i = 0; i < batch_size; i++) { + cache_seqlens[i] += q_len; + } + int init_block_num = 1; + if (config_.block_len <= 32) { + init_block_num = 64 / config_.block_len; + } + + attn(q_in, output, attn_lse, layer_idx, generate_token_idx, q_len, + batch_size, max_block_num, block_table, cache_seqlens, topk, + init_block_num, local, backend); + + // Timer end + auto end = std::chrono::high_resolution_clock::now(); + std::chrono::duration diff = end - start; + // printf("layer %d time of computing attention with kvcache: %f s\n", + // layer_idx, diff.count()); +} + +void KVCache::quantize_q_(const uint16_t *q_in_data, int batch_size) { + // Timer start + auto start = std::chrono::high_resolution_clock::now(); + for (int batch_idx = 0; batch_idx < batch_size; batch_idx++) { + if (config_.kv_type == ggml_type::GGML_TYPE_F16) { + // quantize q + for (int i = 0; i < config_.kv_head_num; i++) { + for (int j = 0; j < n_gqa_ * config_.head_dim; j++) { + q_fp32_[batch_idx][i][j] = GGML_FP16_TO_FP32( + q_in_data[batch_idx * config_.kv_head_num * n_gqa_ * + config_.head_dim + + i * n_gqa_ * config_.head_dim + j]); + } + } + } else { + // quantize q + for (int i = 0; i < config_.kv_head_num; i++) { + for (int j = 0; j < n_gqa_ * config_.head_dim; j++) { + q_fp32[j] = GGML_FP16_TO_FP32( + q_in_data[batch_idx * config_.kv_head_num * n_gqa_ * + config_.head_dim + + i * n_gqa_ * config_.head_dim + j]); + } + quantize_row_q8_0(q_fp32.data(), q_q8_0_[batch_idx][i].data(), + n_gqa_ * config_.head_dim); + } + } + } + // Timer end + auto end = std::chrono::high_resolution_clock::now(); + // printf("time of quantizing q: %f s\n", + // std::chrono::duration(end - start).count()); +} +void KVCache::attn_initialize_layer_(int batch_size, int layer_idx, + int *block_table, int &max_block_num, + int *cache_seqlens) { + // Timer start + auto start = std::chrono::high_resolution_clock::now(); + for (int batch_idx = 0; batch_idx < batch_size; batch_idx++) { + // initialize output_fp32_ and attn_lse_ + for (int i = 0; i < config_.kv_head_num; i++) { + for (int j = 0; j < n_gqa_ * config_.head_dim; j++) { + output_fp32_[batch_idx][i][j] = 0; + } + for (int j = 0; j < n_gqa_; j++) { + attn_lse_[batch_idx][i][j] = 0; + } + } + // clear top_similar_block_ + + while (!top_similar_block_[batch_idx].empty()) + top_similar_block_[batch_idx].pop(); + } + + // get block_table_before_retrieval_ and cache_seqlens_ + if (block_table == nullptr) { + max_block_num = past_block_num_[layer_idx]; + for (int batch_idx = 0; batch_idx < batch_size; batch_idx++) { + if (cache_total_len_ != 0) + cache_seqlens_[batch_idx] = cache_total_len_; + else + cache_seqlens_[batch_idx] = max_block_num * config_.block_len; + for (int i = 0; i < max_block_num; i++) { + block_table_before_retrieval_[batch_idx][i] = i; + block_similar_[batch_idx][i] = 0; + } + } + } else { + for (int batch_idx = 0; batch_idx < batch_size; batch_idx++) { + cache_seqlens_[batch_idx] = cache_seqlens[batch_idx]; + for (int i = 0; i < max_block_num; i++) { + block_table_before_retrieval_[batch_idx][i] = + block_table[batch_idx * max_block_num + i]; + block_similar_[batch_idx][i] = 0; + } + } + } + // Timer end + auto end = std::chrono::high_resolution_clock::now(); + // printf("layer %d time of initializing attention: %f s\n", layer_idx, + // std::chrono::duration(end - start).count()); +} + +void KVCache::calculate_block_similarity_layer_( + const uint16_t *q_in_data, int batch_size, int layer_idx, int q_len, + int max_block_num, int *cache_seqlens, int init_block_num, + int local_block_num, int pick_block_num, Backend *backend) { + // Timer start + auto start = std::chrono::high_resolution_clock::now(); + + if (batch_size == 1 && + config_.anchor_num == 1) { // TODO: improve batch_size > 1 + for (int batch_id = 0; batch_id < batch_size; batch_id++) { + if (q_len == 1) { + for (int j = 0; j < config_.head_dim * config_.q_head_num; + j++) { + avg_q[batch_id][j] = GGML_FP16_TO_FP32( + q_in_data[batch_id * q_len * config_.q_head_num * + config_.head_dim + + j]); + avg_q_fp16[batch_id][j] = + q_in_data[batch_id * q_len * config_.q_head_num * + config_.head_dim + + j]; + } + } else { + for (int j = 0; j < config_.head_dim * config_.q_head_num; + j++) { + avg_q[batch_id][j] = 0; + } + for (int i = 0; i < q_len; i++) { + for (int j = 0; j < config_.head_dim; j++) { + avg_q[batch_id][j] += GGML_FP16_TO_FP32( + q_in_data[batch_id * q_len * config_.q_head_num * + config_.head_dim + + i * config_.q_head_num * + config_.head_dim + + j]); + } + } + for (int j = 0; j < config_.head_dim * config_.q_head_num; + j++) { + avg_q[batch_id][j] /= q_len; + avg_q_fp16[batch_id][j] = + GGML_FP32_TO_FP16(avg_q[batch_id][j]); + } + } + int seq_len = cache_seqlens_[batch_id]; + int block_num = (seq_len / config_.block_len) - local_block_num - + init_block_num; + if (block_num <= 0) { + continue; + } + bool is_seq = true; + for (int i = init_block_num + 1; + i < (seq_len / config_.block_len) - local_block_num; i++) { + if (block_table_before_retrieval_[batch_id][i] != + block_table_before_retrieval_[batch_id][i - 1] + 1) { + is_seq = false; + break; + } + } + if (is_seq) { + int nth = backend->get_thread_num(); + backend->do_work_stealing_job( + nth, nullptr, + [&](int task_id) { + int ith = task_id; + bool ok = llamafile_sgemm( + block_num, 1, config_.q_head_num * config_.head_dim, + anchor_.data() + + (layer_idx * config_.max_block_num + + block_table_before_retrieval_ + [batch_id][init_block_num]) * + config_.anchor_num * config_.q_head_num * + config_.head_dim, + config_.q_head_num * config_.head_dim, + avg_q_fp16[batch_id].data(), + config_.q_head_num * config_.head_dim, + block_similar_[batch_id].data() + init_block_num, + block_num, ith, nth, GGML_TASK_TYPE_COMPUTE, + GGML_TYPE_F16, GGML_TYPE_F16, GGML_TYPE_F32, + GGML_PREC_DEFAULT); + if (!ok) { + printf("llamafile_sgemm failed\n"); + } + }, + nullptr); + } else { + backend->do_work_stealing_job( + block_num, nullptr, + [&](int task_id) { + int block_id = task_id + init_block_num; + int block_idx = + block_table_before_retrieval_[batch_id][block_id]; + bool ok = llamafile_sgemm( + 1, 1, config_.q_head_num * config_.head_dim, + anchor_.data() + + (layer_idx * config_.max_block_num + + block_table_before_retrieval_[batch_id] + [block_idx]) * + config_.anchor_num * config_.q_head_num * + config_.head_dim, + config_.q_head_num * config_.head_dim, + avg_q_fp16[batch_id].data(), + config_.q_head_num * config_.head_dim, + block_similar_[batch_id].data() + block_id, 1, 0, 1, + GGML_TASK_TYPE_COMPUTE, GGML_TYPE_F16, + GGML_TYPE_F16, GGML_TYPE_F32, GGML_PREC_DEFAULT); + if (!ok) { + printf("llamafile_sgemm failed\n"); + } + }, + nullptr); + } + } + } else { + backend->do_work_stealing_job( + batch_size * max_block_num, nullptr, + [&](int task_id) { + int batch_id = task_id / max_block_num; + int block_id = task_id % max_block_num; + int seq_len = cache_seqlens_[batch_id]; + + if (block_id < init_block_num || + block_id >= + (seq_len / config_.block_len) - local_block_num) { + return; + } + + int block_idx = + block_table_before_retrieval_[batch_id][block_id]; + float sim = 0; + + for (int head_id = 0; head_id < config_.q_head_num; head_id++) { + for (int i = 0; i < config_.head_dim; i++) { + float q_i = 0, + qa_i = std::numeric_limits::lowest(); + for (int q_id = 0; q_id < q_len; q_id++) { + q_i += GGML_FP16_TO_FP32( + q_in_data[batch_id * q_len * + config_.q_head_num * + config_.head_dim + + q_id * config_.q_head_num * + config_.head_dim + + head_id * config_.head_dim + i]); + } + q_i /= q_len; + for (int anchor_id = 0; anchor_id < config_.anchor_num; + anchor_id++) { + qa_i = std::max( + qa_i, + GGML_FP16_TO_FP32( + anchor_[(long long)layer_idx * + config_.max_block_num * + config_.anchor_num * + config_.q_head_num * + config_.head_dim + + block_idx * config_.anchor_num * + config_.q_head_num * + config_.head_dim + + anchor_id * config_.q_head_num * + config_.head_dim + + head_id * config_.head_dim + i]) * + q_i); + } + sim += qa_i; + } + } + block_similar_[batch_id][block_id] = sim; + }, + nullptr); + } + // Timer end + auto end = std::chrono::high_resolution_clock::now(); + std::chrono::duration diff = end - start; + // printf("layer %d time of calculating similarity: %f s\n", layer_idx, + // diff.count()); +} + +void KVCache::select_block_layer_(int batch_size, int layer_idx, + int max_block_num, int init_block_num, + int local_block_num, int pick_block_num) { + // Timer start + auto start = std::chrono::high_resolution_clock::now(); + + for (int batch_idx = 0; batch_idx < batch_size; batch_idx++) { + + if (cache_seqlens_[batch_idx] / config_.block_len <= + init_block_num + pick_block_num + local_block_num) { + block_table_after_retrieval_[batch_idx].swap( + block_table_before_retrieval_[batch_idx]); + selected_blocks_num_history_[(layer_idx - config_.layer_offset) / + config_.layer_step] = 0; + continue; + } + + for (int block_id = init_block_num; + block_id < + (cache_seqlens_[batch_idx] / config_.block_len) - local_block_num; + block_id++) { + top_similar_block_[batch_idx].push(std::make_pair( + block_similar_[batch_idx][block_id], + block_table_before_retrieval_[batch_idx][block_id])); + if (top_similar_block_[batch_idx].size() > pick_block_num) { + top_similar_block_[batch_idx].pop(); + } + } + + int i = 0; + for (; i < init_block_num; i++) { + block_table_after_retrieval_[batch_idx][i] = + block_table_before_retrieval_[batch_idx][i]; + } + while (!top_similar_block_[batch_idx].empty()) { + block_table_after_retrieval_[batch_idx][i] = + top_similar_block_[batch_idx].top().second; + top_similar_block_[batch_idx].pop(); + i++; + } + for (; i < init_block_num + pick_block_num + local_block_num; i++) { + block_table_after_retrieval_[batch_idx][i] = + block_table_before_retrieval_[batch_idx] + [(cache_seqlens_[batch_idx] / + config_.block_len) - + local_block_num + i - + init_block_num - pick_block_num]; + } + if (cache_seqlens_[batch_idx] % config_.block_len != 0) { + block_table_after_retrieval_[batch_idx][i] = + block_table_before_retrieval_[batch_idx][( + cache_seqlens_[batch_idx] / config_.block_len)]; + cache_seqlens_[batch_idx] = + (cache_seqlens_[batch_idx] % config_.block_len) + + i * config_.block_len; + i++; + } else { + cache_seqlens_[batch_idx] = + (cache_seqlens_[batch_idx] % config_.block_len) + + i * config_.block_len; + } + for (int j = 0; j < i; j++) { + selected_blocks_history_[(layer_idx - config_.layer_offset) / + config_.layer_step][batch_idx][j] = + block_table_after_retrieval_[batch_idx][j]; + } + selected_blocks_num_history_[(layer_idx - config_.layer_offset) / + config_.layer_step] = i; + } + + // Timer end + auto end = std::chrono::high_resolution_clock::now(); + std::chrono::duration diff = end - start; + // printf("layer %d time of selecting blocks: %f s\n", layer_idx, + // diff.count()); +} + +// retrieval kvcache, get the init_block_num block at beginning, top +// pick_block_num similar and last local_block_num blocks. Each task +// calculates the simlarity of a certain block with the query, then push +// the block into the priority queue. Finally, the required blocks are +// pushed into the block_table_after_retrieval_. +void KVCache::retrieval_kvcache_layer_(const uint16_t *q_in_data, + int init_block_num, int local_block_num, + int pick_block_num, int q_len, + int generate_token_idx, int batch_size, + int layer_idx, int *cache_seqlens, + int &max_block_num, Backend *backend) { + // Timer start + auto start = std::chrono::high_resolution_clock::now(); + max_block_num_after_retrieval_ = 0; + if (pick_block_num != -1 && + (generate_token_idx % config_.token_step != 0 || + (layer_idx % config_.layer_step != config_.layer_offset))) { + + if (selected_blocks_num_history_[(layer_idx - config_.layer_offset) / + config_.layer_step] == 0) { + max_block_num_after_retrieval_ = max_block_num; + block_table_after_retrieval_.swap(block_table_before_retrieval_); + } else { + max_block_num_after_retrieval_ = selected_blocks_num_history_ + [(layer_idx - config_.layer_offset) / config_.layer_step]; + for (int batch_idx = 0; batch_idx < batch_size; batch_idx++) { + for (int i = 0; i < max_block_num_after_retrieval_; i++) { + block_table_after_retrieval_[batch_idx][i] = + selected_blocks_history_[(layer_idx - + config_.layer_offset) / + config_.layer_step][batch_idx] + [i]; + } + + if (cache_seqlens[batch_idx] % config_.block_len == 1) { + selected_blocks_num_history_[(layer_idx - + config_.layer_offset) / + config_.layer_step] += 1; + int x = + selected_blocks_num_history_[(layer_idx - + config_.layer_offset) / + config_.layer_step]; + int last_block_idx = + block_table_before_retrieval_[batch_idx] + [cache_seqlens[batch_idx] / + config_.block_len]; + selected_blocks_history_[(layer_idx - + config_.layer_offset) / + config_.layer_step][batch_idx] + [x - 1] = last_block_idx; + block_table_after_retrieval_[batch_idx][x - 1] = + last_block_idx; + } + cache_seqlens_[batch_idx] = + (cache_seqlens_[batch_idx] % config_.block_len) + + selected_blocks_num_history_[(layer_idx - + config_.layer_offset) / + config_.layer_step] * + config_.block_len - + config_.block_len; + } + } + } else if (pick_block_num != -1) { + max_block_num_after_retrieval_ = + std::min(max_block_num, + init_block_num + pick_block_num + local_block_num + 1); + calculate_block_similarity_layer_(q_in_data, batch_size, layer_idx, + q_len, max_block_num, cache_seqlens, + init_block_num, local_block_num, + pick_block_num, backend); + select_block_layer_(batch_size, layer_idx, max_block_num, + init_block_num, local_block_num, pick_block_num); + } else { + selected_blocks_num_history_[(layer_idx - config_.layer_offset) / + config_.layer_step] = 0; + max_block_num_after_retrieval_ = max_block_num; + block_table_after_retrieval_.swap(block_table_before_retrieval_); + } + + // Timer end + auto end = std::chrono::high_resolution_clock::now(); + // printf("layer %d time of retrieval kvcache: %f s\n", layer_idx, + // std::chrono::duration(end - start).count()); +} +void KVCache::calculate_sparsity_layer_(const uint16_t *q_in_data, + float *attn_sparsity, int batch_size, + int max_block_num, int *block_table, + int *cache_seqlens, Backend *backend + +) { + // Timer start + auto start = std::chrono::high_resolution_clock::now(); + seq_len_ = config_.block_len; + backend->do_work_stealing_job( + batch_size * config_.kv_head_num * max_block_num, + [&](int thread_id) { + thread_cur_head_idx_[thread_id].first = -1; + thread_cur_head_idx_[thread_id].second = -1; + }, + [&](int task_id) { + int batch_id = task_id / (config_.kv_head_num * max_block_num); + int head_id = (task_id % (config_.kv_head_num * max_block_num)) / + max_block_num; + int block_id = task_id % max_block_num; + int thread_id = Backend::thread_local_id; + // If the block is out of the sequence length, skip it. + if (cache_seqlens[batch_id] / config_.block_len < block_id) { + return; + } + int block_idx = block_table[batch_id * max_block_num + block_id]; + if (cache_seqlens_[batch_id] / config_.block_len == block_id) { + int seq_len = cache_seqlens_[batch_id] % config_.block_len; + if (seq_len == 0) + return; + + // Prepare the attention mask for the last block. + int full_blocks = seq_len / 8; + int remaining_bits = seq_len % 8; + // Fill full blocks with 1s + for (int i = 0; i < full_blocks; ++i) { + thread_local_attn_mask_[thread_id][i] = 0xFF; + } + // Fill the remaining bits in the next block + if (remaining_bits > 0 && full_blocks < seq_len_ / 8) { + thread_local_attn_mask_[thread_id][full_blocks] = + (1 << remaining_bits) - 1; + } else { + thread_local_attn_mask_[thread_id][full_blocks] = 0; + } + + for (int i = full_blocks + 1; i < seq_len_ / 8; ++i) { + thread_local_attn_mask_[thread_id][i] = 0; + } + if (config_.kv_type == ggml_type::GGML_TYPE_F16) { + attn_with_kvcache_one_block_( + config_.head_dim, + config_.q_head_num / config_.kv_head_num, GGML_TYPE_F16, + (void *)&q_in_data[batch_id * config_.kv_head_num * + n_gqa_ * config_.head_dim + + head_id * n_gqa_ * config_.head_dim], + seq_len_, 0, false, + thread_local_attn_mask_[thread_id].data(), + GGML_TYPE_F16, 0, + k_cache_fp16_[layer_id_][head_id][block_idx].data(), 0, + nullptr, nullptr, GGML_TYPE_F16, 1, + v_cache_fp16_[layer_id_][head_id][block_idx].data(), 0, + nullptr, nullptr, + thread_local_attn_score_[thread_id].data(), + thread_local_output_fp32_[thread_id].data(), + thread_local_attn_lse_[thread_id].data(), + thread_local_draft_[thread_id].data(), nullptr, + cos_.data(), sin_.data()); + } else if (config_.kv_type == ggml_type::GGML_TYPE_Q4_0) { + attn_with_kvcache_one_block_( + config_.head_dim, + config_.q_head_num / config_.kv_head_num, + GGML_TYPE_Q8_0, q_q8_0_[batch_id][head_id].data(), + seq_len_, 0, false, + thread_local_attn_mask_[thread_id].data(), + GGML_TYPE_Q4_0, 0, + k_cache_q4[layer_id_][head_id][block_idx].data(), 0, + nullptr, nullptr, GGML_TYPE_Q4_0, 1, + v_cache_q4[layer_id_][head_id][block_idx].data(), 0, + nullptr, nullptr, + thread_local_attn_score_[thread_id].data(), + thread_local_output_q8_0_[thread_id].data(), + thread_local_attn_lse_[thread_id].data(), + thread_local_draft_[thread_id].data(), nullptr, + cos_.data(), sin_.data()); + dequantize_row_q8_0( + thread_local_output_q8_0_[thread_id].data(), + thread_local_output_fp32_[thread_id].data(), + n_gqa_ * config_.head_dim); + } else if (config_.kv_type == ggml_type::GGML_TYPE_Q8_0) { + attn_with_kvcache_one_block_( + config_.head_dim, + config_.q_head_num / config_.kv_head_num, + GGML_TYPE_Q8_0, q_q8_0_[batch_id][head_id].data(), + seq_len_, 0, false, + thread_local_attn_mask_[thread_id].data(), + GGML_TYPE_Q8_0, 0, + k_cache_q8[layer_id_][head_id][block_idx].data(), 0, + nullptr, nullptr, GGML_TYPE_Q8_0, 1, + v_cache_q8[layer_id_][head_id][block_idx].data(), 0, + nullptr, nullptr, + thread_local_attn_score_[thread_id].data(), + thread_local_output_q8_0_[thread_id].data(), + thread_local_attn_lse_[thread_id].data(), + thread_local_draft_[thread_id].data(), nullptr, + cos_.data(), sin_.data()); + dequantize_row_q8_0( + thread_local_output_q8_0_[thread_id].data(), + thread_local_output_fp32_[thread_id].data(), + n_gqa_ * config_.head_dim); + } + } else { + if (config_.kv_type == ggml_type::GGML_TYPE_F16) { + attn_with_kvcache_one_block_( + config_.head_dim, + config_.q_head_num / config_.kv_head_num, GGML_TYPE_F16, + (void *)&q_in_data[batch_id * config_.kv_head_num * + n_gqa_ * config_.head_dim + + head_id * n_gqa_ * config_.head_dim], + seq_len_, 0, true, nullptr, GGML_TYPE_F16, 0, + k_cache_fp16_[layer_id_][head_id][block_idx].data(), 0, + nullptr, nullptr, GGML_TYPE_F16, 1, + v_cache_fp16_[layer_id_][head_id][block_idx].data(), 0, + nullptr, nullptr, + thread_local_attn_score_[thread_id].data(), + thread_local_output_fp32_[thread_id].data(), + thread_local_attn_lse_[thread_id].data(), + thread_local_draft_[thread_id].data(), nullptr, + cos_.data(), sin_.data()); + + } else if (config_.kv_type == ggml_type::GGML_TYPE_Q4_0) { + attn_with_kvcache_one_block_( + config_.head_dim, + config_.q_head_num / config_.kv_head_num, + GGML_TYPE_Q8_0, q_q8_0_[batch_id][head_id].data(), + seq_len_, 0, true, nullptr, GGML_TYPE_Q4_0, 0, + k_cache_q4[layer_id_][head_id][block_idx].data(), 0, + nullptr, nullptr, GGML_TYPE_Q4_0, 1, + v_cache_q4[layer_id_][head_id][block_idx].data(), 0, + nullptr, nullptr, + thread_local_attn_score_[thread_id].data(), + thread_local_output_q8_0_[thread_id].data(), + thread_local_attn_lse_[thread_id].data(), + thread_local_draft_[thread_id].data(), nullptr, + cos_.data(), sin_.data()); + dequantize_row_q8_0( + thread_local_output_q8_0_[thread_id].data(), + thread_local_output_fp32_[thread_id].data(), + n_gqa_ * config_.head_dim); + } else if (config_.kv_type == ggml_type::GGML_TYPE_Q8_0) { + attn_with_kvcache_one_block_( + config_.head_dim, + config_.q_head_num / config_.kv_head_num, + GGML_TYPE_Q8_0, q_q8_0_[batch_id][head_id].data(), + seq_len_, 0, true, nullptr, GGML_TYPE_Q8_0, 0, + k_cache_q8[layer_id_][head_id][block_idx].data(), 0, + nullptr, nullptr, GGML_TYPE_Q8_0, 1, + v_cache_q8[layer_id_][head_id][block_idx].data(), 0, + nullptr, nullptr, + thread_local_attn_score_[thread_id].data(), + thread_local_output_q8_0_[thread_id].data(), + thread_local_attn_lse_[thread_id].data(), + thread_local_draft_[thread_id].data(), nullptr, + cos_.data(), sin_.data()); + dequantize_row_q8_0( + thread_local_output_q8_0_[thread_id].data(), + thread_local_output_fp32_[thread_id].data(), + n_gqa_ * config_.head_dim); + } + } + for (int i = 0; i < n_gqa_; i++) { + block_lse_[batch_id][block_idx][head_id * n_gqa_ + i] = + thread_local_attn_lse_[thread_id][i]; + } + int cur_batch_idx = thread_cur_head_idx_[thread_id].first; + int cur_head_id = thread_cur_head_idx_[thread_id].second; + if (batch_id == cur_batch_idx && head_id == cur_head_id) { + for (int i = 0; i < n_gqa_; i++) { + float new_attn_lse = + thread_local_cur_attn_lse_[thread_id][i] + + std::log( + 1.0 + + std::exp(thread_local_attn_lse_[thread_id][i] - + thread_local_cur_attn_lse_[thread_id][i])); + ggml_vec_scale_f32( + config_.head_dim, + thread_local_cur_output_fp32_[thread_id].data() + + i * config_.head_dim, + std::exp(thread_local_cur_attn_lse_[thread_id][i] - + new_attn_lse)); + ggml_vec_scale_f32( + config_.head_dim, + thread_local_output_fp32_[thread_id].data() + + i * config_.head_dim, + std::exp(thread_local_attn_lse_[thread_id][i] - + new_attn_lse)); + for (int j = 0; j < config_.head_dim; j++) { + thread_local_cur_output_fp32_[thread_id] + [i * config_.head_dim + + j] += + thread_local_output_fp32_[thread_id] + [i * config_.head_dim + j]; + } + thread_local_cur_attn_lse_[thread_id][i] = new_attn_lse; + } + } else { + if (cur_batch_idx != -1) { + mutex_[cur_batch_idx][cur_head_id]->lock(); + for (int i = 0; i < n_gqa_; i++) { + if (std::abs(attn_lse_[cur_batch_idx][cur_head_id][i]) < + 1e-6) { + attn_lse_[cur_batch_idx][cur_head_id][i] = + thread_local_cur_attn_lse_[thread_id][i]; + for (int j = 0; j < config_.head_dim; j++) { + output_fp32_[cur_batch_idx][cur_head_id] + [i * config_.head_dim + j] = + thread_local_cur_output_fp32_ + [thread_id] + [i * config_.head_dim + j]; + } + continue; + } + float new_attn_lse = + attn_lse_[cur_batch_idx][cur_head_id][i] + + std::log( + 1.0 + + std::exp( + thread_local_cur_attn_lse_[thread_id][i] - + attn_lse_[cur_batch_idx][cur_head_id][i])); + ggml_vec_scale_f32( + config_.head_dim, + output_fp32_[cur_batch_idx][cur_head_id].data() + + i * config_.head_dim, + std::exp(attn_lse_[cur_batch_idx][cur_head_id][i] - + new_attn_lse)); + ggml_vec_scale_f32( + config_.head_dim, + thread_local_cur_output_fp32_[thread_id].data() + + i * config_.head_dim, + std::exp(thread_local_cur_attn_lse_[thread_id][i] - + new_attn_lse)); + for (int j = 0; j < config_.head_dim; j++) { + output_fp32_[cur_batch_idx][cur_head_id] + [i * config_.head_dim + j] += + thread_local_cur_output_fp32_ + [thread_id][i * config_.head_dim + j]; + } + attn_lse_[cur_batch_idx][cur_head_id][i] = new_attn_lse; + } + mutex_[cur_batch_idx][cur_head_id]->unlock(); + } + thread_cur_head_idx_[thread_id].first = batch_id; + thread_cur_head_idx_[thread_id].second = head_id; + for (int i = 0; i < n_gqa_; i++) { + thread_local_cur_attn_lse_[thread_id][i] = + thread_local_attn_lse_[thread_id][i]; + for (int j = 0; j < config_.head_dim; j++) { + thread_local_cur_output_fp32_ + [thread_id][i * config_.head_dim + j] = + thread_local_output_fp32_[thread_id] + [i * config_.head_dim + + j]; + } + } + } + }, + // Merge the results of the remaining blocks. + [&](int thread_id) { + int cur_batch_idx = thread_cur_head_idx_[thread_id].first; + int cur_head_id = thread_cur_head_idx_[thread_id].second; + if (cur_head_id != -1) { + mutex_[cur_batch_idx][cur_head_id]->lock(); + for (int i = 0; i < n_gqa_; i++) { + float new_attn_lse; + if (std::abs(attn_lse_[cur_batch_idx][cur_head_id][i]) < + 1e-6) { + attn_lse_[cur_batch_idx][cur_head_id][i] = + thread_local_cur_attn_lse_[thread_id][i]; + for (int j = 0; j < config_.head_dim; j++) { + output_fp32_[cur_batch_idx][cur_head_id] + [i * config_.head_dim + j] = + thread_local_cur_output_fp32_ + [thread_id] + [i * config_.head_dim + j]; + } + continue; + } + new_attn_lse = + attn_lse_[cur_batch_idx][cur_head_id][i] + + std::log( + 1.0 + + std::exp(thread_local_cur_attn_lse_[thread_id][i] - + attn_lse_[cur_batch_idx][cur_head_id][i])); + ggml_vec_scale_f32( + config_.head_dim, + output_fp32_[cur_batch_idx][cur_head_id].data() + + i * config_.head_dim, + std::exp(attn_lse_[cur_batch_idx][cur_head_id][i] - + new_attn_lse)); + ggml_vec_scale_f32( + config_.head_dim, + thread_local_cur_output_fp32_[thread_id].data() + + i * config_.head_dim, + std::exp(thread_local_cur_attn_lse_[thread_id][i] - + new_attn_lse)); + for (int j = 0; j < config_.head_dim; j++) { + output_fp32_[cur_batch_idx][cur_head_id] + [i * config_.head_dim + j] += + thread_local_cur_output_fp32_[thread_id] + [i * config_.head_dim + + j]; + } + attn_lse_[cur_batch_idx][cur_head_id][i] = new_attn_lse; + } + mutex_[cur_batch_idx][cur_head_id]->unlock(); + } + }); + + for (int i = 0; i < batch_size; i++) { + for (int j = 0; j < max_block_num_after_retrieval_; j++) { + int block_idx = block_table_after_retrieval_[i][j]; + for (int k = 0; k < config_.q_head_num; k++) { + attn_sparsity[i * config_.q_head_num + k] += + std::exp(block_lse_[i][block_idx][k] - + attn_lse_[i][k / n_gqa_][k % n_gqa_]); + } + } + } + + // Timer end + auto end = std::chrono::high_resolution_clock::now(); + std::chrono::duration diff = end - start; + // printf("layer %d time of calculating sparsity: %f s\n", layer_id_, + // diff.count()); +} + +void KVCache::attn_initialize_kvhead_(int batch_size, int layer_idx, + int *block_table, int &max_block_num, + int *cache_seqlens) { + // Timer start + auto start = std::chrono::high_resolution_clock::now(); + for (int batch_idx = 0; batch_idx < batch_size; batch_idx++) { + // initialize output_fp32_ and attn_lse_ + for (int i = 0; i < config_.kv_head_num; i++) { + for (int j = 0; j < n_gqa_ * config_.head_dim; j++) { + output_fp32_[batch_idx][i][j] = 0; + } + for (int j = 0; j < n_gqa_; j++) { + attn_lse_[batch_idx][i][j] = 0; + } + } + + // clear top_similar_block_ + while (!top_similar_block_[batch_idx].empty()) + top_similar_block_[batch_idx].pop(); + } + + for (int batch_idx = 0; batch_idx < batch_size; batch_idx++) { + cache_seqlens_[batch_idx] = cache_seqlens[batch_idx]; + for (int i = 0; i < max_block_num; i++) { + for (int j = 0; j < config_.kv_head_num; j++) { + block_table_before_retrieval_kvhead_[batch_idx][i][j] = + block_table[batch_idx * max_block_num + i]; + block_similar_kv_head_[batch_idx][i][j] = 0; + } + } + } + + // Timer end + auto end = std::chrono::high_resolution_clock::now(); + // printf("layer %d time of initializing attn: %f s\n", layer_idx, + // std::chrono::duration(end - start).count()); +} +void KVCache::retrieval_kvcache_kvhead_(const uint16_t *q_in_data, + int init_block_num, int local_block_num, + int pick_block_num, int q_len, + int generate_token_idx, int batch_size, + int layer_idx, int *cache_seqlens, + int &max_block_num, Backend *backend) { + // Timer start + auto start = std::chrono::high_resolution_clock::now(); + max_block_num_after_retrieval_ = 0; + if (pick_block_num != -1 && + (generate_token_idx % config_.token_step != 0 || + (layer_idx % config_.layer_step != config_.layer_offset))) { + + if (selected_blocks_num_history_[(layer_idx - config_.layer_offset) / + config_.layer_step] == 0) { + max_block_num_after_retrieval_ = max_block_num; + for (int batch_idx = 0; batch_idx < batch_size; batch_idx++) { + for (int i = 0; i < max_block_num; i++) { + for (int j = 0; j < config_.kv_head_num; j++) { + block_table_after_retrieval_kvhead_[batch_idx][i][j] = + block_table_before_retrieval_kvhead_[batch_idx][i] + [j]; + } + } + } + } else { + + max_block_num_after_retrieval_ = selected_blocks_num_history_ + [(layer_idx - config_.layer_offset) / config_.layer_step]; + + for (int batch_idx = 0; batch_idx < batch_size; batch_idx++) { + for (int i = 0; i < max_block_num_after_retrieval_; i++) { + for (int j = 0; j < config_.kv_head_num; j++) { + block_table_after_retrieval_kvhead_[batch_idx][i][j] = + selected_blocks_history_kvhead_ + [(layer_idx - config_.layer_offset) / + config_.layer_step][batch_idx][i][j]; + } + } + + if (cache_seqlens[batch_idx] % config_.block_len == 1) { + selected_blocks_num_history_[(layer_idx - + config_.layer_offset) / + config_.layer_step] += 1; + int x = + selected_blocks_num_history_[(layer_idx - + config_.layer_offset) / + config_.layer_step]; + for (int i = 0; i < config_.kv_head_num; i++) { + int last_block_idx = + block_table_before_retrieval_kvhead_ + [batch_idx][cache_seqlens[batch_idx] / + config_.block_len][i]; + selected_blocks_history_kvhead_[(layer_idx - + config_.layer_offset) / + config_.layer_step] + [batch_idx][x - 1][i] = + last_block_idx; + block_table_after_retrieval_kvhead_[batch_idx][x - 1] + [i] = last_block_idx; + } + } + cache_seqlens_[batch_idx] = std::min( + cache_seqlens_[batch_idx], + (cache_seqlens_[batch_idx] % config_.block_len) + + (init_block_num + pick_block_num + local_block_num) * + config_.block_len); + } + } + } else if (pick_block_num != -1) { + max_block_num_after_retrieval_ = + std::min(max_block_num, + init_block_num + pick_block_num + local_block_num + 1); + calculate_block_similarity_kvhead_(q_in_data, batch_size, layer_idx, + q_len, max_block_num, cache_seqlens, + init_block_num, local_block_num, + pick_block_num, backend); + select_block_kvhead_(batch_size, layer_idx, max_block_num, + init_block_num, local_block_num, pick_block_num); + } else { + selected_blocks_num_history_[(layer_idx - config_.layer_offset) / + config_.layer_step] = 0; + max_block_num_after_retrieval_ = max_block_num; + for (int batch_idx = 0; batch_idx < batch_size; batch_idx++) { + for (int i = 0; i < max_block_num; i++) { + for (int j = 0; j < config_.kv_head_num; j++) { + block_table_after_retrieval_kvhead_[batch_idx][i][j] = + block_table_before_retrieval_kvhead_[batch_idx][i][j]; + } + } + } + } + + // Timer end + auto end = std::chrono::high_resolution_clock::now(); + // printf("layer %d time of retrieval kvcache: %f s\n", layer_idx, + // std::chrono::duration(end - start).count()); +} +void KVCache::calculate_sparsity_kvhead_(const uint16_t *q_in_data, + float *attn_sparsity, int batch_size, + int max_block_num, int *block_table, + int *cache_seqlens, Backend *backend) { + // Timer start + auto start = std::chrono::high_resolution_clock::now(); + seq_len_ = config_.block_len; + backend->do_work_stealing_job( + batch_size * config_.kv_head_num * max_block_num, + [&](int thread_id) { + thread_cur_head_idx_[thread_id].first = -1; + thread_cur_head_idx_[thread_id].second = -1; + }, + [&](int task_id) { + int batch_id = task_id / (config_.kv_head_num * max_block_num); + int head_id = (task_id % (config_.kv_head_num * max_block_num)) / + max_block_num; + int block_id = task_id % max_block_num; + int thread_id = Backend::thread_local_id; + // If the block is out of the sequence length, skip it. + if (cache_seqlens[batch_id] / config_.block_len < block_id) { + return; + } + int block_idx = block_table[batch_id * max_block_num + block_id]; + if (cache_seqlens_[batch_id] / config_.block_len == block_id) { + int seq_len = cache_seqlens_[batch_id] % config_.block_len; + if (seq_len == 0) + return; + + // Prepare the attention mask for the last block. + int full_blocks = seq_len / 8; + int remaining_bits = seq_len % 8; + + // Fill full blocks with 1s + for (int i = 0; i < full_blocks; ++i) { + thread_local_attn_mask_[thread_id][i] = 0xFF; + } + // Fill the remaining bits in the next block + if (remaining_bits > 0 && full_blocks < seq_len_ / 8) { + thread_local_attn_mask_[thread_id][full_blocks] = + (1 << remaining_bits) - 1; + } else { + thread_local_attn_mask_[thread_id][full_blocks] = 0; + } + + for (int i = full_blocks + 1; i < seq_len_ / 8; ++i) { + thread_local_attn_mask_[thread_id][i] = 0; + } + if (config_.kv_type == ggml_type::GGML_TYPE_F16) { + attn_with_kvcache_one_block_( + config_.head_dim, + config_.q_head_num / config_.kv_head_num, GGML_TYPE_F16, + (void *)&q_in_data[batch_id * config_.kv_head_num * + n_gqa_ * config_.head_dim + + head_id * n_gqa_ * config_.head_dim], + seq_len_, 0, false, + thread_local_attn_mask_[thread_id].data(), + GGML_TYPE_F16, 0, + k_cache_fp16_[layer_id_][head_id][block_idx].data(), 0, + nullptr, nullptr, GGML_TYPE_F16, 1, + v_cache_fp16_[layer_id_][head_id][block_idx].data(), 0, + nullptr, nullptr, + thread_local_attn_score_[thread_id].data(), + thread_local_output_fp32_[thread_id].data(), + thread_local_attn_lse_[thread_id].data(), + thread_local_draft_[thread_id].data(), nullptr, + cos_.data(), sin_.data()); + } else if (config_.kv_type == ggml_type::GGML_TYPE_Q4_0) { + attn_with_kvcache_one_block_( + config_.head_dim, + config_.q_head_num / config_.kv_head_num, + GGML_TYPE_Q8_0, q_q8_0_[batch_id][head_id].data(), + seq_len_, 0, false, + thread_local_attn_mask_[thread_id].data(), + GGML_TYPE_Q4_0, 0, + k_cache_q4[layer_id_][head_id][block_idx].data(), 0, + nullptr, nullptr, GGML_TYPE_Q4_0, 1, + v_cache_q4[layer_id_][head_id][block_idx].data(), 0, + nullptr, nullptr, + thread_local_attn_score_[thread_id].data(), + thread_local_output_q8_0_[thread_id].data(), + thread_local_attn_lse_[thread_id].data(), + thread_local_draft_[thread_id].data(), nullptr, + cos_.data(), sin_.data()); + dequantize_row_q8_0( + thread_local_output_q8_0_[thread_id].data(), + thread_local_output_fp32_[thread_id].data(), + n_gqa_ * config_.head_dim); + } else if (config_.kv_type == ggml_type::GGML_TYPE_Q8_0) { + attn_with_kvcache_one_block_( + config_.head_dim, + config_.q_head_num / config_.kv_head_num, + GGML_TYPE_Q8_0, q_q8_0_[batch_id][head_id].data(), + seq_len_, 0, false, + thread_local_attn_mask_[thread_id].data(), + GGML_TYPE_Q8_0, 0, + k_cache_q8[layer_id_][head_id][block_idx].data(), 0, + nullptr, nullptr, GGML_TYPE_Q8_0, 1, + v_cache_q8[layer_id_][head_id][block_idx].data(), 0, + nullptr, nullptr, + thread_local_attn_score_[thread_id].data(), + thread_local_output_q8_0_[thread_id].data(), + thread_local_attn_lse_[thread_id].data(), + thread_local_draft_[thread_id].data(), nullptr, + cos_.data(), sin_.data()); + dequantize_row_q8_0( + thread_local_output_q8_0_[thread_id].data(), + thread_local_output_fp32_[thread_id].data(), + n_gqa_ * config_.head_dim); + } + } else { + if (config_.kv_type == ggml_type::GGML_TYPE_F16) { + attn_with_kvcache_one_block_( + config_.head_dim, + config_.q_head_num / config_.kv_head_num, GGML_TYPE_F16, + (void *)&q_in_data[batch_id * config_.kv_head_num * + n_gqa_ * config_.head_dim + + head_id * n_gqa_ * config_.head_dim], + seq_len_, 0, true, nullptr, GGML_TYPE_F16, 0, + k_cache_fp16_[layer_id_][head_id][block_idx].data(), 0, + nullptr, nullptr, GGML_TYPE_F16, 1, + v_cache_fp16_[layer_id_][head_id][block_idx].data(), 0, + nullptr, nullptr, + thread_local_attn_score_[thread_id].data(), + thread_local_output_fp32_[thread_id].data(), + thread_local_attn_lse_[thread_id].data(), + thread_local_draft_[thread_id].data(), nullptr, + cos_.data(), sin_.data()); + + } else if (config_.kv_type == ggml_type::GGML_TYPE_Q4_0) { + attn_with_kvcache_one_block_( + config_.head_dim, + config_.q_head_num / config_.kv_head_num, + GGML_TYPE_Q8_0, q_q8_0_[batch_id][head_id].data(), + seq_len_, 0, true, nullptr, GGML_TYPE_Q4_0, 0, + k_cache_q4[layer_id_][head_id][block_idx].data(), 0, + nullptr, nullptr, GGML_TYPE_Q4_0, 1, + v_cache_q4[layer_id_][head_id][block_idx].data(), 0, + nullptr, nullptr, + thread_local_attn_score_[thread_id].data(), + thread_local_output_q8_0_[thread_id].data(), + thread_local_attn_lse_[thread_id].data(), + thread_local_draft_[thread_id].data(), nullptr, + cos_.data(), sin_.data()); + dequantize_row_q8_0( + thread_local_output_q8_0_[thread_id].data(), + thread_local_output_fp32_[thread_id].data(), + n_gqa_ * config_.head_dim); + } else if (config_.kv_type == ggml_type::GGML_TYPE_Q8_0) { + attn_with_kvcache_one_block_( + config_.head_dim, + config_.q_head_num / config_.kv_head_num, + GGML_TYPE_Q8_0, q_q8_0_[batch_id][head_id].data(), + seq_len_, 0, true, nullptr, GGML_TYPE_Q8_0, 0, + k_cache_q8[layer_id_][head_id][block_idx].data(), 0, + nullptr, nullptr, GGML_TYPE_Q8_0, 1, + v_cache_q8[layer_id_][head_id][block_idx].data(), 0, + nullptr, nullptr, + thread_local_attn_score_[thread_id].data(), + thread_local_output_q8_0_[thread_id].data(), + thread_local_attn_lse_[thread_id].data(), + thread_local_draft_[thread_id].data(), nullptr, + cos_.data(), sin_.data()); + dequantize_row_q8_0( + thread_local_output_q8_0_[thread_id].data(), + thread_local_output_fp32_[thread_id].data(), + n_gqa_ * config_.head_dim); + } + } + for (int i = 0; i < n_gqa_; i++) { + block_lse_[batch_id][block_idx][head_id * n_gqa_ + i] = + thread_local_attn_lse_[thread_id][i]; + } + int cur_batch_idx = thread_cur_head_idx_[thread_id].first; + int cur_head_id = thread_cur_head_idx_[thread_id].second; + if (batch_id == cur_batch_idx && head_id == cur_head_id) { + for (int i = 0; i < n_gqa_; i++) { + float new_attn_lse = + thread_local_cur_attn_lse_[thread_id][i] + + std::log( + 1.0 + + std::exp(thread_local_attn_lse_[thread_id][i] - + thread_local_cur_attn_lse_[thread_id][i])); + ggml_vec_scale_f32( + config_.head_dim, + thread_local_cur_output_fp32_[thread_id].data() + + i * config_.head_dim, + std::exp(thread_local_cur_attn_lse_[thread_id][i] - + new_attn_lse)); + ggml_vec_scale_f32( + config_.head_dim, + thread_local_output_fp32_[thread_id].data() + + i * config_.head_dim, + std::exp(thread_local_attn_lse_[thread_id][i] - + new_attn_lse)); + for (int j = 0; j < config_.head_dim; j++) { + thread_local_cur_output_fp32_[thread_id] + [i * config_.head_dim + + j] += + thread_local_output_fp32_[thread_id] + [i * config_.head_dim + j]; + } + thread_local_cur_attn_lse_[thread_id][i] = new_attn_lse; + } + } else { + if (cur_batch_idx != -1) { + mutex_[cur_batch_idx][cur_head_id]->lock(); + for (int i = 0; i < n_gqa_; i++) { + if (std::abs(attn_lse_[cur_batch_idx][cur_head_id][i]) < + 1e-6) { + attn_lse_[cur_batch_idx][cur_head_id][i] = + thread_local_cur_attn_lse_[thread_id][i]; + for (int j = 0; j < config_.head_dim; j++) { + output_fp32_[cur_batch_idx][cur_head_id] + [i * config_.head_dim + j] = + thread_local_cur_output_fp32_ + [thread_id] + [i * config_.head_dim + j]; + } + continue; + } + float new_attn_lse = + attn_lse_[cur_batch_idx][cur_head_id][i] + + std::log( + 1.0 + + std::exp( + thread_local_cur_attn_lse_[thread_id][i] - + attn_lse_[cur_batch_idx][cur_head_id][i])); + ggml_vec_scale_f32( + config_.head_dim, + output_fp32_[cur_batch_idx][cur_head_id].data() + + i * config_.head_dim, + std::exp(attn_lse_[cur_batch_idx][cur_head_id][i] - + new_attn_lse)); + ggml_vec_scale_f32( + config_.head_dim, + thread_local_cur_output_fp32_[thread_id].data() + + i * config_.head_dim, + std::exp(thread_local_cur_attn_lse_[thread_id][i] - + new_attn_lse)); + for (int j = 0; j < config_.head_dim; j++) { + output_fp32_[cur_batch_idx][cur_head_id] + [i * config_.head_dim + j] += + thread_local_cur_output_fp32_ + [thread_id][i * config_.head_dim + j]; + } + attn_lse_[cur_batch_idx][cur_head_id][i] = new_attn_lse; + } + mutex_[cur_batch_idx][cur_head_id]->unlock(); + } + thread_cur_head_idx_[thread_id].first = batch_id; + thread_cur_head_idx_[thread_id].second = head_id; + for (int i = 0; i < n_gqa_; i++) { + thread_local_cur_attn_lse_[thread_id][i] = + thread_local_attn_lse_[thread_id][i]; + for (int j = 0; j < config_.head_dim; j++) { + thread_local_cur_output_fp32_ + [thread_id][i * config_.head_dim + j] = + thread_local_output_fp32_[thread_id] + [i * config_.head_dim + + j]; + } + } + } + }, + // Merge the results of the remaining blocks. + [&](int thread_id) { + int cur_batch_idx = thread_cur_head_idx_[thread_id].first; + int cur_head_id = thread_cur_head_idx_[thread_id].second; + if (cur_head_id != -1) { + mutex_[cur_batch_idx][cur_head_id]->lock(); + for (int i = 0; i < n_gqa_; i++) { + float new_attn_lse; + if (std::abs(attn_lse_[cur_batch_idx][cur_head_id][i]) < + 1e-6) { + attn_lse_[cur_batch_idx][cur_head_id][i] = + thread_local_cur_attn_lse_[thread_id][i]; + for (int j = 0; j < config_.head_dim; j++) { + output_fp32_[cur_batch_idx][cur_head_id] + [i * config_.head_dim + j] = + thread_local_cur_output_fp32_ + [thread_id] + [i * config_.head_dim + j]; + } + continue; + } + new_attn_lse = + attn_lse_[cur_batch_idx][cur_head_id][i] + + std::log( + 1.0 + + std::exp(thread_local_cur_attn_lse_[thread_id][i] - + attn_lse_[cur_batch_idx][cur_head_id][i])); + ggml_vec_scale_f32( + config_.head_dim, + output_fp32_[cur_batch_idx][cur_head_id].data() + + i * config_.head_dim, + std::exp(attn_lse_[cur_batch_idx][cur_head_id][i] - + new_attn_lse)); + ggml_vec_scale_f32( + config_.head_dim, + thread_local_cur_output_fp32_[thread_id].data() + + i * config_.head_dim, + std::exp(thread_local_cur_attn_lse_[thread_id][i] - + new_attn_lse)); + for (int j = 0; j < config_.head_dim; j++) { + output_fp32_[cur_batch_idx][cur_head_id] + [i * config_.head_dim + j] += + thread_local_cur_output_fp32_[thread_id] + [i * config_.head_dim + + j]; + } + attn_lse_[cur_batch_idx][cur_head_id][i] = new_attn_lse; + } + mutex_[cur_batch_idx][cur_head_id]->unlock(); + } + }); + + for (int i = 0; i < batch_size; i++) { + for (int j = 0; j < max_block_num_after_retrieval_; j++) { + for (int k = 0; k < config_.q_head_num; k++) { + int block_idx = + block_table_after_retrieval_kvhead_[i][j][k / n_gqa_]; + attn_sparsity[i * config_.q_head_num + k] += + std::exp(block_lse_[i][block_idx][k] - + attn_lse_[i][k / n_gqa_][k % n_gqa_]); + } + } + } + + // Timer end + auto end = std::chrono::high_resolution_clock::now(); + std::chrono::duration diff = end - start; + // printf("layer %d time of calculating sparsity: %f s\n", layer_id_, + // diff.count()); +} +void KVCache::calculate_block_similarity_kvhead_( + const uint16_t *q_in_data, int batch_size, int layer_idx, int q_len, + int max_block_num, int *cache_seqlens, int init_block_num, + int local_block_num, int pick_block_num, Backend *backend) { + // Timer start + auto start = std::chrono::high_resolution_clock::now(); + backend->do_work_stealing_job( + batch_size * max_block_num, nullptr, + [&](int task_id) { + int batch_id = task_id / max_block_num; + int block_id = task_id % max_block_num; + int seq_len = cache_seqlens_[batch_id]; + + if (block_id < init_block_num || + block_id >= (seq_len / config_.block_len) - local_block_num) { + return; + } + int block_idx = + block_table_before_retrieval_kvhead_[batch_id][block_id][0]; + + for (int head_id = 0; head_id < config_.q_head_num; head_id++) { + for (int i = 0; i < config_.head_dim; i++) { + float q_i = 0, qa_i = std::numeric_limits::lowest(); + for (int q_id = 0; q_id < q_len; q_id++) { + q_i += GGML_FP16_TO_FP32( + q_in_data[batch_id * q_len * config_.q_head_num * + config_.head_dim + + q_id * config_.q_head_num * + config_.head_dim + + head_id * config_.head_dim + i]); + } + q_i /= q_len; + for (int anchor_id = 0; anchor_id < config_.anchor_num; + anchor_id++) { + qa_i = std::max( + qa_i, + GGML_FP16_TO_FP32( + anchor_[layer_idx * config_.max_block_num * + config_.anchor_num * + config_.q_head_num * + config_.head_dim + + block_idx * config_.anchor_num * + config_.q_head_num * + config_.head_dim + + anchor_id * config_.q_head_num * + config_.head_dim + + head_id * config_.head_dim + i]) * + q_i); + } + block_similar_kv_head_[batch_id][block_id] + [head_id / n_gqa_] += qa_i; + } + } + }, + nullptr); + + // Timer end + auto end = std::chrono::high_resolution_clock::now(); + std::chrono::duration diff = end - start; + // printf("layer %d time of calculating similarity: %f s\n", layer_idx, + // diff.count()); +} +void KVCache::select_block_kvhead_(int batch_size, int layer_idx, + int max_block_num, int init_block_num, + int local_block_num, int pick_block_num) { + // Timer start + auto start = std::chrono::high_resolution_clock::now(); + + for (int batch_idx = 0; batch_idx < batch_size; batch_idx++) { + int cache_len_after_retrieval = 0; + if (cache_seqlens_[batch_idx] / config_.block_len <= + init_block_num + pick_block_num + local_block_num) { + selected_blocks_num_history_[(layer_idx - config_.layer_offset) / + config_.layer_step] = 0; + for (int i = 0; i < max_block_num; i++) { + for (int j = 0; j < config_.kv_head_num; j++) { + block_table_after_retrieval_kvhead_[batch_idx][i][j] = + block_table_before_retrieval_kvhead_[batch_idx][i][j]; + } + } + continue; + } + for (int head_id = 0; head_id < config_.kv_head_num; head_id++) { + + for (int block_id = init_block_num; + block_id < (cache_seqlens_[batch_idx] / config_.block_len) - + local_block_num; + block_id++) { + + top_similar_block_[batch_idx].push(std::make_pair( + block_similar_kv_head_[batch_idx][block_id][head_id], + block_table_before_retrieval_kvhead_[batch_idx][block_id] + [head_id])); + if (top_similar_block_[batch_idx].size() > pick_block_num) { + top_similar_block_[batch_idx].pop(); + } + } + + int i = 0; + for (; i < init_block_num; i++) { + block_table_after_retrieval_kvhead_[batch_idx][i][head_id] = + block_table_before_retrieval_kvhead_[batch_idx][i][head_id]; + } + while (!top_similar_block_[batch_idx].empty()) { + block_table_after_retrieval_kvhead_[batch_idx][i][head_id] = + top_similar_block_[batch_idx].top().second; + top_similar_block_[batch_idx].pop(); + i++; + } + for (; i < init_block_num + pick_block_num + local_block_num; i++) { + block_table_after_retrieval_kvhead_[batch_idx][i][head_id] = + block_table_before_retrieval_kvhead_ + [batch_idx] + [(cache_seqlens_[batch_idx] / config_.block_len) - + local_block_num + i - init_block_num - pick_block_num] + [head_id]; + } + if (cache_seqlens_[batch_idx] % config_.block_len != 0) { + block_table_after_retrieval_kvhead_[batch_idx][i][head_id] = + block_table_before_retrieval_kvhead_[batch_idx][( + cache_seqlens_[batch_idx] / config_.block_len)] + [head_id]; + cache_len_after_retrieval = + (cache_seqlens_[batch_idx] % config_.block_len) + + i * config_.block_len; + i++; + } else { + cache_len_after_retrieval = + (cache_seqlens_[batch_idx] % config_.block_len) + + i * config_.block_len; + } + for (int j = 0; j < i; j++) { + selected_blocks_history_kvhead_ + [(layer_idx - config_.layer_offset) / config_.layer_step] + [batch_idx][j][head_id] = + block_table_after_retrieval_kvhead_[batch_idx][j] + [head_id]; + } + } + cache_seqlens_[batch_idx] = cache_len_after_retrieval; + selected_blocks_num_history_[(layer_idx - config_.layer_offset) / + config_.layer_step] = + (cache_len_after_retrieval + config_.block_len - 1) / + config_.block_len; + } + + // Timer end + auto end = std::chrono::high_resolution_clock::now(); + std::chrono::duration diff = end - start; + // printf("layer %d time of selecting block: %f s\n", layer_idx, + // diff.count()) +} + +void KVCache::get_attn_sparsity(const ggml_fp16_t *q_in, float *attn_sparsity, + int layer_idx, int generate_token_idx, + int q_len, int batch_size, int max_block_num, + int *block_table, int *cache_seqlens, + int *block_table_origin, + int *cache_seqlens_origin, + int max_block_num_origin, int topk, int local, + Backend *backend) { + // Timer start + auto start = std::chrono::high_resolution_clock::now(); + layer_id_ = layer_idx; + int thread_num = backend->get_thread_num(); + batch_size = 1; + + const uint16_t *q_in_data = const_cast(q_in); + + quantize_q_(q_in_data, batch_size); + if (config_.retrieval_type == RetrievalType::LAYER) { + attn_initialize_layer_(batch_size, layer_idx, block_table, + max_block_num, cache_seqlens); + retrieval_kvcache_layer_(q_in_data, 1, local, topk, q_len, + generate_token_idx, batch_size, layer_idx, + cache_seqlens, max_block_num, backend); + calculate_sparsity_layer_(q_in_data, attn_sparsity, batch_size, + max_block_num_origin, block_table_origin, + cache_seqlens_origin, backend); + } else if (config_.retrieval_type == RetrievalType::KVHEAD) { + attn_initialize_kvhead_(batch_size, layer_idx, block_table, + max_block_num, cache_seqlens); + retrieval_kvcache_kvhead_(q_in_data, 1, local, topk, q_len, + generate_token_idx, batch_size, layer_idx, + cache_seqlens, max_block_num, backend); + calculate_sparsity_kvhead_(q_in_data, attn_sparsity, batch_size, + max_block_num_origin, block_table_origin, + cache_seqlens_origin, backend); + } +} + +void KVCache::attn_with_kvcache_one_block_( + int head_dim, int bsz, + ggml_type q_type, // GGML data type of `Q`, only supports fp16 and q8_0 + // [bsz, head_dim] + // Quantization is always on the head_dim dimension (per_token). If + // head_dim % 32 != 0, an error will be raised. The size must be bsz * + // head_dim/32 * qtype_size. + const void *q, + + int past_kv_len, int past_kv_offset, + bool is_full_attn, // true indicates a full 1 mask + // If is_full_attn = false, a bit matrix representing the mask is + // passed. [bsz, past_kv_len] + const uint8_t *attn_mask, + + ggml_type k_type, // GGML data type of `K Cache`, only supports fp16, + // q4_0, q8_0 + int k_quant_type, // 0 for per_token, 1 for per_channel, others raise an + // error + // [seq_len, head_dim] + // If quant_type == 0, head_dim % 32 must be 0. + // If quant_type == 1, seq_len % 32 must be 0. + const void *k_cache, + + // k_anchor_type must be fp16 + int num_k_anchor, // num_k_anchor == 0 indicates no anchor + // [num_k_anchor, head_dim] + const void *k_cache_anchors, + // Each token is associated with the nearest previous position's anchor, + // with the same distance. + const int *k_cache_anchor_pos, + + // v_cache similar to k_cache + ggml_type v_type, int v_quant_type, + // [head_dim, seq_len] + const void *v_cache, int num_v_anchor, const void *v_cache_anchors, + const int *v_cache_anchor_pos, + + // Pre-allocated buffer for intermediate calculations [bsz, + // past_kv_len]. No malloc is performed inside this function. + float *attn_score, + + // Output: [bsz, head_dim], with the same type as q_type + void *output, + // [bsz] + float *lse, + + // Pre-allocated temporary buffer with sufficient size: + // (2 * bsz * past_kv_len + 6 * bsz * head_dim + 2 * past_kv_len * + // head_dim + past_kv_len * head_dim / 32) bytes. + void *draft, + + // Apply rotary embedding online + const int *rotary_angle, const void *rotary_cos, const void *rotary_sin + // rotary_cos=None, + // rotary_sin=None, + // cache_seqlens: Optional[Union[(int, torch.Tensor)]] = None, + // cache_batch_idx: Optional[torch.Tensor] = None, + // rotary_interleaved=True, + + // // Not supported for now + // window_size=(-1, -1), # -1 means infinite context window + // alibi_slopes=None, +) { + assert(head_dim % 32 == 0); + assert(k_quant_type == 0); + assert(v_quant_type == 1); + assert(q_type == GGML_TYPE_F16 || q_type == GGML_TYPE_Q8_0); + if (q_type == GGML_TYPE_F16) { + assert(k_type == GGML_TYPE_F16); + assert(v_type == GGML_TYPE_F16); + + // attn = q * k + q * k_anchor + // TODO: anchor + assert(num_k_anchor == 0); + + if (rotary_angle != nullptr) { + ggml_fp16_t *k_cache_with_rope_fp16 = + (reinterpret_cast(draft) + + sizeof(block_q8_0) * bsz * past_kv_len / QK8_0 + + sizeof(float) * bsz * head_dim); + // dequant k_cache and apply rope + // k_rope(i) = k(i) * cos(i) - k(i+l) * sin(i) + // k_rope(i+l) = k(i+l) * cos(i+l) + k(i) * sin(i) + + // k(i)cos(i) -> k_rope(i) + // k(i)sin(i+l) -> k_rope(i+l) + + // k(i)cos(i) -> k_rope(i) + // -k(i)sin(i-l) -> k_rope(i-l) + + std::vector block_fp32(32); + for (int k = 0; k < past_kv_len; k++) { + int angle = rotary_angle[k]; + for (int l = 0; l < head_dim / 32; l++) { + for (int m = 0; m < 32; m++) { + float x = GGML_FP16_TO_FP32(( + (ggml_fp16_t *)k_cache)[k * head_dim + l * 32 + m]); + float sin_val = GGML_FP16_TO_FP32( + ((ggml_fp16_t *) + rotary_sin)[angle * head_dim + l * 32 + m]); + float cos_val = GGML_FP16_TO_FP32( + ((ggml_fp16_t *) + rotary_cos)[angle * head_dim + l * 32 + m]); + + if (l * 32 + m < head_dim / 2) { + k_cache_with_rope_fp16[k * head_dim + l * 32 + m] = + GGML_FP32_TO_FP16(x * cos_val); + k_cache_with_rope_fp16[k * head_dim + l * 32 + m + + head_dim / 2] = + GGML_FP32_TO_FP16(-x * sin_val); + } else { + k_cache_with_rope_fp16[k * head_dim + l * 32 + m] = + GGML_FP32_TO_FP16( + GGML_FP16_TO_FP32( + k_cache_with_rope_fp16[k * head_dim + + l * 32 + m]) + + x * sin_val); + k_cache_with_rope_fp16[k * head_dim + l * 32 + m - + head_dim / 2] = + GGML_FP32_TO_FP16( + GGML_FP16_TO_FP32( + k_cache_with_rope_fp16[k * head_dim + + l * 32 + m - + head_dim / 2]) - + x * cos_val); + } + } + } + } + + llamafile_sgemm(past_kv_len, bsz, head_dim, + (ggml_fp16_t *)k_cache_with_rope_fp16, head_dim, + (ggml_fp16_t *)q, head_dim, attn_score, past_kv_len, + 0, 1, GGML_TASK_TYPE_COMPUTE, k_type, GGML_TYPE_F16, + GGML_TYPE_F32, GGML_PREC_DEFAULT); + } else { + bool ok = llamafile_sgemm( + past_kv_len, bsz, head_dim, (ggml_fp16_t *)k_cache, head_dim, + (ggml_fp16_t *)q, head_dim, attn_score, past_kv_len, 0, 1, + GGML_TASK_TYPE_COMPUTE, k_type, GGML_TYPE_F16, GGML_TYPE_F32, + GGML_PREC_DEFAULT); + + if (!ok) { + printf("llamafile_sgemm failed\n"); + } + } + // attn = attn * scale + float scale_factor = 1.0 / std::sqrt(float(head_dim)); + ggml_vec_scale_f32(bsz * past_kv_len, attn_score, scale_factor); + + // attn = attn & mask + if (!is_full_attn) { + for (int i = 0; i < bsz; i++) { + for (int j = 0; j < past_kv_len; j++) { + int index = i * past_kv_len + j; + if (!(attn_mask[j / 8] & (1 << (j % 8)))) { + attn_score[index] = + std::numeric_limits::lowest(); + } + } + } + } + + // attn = softmax(attn) + for (int i = 0; i < bsz; i++) { + float sum_exp = 0; + for (int j = 0; j < past_kv_len; j++) { + attn_score[i * past_kv_len + j] = + std::exp(attn_score[i * past_kv_len + j]); + sum_exp += attn_score[i * past_kv_len + j]; + } + for (int j = 0; j < past_kv_len; j++) { + attn_score[i * past_kv_len + j] /= sum_exp; + } + if (lse != nullptr) { + lse[i] = std::log(sum_exp); + } + } + + // output = attn * v + attn * v_anchor + // std::vector sum(bsz * head_dim); + float *sum = reinterpret_cast(reinterpret_cast(draft) + + sizeof(block_q8_0) * bsz * + past_kv_len / QK8_0); + + // float* attn_score_fp16(bsz, past_kv_len) + ggml_fp16_t *attn_score_fp16 = (reinterpret_cast( + reinterpret_cast(draft) + + sizeof(block_q8_0) * bsz * past_kv_len / QK8_0 + + sizeof(float) * bsz * head_dim)); + + for (int i = 0; i < bsz * past_kv_len; i++) { + attn_score_fp16[i] = GGML_FP32_TO_FP16(attn_score[i]); + } + + // TODO: anchor + assert(num_v_anchor == 0); + bool ok = llamafile_sgemm( + head_dim, bsz, past_kv_len, (ggml_fp16_t *)v_cache, past_kv_len, + (ggml_fp16_t *)attn_score_fp16, past_kv_len, sum, head_dim, 0, 1, + GGML_TASK_TYPE_COMPUTE, v_type, GGML_TYPE_F16, GGML_TYPE_F32, + GGML_PREC_DEFAULT); + if (!ok) { + printf("llamafile_sgemm failed\n"); + } + + // copy to output + for (int i = 0; i < bsz; i++) { + for (int j = 0; j < head_dim; j++) { + ((float *)output)[i * head_dim + j] = sum[i * head_dim + j]; + } + } + } else { + assert(k_type == GGML_TYPE_Q4_0 || k_type == GGML_TYPE_Q8_0); + assert(v_type == GGML_TYPE_Q4_0 || v_type == GGML_TYPE_Q8_0); + + // attn = q * k + q * k_anchor + // TODO: anchor + assert(num_k_anchor == 0); + + if (rotary_angle != nullptr) { + ggml_fp16_t *k_cache_with_rope_fp16 = + (reinterpret_cast(draft) + + sizeof(block_q8_0) * bsz * past_kv_len / QK8_0 + + sizeof(float) * bsz * head_dim); + block_q4_0 *k_cache_with_rope_q4 = + (reinterpret_cast(draft) + + sizeof(block_q8_0) * bsz * past_kv_len / QK8_0 + + sizeof(float) * bsz * head_dim) + + sizeof(ggml_fp16_t) * bsz * head_dim; + // dequant k_cache and apply rope + // k_rope(i) = k(i) * cos(i) - k(i+l) * sin(i) + // k_rope(i+l) = k(i+l) * cos(i+l) + k(i) * sin(i) + + // k(i)cos(i) -> k_rope(i) + // k(i)sin(i+l) -> k_rope(i+l) + + // k(i)cos(i) -> k_rope(i) + // -k(i)sin(i-l) -> k_rope(i-l) + + std::vector block_fp32(32); + for (int k = 0; k < past_kv_len; k++) { + int angle = rotary_angle[k]; + for (int l = 0; l < head_dim / 32; l++) { + block_q4_0 block = + ((block_q4_0 *)k_cache)[k * head_dim / 32 + l]; + dequantize_row_q4_0(&block, block_fp32.data(), 32); + for (int m = 0; m < 32; m++) { + float sin_val = GGML_FP16_TO_FP32( + ((ggml_fp16_t *) + rotary_sin)[angle * head_dim + l * 32 + m]); + float cos_val = GGML_FP16_TO_FP32( + ((ggml_fp16_t *) + rotary_cos)[angle * head_dim + l * 32 + m]); + + if (l * 32 + m < head_dim / 2) { + k_cache_with_rope_fp16[k * head_dim + l * 32 + m] = + GGML_FP32_TO_FP16(block_fp32[m] * cos_val); + k_cache_with_rope_fp16[k * head_dim + l * 32 + m + + head_dim / 2] = + GGML_FP32_TO_FP16(-block_fp32[m] * sin_val); + } else { + k_cache_with_rope_fp16[k * head_dim + l * 32 + m] += + GGML_FP32_TO_FP16(block_fp32[m] * sin_val); + k_cache_with_rope_fp16[k * head_dim + l * 32 + m - + head_dim / 2] -= + GGML_FP32_TO_FP16(block_fp32[m] * cos_val); + } + } + } + } + // quantize k_cache_with_rope_fp16 + for (int k = 0; k < past_kv_len; k++) { + for (int l = 0; l < head_dim / 32; l++) { + for (int m = 0; m < 32; m++) { + block_fp32[m] = GGML_FP16_TO_FP32( + k_cache_with_rope_fp16[k * head_dim + l * 32 + m]); + } + quantize_row_q4_0( + block_fp32.data(), + &k_cache_with_rope_q4[k * head_dim / 32 + l], 32); + } + } + + llamafile_sgemm(past_kv_len, bsz, head_dim / 32, + (block_q4_0 *)k_cache_with_rope_q4, head_dim / 32, + (block_q8_0 *)q, head_dim / 32, attn_score, + past_kv_len, 0, 1, GGML_TASK_TYPE_COMPUTE, k_type, + GGML_TYPE_Q8_0, GGML_TYPE_F32, GGML_PREC_DEFAULT); + } else { + llamafile_sgemm(past_kv_len, bsz, head_dim / 32, + (block_q4_0 *)k_cache, head_dim / 32, + (block_q8_0 *)q, head_dim / 32, attn_score, + past_kv_len, 0, 1, GGML_TASK_TYPE_COMPUTE, k_type, + GGML_TYPE_Q8_0, GGML_TYPE_F32, GGML_PREC_DEFAULT); + } + + // attn = attn * scale + float scale_factor = 1.0 / std::sqrt(float(head_dim)); + ggml_vec_scale_f32(bsz * past_kv_len, attn_score, scale_factor); + + // attn = attn & mask + if (!is_full_attn) { + for (int i = 0; i < bsz; i++) { + for (int j = 0; j < past_kv_len; j++) { + int index = i * past_kv_len + j; + if (!(attn_mask[j / 8] & (1 << (j % 8)))) { + attn_score[index] = + std::numeric_limits::lowest(); + } + } + } + } + + // attn = softmax(attn) + for (int i = 0; i < bsz; i++) { + float sum_exp = 0; + for (int j = 0; j < past_kv_len; j++) { + attn_score[i * past_kv_len + j] = + std::exp(attn_score[i * past_kv_len + j]); + sum_exp += attn_score[i * past_kv_len + j]; + } + for (int j = 0; j < past_kv_len; j++) { + attn_score[i * past_kv_len + j] /= sum_exp; + } + if (lse != nullptr) { + lse[i] = std::log(sum_exp); + } + } + + // output = attn * v + attn * v_anchor + // std::vector attn_q8_0(bsz * past_kv_len / QK8_0); + block_q8_0 *attn_q8_0 = reinterpret_cast(draft); + quantize_row_q8_0(attn_score, attn_q8_0, bsz * past_kv_len); + // std::vector sum(bsz * head_dim); + float *sum = reinterpret_cast(reinterpret_cast(draft) + + sizeof(block_q8_0) * bsz * + past_kv_len / QK8_0); + // TODO: anchor + assert(num_v_anchor == 0); + llamafile_sgemm(head_dim, bsz, past_kv_len / 32, (block_q4_0 *)v_cache, + past_kv_len / 32, attn_q8_0, past_kv_len / 32, sum, + head_dim, 0, 1, GGML_TASK_TYPE_COMPUTE, v_type, + GGML_TYPE_Q8_0, GGML_TYPE_F32, GGML_PREC_DEFAULT); + + quantize_row_q8_0(sum, (block_q8_0 *)output, bsz * head_dim); + } +} diff --git a/archive/kt-sft/csrc/ktransformers_ext/operators/kvcache/kvcache_load_dump.cpp b/archive/kt-sft/csrc/ktransformers_ext/operators/kvcache/kvcache_load_dump.cpp new file mode 100644 index 0000000..4de217f --- /dev/null +++ b/archive/kt-sft/csrc/ktransformers_ext/operators/kvcache/kvcache_load_dump.cpp @@ -0,0 +1,126 @@ +/** + * @Description : + * @Author : Jianwei Dong + * @Date : 2024-08-26 22:47:06 + * @Version : 1.0.0 + * @LastEditors : Jianwei Dong + * @LastEditTime : 2024-08-26 22:47:06 + * @Copyright (c) 2024 by KVCache.AI, All Rights Reserved. + **/ + +#include "kvcache.h" + +#include + +void KVCache::load_kvcache(std::string tensor_file_path, Backend *backend) { + // Timer start + auto start = std::chrono::high_resolution_clock::now(); + std::ifstream ifs_tensor(tensor_file_path, std::ios::binary); + if (!ifs_tensor) { + throw std::runtime_error("Failed to open tensor file"); + } + ifs_tensor.read(reinterpret_cast(&cache_total_len_), + sizeof(cache_total_len_)); + int past_block_num = + (cache_total_len_ + config_.block_len - 1) / config_.block_len; + printf("cache_total_len: %d, past_block_num: %d\n", cache_total_len_, + past_block_num); + for (int i = 0; i < config_.layer_num; ++i) { + past_block_num_[i] = past_block_num; + } + ifs_tensor.read(reinterpret_cast(anchor_.data()), + anchor_.size() * sizeof(ggml_fp16_t)); + for (int i = 0; i < config_.layer_num; ++i) { + for (int j = 0; j < config_.kv_head_num; ++j) { + for (int k = 0; k < past_block_num_[i]; ++k) { + if (config_.kv_type == GGML_TYPE_F16) { + ifs_tensor.read( + reinterpret_cast(k_cache_fp16_[i][j][k].data()), + k_cache_fp16_[i][j][k].size() * sizeof(ggml_fp16_t)); + ifs_tensor.read( + reinterpret_cast(v_cache_fp16_[i][j][k].data()), + v_cache_fp16_[i][j][k].size() * sizeof(ggml_fp16_t)); + } else if (config_.kv_type == GGML_TYPE_Q4_0) { + ifs_tensor.read( + reinterpret_cast(k_cache_q4[i][j][k].data()), + k_cache_q4[i][j][k].size() * sizeof(block_q4_0)); + ifs_tensor.read( + reinterpret_cast(v_cache_q4[i][j][k].data()), + v_cache_q4[i][j][k].size() * sizeof(block_q4_0)); + } + } + } + for (int k = 0; k < past_block_num_[i]; ++k) { + for (int l = 0; l < config_.block_len; l++) { + ifs_tensor.read( + reinterpret_cast(importance_[i][k][l].data()), + importance_[i][k][l].size() * sizeof(ggml_fp16_t)); + } + } + } + ifs_tensor.close(); + // Timer end + auto end = std::chrono::high_resolution_clock::now(); + std::chrono::duration diff = end - start; + printf("time of load: %f s\n", diff.count()); +} +void KVCache::dump_kvcache(int *block_table, int cache_total_len, + std::string tensor_file_path, Backend *backend) { + // Timer start + auto start = std::chrono::high_resolution_clock::now(); + std::ofstream ofs(tensor_file_path, std::ios::binary); + printf("dump_kvcache: %s\n", tensor_file_path.c_str()); + if (!ofs.is_open()) { + std::cerr << "Cannot open file " << tensor_file_path << std::endl; + return; + } + ofs.write(reinterpret_cast(&cache_total_len), + sizeof(cache_total_len)); + int past_block_num = + (cache_total_len + config_.block_len - 1) / config_.block_len; + printf("cache_total_len: %d, past_block_num: %d\n", cache_total_len, + past_block_num); + ofs.write(reinterpret_cast(anchor_.data()), + anchor_.size() * sizeof(ggml_fp16_t)); + for (int i = 0; i < config_.layer_num; ++i) { + for (int j = 0; j < config_.kv_head_num; ++j) { + for (int k = 0; k < past_block_num; ++k) { + int block_idx = block_table[k]; + if (config_.kv_type == GGML_TYPE_F16) { + ofs.write(reinterpret_cast( + k_cache_fp16_[i][j][block_idx].data()), + k_cache_fp16_[i][j][block_idx].size() * + sizeof(ggml_fp16_t)); + ofs.write(reinterpret_cast( + v_cache_fp16_[i][j][block_idx].data()), + v_cache_fp16_[i][j][block_idx].size() * + sizeof(ggml_fp16_t)); + + } else if (config_.kv_type == GGML_TYPE_Q4_0) { + ofs.write(reinterpret_cast( + k_cache_q4[i][j][block_idx].data()), + k_cache_q4[i][j][block_idx].size() * + sizeof(block_q4_0)); + ofs.write(reinterpret_cast( + v_cache_q4[i][j][block_idx].data()), + v_cache_q4[i][j][block_idx].size() * + sizeof(block_q4_0)); + } + } + } + for (int k = 0; k < past_block_num; ++k) { + int block_idx = block_table[k]; + for (int l = 0; l < config_.block_len; l++) { + ofs.write(reinterpret_cast( + importance_[i][block_idx][l].data()), + importance_[i][block_idx][l].size() * + sizeof(ggml_fp16_t)); + } + } + } + ofs.close(); + // Timer end + auto end = std::chrono::high_resolution_clock::now(); + std::chrono::duration diff = end - start; + printf("time of dump: %f s\n", diff.count()); +} \ No newline at end of file diff --git a/archive/kt-sft/csrc/ktransformers_ext/operators/kvcache/kvcache_read_write.cpp b/archive/kt-sft/csrc/ktransformers_ext/operators/kvcache/kvcache_read_write.cpp new file mode 100644 index 0000000..0104905 --- /dev/null +++ b/archive/kt-sft/csrc/ktransformers_ext/operators/kvcache/kvcache_read_write.cpp @@ -0,0 +1,1021 @@ +/** + * @Description : + * @Author : Jianwei Dong + * @Date : 2024-08-26 22:47:06 + * @Version : 1.0.0 + * @LastEditors : Jianwei Dong + * @LastEditTime : 2024-08-26 22:47:06 + * @Copyright (c) 2024 by KVCache.AI, All Rights Reserved. + **/ + +#include "kvcache.h" + +#include + +void KVCache::get_anchor_one_block(ggml_fp16_t *anchor, int layer_id, + int block_idx, Backend *backend) { + // Timer start + auto start = std::chrono::high_resolution_clock::now(); + + layer_id_ = layer_id; + block_idx = block_idx; + seq_len_ = config_.block_len; + anchor_data_ = const_cast(anchor); + + // Timer end + auto end = std::chrono::high_resolution_clock::now(); + std::chrono::duration duration = end - start; + printf("layer %d block %d time of reading anchor: %f s\n", layer_id, + block_idx, duration.count()); +} + +void KVCache::update_anchor_one_block(const ggml_fp16_t *anchor, int layer_id, + int block_idx, Backend *backend) { + // Timer start + auto start = std::chrono::high_resolution_clock::now(); + + layer_id_ = layer_id; + block_idx = block_idx; + seq_len_ = config_.block_len; + anchor_data_ = const_cast(anchor); + + // Each task updates the anchor of a certain position + // backend->do_work_stealing_job(config_.anchor_num, [&](int task_id) { + // int k = task_id % config_.anchor_num; + // int head_id = task_id / config_.anchor_num; + // memcpy(anchor_[layer_id_][head_id][block_idx].data() + + // k * config_.head_dim, + // anchor_data_ + k * config_.head_dim, + // sizeof(uint16_t) * config_.head_dim); + // }); + + // Timer end + auto end = std::chrono::high_resolution_clock::now(); + std::chrono::duration duration = end - start; + printf("layer %d block %d time of writting anchor: %f s\n", layer_id, + block_idx, duration.count()); +} + +void KVCache::update_importance_one_block(const ggml_fp16_t *importance, + int layer_id, int block_idx, + Backend *backend) { + // Timer start + auto start = std::chrono::high_resolution_clock::now(); + + layer_id_ = layer_id; + block_idx = block_idx; + seq_len_ = config_.block_len; + importance_data_ = const_cast(importance); + + // Each task updates the importance of a certain position + backend->do_work_stealing_job( + config_.block_len, nullptr, + [&](int task_id) { + int k = task_id; + memcpy(importance_[layer_id_][block_idx].data() + k, + importance_data_ + k, sizeof(uint16_t)); + }, + nullptr); + + // Timer end + auto end = std::chrono::high_resolution_clock::now(); + std::chrono::duration duration = end - start; + printf("layer %d block %d time of writting importance: %f s\n", layer_id, + block_idx, duration.count()); +} + +void KVCache::get_importance_one_block(ggml_fp16_t *importance, int layer_id, + int block_idx, Backend *backend) { + // Timer start + auto start = std::chrono::high_resolution_clock::now(); + + layer_id_ = layer_id; + block_idx = block_idx; + seq_len_ = config_.block_len; + importance_data_ = const_cast(importance); + + // Each task updates the importance of a certain position + backend->do_work_stealing_job( + config_.block_len, nullptr, + [&](int task_id) { + int k = task_id; + memcpy(importance_data_ + k, + importance_[layer_id_][block_idx].data() + k, + sizeof(uint16_t)); + }, + nullptr); + + // Timer end + auto end = std::chrono::high_resolution_clock::now(); + std::chrono::duration duration = end - start; + printf("layer %d block %d time of reading importance: %f s\n", layer_id, + block_idx, duration.count()); +} + +void KVCache::update_kvcache_one_block_fp16(const ggml_fp16_t *k_in, + const ggml_fp16_t *v_in, + int layer_id, int block_idx, + Backend *backend) { + // Timer start + auto start = std::chrono::high_resolution_clock::now(); + + layer_id_ = layer_id; + block_idx = block_idx; + seq_len_ = config_.block_len; + k_data_ = const_cast(k_in); + v_data_ = const_cast(v_in); + + int new_block_num = std::max((int)past_block_num_[layer_id], block_idx + 1); + + importance_[layer_id_].resize(new_block_num); + + for (int i = 0; i < config_.kv_head_num; i++) { + k_cache_q4[layer_id][i].resize(new_block_num); + v_cache_q4[layer_id][i].resize(new_block_num); + // anchor_[layer_id][i].resize(new_block_num); + } + + for (int i = 0; i < new_block_num; i++) { + importance_[layer_id][i].resize(config_.block_len); + } + + // Each task updates the k cache or v cache of a certain header + backend->do_work_stealing_job( + config_.kv_head_num * 2, nullptr, + [&](int task_id) { + std::vector block_fp32(32); + int head_id = task_id / 2; + if (task_id & 1) { + // fill k_cache_ + k_cache_q4[layer_id_][head_id][block_idx].resize( + config_.block_len * config_.head_dim / 32); + for (int k = 0; k < config_.block_len; k++) { + for (int l = 0; l < config_.head_dim / 32; l++) { + block_q4_0 block; + for (int m = 0; m < 32; m++) { + + block_fp32[m] = GGML_FP16_TO_FP32( + k_data_[((0 * config_.kv_head_num + head_id) * + seq_len_ + + 0 * config_.block_len + k) * + config_.head_dim + + l * 32 + m]); + } + quantize_row_q4_0(block_fp32.data(), &block, 32); + k_cache_q4[layer_id_][head_id][block_idx] + [k * config_.head_dim / 32 + l] = block; + } + } + } else { + // fill v_cache_ + v_cache_q4[layer_id_][head_id][block_idx].resize( + config_.head_dim * config_.block_len / 32); + for (int k = 0; k < config_.block_len / 32; k++) { + for (int l = 0; l < config_.head_dim; l++) { + block_q4_0 block; + for (int m = 0; m < 32; m++) { + + block_fp32[m] = GGML_FP16_TO_FP32( + v_data_[((0 * config_.kv_head_num + head_id) * + seq_len_ + + 0 * config_.block_len + k * 32 + m) * + config_.head_dim + + l]); + } + quantize_row_q4_0(block_fp32.data(), &block, 32); + v_cache_q4[layer_id_][head_id][block_idx] + [l * config_.block_len / 32 + k] = block; + } + } + } + }, + nullptr); + past_block_num_[layer_id] = new_block_num; + + // Timer end + auto end = std::chrono::high_resolution_clock::now(); + std::chrono::duration duration = end - start; + printf("layer %d block %d time of writting KV Cache: %f s\n", layer_id, + block_idx, duration.count()); + // printf("get_one_block_fp16 duration: %ld\n", duration); +} + +void KVCache::get_kvcache_one_block_fp16(ggml_fp16_t *k_in, ggml_fp16_t *v_in, + int layer_id, int block_idx, + Backend *backend) { + // Timer start + auto start = std::chrono::high_resolution_clock::now(); + + layer_id_ = layer_id; + seq_len_ = config_.block_len; + k_data_ = reinterpret_cast(k_in); + v_data_ = reinterpret_cast(v_in); + + // printf("layer_id: %d, block_idx: %d\n", layer_id, block_idx); + // Each task gets the k cache or v cache of a certain header + backend->do_work_stealing_job( + config_.kv_head_num * 2, nullptr, + [&](int task_id) { + std::vector block_fp32(32); + int head_id = task_id / 2; + if (task_id & 1) { + // get k_cache_ + for (int k = 0; k < config_.block_len; k++) { + for (int l = 0; l < config_.head_dim / 32; l++) { + block_q4_0 block = + k_cache_q4[layer_id_][head_id][block_idx] + [k * config_.head_dim / 32 + l]; + dequantize_row_q4_0(&block, block_fp32.data(), 32); + for (int m = 0; m < 32; m++) { + + k_data_[((0 * config_.kv_head_num + head_id) * + seq_len_ + + 0 * config_.block_len + k) * + config_.head_dim + + l * 32 + m] = + GGML_FP32_TO_FP16(block_fp32[m]); + } + } + } + } else { + // get v_cache_ + for (int k = 0; k < config_.block_len / 32; k++) { + for (int l = 0; l < config_.head_dim; l++) { + block_q4_0 block = + v_cache_q4[layer_id_][head_id][block_idx] + [l * config_.block_len / 32 + k]; + dequantize_row_q4_0(&block, block_fp32.data(), 32); + for (int m = 0; m < 32; m++) { + + v_data_[((0 * config_.kv_head_num + head_id) * + seq_len_ + + 0 * config_.block_len + k * 32 + m) * + config_.head_dim + + l] = GGML_FP32_TO_FP16(block_fp32[m]); + } + } + } + } + }, + nullptr); + + // Timer end + auto end = std::chrono::high_resolution_clock::now(); + std::chrono::duration duration = end - start; + printf("layer %d block %d time of reading KV Cache: %f s\n", layer_id, + block_idx, duration.count()); + // printf("get_one_block_fp16 duration: %ld\n", duration); +} + +// k_in: (batch_size, seq_len, head_num, head_dim) +// v_in: (batch_size, seq_len, head_num, head_dim) +void KVCache::get_and_update_kvcache_fp16(ggml_fp16_t *k_in, ggml_fp16_t *v_in, + int layer_id, int *block_table, + int batch_size, int max_block_num, + int *cache_seqlens, int q_len, + Backend *backend) { + // Timer start + auto start = std::chrono::high_resolution_clock::now(); + + layer_id_ = layer_id; + k_data_ = const_cast(k_in); + v_data_ = const_cast(v_in); + + // Each task updates the k cache and v cache of a certain header + backend->do_work_stealing_job( + config_.kv_head_num * max_block_num * batch_size, nullptr, + [&](int task_id) { + // printf("block_idx: %d, task_id: %d\n", block_idx, task_id); + std::vector block_fp32(32); + int batch_id = task_id / (config_.kv_head_num * max_block_num); + int block_id = (task_id / config_.kv_head_num) % max_block_num; + int head_id = task_id % config_.kv_head_num; + int block_idx = block_table[batch_id * max_block_num + block_id]; + int seq_len = cache_seqlens[batch_id]; + int block_l = block_id * config_.block_len; + int block_r = block_id * config_.block_len + config_.block_len; + + if (block_l < seq_len) { + if (config_.kv_type == ggml_type::GGML_TYPE_F16) { + for (int k = 0; k < config_.block_len; k++) { + if (block_id * config_.block_len + k >= seq_len) + break; + for (int l = 0; l < config_.head_dim; l++) { + k_data_ + [batch_id * + (max_block_num * config_.block_len * + config_.kv_head_num * config_.head_dim) + + block_id * + (config_.block_len * config_.kv_head_num * + config_.head_dim) + + k * (config_.kv_head_num * config_.head_dim) + + head_id * config_.head_dim + l] = + k_cache_fp16_[layer_id_][head_id][block_idx] + [k * config_.head_dim + l]; + v_data_ + [batch_id * + (max_block_num * config_.block_len * + config_.kv_head_num * config_.head_dim) + + block_id * + (config_.block_len * config_.kv_head_num * + config_.head_dim) + + k * (config_.kv_head_num * config_.head_dim) + + head_id * config_.head_dim + l] = + v_cache_fp16_[layer_id_][head_id][block_idx] + [l * config_.block_len + k]; + } + } + } else if (config_.kv_type == ggml_type::GGML_TYPE_Q4_0) { + // get k_cache_ + for (int k = 0; k < config_.block_len; k++) { + if (block_id * config_.block_len + k >= seq_len) + break; + for (int l = 0; l < config_.head_dim / 32; l++) { + block_q4_0 block = + k_cache_q4[layer_id_][head_id][block_idx] + [k * config_.head_dim / 32 + l]; + dequantize_row_q4_0(&block, block_fp32.data(), 32); + for (int m = 0; m < 32; m++) { + + k_data_[batch_id * + (max_block_num * config_.block_len * + config_.kv_head_num * + config_.head_dim) + + block_id * (config_.block_len * + config_.kv_head_num * + config_.head_dim) + + k * (config_.kv_head_num * + config_.head_dim) + + head_id * config_.head_dim + l * 32 + + m] = GGML_FP32_TO_FP16(block_fp32[m]); + } + } + } + // get v_cache_ + for (int k = 0; k < config_.block_len / 32; k++) { + for (int l = 0; l < config_.head_dim; l++) { + block_q4_0 block = + v_cache_q4[layer_id_][head_id][block_idx] + [l * config_.block_len / 32 + k]; + dequantize_row_q4_0(&block, block_fp32.data(), 32); + for (int m = 0; m < 32; m++) { + + if (block_id * config_.block_len + k * 32 + m >= + seq_len) + break; + v_data_[batch_id * + (max_block_num * config_.block_len * + config_.kv_head_num * + config_.head_dim) + + block_id * (config_.block_len * + config_.kv_head_num * + config_.head_dim) + + (k * 32 + m) * config_.kv_head_num * + config_.head_dim + + head_id * config_.head_dim + l] = + GGML_FP32_TO_FP16(block_fp32[m]); + } + } + } + } else if (config_.kv_type == ggml_type::GGML_TYPE_Q8_0) { + // get k_cache_ + for (int k = 0; k < config_.block_len; k++) { + if (block_id * config_.block_len + k >= seq_len) + break; + for (int l = 0; l < config_.head_dim / 32; l++) { + block_q8_0 block = + k_cache_q8[layer_id_][head_id][block_idx] + [k * config_.head_dim / 32 + l]; + dequantize_row_q8_0(&block, block_fp32.data(), 32); + for (int m = 0; m < 32; m++) { + + k_data_[batch_id * + (max_block_num * config_.block_len * + config_.kv_head_num * + config_.head_dim) + + block_id * (config_.block_len * + config_.kv_head_num * + config_.head_dim) + + k * (config_.kv_head_num * + config_.head_dim) + + head_id * config_.head_dim + l * 32 + + m] = GGML_FP32_TO_FP16(block_fp32[m]); + } + } + } + // get v_cache_ + for (int k = 0; k < config_.block_len / 32; k++) { + for (int l = 0; l < config_.head_dim; l++) { + block_q8_0 block = + v_cache_q8[layer_id_][head_id][block_idx] + [l * config_.block_len / 32 + k]; + dequantize_row_q8_0(&block, block_fp32.data(), 32); + for (int m = 0; m < 32; m++) { + + if (block_id * config_.block_len + k * 32 + m >= + seq_len) + break; + v_data_[batch_id * + (max_block_num * config_.block_len * + config_.kv_head_num * + config_.head_dim) + + block_id * (config_.block_len * + config_.kv_head_num * + config_.head_dim) + + (k * 32 + m) * config_.kv_head_num * + config_.head_dim + + head_id * config_.head_dim + l] = + GGML_FP32_TO_FP16(block_fp32[m]); + } + } + } + } + } + if (block_r > seq_len && block_l < seq_len + q_len) { + if (config_.kv_type == ggml_type::GGML_TYPE_F16) { + for (int k = 0; k < config_.block_len; k++) { + if (block_id * config_.block_len + k >= + seq_len + q_len || + block_id * config_.block_len + k < seq_len) + continue; + for (int l = 0; l < config_.head_dim; l++) { + k_cache_fp16_[layer_id_][head_id][block_idx] + [k * config_.head_dim + l] = k_data_ + [batch_id * (max_block_num * + config_.block_len * + config_.kv_head_num * + config_.head_dim) + + block_id * (config_.block_len * + config_.kv_head_num * + config_.head_dim) + + k * (config_.kv_head_num * + config_.head_dim) + + head_id * config_.head_dim + l]; + v_cache_fp16_[layer_id_][head_id][block_idx] + [l * config_.block_len + k] = v_data_ + [batch_id * (max_block_num * + config_.block_len * + config_.kv_head_num * + config_.head_dim) + + block_id * (config_.block_len * + config_.kv_head_num * + config_.head_dim) + + k * (config_.kv_head_num * + config_.head_dim) + + head_id * config_.head_dim + l]; + } + } + } else if (config_.kv_type == ggml_type::GGML_TYPE_Q4_0) { + // fill k_cache_ + for (int k = 0; k < config_.block_len; k++) { + if (block_id * config_.block_len + k >= + seq_len + q_len || + block_id * config_.block_len + k < seq_len) + continue; + for (int l = 0; l < config_.head_dim / 32; l++) { + block_q4_0 block; + for (int m = 0; m < 32; m++) { + + block_fp32[m] = GGML_FP16_TO_FP32( + k_data_[batch_id * (max_block_num * + config_.block_len * + config_.kv_head_num * + config_.head_dim) + + block_id * (config_.block_len * + config_.kv_head_num * + config_.head_dim) + + k * (config_.kv_head_num * + config_.head_dim) + + head_id * config_.head_dim + + l * 32 + m]); + } + quantize_row_q4_0(block_fp32.data(), &block, 32); + k_cache_q4[layer_id_][head_id][block_idx] + [k * config_.head_dim / 32 + l] = block; + } + } + + // fill v_cache_ + for (int k = 0; k < config_.block_len / 32; k++) { + for (int l = 0; l < config_.head_dim; l++) { + block_q4_0 block; + for (int m = 0; m < 32; m++) { + + if (block_id * config_.block_len + k * 32 + m >= + seq_len + q_len) { + block_fp32[m] = 0; + continue; + } + block_fp32[m] = GGML_FP16_TO_FP32( + v_data_[batch_id * (max_block_num * + config_.block_len * + config_.kv_head_num * + config_.head_dim) + + block_id * (config_.block_len * + config_.kv_head_num * + config_.head_dim) + + (k * 32 + m) * config_.kv_head_num * + config_.head_dim + + head_id * config_.head_dim + l]); + } + quantize_row_q4_0(block_fp32.data(), &block, 32); + v_cache_q4[layer_id_][head_id][block_idx] + [l * config_.block_len / 32 + k] = block; + } + } + } else if (config_.kv_type == ggml_type::GGML_TYPE_Q8_0) { + // fill k_cache_ + for (int k = 0; k < config_.block_len; k++) { + if (block_id * config_.block_len + k >= + seq_len + q_len || + block_id * config_.block_len + k < seq_len) + continue; + for (int l = 0; l < config_.head_dim / 32; l++) { + block_q8_0 block; + for (int m = 0; m < 32; m++) { + + block_fp32[m] = GGML_FP16_TO_FP32( + k_data_[batch_id * (max_block_num * + config_.block_len * + config_.kv_head_num * + config_.head_dim) + + block_id * (config_.block_len * + config_.kv_head_num * + config_.head_dim) + + k * (config_.kv_head_num * + config_.head_dim) + + head_id * config_.head_dim + + l * 32 + m]); + } + quantize_row_q8_0(block_fp32.data(), &block, 32); + k_cache_q8[layer_id_][head_id][block_idx] + [k * config_.head_dim / 32 + l] = block; + } + } + + // fill v_cache_ + for (int k = 0; k < config_.block_len / 32; k++) { + for (int l = 0; l < config_.head_dim; l++) { + block_q8_0 block; + for (int m = 0; m < 32; m++) { + + if (block_id * config_.block_len + k * 32 + m >= + seq_len + q_len) { + block_fp32[m] = 0; + continue; + } + block_fp32[m] = GGML_FP16_TO_FP32( + v_data_[batch_id * (max_block_num * + config_.block_len * + config_.kv_head_num * + config_.head_dim) + + block_id * (config_.block_len * + config_.kv_head_num * + config_.head_dim) + + (k * 32 + m) * config_.kv_head_num * + config_.head_dim + + head_id * config_.head_dim + l]); + } + quantize_row_q8_0(block_fp32.data(), &block, 32); + v_cache_q8[layer_id_][head_id][block_idx] + [l * config_.block_len / 32 + k] = block; + } + } + } + } + }, + nullptr); + + // Timer end + auto end = std::chrono::high_resolution_clock::now(); + std::chrono::duration duration = end - start; + + // printf("layer %d time of reading and updating KV Cache: %f s\n", + // layer_id, + // duration.count()); +} + +void KVCache::update_importance(const ggml_fp16_t *importance, int layer_id, + int *block_table, int batch_size, + int max_block_num, int *offset, int width, + Backend *backend) { + // Timer start + auto start = std::chrono::high_resolution_clock::now(); + + layer_id_ = layer_id; + importance_data_ = const_cast(importance); + + // Each task updates the importance of a certain position + backend->do_work_stealing_job( + max_block_num * batch_size, nullptr, + [&](int task_id) { + int block_id = task_id % max_block_num; + int batch_id = task_id / max_block_num; + int block_idx = block_table[batch_id * max_block_num + block_id]; + if (block_id > (offset[batch_id] + width) / config_.block_len) { + return; + } + for (int k = 0; k < config_.block_len; k++) { + for (int head_id = 0; head_id < config_.q_head_num; head_id++) { + importance_[layer_id_][block_idx][k][head_id] = + GGML_FP32_TO_FP16( + GGML_FP16_TO_FP32( + importance_data_[batch_id * max_block_num * + config_.block_len * + config_.q_head_num + + (block_id * config_.block_len + + k) * + config_.q_head_num + + head_id]) + + GGML_FP16_TO_FP32( + importance_[layer_id_][block_idx][k][head_id])); + } + } + }, + nullptr); + + // Timer end + auto end = std::chrono::high_resolution_clock::now(); + std::chrono::duration duration = end - start; + + // printf("layer %d time of updating importance: %f s\n", layer_id, + // duration.count()); +} + +void KVCache::get_kvcache_fp16(ggml_fp16_t *k_in, ggml_fp16_t *v_in, + int layer_id, int *block_table, int batch_size, + int max_block_num, int *cache_seqlens, + Backend *backend) { + // Timer start + auto start = std::chrono::high_resolution_clock::now(); + + layer_id_ = layer_id; + k_data_ = const_cast(k_in); + v_data_ = const_cast(v_in); + + // Each task updates the k cache and v cache of a certain header + backend->do_work_stealing_job( + config_.kv_head_num * max_block_num * batch_size, nullptr, + [&](int task_id) { + // printf("block_idx: %d, task_id: %d\n", block_idx, task_id); + std::vector block_fp32(32); + int batch_id = task_id / (config_.kv_head_num * max_block_num); + int block_id = (task_id / config_.kv_head_num) % max_block_num; + int head_id = task_id % config_.kv_head_num; + int block_idx = block_table[batch_id * max_block_num + block_id]; + int seq_len = cache_seqlens[batch_id]; + int block_l = block_id * config_.block_len; + int block_r = block_id * config_.block_len + config_.block_len; + + if (block_l < seq_len) { + if (config_.kv_type == ggml_type::GGML_TYPE_F16) { + for (int k = 0; k < config_.block_len; k++) { + if (block_id * config_.block_len + k >= seq_len) + break; + for (int l = 0; l < config_.head_dim; l++) { + k_data_ + [batch_id * + (max_block_num * config_.block_len * + config_.kv_head_num * config_.head_dim) + + block_id * + (config_.block_len * config_.kv_head_num * + config_.head_dim) + + k * (config_.kv_head_num * config_.head_dim) + + head_id * config_.head_dim + l] = + k_cache_fp16_[layer_id_][head_id][block_idx] + [k * config_.head_dim + l]; + v_data_ + [batch_id * + (max_block_num * config_.block_len * + config_.kv_head_num * config_.head_dim) + + block_id * + (config_.block_len * config_.kv_head_num * + config_.head_dim) + + k * (config_.kv_head_num * config_.head_dim) + + head_id * config_.head_dim + l] = + v_cache_fp16_[layer_id_][head_id][block_idx] + [l * config_.block_len + k]; + } + } + } else if (config_.kv_type == ggml_type::GGML_TYPE_Q4_0) { + // get k_cache_ + for (int k = 0; k < config_.block_len; k++) { + if (block_id * config_.block_len + k >= seq_len) + break; + for (int l = 0; l < config_.head_dim / 32; l++) { + block_q4_0 block = + k_cache_q4[layer_id_][head_id][block_idx] + [k * config_.head_dim / 32 + l]; + dequantize_row_q4_0(&block, block_fp32.data(), 32); + for (int m = 0; m < 32; m++) { + + k_data_[batch_id * + (max_block_num * config_.block_len * + config_.kv_head_num * + config_.head_dim) + + block_id * (config_.block_len * + config_.kv_head_num * + config_.head_dim) + + k * (config_.kv_head_num * + config_.head_dim) + + head_id * config_.head_dim + l * 32 + + m] = GGML_FP32_TO_FP16(block_fp32[m]); + } + } + } + // get v_cache_ + for (int k = 0; k < config_.block_len / 32; k++) { + for (int l = 0; l < config_.head_dim; l++) { + block_q4_0 block = + v_cache_q4[layer_id_][head_id][block_idx] + [l * config_.block_len / 32 + k]; + dequantize_row_q4_0(&block, block_fp32.data(), 32); + for (int m = 0; m < 32; m++) { + + if (block_id * config_.block_len + k * 32 + m >= + seq_len) + break; + v_data_[batch_id * + (max_block_num * config_.block_len * + config_.kv_head_num * + config_.head_dim) + + block_id * (config_.block_len * + config_.kv_head_num * + config_.head_dim) + + (k * 32 + m) * config_.kv_head_num * + config_.head_dim + + head_id * config_.head_dim + l] = + GGML_FP32_TO_FP16(block_fp32[m]); + } + } + } + } else if (config_.kv_type == ggml_type::GGML_TYPE_Q8_0) { + // get k_cache_ + for (int k = 0; k < config_.block_len; k++) { + if (block_id * config_.block_len + k >= seq_len) + break; + for (int l = 0; l < config_.head_dim / 32; l++) { + block_q8_0 block = + k_cache_q8[layer_id_][head_id][block_idx] + [k * config_.head_dim / 32 + l]; + dequantize_row_q8_0(&block, block_fp32.data(), 32); + for (int m = 0; m < 32; m++) { + + k_data_[batch_id * + (max_block_num * config_.block_len * + config_.kv_head_num * + config_.head_dim) + + block_id * (config_.block_len * + config_.kv_head_num * + config_.head_dim) + + k * (config_.kv_head_num * + config_.head_dim) + + head_id * config_.head_dim + l * 32 + + m] = GGML_FP32_TO_FP16(block_fp32[m]); + } + } + } + // get v_cache_ + for (int k = 0; k < config_.block_len / 32; k++) { + for (int l = 0; l < config_.head_dim; l++) { + block_q8_0 block = + v_cache_q8[layer_id_][head_id][block_idx] + [l * config_.block_len / 32 + k]; + dequantize_row_q8_0(&block, block_fp32.data(), 32); + for (int m = 0; m < 32; m++) { + + if (block_id * config_.block_len + k * 32 + m >= + seq_len) + break; + v_data_[batch_id * + (max_block_num * config_.block_len * + config_.kv_head_num * + config_.head_dim) + + block_id * (config_.block_len * + config_.kv_head_num * + config_.head_dim) + + (k * 32 + m) * config_.kv_head_num * + config_.head_dim + + head_id * config_.head_dim + l] = + GGML_FP32_TO_FP16(block_fp32[m]); + } + } + } + } + } + }, + nullptr); + + // Timer end + auto end = std::chrono::high_resolution_clock::now(); + std::chrono::duration duration = end - start; +} + +void KVCache::update_kvcache_fp16(const ggml_fp16_t *k_in, + const ggml_fp16_t *v_in, int layer_id, + int *block_table, int batch_size, + int max_block_num, int *cache_seqlens, + int q_len, Backend *backend) { + // Timer start + auto start = std::chrono::high_resolution_clock::now(); + + layer_id_ = layer_id; + k_data_ = const_cast(k_in); + v_data_ = const_cast(v_in); + // Each task updates the k cache and v cache of a certain header + backend->do_work_stealing_job( + batch_size * config_.kv_head_num * q_len, nullptr, + [&](int task_id) { + int batch_id = task_id / (config_.kv_head_num * q_len); + int head_id = task_id / q_len % config_.kv_head_num; + int seq_len = cache_seqlens[batch_id] + task_id % q_len; + int q_offset = task_id % q_len; + + int block_id = seq_len / config_.block_len; + int block_idx = block_table[batch_id * max_block_num + block_id]; + int pos_in_block = seq_len % config_.block_len; + + if (config_.kv_type == ggml_type::GGML_TYPE_F16) { + for (int l = 0; l < config_.head_dim; l++) { + k_cache_fp16_[layer_id_][head_id][block_idx] + [pos_in_block * config_.head_dim + l] = + k_data_[batch_id * + (q_len * config_.kv_head_num * + config_.head_dim) + + q_offset * config_.kv_head_num * + config_.head_dim + + head_id * config_.head_dim + l]; + v_cache_fp16_[layer_id_][head_id][block_idx] + [l * config_.block_len + pos_in_block] = + v_data_[batch_id * + (q_len * config_.kv_head_num * + config_.head_dim) + + q_offset * config_.kv_head_num * + config_.head_dim + + head_id * config_.head_dim + l]; + } + } else if (config_.kv_type == ggml_type::GGML_TYPE_Q4_0) { + std::vector block_fp32(32); + // fill k_cache_ + for (int l = 0; l < config_.head_dim / 32; l++) { + block_q4_0 block; + for (int m = 0; m < 32; m++) { + + block_fp32[m] = GGML_FP16_TO_FP32( + k_data_[batch_id * (q_len * config_.kv_head_num * + config_.head_dim) + + head_id * config_.head_dim + l * 32 + m]); + } + quantize_row_q4_0(block_fp32.data(), &block, 32); + + k_cache_q4[layer_id_][head_id][block_idx] + [pos_in_block * config_.head_dim / 32 + l] = + block; + } + + // fill v_cache_ + for (int l = 0; l < config_.head_dim; l++) { + block_q4_0 block = v_cache_q4[layer_id_][head_id][block_idx] + [l * config_.block_len / 32 + + pos_in_block / 32]; + dequantize_row_q4_0(&block, block_fp32.data(), 32); + block_fp32[pos_in_block % 32] = GGML_FP16_TO_FP32( + v_data_[batch_id * (q_len * config_.kv_head_num * + config_.head_dim) + + head_id * config_.head_dim + l]); + quantize_row_q4_0(block_fp32.data(), &block, 32); + v_cache_q4[layer_id_][head_id][block_idx] + [l * config_.block_len / 32 + pos_in_block / 32] = + block; + } + } else if (config_.kv_type == ggml_type::GGML_TYPE_Q8_0) { + std::vector block_fp32(32); + // fill k_cache_ + for (int l = 0; l < config_.head_dim / 32; l++) { + block_q8_0 block; + for (int m = 0; m < 32; m++) { + + block_fp32[m] = GGML_FP16_TO_FP32( + k_data_[batch_id * (q_len * config_.kv_head_num * + config_.head_dim) + + head_id * config_.head_dim + l * 32 + m]); + } + quantize_row_q8_0(block_fp32.data(), &block, 32); + + k_cache_q8[layer_id_][head_id][block_idx] + [pos_in_block * config_.head_dim / 32 + l] = + block; + } + + // fill v_cache_ + for (int l = 0; l < config_.head_dim; l++) { + block_q8_0 block = v_cache_q8[layer_id_][head_id][block_idx] + [l * config_.block_len / 32 + + pos_in_block / 32]; + dequantize_row_q8_0(&block, block_fp32.data(), 32); + block_fp32[pos_in_block % 32] = GGML_FP16_TO_FP32( + v_data_[batch_id * (q_len * config_.kv_head_num * + config_.head_dim) + + head_id * config_.head_dim + l]); + quantize_row_q8_0(block_fp32.data(), &block, 32); + v_cache_q8[layer_id_][head_id][block_idx] + [l * config_.block_len / 32 + pos_in_block / 32] = + block; + } + } + }, + nullptr); + + // Timer end + auto end = std::chrono::high_resolution_clock::now(); + std::chrono::duration duration = end - start; + // printf("layer %d time of reading KV Cache: %f s\n", layer_id, + // duration.count()); +} + +void KVCache::get_all_kvcache_one_layer(int layer_id, ggml_fp16_t *k_in, + ggml_fp16_t *v_in, Backend *backend) { + // Timer start + auto start = std::chrono::high_resolution_clock::now(); + + layer_id_ = layer_id; + seq_len_ = config_.block_len; + block_num_ = get_cache_total_block_num(); + k_data_ = reinterpret_cast(k_in); + v_data_ = reinterpret_cast(v_in); + + // Each task gets the k cache or v cache of a certain header + backend->do_work_stealing_job( + config_.kv_head_num * past_block_num_[layer_id] * 2, nullptr, + [&](int task_id) { + std::vector block_fp32(32); + int head_id = task_id / 2 / past_block_num_[layer_id]; + int block_idx = task_id / 2 % past_block_num_[layer_id]; + if (block_idx >= block_num_) + return; + + int max_offset = 0; + if (task_id & 1) { + // get k_cache_ + for (int k = 0; k < config_.block_len; k++) { + if (block_idx * seq_len_ + k >= cache_total_len_) + break; + for (int l = 0; l < config_.head_dim / 32; l++) { + block_q4_0 block = + k_cache_q4[layer_id_][head_id][block_idx] + [k * config_.head_dim / 32 + l]; + dequantize_row_q4_0(&block, block_fp32.data(), 32); + for (int m = 0; m < 32; m++) { + + k_data_[(head_id * cache_total_len_ + + block_idx * config_.block_len + k) * + config_.head_dim + + l * 32 + m] = + GGML_FP32_TO_FP16(block_fp32[m]); + max_offset = std::max( + max_offset, + (int)(head_id * cache_total_len_ + + block_idx * config_.block_len + k) * + config_.head_dim + + l * 32 + m); + } + } + } + } else { + // get v_cache_ + for (int k = 0; k < config_.block_len / 32; k++) { + for (int l = 0; l < config_.head_dim; l++) { + block_q4_0 block = + v_cache_q4[layer_id_][head_id][block_idx] + [l * config_.block_len / 32 + k]; + dequantize_row_q4_0(&block, block_fp32.data(), 32); + for (int m = 0; m < 32; m++) { + + if (block_idx * seq_len_ + k * 32 + m >= + cache_total_len_) + break; + v_data_[(head_id * cache_total_len_ + + block_idx * config_.block_len + k * 32 + + m) * + config_.head_dim + + l] = GGML_FP32_TO_FP16(block_fp32[m]); + max_offset = + std::max(max_offset, + (int)((head_id * cache_total_len_ + + block_idx * config_.block_len + + k * 32 + m) * + config_.head_dim + + l)); + } + } + } + } + }, + nullptr); + + // Timer end + auto end = std::chrono::high_resolution_clock::now(); + std::chrono::duration duration = end - start; + // printf("layer %d block num %d time of reading all KV Cache: %f s\n", + // layer_id, block_num_, duration.count()); +} diff --git a/archive/kt-sft/csrc/ktransformers_ext/operators/kvcache/kvcache_utils.cpp b/archive/kt-sft/csrc/ktransformers_ext/operators/kvcache/kvcache_utils.cpp new file mode 100644 index 0000000..c57d475 --- /dev/null +++ b/archive/kt-sft/csrc/ktransformers_ext/operators/kvcache/kvcache_utils.cpp @@ -0,0 +1,1159 @@ +/** + * @Description : + * @Author : Jianwei Dong + * @Date : 2024-08-26 22:47:06 + * @Version : 1.0.0 + * @LastEditors : Jianwei Dong + * @LastEditTime : 2024-08-26 22:47:06 + * @Copyright (c) 2024 by KVCache.AI, All Rights Reserved. + **/ + +#include "kvcache.h" + +#include + +std::string ggml_type_to_string(ggml_type type) { + switch (type) { + case GGML_TYPE_F32: + return "GGML_TYPE_F32"; + case GGML_TYPE_F16: + return "GGML_TYPE_F16"; + case GGML_TYPE_Q4_0: + return "GGML_TYPE_Q4_0"; + case GGML_TYPE_Q8_0: + return "GGML_TYPE_Q8_0"; + } + return "UNDIFINED"; +} +std::string AnchorTypeToString(AnchorType type) { + switch (type) { + case AnchorType::DYNAMIC: + return "DYNAMIC"; + case AnchorType::BLOCK_MEAN: + return "BLOCK_MEAN"; + case AnchorType::BLOCK_MAX: + return "BLOCK_MAX"; + case AnchorType::FIXED_ANCHOR: + return "FIXED_ANCHOR"; + case AnchorType::QUEST: + return "QUEST"; + } + return "UNDIFINED"; +} +std::string RetrievalTypeToString(RetrievalType type) { + switch (type) { + case RetrievalType::LAYER: + return "SHARED"; + case RetrievalType::KVHEAD: + return "SEPARATE"; + case RetrievalType::QHEAD: + return "INDIVIDUAL"; + } + return "UNDIFINED"; +} +KVCacheConfig::KVCacheConfig(int layer_num, int kv_head_num, int q_head_num, + int head_dim, int block_len, int anchor_num, + AnchorType anchor_type, ggml_type kv_type, + RetrievalType retrieval_type, int layer_step, + int token_step, int layer_offset, + int max_block_num, int max_batch_size, + int max_thread_num) + : layer_num(layer_num), kv_head_num(kv_head_num), q_head_num(q_head_num), + head_dim(head_dim), block_len(block_len), anchor_num(anchor_num), + anchor_type(anchor_type), kv_type(kv_type), + retrieval_type(retrieval_type), layer_step(layer_step), + token_step(token_step), layer_offset(layer_offset), + max_block_num(max_block_num), max_batch_size(max_batch_size), + max_thread_num(max_thread_num) { + printf( + "layer_num: %d, kv_head_num: %d, q_head_num: %d, head_dim: %d, " + "block_len: %d, anchor_num: %d, anchor_type: %s, kv_type: %s, " + "retrieval_type: %s, layer_step: %d, token_step: %d, layer_offset: %d," + "max_block_num: %d, max_batch_size: %d, max_thread_num: %d\n", + layer_num, kv_head_num, q_head_num, head_dim, block_len, anchor_num, + AnchorTypeToString(anchor_type).c_str(), + ggml_type_to_string(kv_type).c_str(), + RetrievalTypeToString(retrieval_type).c_str(), layer_step, token_step, + layer_offset, max_block_num, max_batch_size, max_thread_num); + assert(q_head_num % kv_head_num == 0); +} +KVCache::KVCache(KVCacheConfig config) { + this->config_ = config; + + n_gqa_ = config_.q_head_num / config_.kv_head_num; + if (config_.kv_type == ggml_type::GGML_TYPE_F16) { + // TODO: Elegant implement + k_cache_fp16_.resize(config_.layer_num); + v_cache_fp16_.resize(config_.layer_num); + selected_blocks_num_history_.resize(config_.layer_num / + config_.layer_step); + if (config_.retrieval_type == RetrievalType::LAYER) { + selected_blocks_history_.resize(config_.layer_num / + config_.layer_step); + } else if (config_.retrieval_type == RetrievalType::KVHEAD) { + selected_blocks_history_kvhead_.resize(config_.layer_num / + config_.layer_step); + } else if (config_.retrieval_type == RetrievalType::QHEAD) { + } + } else if (config_.kv_type == ggml_type::GGML_TYPE_Q4_0) { + k_cache_q4.resize(config.layer_num); + v_cache_q4.resize(config.layer_num); + } else if (config_.kv_type == ggml_type::GGML_TYPE_Q8_0) { + k_cache_q8.resize(config.layer_num); + v_cache_q8.resize(config.layer_num); + } else { + assert(false); + } + anchor_.resize(config.layer_num * config.max_block_num * config.anchor_num * + config.q_head_num * config.head_dim); + importance_.resize(config.layer_num); + past_block_num_.resize(config.layer_num); + for (int i = 0; i < config.layer_num; i++) { + past_block_num_[i] = 0; + } + + ThreadResize(config.max_thread_num); + BatchResize(config.max_batch_size); + BlockResize(config.max_block_num); + q_fp32.resize(n_gqa_ * config.head_dim); +} + +void KVCache::ThreadResize(int thread_num) { + thread_local_output_q8_0_.resize(thread_num); + thread_local_attn_score_.resize(thread_num); + thread_local_output_fp32_.resize(thread_num); + thread_local_attn_lse_.resize(thread_num); + thread_local_cur_output_fp32_.resize(thread_num); + thread_local_cur_attn_lse_.resize(thread_num); + thread_local_draft_.resize(thread_num); + thread_cur_head_idx_.resize(thread_num); + thread_local_attn_mask_.resize(thread_num); + for (int i = 0; i < thread_num; i++) { + thread_local_output_q8_0_[i].resize(n_gqa_ * config_.head_dim / QK8_0); + thread_local_attn_score_[i].resize(n_gqa_ * config_.block_len); + thread_local_output_fp32_[i].resize(n_gqa_ * config_.head_dim); + thread_local_attn_lse_[i].resize(n_gqa_); + thread_local_cur_output_fp32_[i].resize(n_gqa_ * config_.head_dim); + thread_local_cur_attn_lse_[i].resize(n_gqa_); + thread_local_draft_[i].resize( + 2 * n_gqa_ * config_.block_len + 6 * n_gqa_ * config_.head_dim + + 2 * config_.block_len * config_.head_dim + + config_.block_len * config_.head_dim / QK4_0); + thread_local_attn_mask_[i].resize(config_.block_len / 8); + } +} +void KVCache::BatchResize(int batch_size) { + mutex_.resize(batch_size); + q_q8_0_.resize(batch_size); + q_fp32_.resize(batch_size); + output_fp32_.resize(batch_size); + attn_lse_.resize(batch_size); + block_lse_.resize(batch_size); + attn_sparsity_.resize(batch_size); + + if (config_.retrieval_type == RetrievalType::LAYER) { + block_table_before_retrieval_.resize(batch_size); + block_table_after_retrieval_.resize(batch_size); + + for (int i = 0; i < config_.layer_num / config_.layer_step; i++) { + selected_blocks_history_[i].resize(batch_size); + } + + } else if (config_.retrieval_type == RetrievalType::KVHEAD) { + block_table_before_retrieval_kvhead_.resize(batch_size); + block_table_after_retrieval_kvhead_.resize(batch_size); + for (int i = 0; i < config_.layer_num / config_.layer_step; i++) { + selected_blocks_history_kvhead_[i].resize(batch_size); + } + } else if (config_.retrieval_type == RetrievalType::QHEAD) { + block_table_before_retrieval_qhead_.resize(batch_size); + block_table_after_retrieval_qhead_.resize(batch_size); + } + cache_seqlens_.resize(batch_size); + if (config_.retrieval_type == RetrievalType::LAYER) { + block_similar_.resize(batch_size); + } else if (config_.retrieval_type == RetrievalType::KVHEAD) { + block_similar_kv_head_.resize(batch_size); + } else if (config_.retrieval_type == RetrievalType::QHEAD) { + block_similar_q_head_.resize(batch_size); + } + for (int i = 0; i < batch_size; i++) { + top_similar_block_.resize(batch_size); + + mutex_[i].resize(config_.kv_head_num); + q_q8_0_[i].resize(config_.kv_head_num); + q_fp32_[i].resize(config_.kv_head_num); + output_fp32_[i].resize(config_.kv_head_num); + attn_lse_[i].resize(config_.kv_head_num); + + for (int j = 0; j < config_.kv_head_num; j++) { + if (!mutex_[i][j]) { + mutex_[i][j] = std::make_unique(); + } + q_q8_0_[i][j].resize(n_gqa_ * config_.head_dim / QK8_0); + q_fp32_[i][j].resize(n_gqa_ * config_.head_dim); + output_fp32_[i][j].resize(n_gqa_ * config_.head_dim); + attn_lse_[i][j].resize(n_gqa_); + } + } + avg_q.resize(batch_size); + avg_q_fp16.resize(batch_size); + for (int i = 0; i < batch_size; i++) { + attn_sparsity_[i].resize(config_.q_head_num); + avg_q[i].resize(config_.q_head_num * config_.head_dim); + avg_q_fp16[i].resize(config_.q_head_num * config_.head_dim); + } +} + +void KVCache::BlockResize(int max_block_num) { + sin_.resize(max_block_num * config_.block_len); + cos_.resize(max_block_num * config_.block_len); + for (int i = 0; i < max_block_num * config_.block_len; i++) { + sin_[i].resize(config_.head_dim); + cos_[i].resize(config_.head_dim); + } + + for (int i = 0; i < config_.layer_num / config_.layer_step; i++) { + for (int j = 0; j < config_.max_batch_size; j++) { + if (config_.retrieval_type == RetrievalType::LAYER) { + selected_blocks_history_[i][j].resize(max_block_num); + } else if (config_.retrieval_type == RetrievalType::KVHEAD) { + selected_blocks_history_kvhead_[i][j].resize(max_block_num); + for (int k = 0; k < config_.max_block_num; k++) { + selected_blocks_history_kvhead_[i][j][k].resize( + config_.kv_head_num); + } + } else if (config_.retrieval_type == RetrievalType::QHEAD) { + } + } + } + + for (int layer_id = 0; layer_id < config_.layer_num; layer_id++) { + importance_[layer_id].resize(max_block_num); + + if (config_.kv_type == ggml_type::GGML_TYPE_F16) { + // TODO: Elegant implement + k_cache_fp16_[layer_id].resize(config_.kv_head_num); + v_cache_fp16_[layer_id].resize(config_.kv_head_num); + + for (int i = 0; i < config_.kv_head_num; i++) { + k_cache_fp16_[layer_id][i].resize(max_block_num); + v_cache_fp16_[layer_id][i].resize(max_block_num); + + for (int j = 0; j < max_block_num; j++) { + k_cache_fp16_[layer_id][i][j].resize(config_.block_len * + config_.head_dim); + v_cache_fp16_[layer_id][i][j].resize(config_.block_len * + config_.head_dim); + } + } + + } else if (config_.kv_type == ggml_type::GGML_TYPE_Q4_0) { + k_cache_q4[layer_id].resize(config_.kv_head_num); + v_cache_q4[layer_id].resize(config_.kv_head_num); + for (int i = 0; i < config_.kv_head_num; i++) { + k_cache_q4[layer_id][i].resize(max_block_num); + v_cache_q4[layer_id][i].resize(max_block_num); + + for (int j = 0; j < max_block_num; j++) { + k_cache_q4[layer_id][i][j].resize(config_.block_len * + config_.head_dim / 32); + v_cache_q4[layer_id][i][j].resize(config_.block_len * + config_.head_dim / 32); + } + } + } else if (config_.kv_type == ggml_type::GGML_TYPE_Q8_0) { + k_cache_q8[layer_id].resize(config_.kv_head_num); + v_cache_q8[layer_id].resize(config_.kv_head_num); + for (int i = 0; i < config_.kv_head_num; i++) { + k_cache_q8[layer_id][i].resize(max_block_num); + v_cache_q8[layer_id][i].resize(max_block_num); + + for (int j = 0; j < max_block_num; j++) { + k_cache_q8[layer_id][i][j].resize(config_.block_len * + config_.head_dim / 32); + v_cache_q8[layer_id][i][j].resize(config_.block_len * + config_.head_dim / 32); + } + } + } else { + assert(false); + } + for (int i = 0; i < config_.max_batch_size; i++) { + if (config_.retrieval_type == RetrievalType::LAYER) { + block_similar_[i].resize(max_block_num); + block_table_before_retrieval_[i].resize(max_block_num); + block_table_after_retrieval_[i].resize(max_block_num); + } else if (config_.retrieval_type == RetrievalType::KVHEAD) { + block_similar_kv_head_[i].resize(max_block_num); + block_table_before_retrieval_kvhead_[i].resize(max_block_num); + block_table_after_retrieval_kvhead_[i].resize(max_block_num); + for (int j = 0; j < max_block_num; j++) { + block_similar_kv_head_[i][j].resize(config_.kv_head_num); + block_table_before_retrieval_kvhead_[i][j].resize( + config_.kv_head_num); + block_table_after_retrieval_kvhead_[i][j].resize( + config_.kv_head_num); + } + } else if (config_.retrieval_type == RetrievalType::QHEAD) { + block_similar_q_head_[i].resize(max_block_num); + block_table_before_retrieval_qhead_[i].resize(max_block_num); + block_table_after_retrieval_qhead_[i].resize(max_block_num); + for (int j = 0; j < max_block_num; j++) { + block_similar_q_head_[i][j].resize(config_.q_head_num); + block_table_before_retrieval_qhead_[i][j].resize( + config_.q_head_num); + block_table_after_retrieval_qhead_[i][j].resize( + config_.q_head_num); + } + } + block_lse_[i].resize(max_block_num); + for (int j = 0; j < max_block_num; j++) { + block_lse_[i][j].resize(config_.q_head_num); + } + } + + for (int i = 0; i < max_block_num; i++) { + importance_[layer_id][i].resize(config_.block_len); + for (int j = 0; j < config_.block_len; j++) { + importance_[layer_id][i][j].resize(config_.q_head_num); + } + } + } +} + +void KVCache::calc_anchor_all_layers(int *block_table, int *cache_seqlens, + int batch_size, int max_block_num, + Backend *backend) { + // Timer start + auto start = std::chrono::high_resolution_clock::now(); + + // Each task updates the importance of a certain block + seq_len_ = config_.block_len; + backend->do_work_stealing_job( + config_.layer_num * batch_size * max_block_num, nullptr, + [&](int task_id) { + int layer_id = task_id / (batch_size * max_block_num); + int batch_id = (task_id / max_block_num) % batch_size; + int block_id = task_id % max_block_num; + // If the block is out of the sequence length, skip it. In + // particular, the last block of the sequence that is shorter than + // the block length should be skipped. + + if (cache_seqlens[batch_id] / config_.block_len < block_id) { + return; + } + int block_idx = block_table[batch_id * max_block_num + block_id]; + + std::vector block_fp32(32); + if (config_.anchor_type == AnchorType::DYNAMIC) { + + // clear anchor_ + for (int anchor_id = 0; anchor_id < 1; anchor_id++) { + for (int head_id = 0; head_id < config_.q_head_num; + head_id++) { + for (int l = 0; l < config_.head_dim; l++) { + anchor_[layer_id * config_.max_block_num * + config_.anchor_num * + config_.q_head_num * config_.head_dim + + block_idx * config_.anchor_num * + config_.q_head_num * config_.head_dim + + anchor_id * config_.q_head_num * + config_.head_dim + + head_id * config_.head_dim + l] = 0; + } + } + } + + // find top anchor_num importances and their corresponding + // positions in the importance_ tensor + // TODO: Move top_importances to the class member to avoid + // repeated memory allocation + std::priority_queue< + std::pair>, + std::vector>>, + std::greater<>> + top_importances; + for (int head_id = 0; head_id < config_.q_head_num; head_id++) { + for (int k = 0; k < seq_len_; k++) { + top_importances.push(std::make_pair( + GGML_FP16_TO_FP32( + importance_[layer_id][block_idx][k][head_id]), + std::make_pair(block_idx, k))); + // TODO: change to config_ item + if (top_importances.size() > config_.anchor_num) { + top_importances.pop(); + } + } + + // fill anchor_ + + for (int l = 0; l < config_.head_dim; l++) { + anchor_[layer_id * config_.max_block_num * + config_.anchor_num * config_.q_head_num * + config_.head_dim + + block_idx * config_.anchor_num * + config_.q_head_num * config_.head_dim + + 0 * config_.q_head_num * config_.head_dim + + head_id * config_.head_dim + l] = 0; + } + for (int k = 0; k < config_.anchor_num; k++) { + int top_indice = top_importances.top().second.second; + int top_block_idx = top_importances.top().second.first; + + if (config_.kv_type == ggml_type::GGML_TYPE_F16) { + + for (int l = 0; l < config_.head_dim; l++) { + anchor_[layer_id * config_.max_block_num * + config_.anchor_num * + config_.q_head_num * + config_.head_dim + + top_block_idx * config_.anchor_num * + config_.q_head_num * + config_.head_dim + + 0 * config_.q_head_num * + config_.head_dim + + head_id * config_.head_dim + l] = + GGML_FP32_TO_FP16( + GGML_FP16_TO_FP32( + anchor_[layer_id * + config_.max_block_num * + config_.anchor_num * + config_.q_head_num * + config_.head_dim + + top_block_idx * + config_.anchor_num * + config_.q_head_num * + config_.head_dim + + 0 * config_.q_head_num * + config_.head_dim + + head_id * config_.head_dim + + l]) + + GGML_FP16_TO_FP32( + k_cache_fp16_[layer_id] + [head_id / n_gqa_] + [top_block_idx] + [top_indice * + config_.head_dim + + l])); + } + + } else if (config_.kv_type == + ggml_type::GGML_TYPE_Q4_0) { + for (int l = 0; l < config_.head_dim / 32; l++) { + block_q4_0 block = k_cache_q4 + [layer_id][head_id / n_gqa_][top_block_idx] + [top_indice * config_.head_dim / 32 + l]; + dequantize_row_q4_0(&block, block_fp32.data(), + 32); + for (int m = 0; m < 32; m++) { + anchor_[layer_id * config_.max_block_num * + config_.anchor_num * + config_.q_head_num * + config_.head_dim + + top_block_idx * config_.anchor_num * + config_.q_head_num * + config_.head_dim + + 0 * config_.q_head_num * + config_.head_dim + + head_id * config_.head_dim + + l * 32 + m] = + GGML_FP32_TO_FP16( + block_fp32[m] / 4 + + GGML_FP16_TO_FP32( + anchor_[layer_id * + config_ + .max_block_num * + config_.anchor_num * + config_.q_head_num * + config_.head_dim + + top_block_idx * + config_.anchor_num * + config_.q_head_num * + config_.head_dim + + 0 * config_.q_head_num * + config_.head_dim + + head_id * + config_.head_dim + + l * 32 + m])); + } + } + } else if (config_.kv_type == + ggml_type::GGML_TYPE_Q8_0) { + for (int l = 0; l < config_.head_dim / 32; l++) { + block_q8_0 block = k_cache_q8 + [layer_id][head_id / n_gqa_][top_block_idx] + [top_indice * config_.head_dim / 32 + l]; + dequantize_row_q8_0(&block, block_fp32.data(), + 32); + for (int m = 0; m < 32; m++) { + anchor_[layer_id * config_.max_block_num * + config_.anchor_num * + config_.q_head_num * + config_.head_dim + + top_block_idx * config_.anchor_num * + config_.q_head_num * + config_.head_dim + + 0 * config_.q_head_num * + config_.head_dim + + head_id * config_.head_dim + + l * 32 + m] = + GGML_FP32_TO_FP16( + block_fp32[m] / 4 + + GGML_FP16_TO_FP32( + anchor_[layer_id * + config_ + .max_block_num * + config_.anchor_num * + config_.q_head_num * + config_.head_dim + + top_block_idx * + config_.anchor_num * + config_.q_head_num * + config_.head_dim + + 0 * config_.q_head_num * + config_.head_dim + + head_id * + config_.head_dim + + l * 32 + m])); + } + } + } + top_importances.pop(); + } + } + } else if (config_.anchor_type == AnchorType::BLOCK_MEAN) { + // clear anchor_ + for (int anchor_id = 0; anchor_id < config_.anchor_num; + anchor_id++) { + for (int head_id = 0; head_id < config_.q_head_num; + head_id++) { + for (int l = 0; l < config_.head_dim; l++) { + anchor_[layer_id * config_.max_block_num * + config_.anchor_num * + config_.q_head_num * config_.head_dim + + block_idx * config_.anchor_num * + config_.q_head_num * config_.head_dim + + anchor_id * config_.q_head_num * + config_.head_dim + + head_id * config_.head_dim + l] = 0; + } + } + } + + // fill anchor_ + if (config_.kv_type == ggml_type::GGML_TYPE_F16) { + + for (int head_id = 0; head_id < config_.q_head_num; + head_id++) { + for (int k = 0; k < config_.block_len; k++) { + for (int l = 0; l < config_.head_dim; l++) { + anchor_[layer_id * config_.max_block_num * + config_.anchor_num * + config_.q_head_num * + config_.head_dim + + block_idx * config_.anchor_num * + config_.q_head_num * + config_.head_dim + + 0 * config_.q_head_num * + config_.head_dim + + head_id * config_.head_dim + l] = + GGML_FP32_TO_FP16( + GGML_FP16_TO_FP32( + anchor_[layer_id * + config_.max_block_num * + config_.anchor_num * + config_.q_head_num * + config_.head_dim + + block_idx * + config_.anchor_num * + config_.q_head_num * + config_.head_dim + + 0 * config_.q_head_num * + config_.head_dim + + head_id * config_.head_dim + + l]) + + GGML_FP16_TO_FP32( + k_cache_fp16_[layer_id] + [head_id / n_gqa_] + [block_idx] + [k * config_.head_dim + + l]) / + config_.block_len); + } + } + } + } + } else if (config_.anchor_type == AnchorType::BLOCK_MAX) { + // clear anchor_ + for (int anchor_id = 0; anchor_id < config_.anchor_num; + anchor_id++) { + for (int head_id = 0; head_id < config_.q_head_num; + head_id++) { + for (int l = 0; l < config_.head_dim; l++) { + anchor_[layer_id * config_.max_block_num * + config_.anchor_num * + config_.q_head_num * config_.head_dim + + block_idx * config_.anchor_num * + config_.q_head_num * config_.head_dim + + anchor_id * config_.q_head_num * + config_.head_dim + + head_id * config_.head_dim + l] = 0; + } + } + } + + // fill anchor_ + if (config_.kv_type == ggml_type::GGML_TYPE_F16) { + + for (int head_id = 0; head_id < config_.q_head_num; + head_id++) { + for (int k = 0; k < config_.block_len; k++) { + for (int l = 0; l < config_.head_dim; l++) { + anchor_[layer_id * config_.max_block_num * + config_.anchor_num * + config_.q_head_num * + config_.head_dim + + block_idx * config_.anchor_num * + config_.q_head_num * + config_.head_dim + + 0 * config_.q_head_num * + config_.head_dim + + head_id * config_.head_dim + l] = + GGML_FP32_TO_FP16(std::max( + GGML_FP16_TO_FP32( + anchor_[layer_id * + config_.max_block_num * + config_.anchor_num * + config_.q_head_num * + config_.head_dim + + block_idx * + config_.anchor_num * + config_.q_head_num * + config_.head_dim + + 0 * config_.q_head_num * + config_.head_dim + + head_id * config_.head_dim + + l]), + GGML_FP16_TO_FP32( + k_cache_fp16_ + [layer_id][head_id / n_gqa_] + [block_idx] + [k * config_.head_dim + l]))); + } + } + } + } + } else if (config_.anchor_type == AnchorType::FIXED_ANCHOR) { + // clear anchor_ + for (int anchor_id = 0; anchor_id < 1; anchor_id++) { + for (int head_id = 0; head_id < config_.q_head_num; + head_id++) { + for (int l = 0; l < config_.head_dim; l++) { + anchor_[layer_id * config_.max_block_num * + config_.anchor_num * + config_.q_head_num * config_.head_dim + + block_idx * config_.anchor_num * + config_.q_head_num * config_.head_dim + + anchor_id * config_.q_head_num * + config_.head_dim + + head_id * config_.head_dim + l] = 0; + } + } + } + + // fill anchor_ + if (config_.kv_type == ggml_type::GGML_TYPE_F16) { + + int stride = config_.block_len / config_.anchor_num; + for (int head_id = 0; head_id < config_.q_head_num; + head_id++) { + for (int k = 0, tot = 0; + k < config_.block_len, tot < config_.anchor_num; + k += stride, tot++) { + for (int l = 0; l < config_.head_dim; l++) { + anchor_[layer_id * config_.max_block_num * + config_.anchor_num * + config_.q_head_num * + config_.head_dim + + block_idx * config_.anchor_num * + config_.q_head_num * + config_.head_dim + + 0 * config_.q_head_num * + config_.head_dim + + head_id * config_.head_dim + l] = + GGML_FP32_TO_FP16( + GGML_FP16_TO_FP32( + anchor_[layer_id * + config_.max_block_num * + config_.anchor_num * + config_.q_head_num * + config_.head_dim + + block_idx * + config_.anchor_num * + config_.q_head_num * + config_.head_dim + + 0 * config_.q_head_num * + config_.head_dim + + head_id * config_.head_dim + + l]) + + GGML_FP16_TO_FP32( + k_cache_fp16_[layer_id] + [head_id / n_gqa_] + [block_idx] + [k * config_.head_dim + + l]) / + config_.anchor_num); + } + } + } + } + + } else if (config_.anchor_type == AnchorType::QUEST) { + // clear anchor_ + for (int head_id = 0; head_id < config_.q_head_num; head_id++) { + for (int l = 0; l < config_.head_dim; l++) { + anchor_[layer_id * config_.max_block_num * + config_.anchor_num * config_.q_head_num * + config_.head_dim + + block_idx * config_.anchor_num * + config_.q_head_num * config_.head_dim + + 1 * config_.q_head_num * config_.head_dim + + head_id * config_.head_dim + l] = + GGML_FP32_TO_FP16( + std::numeric_limits::max()); + + anchor_[layer_id * config_.max_block_num * + config_.anchor_num * config_.q_head_num * + config_.head_dim + + block_idx * config_.anchor_num * + config_.q_head_num * config_.head_dim + + 0 * config_.q_head_num * config_.head_dim + + head_id * config_.head_dim + l] = + GGML_FP32_TO_FP16( + std::numeric_limits::min()); + } + } + + // fill anchor_ + + if (config_.kv_type == ggml_type::GGML_TYPE_F16) { + for (int indice = 0; indice < seq_len_; indice++) { + for (int head_id = 0; head_id < config_.kv_head_num; + head_id++) { + for (int l = 0; l < config_.head_dim; l++) { + anchor_[layer_id * config_.max_block_num * + config_.anchor_num * + config_.q_head_num * + config_.head_dim + + block_idx * config_.anchor_num * + config_.q_head_num * + config_.head_dim + + 0 * config_.q_head_num * + config_.head_dim + + head_id * config_.head_dim + l] = + GGML_FP32_TO_FP16(std::max( + GGML_FP16_TO_FP32( + k_cache_fp16_ + [layer_id][head_id][block_idx] + [indice * config_.head_dim + + l]), + GGML_FP16_TO_FP32( + anchor_[layer_id * + config_.max_block_num * + config_.anchor_num * + config_.q_head_num * + config_.head_dim + + block_idx * + config_.anchor_num * + config_.q_head_num * + config_.head_dim + + 0 * config_.q_head_num * + config_.head_dim + + head_id * config_.head_dim + + l]))); + + anchor_[layer_id * config_.max_block_num * + config_.anchor_num * + config_.q_head_num * + config_.head_dim + + block_idx * config_.anchor_num * + config_.q_head_num * + config_.head_dim + + 1 * config_.q_head_num * + config_.head_dim + + head_id * config_.head_dim + l] = + GGML_FP32_TO_FP16(std::min( + GGML_FP16_TO_FP32( + k_cache_fp16_ + [layer_id][head_id][block_idx] + [indice * config_.head_dim + + l]), + GGML_FP16_TO_FP32( + anchor_[layer_id * + config_.max_block_num * + config_.anchor_num * + config_.q_head_num * + config_.head_dim + + block_idx * + config_.anchor_num * + config_.q_head_num * + config_.head_dim + + 1 * config_.q_head_num * + config_.head_dim + + head_id * config_.head_dim + + l]))); + } + } + } + + } else if (config_.kv_type == ggml_type::GGML_TYPE_Q4_0) { + for (int indice = 0; indice < seq_len_; indice++) { + for (int head_id = 0; head_id < config_.kv_head_num; + head_id++) { + for (int l = 0; l < config_.head_dim / 32; l++) { + block_q4_0 block = + k_cache_q4[layer_id][head_id][block_idx] + [indice * config_.head_dim / 32 + + l]; + dequantize_row_q4_0(&block, block_fp32.data(), + 32); + + for (int m = 0; m < 32; m++) { + for (int gqa_idx = 0; gqa_idx < n_gqa_; + gqa_idx++) { + + anchor_[layer_id * + config_.max_block_num * + config_.anchor_num * + config_.q_head_num * + config_.head_dim + + block_idx * config_.anchor_num * + config_.q_head_num * + config_.head_dim + + 0 * config_.q_head_num * + config_.head_dim + + head_id * config_.head_dim + + l * 32 + m] = + GGML_FP32_TO_FP16(std::max( + block_fp32[m], + GGML_FP16_TO_FP32( + anchor_ + [layer_id * + config_ + .max_block_num * + config_ + .anchor_num * + config_ + .q_head_num * + config_.head_dim + + block_idx * + config_ + .anchor_num * + config_ + .q_head_num * + config_.head_dim + + 0 * + config_ + .q_head_num * + config_.head_dim + + head_id * + config_.head_dim + + l * 32 + m]))); + + anchor_[layer_id * + config_.max_block_num * + config_.anchor_num * + config_.q_head_num * + config_.head_dim + + block_idx * config_.anchor_num * + config_.q_head_num * + config_.head_dim + + 1 * config_.q_head_num * + config_.head_dim + + head_id * config_.head_dim + + l * 32 + m] = + GGML_FP32_TO_FP16(std::min( + block_fp32[m], + GGML_FP16_TO_FP32( + anchor_ + [layer_id * + config_ + .max_block_num * + config_ + .anchor_num * + config_ + .q_head_num * + config_.head_dim + + block_idx * + config_ + .anchor_num * + config_ + .q_head_num * + config_.head_dim + + 1 * + config_ + .q_head_num * + config_.head_dim + + head_id * + config_.head_dim + + l * 32 + m]))); + } + } + } + } + } + } else if (config_.kv_type == ggml_type::GGML_TYPE_Q8_0) { + for (int indice = 0; indice < seq_len_; indice++) { + for (int head_id = 0; head_id < config_.kv_head_num; + head_id++) { + for (int l = 0; l < config_.head_dim / 32; l++) { + block_q8_0 block = + k_cache_q8[layer_id][head_id][block_idx] + [indice * config_.head_dim / 32 + + l]; + dequantize_row_q8_0(&block, block_fp32.data(), + 32); + + for (int m = 0; m < 32; m++) { + for (int gqa_idx = 0; gqa_idx < n_gqa_; + gqa_idx++) { + + anchor_[layer_id * + config_.max_block_num * + config_.anchor_num * + config_.q_head_num * + config_.head_dim + + block_idx * config_.anchor_num * + config_.q_head_num * + config_.head_dim + + 0 * config_.q_head_num * + config_.head_dim + + head_id * config_.head_dim + + l * 32 + m] = + GGML_FP32_TO_FP16(std::max( + block_fp32[m], + GGML_FP16_TO_FP32( + anchor_ + [layer_id * + config_ + .max_block_num * + config_ + .anchor_num * + config_ + .q_head_num * + config_.head_dim + + block_idx * + config_ + .anchor_num * + config_ + .q_head_num * + config_.head_dim + + 0 * + config_ + .q_head_num * + config_.head_dim + + head_id * + config_.head_dim + + l * 32 + m]))); + + anchor_[layer_id * + config_.max_block_num * + config_.anchor_num * + config_.q_head_num * + config_.head_dim + + block_idx * config_.anchor_num * + config_.q_head_num * + config_.head_dim + + 1 * config_.q_head_num * + config_.head_dim + + head_id * config_.head_dim + + l * 32 + m] = + GGML_FP32_TO_FP16(std::min( + block_fp32[m], + GGML_FP16_TO_FP32( + anchor_ + [layer_id * + config_ + .max_block_num * + config_ + .anchor_num * + config_ + .q_head_num * + config_.head_dim + + block_idx * + config_ + .anchor_num * + config_ + .q_head_num * + config_.head_dim + + 1 * + config_ + .q_head_num * + config_.head_dim + + head_id * + config_.head_dim + + l * 32 + m]))); + } + } + } + } + } + } + } else { + assert(false); + } + }, + nullptr); + + // Timer end + auto end = std::chrono::high_resolution_clock::now(); + std::chrono::duration duration = end - start; + // printf("time of calc_anchor_all_layers: %f s\n", duration.count()); +} + +void KVCache::clear_importance_all_layers(int *block_table, int *cache_seqlens, + int batch_size, int max_block_num, + Backend *backend) { + // Timer start + auto start = std::chrono::high_resolution_clock::now(); + + // Each task updates the importance of a certain block + seq_len_ = config_.block_len; + backend->do_work_stealing_job( + config_.layer_num * batch_size * max_block_num, nullptr, + [&](int task_id) { + int layer_id = task_id / (batch_size * max_block_num); + int batch_id = (task_id / max_block_num) % batch_size; + int block_id = task_id % max_block_num; + // If the block is out of the sequence length, skip it. In + // particular, the last block of the sequence that is shorter than + // the block length should be skipped. + + if (cache_seqlens[batch_id] / config_.block_len < block_id) { + return; + } + int block_idx = block_table[batch_id * max_block_num + block_id]; + + if (config_.anchor_type == AnchorType::DYNAMIC) { + + // clear anchor_ + for (int head_id = 0; head_id < config_.q_head_num; head_id++) { + for (int l = 0; l < config_.block_len; l++) { + importance_[layer_id][block_idx][l][head_id] = 0; + } + } + } + }, + nullptr); + + // Timer end + auto end = std::chrono::high_resolution_clock::now(); + std::chrono::duration duration = end - start; + // printf("time of clear_importance_all_layerssssss: %f s\n", + // duration.count()); +} + +void KVCache::clear_kvcache_all_layers(int *block_table, int *cache_seqlens, + int batch_size, int max_block_num, + Backend *backend) { + // Timer start + auto start = std::chrono::high_resolution_clock::now(); + + // Each task updates the importance of a certain block + seq_len_ = config_.block_len; + backend->do_work_stealing_job( + config_.layer_num * batch_size * max_block_num * config_.kv_head_num, + nullptr, + [&](int task_id) { + int layer_id = + task_id / (batch_size * max_block_num * config_.kv_head_num); + int batch_id = + (task_id / (max_block_num * config_.kv_head_num)) % batch_size; + int block_id = task_id / config_.kv_head_num % max_block_num; + int head_id = task_id % config_.kv_head_num; + // If the block is out of the sequence length, skip it. In + // particular, the last block of the sequence that is shorter than + // the block length should be skipped. + if (cache_seqlens[batch_id] / config_.block_len < block_id) { + return; + } + int block_idx = block_table[batch_id * max_block_num + block_id]; + + if (config_.kv_type == ggml_type::GGML_TYPE_F16) { + for (int l = 0; l < config_.block_len * config_.head_dim; l++) { + k_cache_fp16_[layer_id][head_id][block_idx][l] = 0; + v_cache_fp16_[layer_id][head_id][block_idx][l] = 0; + } + } else if (config_.kv_type == ggml_type::GGML_TYPE_Q4_0) { + for (int l = 0; l < config_.block_len * config_.head_dim / 32; + l++) { + k_cache_q4[layer_id][head_id][block_idx][l].d = 0; + v_cache_q4[layer_id][head_id][block_idx][l].d = 0; + } + } else if (config_.kv_type == ggml_type::GGML_TYPE_Q8_0) { + for (int l = 0; l < config_.block_len * config_.head_dim / 32; + l++) { + k_cache_q8[layer_id][head_id][block_idx][l].d = 0; + v_cache_q8[layer_id][head_id][block_idx][l].d = 0; + } + } + }, + nullptr); + + // Timer end + auto end = std::chrono::high_resolution_clock::now(); + std::chrono::duration duration = end - start; + // printf("time of clear_kvcache_all_layers: %f s\n", duration.count()); +} + +void KVCache::get_sincos(ggml_fp16_t *sin, ggml_fp16_t *cos, int seqlen) { + // Timer start + auto start = std::chrono::high_resolution_clock::now(); + + const uint16_t *sin_data = const_cast(sin); + const uint16_t *cos_data = const_cast(cos); + + for (int i = 0; i < seqlen; i++) { + for (int j = 0; j < config_.head_dim; j++) { + sin_[i][j] = sin_data[i * config_.head_dim + j]; + cos_[i][j] = cos_data[i * config_.head_dim + j]; + } + } + + // Timer end + auto end = std::chrono::high_resolution_clock::now(); + std::chrono::duration duration = end - start; + printf("time of get_sincos: %f s\n", duration.count()); +} + +void ggml_vec_scale_f32(const int n, float *y, const float v) { +#if defined(GGML_USE_ACCELERATE) + vDSP_vsmul(y, 1, &v, y, 1, n); +#elif defined(GGML_SIMD) + const int np = (n & ~(GGML_F32_STEP - 1)); + + GGML_F32_VEC vx = GGML_F32_VEC_SET1(v); + + GGML_F32_VEC ay[GGML_F32_ARR]; + + for (int i = 0; i < np; i += GGML_F32_STEP) { + for (int j = 0; j < GGML_F32_ARR; j++) { + ay[j] = GGML_F32_VEC_LOAD(y + i + j * GGML_F32_EPR); + ay[j] = GGML_F32_VEC_MUL(ay[j], vx); + + GGML_F32_VEC_STORE(y + i + j * GGML_F32_EPR, ay[j]); + } + } + + // leftovers + for (int i = np; i < n; ++i) { + y[i] *= v; + } +#else + // scalar + for (int i = 0; i < n; ++i) { + y[i] *= v; + } +#endif +} \ No newline at end of file diff --git a/archive/kt-sft/csrc/ktransformers_ext/operators/llamafile/conversion.h b/archive/kt-sft/csrc/ktransformers_ext/operators/llamafile/conversion.h new file mode 100644 index 0000000..2a194d6 --- /dev/null +++ b/archive/kt-sft/csrc/ktransformers_ext/operators/llamafile/conversion.h @@ -0,0 +1,32 @@ +/** + * @Description : + * @Author : chenht2022 + * @Date : 2024-07-12 10:07:58 + * @Version : 1.0.0 + * @LastEditors : chenht2022 + * @LastEditTime : 2024-07-25 10:34:55 + * @Copyright (c) 2024 by KVCache.AI, All Rights Reserved. + **/ +#ifndef CPUINFER_CONVERSION_H +#define CPUINFER_CONVERSION_H + +#include +#include "llama.cpp/ggml.h" + +inline void to_float(const void* input, float* output, int size, ggml_type type) { + if (type == ggml_type::GGML_TYPE_F32) { + memcpy(output, input, size * sizeof(float)); + } else { + ggml_internal_get_type_traits(type).to_float(input, output, size); + } +} + +inline void from_float(const float* input, void* output, int size, ggml_type type) { + if (type == ggml_type::GGML_TYPE_F32) { + memcpy(output, input, size * sizeof(float)); + } else { + ggml_internal_get_type_traits(type).from_float(input, output, size); + } +} + +#endif \ No newline at end of file diff --git a/archive/kt-sft/csrc/ktransformers_ext/operators/llamafile/linear.cpp b/archive/kt-sft/csrc/ktransformers_ext/operators/llamafile/linear.cpp new file mode 100644 index 0000000..d1e7967 --- /dev/null +++ b/archive/kt-sft/csrc/ktransformers_ext/operators/llamafile/linear.cpp @@ -0,0 +1,77 @@ +/** + * @Description : + * @Author : chenht2022 + * @Date : 2024-07-12 10:07:58 + * @Version : 1.0.0 + * @LastEditors : kkk1nak0 + * @LastEditTime : 2024-08-15 07:45:18 + * @Copyright (c) 2024 by KVCache.AI, All Rights Reserved. + **/ +#include "linear.h" + +Linear::Linear(LinearConfig config) { + config_ = config; + proj_ = config_.proj; + + std::vector> mem_requests; + mem_requests.push_back({(void**)&input_fp32_, sizeof(float) * config_.group_max_len * config_.input_size}); + mem_requests.push_back({(void**)&proj_input_, config_.group_max_len * config_.input_size * ggml_type_size(ggml_internal_get_type_traits(config_.proj_type).vec_dot_type) / ggml_blck_size(ggml_internal_get_type_traits(config_.proj_type).vec_dot_type)}); + mem_requests.push_back({(void**)&proj_output_, sizeof(float) * config_.group_max_len * config_.output_size}); + shared_mem_buffer.alloc(this, mem_requests); +} + +Linear::~Linear() { + shared_mem_buffer.dealloc(this); +} + +void Linear::warm_up(Backend *backend) { + std::vector input_fp32(config_.input_size); + std::vector input(config_.input_size * + ggml_type_size(config_.hidden_type) / + ggml_blck_size(config_.hidden_type)); + std::vector output(config_.output_size * + ggml_type_size(config_.hidden_type) / + ggml_blck_size(config_.hidden_type)); + for (int i = 0; i < config_.input_size; i++) { + input_fp32[i] = 0; + } + from_float(input_fp32.data(), input.data(), config_.input_size, config_.hidden_type); + forward_many(1, input.data(), output.data(), backend); +} + +void Linear::forward_many(int qlen, const void* input, void* output, Backend* backend) { + const void* proj_input_ptr; + if (config_.hidden_type == ggml_internal_get_type_traits(config_.proj_type).vec_dot_type) { + proj_input_ptr = input; + } else { + to_float(input, input_fp32_, qlen * config_.input_size, config_.hidden_type); + from_float(input_fp32_, proj_input_, qlen * config_.input_size, ggml_internal_get_type_traits(config_.proj_type).vec_dot_type); + proj_input_ptr = proj_input_; + } + int nth = config_.output_size / config_.stride; + backend->do_work_stealing_job(nth, nullptr, [&](int task_id) { + int ith = task_id; + void* proj_ptr = (uint8_t*)proj_ + ith * config_.stride * config_.input_size * ggml_type_size(config_.proj_type) / ggml_blck_size(config_.proj_type); + float* proj_output_ptr = proj_output_ + ith * config_.stride; + llamafile_sgemm(config_.stride, qlen, config_.input_size / ggml_blck_size(config_.proj_type), proj_ptr, config_.input_size / ggml_blck_size(config_.proj_type), proj_input_ptr, config_.input_size / ggml_blck_size(config_.proj_type), proj_output_ptr, config_.output_size, 0, 1, GGML_TASK_TYPE_COMPUTE, config_.proj_type, ggml_internal_get_type_traits(config_.proj_type).vec_dot_type, GGML_TYPE_F32, GGML_PREC_DEFAULT); + if (config_.stride % ggml_blck_size(config_.hidden_type) == 0) { + for (int i = 0; i < qlen; i++) { + float* output_fp32_ptr = proj_output_ + i * config_.output_size + ith * config_.stride; + void* output_ptr = (uint8_t*)output + i * config_.output_size * ggml_type_size(config_.hidden_type) / ggml_blck_size(config_.hidden_type) + ith * config_.stride * ggml_type_size(config_.hidden_type) / ggml_blck_size(config_.hidden_type); + from_float(output_fp32_ptr, output_ptr, config_.stride, config_.hidden_type); + } + } + }, nullptr); + if (config_.stride % ggml_blck_size(config_.hidden_type) != 0) { + from_float(proj_output_, output, qlen * config_.output_size, config_.hidden_type); + } +} + +void Linear::forward(int qlen, const void* input, void* output, Backend* backend) { + if (qlen <= 0) { + return; + } + int forward_len = std::min(qlen, config_.group_max_len); + forward_many(forward_len, input, output, backend); + forward(qlen - forward_len, (uint8_t*)input + forward_len * config_.input_size * ggml_type_size(config_.hidden_type) / ggml_blck_size(config_.hidden_type), (uint8_t*)output + forward_len * config_.output_size * ggml_type_size(config_.hidden_type) / ggml_blck_size(config_.hidden_type), backend); +} \ No newline at end of file diff --git a/archive/kt-sft/csrc/ktransformers_ext/operators/llamafile/linear.h b/archive/kt-sft/csrc/ktransformers_ext/operators/llamafile/linear.h new file mode 100644 index 0000000..f8ae7ae --- /dev/null +++ b/archive/kt-sft/csrc/ktransformers_ext/operators/llamafile/linear.h @@ -0,0 +1,59 @@ +/** + * @Description : + * @Author : chenht2022 + * @Date : 2024-07-12 10:07:58 + * @Version : 1.0.0 + * @LastEditors : chenht2022 + * @LastEditTime : 2024-07-25 10:35:00 + * @Copyright (c) 2024 by KVCache.AI, All Rights Reserved. + **/ +#ifndef CPUINFER_OPERATOR_LINEAR_H +#define CPUINFER_OPERATOR_LINEAR_H + +#include +#include +#include +#include +#include + +#include "../../cpu_backend/backend.h" +#include "../../cpu_backend/shared_mem_buffer.h" +#include "conversion.h" +#include "llama.cpp/ggml-impl.h" +#include "llama.cpp/ggml-quants.h" +#include "llama.cpp/ggml.h" +#include "llamafile/sgemm.h" + +struct LinearConfig { + int input_size; + int output_size; + int stride; + int group_max_len; + void* proj; + ggml_type proj_type; + ggml_type hidden_type; + + LinearConfig() {} + + LinearConfig(int input_size, int output_size, int stride, int group_max_len, void* proj, ggml_type proj_type, ggml_type hidden_type) + : input_size(input_size), output_size(output_size), stride(stride), group_max_len(group_max_len), proj(proj), proj_type(proj_type), hidden_type(hidden_type) {} +}; + +class Linear { + public: + Linear(LinearConfig); + ~Linear(); + void warm_up(Backend* backend); + void forward_many(int qlen, const void* input, void* output, Backend* backend); + void forward(int qlen, const void* input, void* output, Backend* backend); + + private: + LinearConfig config_; + void* proj_; // [output_size * input_size ( /32 if quantized)] + + float* input_fp32_; // [group_max_len * input_size] + uint8_t* proj_input_; // [group_max_len * input_size * ggml_type_size(ggml_internal_get_type_traits(proj_type).vec_dot_type) / ggml_blck_size(ggml_internal_get_type_traits(proj_type).vec_dot_type)] + float* proj_output_; // [group_max_len * output_size] +}; + +#endif \ No newline at end of file diff --git a/archive/kt-sft/csrc/ktransformers_ext/operators/llamafile/mlp.cpp b/archive/kt-sft/csrc/ktransformers_ext/operators/llamafile/mlp.cpp new file mode 100644 index 0000000..602fdcb --- /dev/null +++ b/archive/kt-sft/csrc/ktransformers_ext/operators/llamafile/mlp.cpp @@ -0,0 +1,125 @@ +/** + * @Description : + * @Author : chenht2022 + * @Date : 2024-07-16 10:43:18 + * @Version : 1.0.0 + * @LastEditors : kkk1nak0 + * @LastEditTime : 2024-08-15 07:44:38 + * @Copyright (c) 2024 by KVCache.AI, All Rights Reserved. + **/ +#include "mlp.h" + +MLP::MLP(MLPConfig config) { + config_ = config; + gate_proj_ = config_.gate_proj; + up_proj_ = config_.up_proj; + down_proj_ = config_.down_proj; + + std::vector> mem_requests; + mem_requests.push_back({(void**)&input_fp32_, sizeof(float) * config_.group_max_len * config_.hidden_size}); + mem_requests.push_back({(void**)&gate_input_, config_.group_max_len * config_.hidden_size * ggml_type_size(ggml_internal_get_type_traits(config_.gate_type).vec_dot_type) / ggml_blck_size(ggml_internal_get_type_traits(config_.gate_type).vec_dot_type)}); + mem_requests.push_back({(void**)&up_input_, config_.group_max_len * config_.hidden_size * ggml_type_size(ggml_internal_get_type_traits(config_.up_type).vec_dot_type) / ggml_blck_size(ggml_internal_get_type_traits(config_.up_type).vec_dot_type)}); + mem_requests.push_back({(void**)&gate_output_, sizeof(float) * config_.group_max_len * config_.intermediate_size}); + mem_requests.push_back({(void**)&up_output_, sizeof(float) * config_.group_max_len * config_.intermediate_size}); + mem_requests.push_back({(void**)&intermediate_fp32_, sizeof(float) * config_.group_max_len * config_.intermediate_size}); + mem_requests.push_back({(void**)&down_input_, config_.group_max_len * config_.intermediate_size * ggml_type_size(ggml_internal_get_type_traits(config_.down_type).vec_dot_type) / ggml_blck_size(ggml_internal_get_type_traits(config_.down_type).vec_dot_type)}); + mem_requests.push_back({(void**)&down_output_, sizeof(float) * config_.group_max_len * config_.hidden_size}); + shared_mem_buffer.alloc(this, mem_requests); +} + +MLP::~MLP() { + shared_mem_buffer.dealloc(this); +} + +void MLP::warm_up(Backend *backend) { + std::vector input_fp32(config_.hidden_size); + std::vector input(config_.hidden_size * + ggml_type_size(config_.hidden_type) / + ggml_blck_size(config_.hidden_type)); + std::vector output(config_.hidden_size * + ggml_type_size(config_.hidden_type) / + ggml_blck_size(config_.hidden_type)); + for (int i = 0; i < config_.hidden_size; i++) { + input_fp32[i] = 0; + } + from_float(input_fp32.data(), input.data(), config_.hidden_size, config_.hidden_type); + forward_many(1, input.data(), output.data(), backend); +} + +static float act_fn(float x) { return x / (1.0f + expf(-x)); } + +void MLP::forward_many(int qlen, const void* input, void* output, Backend* backend) { + const void* gate_input_ptr; + const void* up_input_ptr; + if (config_.hidden_type == ggml_internal_get_type_traits(config_.gate_type).vec_dot_type && config_.hidden_type == ggml_internal_get_type_traits(config_.up_type).vec_dot_type) { + gate_input_ptr = up_input_ptr = input; + } else { + to_float(input, input_fp32_, qlen * config_.hidden_size, config_.hidden_type); + if (ggml_internal_get_type_traits(config_.gate_type).vec_dot_type == ggml_internal_get_type_traits(config_.up_type).vec_dot_type) { + from_float(input_fp32_, gate_input_, qlen * config_.hidden_size, ggml_internal_get_type_traits(config_.gate_type).vec_dot_type); + gate_input_ptr = up_input_ptr = gate_input_; + } else { + if (config_.hidden_type != ggml_internal_get_type_traits(config_.gate_type).vec_dot_type) { + from_float(input_fp32_, gate_input_, qlen * config_.hidden_size, ggml_internal_get_type_traits(config_.gate_type).vec_dot_type); + gate_input_ptr = gate_input_; + } else { + gate_input_ptr = input; + } + if (config_.hidden_type != ggml_internal_get_type_traits(config_.up_type).vec_dot_type) { + from_float(input_fp32_, up_input_, qlen * config_.hidden_size, ggml_internal_get_type_traits(config_.up_type).vec_dot_type); + up_input_ptr = up_input_; + } else { + up_input_ptr = input; + } + } + } + int nth = config_.intermediate_size / config_.stride; + backend->do_work_stealing_job(nth, nullptr, [&](int task_id) { + int ith = task_id; + void* gate_proj_ptr = (uint8_t*)gate_proj_ + ith * config_.stride * config_.hidden_size * ggml_type_size(config_.gate_type) / ggml_blck_size(config_.gate_type); + float* gate_output_ptr = gate_output_ + ith * config_.stride; + llamafile_sgemm(config_.stride, qlen, config_.hidden_size / ggml_blck_size(config_.gate_type), gate_proj_ptr, config_.hidden_size / ggml_blck_size(config_.gate_type), gate_input_ptr, config_.hidden_size / ggml_blck_size(config_.gate_type), gate_output_ptr, config_.intermediate_size, 0, 1, GGML_TASK_TYPE_COMPUTE, config_.gate_type, ggml_internal_get_type_traits(config_.gate_type).vec_dot_type, GGML_TYPE_F32, GGML_PREC_DEFAULT); + void* up_proj_ptr = (uint8_t*)up_proj_ + ith * config_.stride * config_.hidden_size * ggml_type_size(config_.up_type) / ggml_blck_size(config_.up_type); + float* up_output_ptr = up_output_ + ith * config_.stride; + llamafile_sgemm(config_.stride, qlen, config_.hidden_size / ggml_blck_size(config_.up_type), up_proj_ptr, config_.hidden_size / ggml_blck_size(config_.up_type), up_input_ptr, config_.hidden_size / ggml_blck_size(config_.up_type), up_output_ptr, config_.intermediate_size, 0, 1, GGML_TASK_TYPE_COMPUTE, config_.up_type, ggml_internal_get_type_traits(config_.up_type).vec_dot_type, GGML_TYPE_F32, GGML_PREC_DEFAULT); + for (int i = 0; i < qlen; i++) { + for (int j = ith * config_.stride; j < (ith + 1) * config_.stride; j++) { + intermediate_fp32_[i * config_.intermediate_size + j] = act_fn(gate_output_[i * config_.intermediate_size + j]) * up_output_[i * config_.intermediate_size + j]; + } + if (config_.stride % ggml_blck_size(ggml_internal_get_type_traits(config_.down_type).vec_dot_type) == 0) { + float* intermediate_fp32_ptr = intermediate_fp32_ + i * config_.intermediate_size + ith * config_.stride; + void* down_input_ptr = (uint8_t*)down_input_ + i * config_.intermediate_size * ggml_type_size(ggml_internal_get_type_traits(config_.down_type).vec_dot_type) / ggml_blck_size(ggml_internal_get_type_traits(config_.down_type).vec_dot_type) + ith * config_.stride * ggml_type_size(ggml_internal_get_type_traits(config_.down_type).vec_dot_type) / ggml_blck_size(ggml_internal_get_type_traits(config_.down_type).vec_dot_type); + from_float(intermediate_fp32_ptr, down_input_ptr, config_.stride, ggml_internal_get_type_traits(config_.down_type).vec_dot_type); + } + } + }, nullptr); + if (config_.stride % ggml_blck_size(ggml_internal_get_type_traits(config_.down_type).vec_dot_type) != 0) { + from_float(intermediate_fp32_, down_input_, qlen * config_.intermediate_size, ggml_internal_get_type_traits(config_.down_type).vec_dot_type); + } + nth = config_.hidden_size / config_.stride; + backend->do_work_stealing_job(nth, nullptr, [&](int task_id) { + int ith = task_id; + void* down_proj_ptr = (uint8_t*)down_proj_ + ith * config_.stride * config_.intermediate_size * ggml_type_size(config_.down_type) / ggml_blck_size(config_.down_type); + float* down_output_ptr = down_output_ + ith * config_.stride; + llamafile_sgemm(config_.stride, qlen, config_.intermediate_size / ggml_blck_size(config_.down_type), down_proj_ptr, config_.intermediate_size / ggml_blck_size(config_.down_type), down_input_, config_.intermediate_size / ggml_blck_size(config_.down_type), down_output_ptr, config_.hidden_size, 0, 1, GGML_TASK_TYPE_COMPUTE, config_.down_type, ggml_internal_get_type_traits(config_.down_type).vec_dot_type, GGML_TYPE_F32, GGML_PREC_DEFAULT); + if (config_.stride % ggml_blck_size(config_.hidden_type) == 0) { + for (int i = 0; i < qlen; i++) { + float* output_fp32_ptr = down_output_ + i * config_.hidden_size + ith * config_.stride; + void* output_ptr = (uint8_t*)output + i * config_.hidden_size * ggml_type_size(config_.hidden_type) / ggml_blck_size(config_.hidden_type) + ith * config_.stride * ggml_type_size(config_.hidden_type) / ggml_blck_size(config_.hidden_type); + from_float(output_fp32_ptr, output_ptr, config_.stride, config_.hidden_type); + } + } + }, nullptr); + if (config_.stride % ggml_blck_size(config_.hidden_type) != 0) { + from_float(down_output_, output, qlen * config_.hidden_size, config_.hidden_type); + } +} + +void MLP::forward(int qlen, const void* input, void* output, Backend* backend) { + if (qlen <= 0) { + return; + } + int forward_len = std::min(qlen, config_.group_max_len); + forward_many(forward_len, input, output, backend); + forward(qlen - forward_len, (uint8_t*)input + forward_len * config_.hidden_size * ggml_type_size(config_.hidden_type) / ggml_blck_size(config_.hidden_type), (uint8_t*)output + forward_len * config_.hidden_size * ggml_type_size(config_.hidden_type) / ggml_blck_size(config_.hidden_type), backend); +} \ No newline at end of file diff --git a/archive/kt-sft/csrc/ktransformers_ext/operators/llamafile/mlp.h b/archive/kt-sft/csrc/ktransformers_ext/operators/llamafile/mlp.h new file mode 100644 index 0000000..7e6e5cc --- /dev/null +++ b/archive/kt-sft/csrc/ktransformers_ext/operators/llamafile/mlp.h @@ -0,0 +1,70 @@ +/** + * @Description : + * @Author : chenht2022 + * @Date : 2024-07-12 10:07:58 + * @Version : 1.0.0 + * @LastEditors : chenht2022 + * @LastEditTime : 2024-07-25 10:35:06 + * @Copyright (c) 2024 by KVCache.AI, All Rights Reserved. + **/ +#ifndef CPUINFER_OPERATOR_MLP_H +#define CPUINFER_OPERATOR_MLP_H + +#include +#include +#include +#include +#include + +#include "../../cpu_backend/backend.h" +#include "../../cpu_backend/shared_mem_buffer.h" +#include "conversion.h" +#include "llama.cpp/ggml-impl.h" +#include "llama.cpp/ggml-quants.h" +#include "llama.cpp/ggml.h" +#include "llamafile/sgemm.h" + +struct MLPConfig { + int hidden_size; + int intermediate_size; + int stride; + int group_max_len; + void* gate_proj; + void* up_proj; + void* down_proj; + ggml_type gate_type; + ggml_type up_type; + ggml_type down_type; + ggml_type hidden_type; + + MLPConfig() {} + + MLPConfig(int hidden_size, int intermediate_size, int stride, int group_max_len, void* gate_proj, void* up_proj, void* down_proj, ggml_type gate_type, ggml_type up_type, ggml_type down_type, ggml_type hidden_type) + : hidden_size(hidden_size), intermediate_size(intermediate_size), stride(stride), group_max_len(group_max_len), gate_proj(gate_proj), up_proj(up_proj), down_proj(down_proj), gate_type(gate_type), up_type(up_type), down_type(down_type), hidden_type(hidden_type) {} +}; + +class MLP { + public: + MLP(MLPConfig); + ~MLP(); + void warm_up(Backend* backend); + void forward_many(int qlen, const void* input, void* output, Backend* backend); + void forward(int qlen, const void* input, void* output, Backend* backend); + + private: + MLPConfig config_; + void* gate_proj_; // [intermediate_size * hidden_size ( /32 if quantized)] + void* up_proj_; // [intermediate_size * hidden_size ( /32 if quantized)] + void* down_proj_; // [hidden_size * intermediate_size ( /32 if quantized)] + + float* input_fp32_; // [group_max_len * hidden_size] + uint8_t* gate_input_; // [group_max_len * hidden_size * ggml_type_size(ggml_internal_get_type_traits(gate_type).vec_dot_type) / ggml_blck_size(ggml_internal_get_type_traits(gate_type).vec_dot_type)] + uint8_t* up_input_; // [group_max_len * hidden_size * ggml_type_size(ggml_internal_get_type_traits(up_type).vec_dot_type) / ggml_blck_size(ggml_internal_get_type_traits(up_type).vec_dot_type)] + float* gate_output_; // [group_max_len * intermediate_size] + float* up_output_; // [group_max_len * intermediate_size] + float* intermediate_fp32_; // [group_max_len * intermediate_size] + uint8_t* down_input_; // [group_max_len * intermediate_size * ggml_type_size(ggml_internal_get_type_traits(down_type).vec_dot_type) / ggml_blck_size(ggml_internal_get_type_traits(down_type).vec_dot_type)] + float* down_output_; // [group_max_len * hidden_size] +}; + +#endif \ No newline at end of file diff --git a/archive/kt-sft/csrc/ktransformers_ext/operators/llamafile/moe.cpp b/archive/kt-sft/csrc/ktransformers_ext/operators/llamafile/moe.cpp new file mode 100644 index 0000000..cd42691 --- /dev/null +++ b/archive/kt-sft/csrc/ktransformers_ext/operators/llamafile/moe.cpp @@ -0,0 +1,357 @@ +/** + * @Description : + * @Author : chenht2022 + * @Date : 2024-07-22 02:03:22 + * @Version : 1.0.0 + * @LastEditors : kkk1nak0 + * @LastEditTime : 2024-08-15 07:43:41 + * @Copyright (c) 2024 by KVCache.AI, All Rights Reserved. + **/ +#include "moe.h" +#include +#include + +#ifdef USE_NUMA +#include +#include +#endif + +MOE::MOE(MOEConfig config) { + config_ = config; + gate_proj_ = config_.gate_proj; + up_proj_ = config_.up_proj; + down_proj_ = config_.down_proj; + + #ifdef USE_NUMA + int numa_nodes = numa_num_configured_nodes(); + gate_proj_numa_.resize(numa_nodes); + up_proj_numa_.resize(numa_nodes); + down_proj_numa_.resize(numa_nodes); + size_t exp_inter_hidden_mul_ = (size_t)config.expert_num * config.intermediate_size * config.hidden_size; + for (int i = 0; i < numa_nodes; i++) { + gate_proj_numa_[i] = numa_alloc_onnode(exp_inter_hidden_mul_* ggml_type_size(config.gate_type) / ggml_blck_size(config.gate_type), i); + up_proj_numa_[i] = numa_alloc_onnode(exp_inter_hidden_mul_* ggml_type_size(config.up_type) / ggml_blck_size(config.up_type), i); + down_proj_numa_[i] = numa_alloc_onnode(exp_inter_hidden_mul_* ggml_type_size(config.down_type) / ggml_blck_size(config.down_type), i); + if (!gate_proj_numa_[i]) { + std::cout << "Memory allocation failed for gate_proj_numa_ on node " << i << std::endl; + } + if (!up_proj_numa_[i]) { + std::cout << "Memory allocation failed for up_proj_numa_ on node " << i << std::endl; + } + if (!down_proj_numa_[i]) { + std::cout << "Memory allocation failed for down_proj_numa_ on node " << i << std::endl; + } + memcpy(gate_proj_numa_[i], gate_proj_, exp_inter_hidden_mul_* ggml_type_size(config.gate_type) / ggml_blck_size(config.gate_type)); + memcpy(up_proj_numa_[i], up_proj_, exp_inter_hidden_mul_* ggml_type_size(config.up_type) / ggml_blck_size(config.up_type)); + memcpy(down_proj_numa_[i], down_proj_, exp_inter_hidden_mul_* ggml_type_size(config.down_type) / ggml_blck_size(config.down_type)); + } + #endif + + std::vector> s_mem_requests; + s_mem_requests.push_back({(void**)&s_input_fp32_, sizeof(float) * config_.hidden_size}); + s_mem_requests.push_back({(void**)&s_gate_input_, config_.hidden_size * ggml_type_size(ggml_internal_get_type_traits(config_.gate_type).vec_dot_type) / ggml_blck_size(ggml_internal_get_type_traits(config_.gate_type).vec_dot_type)}); + s_mem_requests.push_back({(void**)&s_up_input_, config_.hidden_size * ggml_type_size(ggml_internal_get_type_traits(config_.up_type).vec_dot_type) / ggml_blck_size(ggml_internal_get_type_traits(config_.up_type).vec_dot_type)}); + s_gate_output_.resize(config_.routed_expert_num); + s_up_output_.resize(config_.routed_expert_num); + s_intermediate_fp32_.resize(config_.routed_expert_num); + s_down_input_.resize(config_.routed_expert_num); + s_down_output_.resize(config_.routed_expert_num); + for (int i = 0; i < config_.routed_expert_num; i++) { + s_mem_requests.push_back({(void**)&s_gate_output_[i], sizeof(float) * config_.intermediate_size}); + s_mem_requests.push_back({(void**)&s_up_output_[i], sizeof(float) * config_.intermediate_size}); + s_mem_requests.push_back({(void**)&s_intermediate_fp32_[i], sizeof(float) * config_.intermediate_size}); + s_mem_requests.push_back({(void**)&s_down_input_[i], config_.intermediate_size * ggml_type_size(ggml_internal_get_type_traits(config_.down_type).vec_dot_type) / ggml_blck_size(ggml_internal_get_type_traits(config_.down_type).vec_dot_type)}); + s_mem_requests.push_back({(void**)&s_down_output_[i], sizeof(float) * config_.hidden_size}); + } + s_mem_requests.push_back({(void**)&s_output_fp32_, sizeof(float) * config_.hidden_size}); + shared_mem_buffer.alloc(this, s_mem_requests); + + std::vector> m_mem_requests; + m_input_fp32_.resize(config_.group_max_len); + m_gate_input_.resize(config_.group_max_len); + m_up_input_.resize(config_.group_max_len); + for (int i = 0; i < config_.group_max_len; i++) { + m_mem_requests.push_back({(void**)&m_input_fp32_[i], sizeof(float) * config_.hidden_size}); + m_mem_requests.push_back({(void**)&m_gate_input_[i], config_.hidden_size * ggml_type_size(ggml_internal_get_type_traits(config_.gate_type).vec_dot_type) / ggml_blck_size(ggml_internal_get_type_traits(config_.gate_type).vec_dot_type)}); + m_mem_requests.push_back({(void**)&m_up_input_[i], config_.hidden_size * ggml_type_size(ggml_internal_get_type_traits(config_.up_type).vec_dot_type) / ggml_blck_size(ggml_internal_get_type_traits(config_.up_type).vec_dot_type)}); + } + m_mem_requests.push_back({(void**)&m_local_gate_input_, config_.routed_expert_num * config_.group_max_len * config_.hidden_size * ggml_type_size(ggml_internal_get_type_traits(config_.gate_type).vec_dot_type) / ggml_blck_size(ggml_internal_get_type_traits(config_.gate_type).vec_dot_type)}); + m_mem_requests.push_back({(void**)&m_local_up_input_, config_.routed_expert_num * config_.group_max_len * config_.hidden_size * ggml_type_size(ggml_internal_get_type_traits(config_.up_type).vec_dot_type) / ggml_blck_size(ggml_internal_get_type_traits(config_.up_type).vec_dot_type)}); + m_mem_requests.push_back({(void**)&m_local_gate_output_, sizeof(float) * config_.routed_expert_num * config_.group_max_len * config_.intermediate_size}); + m_mem_requests.push_back({(void**)&m_local_up_output_, sizeof(float) * config_.routed_expert_num * config_.group_max_len * config_.intermediate_size}); + m_mem_requests.push_back({(void**)&m_local_intermediate_fp32_, sizeof(float) * config_.routed_expert_num * config_.group_max_len * config_.intermediate_size}); + m_mem_requests.push_back({(void**)&m_local_down_input_, config_.routed_expert_num * config_.group_max_len * config_.intermediate_size * ggml_type_size(ggml_internal_get_type_traits(config_.down_type).vec_dot_type) / ggml_blck_size(ggml_internal_get_type_traits(config_.down_type).vec_dot_type)}); + m_mem_requests.push_back({(void**)&m_local_down_output_, sizeof(float) * config_.routed_expert_num * config_.group_max_len * config_.hidden_size}); + m_output_fp32_.resize(config_.group_max_len); + for (int i = 0; i < config_.group_max_len; i++) { + m_mem_requests.push_back({(void**)&m_output_fp32_[i], sizeof(float) * config_.hidden_size}); + } + shared_mem_buffer.alloc(this, m_mem_requests); + + m_local_pos_.resize(config_.group_max_len); + for (int i = 0; i < config_.group_max_len; i++) { + m_local_pos_[i].resize(config_.routed_expert_num); + } + m_local_num_.resize(config_.expert_num); + m_local_gate_input_ptr_.resize(config_.expert_num); + m_local_up_input_ptr_.resize(config_.expert_num); + m_local_gate_output_ptr_.resize(config_.expert_num); + m_local_up_output_ptr_.resize(config_.expert_num); + m_local_intermediate_fp32_ptr_.resize(config_.expert_num); + m_local_down_input_ptr_.resize(config_.expert_num); + m_local_down_output_ptr_.resize(config_.expert_num); +} + +MOE::~MOE() { + shared_mem_buffer.dealloc(this); + + #ifdef USE_NUMA + int numa_nodes = numa_num_configured_nodes(); + for (int i = 0; i < numa_nodes; i++) { + numa_free(gate_proj_numa_[i], config_.expert_num * config_.intermediate_size * config_.hidden_size * ggml_type_size(config_.gate_type) / ggml_blck_size(config_.gate_type)); + numa_free(up_proj_numa_[i], config_.expert_num * config_.intermediate_size * config_.hidden_size * ggml_type_size(config_.up_type) / ggml_blck_size(config_.up_type)); + numa_free(down_proj_numa_[i], config_.expert_num * config_.hidden_size * config_.intermediate_size * ggml_type_size(config_.down_type) / ggml_blck_size(config_.down_type)); + } + #endif +} + +void MOE::warm_up(Backend* backend) { + std::vector input_fp32(config_.hidden_size); + std::vector input(config_.hidden_size * ggml_type_size(config_.hidden_type) / ggml_blck_size(config_.hidden_type)); + std::vector output(config_.hidden_size * ggml_type_size(config_.hidden_type) / ggml_blck_size(config_.hidden_type)); + for (int i = 0; i < config_.hidden_size; i++) { + input_fp32[i] = 0; + } + from_float(input_fp32.data(), input.data(), config_.hidden_size, config_.hidden_type); + for (int i = 0; i < config_.expert_num; i++) { + uint64_t expert_ids = i; + float weights = 0; + forward_one(1, &expert_ids, &weights, input.data(), output.data(), backend); + } +} + +static float act_fn(float x) { + return x / (1.0f + expf(-x)); +} + +void MOE::forward_one(int k, const uint64_t* expert_ids, const float* weights, const void* input, void* output, Backend* backend) { + const void* gate_input_ptr; + const void* up_input_ptr; + if (config_.hidden_type == ggml_internal_get_type_traits(config_.gate_type).vec_dot_type && config_.hidden_type == ggml_internal_get_type_traits(config_.up_type).vec_dot_type) { + gate_input_ptr = up_input_ptr = input; + } else { + to_float(input, s_input_fp32_, config_.hidden_size, config_.hidden_type); + if (ggml_internal_get_type_traits(config_.gate_type).vec_dot_type == ggml_internal_get_type_traits(config_.up_type).vec_dot_type) { + from_float(s_input_fp32_, s_gate_input_, config_.hidden_size, ggml_internal_get_type_traits(config_.gate_type).vec_dot_type); + gate_input_ptr = up_input_ptr = s_gate_input_; + } else { + if (config_.hidden_type != ggml_internal_get_type_traits(config_.gate_type).vec_dot_type) { + from_float(s_input_fp32_, s_gate_input_, config_.hidden_size, ggml_internal_get_type_traits(config_.gate_type).vec_dot_type); + gate_input_ptr = s_gate_input_; + } else { + gate_input_ptr = input; + } + if (config_.hidden_type != ggml_internal_get_type_traits(config_.up_type).vec_dot_type) { + from_float(s_input_fp32_, s_up_input_, config_.hidden_size, ggml_internal_get_type_traits(config_.up_type).vec_dot_type); + up_input_ptr = s_up_input_; + } else { + up_input_ptr = input; + } + } + } + int nth = config_.intermediate_size / config_.stride; + backend->do_work_stealing_job(nth * k, nullptr, [&](int task_id) { + int expert_idx = task_id / nth; + uint64_t expert_id = expert_ids[expert_idx]; + int ith = task_id % nth; + + #ifdef USE_NUMA + void* gate_proj_ptr = (uint8_t*)gate_proj_numa_[Backend::numa_node] + (expert_id * config_.intermediate_size + ith * config_.stride) * config_.hidden_size * ggml_type_size(config_.gate_type) / ggml_blck_size(config_.gate_type); + #else + void* gate_proj_ptr = (uint8_t*)gate_proj_ + (expert_id * config_.intermediate_size + ith * config_.stride) * config_.hidden_size * ggml_type_size(config_.gate_type) / ggml_blck_size(config_.gate_type); + #endif + + float* gate_output_ptr = s_gate_output_[expert_idx] + ith * config_.stride; + llamafile_sgemm(config_.stride, 1, config_.hidden_size / ggml_blck_size(config_.gate_type), gate_proj_ptr, config_.hidden_size / ggml_blck_size(config_.gate_type), gate_input_ptr, config_.hidden_size / ggml_blck_size(config_.gate_type), gate_output_ptr, config_.stride, 0, 1, GGML_TASK_TYPE_COMPUTE, config_.gate_type, ggml_internal_get_type_traits(config_.gate_type).vec_dot_type, GGML_TYPE_F32, GGML_PREC_DEFAULT); + + #ifdef USE_NUMA + void* up_proj_ptr = (uint8_t*)up_proj_numa_[Backend::numa_node] + (expert_id * config_.intermediate_size + ith * config_.stride) * config_.hidden_size * ggml_type_size(config_.up_type) / ggml_blck_size(config_.up_type); + #else + void* up_proj_ptr = (uint8_t*)up_proj_ + (expert_id * config_.intermediate_size + ith * config_.stride) * config_.hidden_size * ggml_type_size(config_.up_type) / ggml_blck_size(config_.up_type); + #endif + + float* up_output_ptr = s_up_output_[expert_idx] + ith * config_.stride; + llamafile_sgemm(config_.stride, 1, config_.hidden_size / ggml_blck_size(config_.up_type), up_proj_ptr, config_.hidden_size / ggml_blck_size(config_.up_type), up_input_ptr, config_.hidden_size / ggml_blck_size(config_.up_type), up_output_ptr, config_.stride, 0, 1, GGML_TASK_TYPE_COMPUTE, config_.up_type, ggml_internal_get_type_traits(config_.up_type).vec_dot_type, GGML_TYPE_F32, GGML_PREC_DEFAULT); + for (int i = ith * config_.stride; i < (ith + 1) * config_.stride; i++) { + s_intermediate_fp32_[expert_idx][i] = act_fn(s_gate_output_[expert_idx][i]) * s_up_output_[expert_idx][i]; + } + if (config_.stride % ggml_blck_size(ggml_internal_get_type_traits(config_.down_type).vec_dot_type) == 0) { + float* intermediate_fp32_ptr = s_intermediate_fp32_[expert_idx] + ith * config_.stride; + void* down_input_ptr = s_down_input_[expert_idx] + ith * config_.stride * ggml_type_size(ggml_internal_get_type_traits(config_.down_type).vec_dot_type) / ggml_blck_size(ggml_internal_get_type_traits(config_.down_type).vec_dot_type); + from_float(intermediate_fp32_ptr, down_input_ptr, config_.stride, ggml_internal_get_type_traits(config_.down_type).vec_dot_type); + } + }, nullptr); + if (config_.stride % ggml_blck_size(ggml_internal_get_type_traits(config_.down_type).vec_dot_type) != 0) { + for (int i = 0; i < k; i++) { + from_float(s_intermediate_fp32_[i], s_down_input_[i], config_.intermediate_size, ggml_internal_get_type_traits(config_.down_type).vec_dot_type); + } + } + nth = config_.hidden_size / config_.stride; + backend->do_work_stealing_job(nth, nullptr, [&](int task_id) { + int ith = task_id; + for (int i = ith * config_.stride; i < (ith + 1) * config_.stride; i++) { + s_output_fp32_[i] = 0; + } + for (int expert_idx = 0; expert_idx < k; expert_idx++) { + uint64_t expert_id = expert_ids[expert_idx]; + + #ifdef USE_NUMA + void* down_proj_ptr = (uint8_t*)down_proj_numa_[Backend::numa_node] + (expert_id * config_.hidden_size + ith * config_.stride) * config_.intermediate_size * ggml_type_size(config_.down_type) / ggml_blck_size(config_.down_type); + #else + void* down_proj_ptr = (uint8_t*)down_proj_ + (expert_id * config_.hidden_size + ith * config_.stride) * config_.intermediate_size * ggml_type_size(config_.down_type) / ggml_blck_size(config_.down_type); + #endif + + float* down_output_ptr = s_down_output_[expert_idx] + ith * config_.stride; + llamafile_sgemm(config_.stride, 1, config_.intermediate_size / ggml_blck_size(config_.down_type), down_proj_ptr, config_.intermediate_size / ggml_blck_size(config_.down_type), s_down_input_[expert_idx], config_.intermediate_size / ggml_blck_size(config_.down_type), down_output_ptr, config_.stride, 0, 1, GGML_TASK_TYPE_COMPUTE, config_.down_type, ggml_internal_get_type_traits(config_.down_type).vec_dot_type, GGML_TYPE_F32, GGML_PREC_DEFAULT); + for (int i = ith * config_.stride; i < (ith + 1) * config_.stride; i++) { + s_output_fp32_[i] += s_down_output_[expert_idx][i] * weights[expert_idx]; + } + } + if (config_.stride % ggml_blck_size(config_.hidden_type) == 0) { + float* output_fp32_ptr = s_output_fp32_ + ith * config_.stride; + void* output_ptr = (uint8_t*)output + ith * config_.stride * ggml_type_size(config_.hidden_type) / ggml_blck_size(config_.hidden_type); + from_float(output_fp32_ptr, output_ptr, config_.stride, config_.hidden_type); + } + }, nullptr); + if (config_.stride % ggml_blck_size(config_.hidden_type) != 0) { + from_float(s_output_fp32_, output, config_.hidden_size, config_.hidden_type); + } +} + +void MOE::forward_many(int qlen, int k, const uint64_t* expert_ids, const float* weights, const void* input, void* output, Backend* backend) { + for (int i = 0; i < config_.expert_num; i++) { + m_local_num_[i] = 0; + } + for (int i = 0; i < qlen; i++) { + for (int j = 0; j < k; j++) { + m_local_pos_[i][j] = m_local_num_[expert_ids[i * k + j]]++; + } + } + uint64_t offset = 0; + for (int i = 0; i < config_.expert_num; i++) { + m_local_gate_input_ptr_[i] = m_local_gate_input_ + offset * config_.hidden_size * ggml_type_size(ggml_internal_get_type_traits(config_.gate_type).vec_dot_type) / ggml_blck_size(ggml_internal_get_type_traits(config_.gate_type).vec_dot_type); + m_local_up_input_ptr_[i] = m_local_up_input_ + offset * config_.hidden_size * ggml_type_size(ggml_internal_get_type_traits(config_.up_type).vec_dot_type) / ggml_blck_size(ggml_internal_get_type_traits(config_.up_type).vec_dot_type); + m_local_gate_output_ptr_[i] = m_local_gate_output_ + offset * config_.intermediate_size; + m_local_up_output_ptr_[i] = m_local_up_output_ + offset * config_.intermediate_size; + m_local_intermediate_fp32_ptr_[i] = m_local_intermediate_fp32_ + offset * config_.intermediate_size; + m_local_down_input_ptr_[i] = m_local_down_input_ + offset * config_.intermediate_size * ggml_type_size(ggml_internal_get_type_traits(config_.down_type).vec_dot_type) / ggml_blck_size(ggml_internal_get_type_traits(config_.down_type).vec_dot_type); + m_local_down_output_ptr_[i] = m_local_down_output_ + offset * config_.hidden_size; + offset += m_local_num_[i]; + } + backend->do_work_stealing_job(qlen, nullptr, [&](int i) { + const void* gate_input_ptr; + const void* up_input_ptr; + if (config_.hidden_type == ggml_internal_get_type_traits(config_.gate_type).vec_dot_type && config_.hidden_type == ggml_internal_get_type_traits(config_.up_type).vec_dot_type) { + gate_input_ptr = up_input_ptr = (uint8_t*)input + i * config_.hidden_size * ggml_type_size(config_.hidden_type) / ggml_blck_size(config_.hidden_type); + } else { + to_float((uint8_t*)input + i * config_.hidden_size * ggml_type_size(config_.hidden_type) / ggml_blck_size(config_.hidden_type), m_input_fp32_[i], config_.hidden_size, config_.hidden_type); + if (ggml_internal_get_type_traits(config_.gate_type).vec_dot_type == ggml_internal_get_type_traits(config_.up_type).vec_dot_type) { + from_float(m_input_fp32_[i], m_gate_input_[i], config_.hidden_size, ggml_internal_get_type_traits(config_.gate_type).vec_dot_type); + gate_input_ptr = up_input_ptr = m_gate_input_[i]; + } else { + if (config_.hidden_type != ggml_internal_get_type_traits(config_.gate_type).vec_dot_type) { + from_float(m_input_fp32_[i], m_gate_input_[i], config_.hidden_size, ggml_internal_get_type_traits(config_.gate_type).vec_dot_type); + gate_input_ptr = m_gate_input_[i]; + } else { + gate_input_ptr = (uint8_t*)input + i * config_.hidden_size * ggml_type_size(config_.hidden_type) / ggml_blck_size(config_.hidden_type); + } + if (config_.hidden_type != ggml_internal_get_type_traits(config_.up_type).vec_dot_type) { + from_float(m_input_fp32_[i], m_up_input_[i], config_.hidden_size, ggml_internal_get_type_traits(config_.up_type).vec_dot_type); + up_input_ptr = m_up_input_[i]; + } else { + up_input_ptr = (uint8_t*)input + i * config_.hidden_size * ggml_type_size(config_.hidden_type) / ggml_blck_size(config_.hidden_type); + } + } + } + for (int j = 0; j < k; j++) { + memcpy(m_local_gate_input_ptr_[expert_ids[i * k + j]] + m_local_pos_[i][j] * config_.hidden_size * ggml_type_size(ggml_internal_get_type_traits(config_.gate_type).vec_dot_type) / ggml_blck_size(ggml_internal_get_type_traits(config_.gate_type).vec_dot_type), gate_input_ptr, config_.hidden_size * ggml_type_size(ggml_internal_get_type_traits(config_.gate_type).vec_dot_type) / ggml_blck_size(ggml_internal_get_type_traits(config_.gate_type).vec_dot_type)); + memcpy(m_local_up_input_ptr_[expert_ids[i * k + j]] + m_local_pos_[i][j] * config_.hidden_size * ggml_type_size(ggml_internal_get_type_traits(config_.up_type).vec_dot_type) / ggml_blck_size(ggml_internal_get_type_traits(config_.up_type).vec_dot_type), up_input_ptr, config_.hidden_size * ggml_type_size(ggml_internal_get_type_traits(config_.up_type).vec_dot_type) / ggml_blck_size(ggml_internal_get_type_traits(config_.up_type).vec_dot_type)); + } + }, nullptr); + int stride = QK_K; + int nth = config_.intermediate_size / stride; + backend->do_work_stealing_job(nth * config_.expert_num, nullptr, [&](int task_id) { + uint64_t expert_idx = task_id / nth; + int ith = task_id % nth; + void* gate_input_ptr = m_local_gate_input_ptr_[expert_idx]; + + #ifdef USE_NUMA + void* gate_proj_ptr = (uint8_t*)gate_proj_numa_[Backend::numa_node] + (expert_idx * config_.intermediate_size + ith * stride) * config_.hidden_size * ggml_type_size(config_.gate_type) / ggml_blck_size(config_.gate_type); + #else + void* gate_proj_ptr = (uint8_t*)gate_proj_ + (expert_idx * config_.intermediate_size + ith * stride) * config_.hidden_size * ggml_type_size(config_.gate_type) / ggml_blck_size(config_.gate_type); + #endif + + float* gate_output_ptr = m_local_gate_output_ptr_[expert_idx] + ith * stride; + llamafile_sgemm(stride, m_local_num_[expert_idx], config_.hidden_size / ggml_blck_size(config_.gate_type), gate_proj_ptr, config_.hidden_size / ggml_blck_size(config_.gate_type), gate_input_ptr, config_.hidden_size / ggml_blck_size(config_.gate_type), gate_output_ptr, config_.intermediate_size, 0, 1, GGML_TASK_TYPE_COMPUTE, config_.gate_type, ggml_internal_get_type_traits(config_.gate_type).vec_dot_type, GGML_TYPE_F32, GGML_PREC_DEFAULT); + void* up_input_ptr = m_local_up_input_ptr_[expert_idx]; + + #ifdef USE_NUMA + void* up_proj_ptr = (uint8_t*)up_proj_numa_[Backend::numa_node] + (expert_idx * config_.intermediate_size + ith * stride) * config_.hidden_size * ggml_type_size(config_.up_type) / ggml_blck_size(config_.up_type); + #else + void* up_proj_ptr = (uint8_t*)up_proj_ + (expert_idx * config_.intermediate_size + ith * stride) * config_.hidden_size * ggml_type_size(config_.up_type) / ggml_blck_size(config_.up_type); + #endif + + float* up_output_ptr = m_local_up_output_ptr_[expert_idx] + ith * stride; + llamafile_sgemm(stride, m_local_num_[expert_idx], config_.hidden_size / ggml_blck_size(config_.up_type), up_proj_ptr, config_.hidden_size / ggml_blck_size(config_.up_type), up_input_ptr, config_.hidden_size / ggml_blck_size(config_.up_type), up_output_ptr, config_.intermediate_size, 0, 1, GGML_TASK_TYPE_COMPUTE, config_.up_type, ggml_internal_get_type_traits(config_.up_type).vec_dot_type, GGML_TYPE_F32, GGML_PREC_DEFAULT); + for (int i = 0; i < m_local_num_[expert_idx]; i++) { + for (int j = ith * stride; j < (ith + 1) * stride; j++) { + m_local_intermediate_fp32_ptr_[expert_idx][i * config_.intermediate_size + j] = act_fn(m_local_gate_output_ptr_[expert_idx][i * config_.intermediate_size + j]) * m_local_up_output_ptr_[expert_idx][i * config_.intermediate_size + j]; + } + float* intermediate_fp32_ptr = m_local_intermediate_fp32_ptr_[expert_idx] + i * config_.intermediate_size + ith * stride; + void* down_input_ptr = m_local_down_input_ptr_[expert_idx] + i * config_.intermediate_size * ggml_type_size(ggml_internal_get_type_traits(config_.down_type).vec_dot_type) / ggml_blck_size(ggml_internal_get_type_traits(config_.down_type).vec_dot_type) + ith * stride * ggml_type_size(ggml_internal_get_type_traits(config_.down_type).vec_dot_type) / ggml_blck_size(ggml_internal_get_type_traits(config_.down_type).vec_dot_type); + from_float(intermediate_fp32_ptr, down_input_ptr, stride, ggml_internal_get_type_traits(config_.down_type).vec_dot_type); + } + }, nullptr); + stride = QK_K; + nth = config_.hidden_size / stride; + backend->do_work_stealing_job(nth * config_.expert_num, nullptr, [&](int task_id) { + uint64_t expert_idx = task_id / nth; + int ith = task_id % nth; + void* down_input_ptr = m_local_down_input_ptr_[expert_idx]; + + #ifdef USE_NUMA + void* down_proj_ptr = (uint8_t*)down_proj_numa_[Backend::numa_node] + (expert_idx * config_.hidden_size + ith * stride) * config_.intermediate_size * ggml_type_size(config_.down_type) / ggml_blck_size(config_.down_type); + #else + void* down_proj_ptr = (uint8_t*)down_proj_ + (expert_idx * config_.hidden_size + ith * stride) * config_.intermediate_size * ggml_type_size(config_.down_type) / ggml_blck_size(config_.down_type); + #endif + + float* down_output_ptr = m_local_down_output_ptr_[expert_idx] + ith * stride; + llamafile_sgemm(stride, m_local_num_[expert_idx], config_.intermediate_size / ggml_blck_size(config_.down_type), down_proj_ptr, config_.intermediate_size / ggml_blck_size(config_.down_type), down_input_ptr, config_.intermediate_size / ggml_blck_size(config_.down_type), down_output_ptr, config_.hidden_size, 0, 1, GGML_TASK_TYPE_COMPUTE, config_.down_type, ggml_internal_get_type_traits(config_.down_type).vec_dot_type, GGML_TYPE_F32, GGML_PREC_DEFAULT); + }, nullptr); + backend->do_work_stealing_job(qlen, nullptr, [&](int i) { + for (int e = 0; e < config_.hidden_size; e++) { + m_output_fp32_[i][e] = 0; + } + for (int j = 0; j < k; j++) { + for (int e = 0; e < config_.hidden_size; e++) { + m_output_fp32_[i][e] += m_local_down_output_ptr_[expert_ids[i * k + j]][m_local_pos_[i][j] * config_.hidden_size + e] * weights[i * k + j]; + } + } + from_float(m_output_fp32_[i], (uint8_t*)output + i * config_.hidden_size * ggml_type_size(config_.hidden_type) / ggml_blck_size(config_.hidden_type), config_.hidden_size, config_.hidden_type); + }, nullptr); +} + +void MOE::forward(int qlen, int k, const uint64_t* expert_ids, const float* weights, const void* input, void* output, int* batch_size_tensor, Backend* backend) { + qlen = batch_size_tensor[0]; + if (qlen < config_.group_min_len) { + for (int i = 0; i < qlen; i++) { + forward_one(k, expert_ids + i * k, weights + i * k, (uint8_t*)input + i * config_.hidden_size * ggml_type_size(config_.hidden_type) / ggml_blck_size(config_.hidden_type), (uint8_t*)output + i * config_.hidden_size * ggml_type_size(config_.hidden_type) / ggml_blck_size(config_.hidden_type), backend); + } + return; + } + int forward_len = std::min(config_.group_max_len, qlen); + forward_many(forward_len, k, expert_ids, weights, input, output, backend); + + batch_size_tensor[0] -= forward_len; + forward(qlen - forward_len, k, expert_ids + forward_len * k, weights + forward_len * k, (uint8_t*)input + forward_len * config_.hidden_size * ggml_type_size(config_.hidden_type) / ggml_blck_size(config_.hidden_type), (uint8_t*)output + forward_len * config_.hidden_size * ggml_type_size(config_.hidden_type) / ggml_blck_size(config_.hidden_type), batch_size_tensor, backend); +} \ No newline at end of file diff --git a/archive/kt-sft/csrc/ktransformers_ext/operators/llamafile/moe.h b/archive/kt-sft/csrc/ktransformers_ext/operators/llamafile/moe.h new file mode 100644 index 0000000..28d7ad3 --- /dev/null +++ b/archive/kt-sft/csrc/ktransformers_ext/operators/llamafile/moe.h @@ -0,0 +1,103 @@ +/** + * @Description : + * @Author : chenht2022 + * @Date : 2024-07-22 02:03:22 + * @Version : 1.0.0 + * @LastEditors : chenht2022 + * @LastEditTime : 2024-07-25 10:35:10 + * @Copyright (c) 2024 by KVCache.AI, All Rights Reserved. + **/ +#ifndef CPUINFER_OPERATOR_MOE_H +#define CPUINFER_OPERATOR_MOE_H + +#include +#include +#include +#include +#include + +#include "../../cpu_backend/backend.h" +#include "../../cpu_backend/shared_mem_buffer.h" +#include "conversion.h" +#include "llama.cpp/ggml-impl.h" +#include "llama.cpp/ggml-quants.h" +#include "llama.cpp/ggml.h" +#include "llamafile/sgemm.h" + +struct MOEConfig { + int expert_num; + int routed_expert_num; + int hidden_size; + int intermediate_size; + int stride; + int group_min_len; + int group_max_len; + void* gate_proj; + void* up_proj; + void* down_proj; + ggml_type gate_type; + ggml_type up_type; + ggml_type down_type; + ggml_type hidden_type; + + MOEConfig() {} + + MOEConfig(int expert_num, int routed_expert_num, int hidden_size, int intermediate_size, int stride, int group_min_len, int group_max_len, void* gate_proj, void* up_proj, void* down_proj, ggml_type gate_type, ggml_type up_type, ggml_type down_type, ggml_type hidden_type) + : expert_num(expert_num), routed_expert_num(routed_expert_num), hidden_size(hidden_size), intermediate_size(intermediate_size), stride(stride), group_min_len(group_min_len), group_max_len(group_max_len), gate_proj(gate_proj), up_proj(up_proj), down_proj(down_proj), gate_type(gate_type), up_type(up_type), down_type(down_type), hidden_type(hidden_type) {} +}; + +class MOE { + public: + MOE(MOEConfig); + ~MOE(); + void warm_up(Backend* backend); + void forward_one(int k, const uint64_t* expert_ids, const float* weights, const void* input, void* output, Backend* backend); + void forward_many(int qlen, int k, const uint64_t* expert_ids, const float* weights, const void* input, void* output, Backend* backend); + void forward(int qlen, int k, const uint64_t* expert_ids, const float* weights, const void* input, void* output, int* batch_size_tensor, Backend* backend); + + private: + MOEConfig config_; + void* gate_proj_; // [expert_num * intermediate_size * hidden_size ( /32 if quantized)] + void* up_proj_; // [expert_num * intermediate_size * hidden_size ( /32 if quantized)] + void* down_proj_; // [expert_num * hidden_size * intermediate_size ( /32 if quantized)] + + #ifdef USE_NUMA + std::vector gate_proj_numa_; // [numa_num, expert_num * intermediate_size * hidden_size ( /32 if quantized)] + std::vector up_proj_numa_; // [numa_num, expert_num * intermediate_size * hidden_size ( /32 if quantized)] + std::vector down_proj_numa_; // [numa_num, expert_num * hidden_size * intermediate_size ( /32 if quantized)] + #endif + + float* s_input_fp32_; // [hidden_size] + uint8_t* s_gate_input_; // [hidden_size * ggml_type_size(ggml_internal_get_type_traits(gate_type).vec_dot_type) / ggml_blck_size(ggml_internal_get_type_traits(gate_type).vec_dot_type)] + uint8_t* s_up_input_; // [hidden_size * ggml_type_size(ggml_internal_get_type_traits(up_type).vec_dot_type) / ggml_blck_size(ggml_internal_get_type_traits(up_type).vec_dot_type)] + std::vector s_gate_output_; // [routed_expert_num, intermediate_size] + std::vector s_up_output_; // [routed_expert_num, intermediate_size] + std::vector s_intermediate_fp32_; // [routed_expert_num, intermediate_size] + std::vector s_down_input_; // [routed_expert_num, intermediate_size * ggml_type_size(ggml_internal_get_type_traits(down_type).vec_dot_type) / ggml_blck_size(ggml_internal_get_type_traits(down_type).vec_dot_type)] + std::vector s_down_output_; // [routed_expert_num, hidden_size] + float* s_output_fp32_; // [hidden_size] + + std::vector m_input_fp32_; // [group_max_len, hidden_size] + std::vector m_gate_input_; // [group_max_len, hidden_size * ggml_type_size(ggml_internal_get_type_traits(gate_type).vec_dot_type) / ggml_blck_size(ggml_internal_get_type_traits(gate_type).vec_dot_type)] + std::vector m_up_input_; // [group_max_len, hidden_size * ggml_type_size(ggml_internal_get_type_traits(up_type).vec_dot_type) / ggml_blck_size(ggml_internal_get_type_traits(up_type).vec_dot_type)] + uint8_t* m_local_gate_input_; // [routed_expert_num * group_max_len * hidden_size * ggml_type_size(ggml_internal_get_type_traits(gate_type).vec_dot_type) / ggml_blck_size(ggml_internal_get_type_traits(gate_type).vec_dot_type)] + uint8_t* m_local_up_input_; // [routed_expert_num * group_max_len * hidden_size * ggml_type_size(ggml_internal_get_type_traits(up_type).vec_dot_type) / ggml_blck_size(ggml_internal_get_type_traits(up_type).vec_dot_type)] + float* m_local_gate_output_; // [routed_expert_num * group_max_len * intermediate_size] + float* m_local_up_output_; // [routed_expert_num * group_max_len * intermediate_size] + float* m_local_intermediate_fp32_; // [routed_expert_num * group_max_len * intermediate_size] + uint8_t* m_local_down_input_; // [routed_expert_num * group_max_len * intermediate_size * ggml_type_size(ggml_internal_get_type_traits(down_type).vec_dot_type) / ggml_blck_size(ggml_internal_get_type_traits(down_type).vec_dot_type)] + float* m_local_down_output_; // [routed_expert_num * group_max_len * hidden_size] + std::vector m_output_fp32_; // [group_max_len, hidden_size] + + std::vector> m_local_pos_; // [group_max_len, routed_expert_num] + std::vector m_local_num_; // [expert_num] + std::vector m_local_gate_input_ptr_; // [expert_num] + std::vector m_local_up_input_ptr_; // [expert_num] + std::vector m_local_gate_output_ptr_; // [expert_num] + std::vector m_local_up_output_ptr_; // [expert_num] + std::vector m_local_intermediate_fp32_ptr_; // [expert_num] + std::vector m_local_down_input_ptr_; // [expert_num] + std::vector m_local_down_output_ptr_; // [expert_num] +}; + +#endif \ No newline at end of file diff --git a/archive/kt-sft/csrc/ktransformers_ext/operators/llamafile/sft_moe.cpp b/archive/kt-sft/csrc/ktransformers_ext/operators/llamafile/sft_moe.cpp new file mode 100644 index 0000000..88fa657 --- /dev/null +++ b/archive/kt-sft/csrc/ktransformers_ext/operators/llamafile/sft_moe.cpp @@ -0,0 +1,737 @@ +/** + * @Description : + * @Author : chenht2022 + * @Date : 2024-07-22 02:03:22 + * @Version : 1.0.0 + * @LastEditors : kkk1nak0 + * @LastEditTime : 2024-08-15 07:43:41 + * @Copyright (c) 2024 by KVCache.AI, All Rights Reserved. + **/ +#include "sft_moe.h" +#include +#include +#include +#include + +#ifdef USE_NUMA +#include +#include +#endif + +SFT_MOE::SFT_MOE(SFT_MOEConfig config) { + config_ = config; + gate_proj_ = config_.gate_proj; + up_proj_ = config_.up_proj; + down_proj_ = config_.down_proj; + + #ifdef USE_NUMA + int numa_nodes = numa_num_configured_nodes(); + gate_proj_numa_.resize(numa_nodes); + up_proj_numa_.resize(numa_nodes); + down_proj_numa_.resize(numa_nodes); + size_t exp_inter_hidden_mul_ = (size_t)config.expert_num * config.intermediate_size * config.hidden_size; + for (int i = 0; i < numa_nodes; i++) { + gate_proj_numa_[i] = numa_alloc_onnode(exp_inter_hidden_mul_* ggml_type_size(config.gate_type) / ggml_blck_size(config.gate_type), i); + up_proj_numa_[i] = numa_alloc_onnode(exp_inter_hidden_mul_* ggml_type_size(config.up_type) / ggml_blck_size(config.up_type), i); + down_proj_numa_[i] = numa_alloc_onnode(exp_inter_hidden_mul_* ggml_type_size(config.down_type) / ggml_blck_size(config.down_type), i); + if (!gate_proj_numa_[i]) { + std::cout << "Memory allocation failed for gate_proj_numa_ on node " << i << std::endl; + } + if (!up_proj_numa_[i]) { + std::cout << "Memory allocation failed for up_proj_numa_ on node " << i << std::endl; + } + if (!down_proj_numa_[i]) { + std::cout << "Memory allocation failed for down_proj_numa_ on node " << i << std::endl; + } + memcpy(gate_proj_numa_[i], gate_proj_, exp_inter_hidden_mul_* ggml_type_size(config.gate_type) / ggml_blck_size(config.gate_type)); + memcpy(up_proj_numa_[i], up_proj_, exp_inter_hidden_mul_* ggml_type_size(config.up_type) / ggml_blck_size(config.up_type)); + memcpy(down_proj_numa_[i], down_proj_, exp_inter_hidden_mul_* ggml_type_size(config.down_type) / ggml_blck_size(config.down_type)); + } + #endif + + std::vector> s_mem_requests; + s_mem_requests.push_back({(void**)&gate_proj_t_, config_.expert_num * config_.hidden_size * config_.intermediate_size * ggml_type_size(config_.grad_type)}); + s_mem_requests.push_back({(void**)&up_proj_t_, config_.expert_num * config_.hidden_size * config_.intermediate_size * ggml_type_size(config_.grad_type)}); + s_mem_requests.push_back({(void**)&down_proj_t_, config_.expert_num * config_.hidden_size * config_.intermediate_size * ggml_type_size(config_.grad_type)}); + s_mem_requests.push_back({(void**)&transpose_buffer_fp32_, config_.expert_num * config_.intermediate_size * config_.hidden_size * sizeof(float)}); + s_mem_requests.push_back({(void**)&transpose_buffer_, config_.expert_num * config_.intermediate_size * config_.hidden_size * ggml_type_size(config_.grad_type)}); + + s_mem_requests.push_back({(void**)&s_input_fp32_, sizeof(float) * config_.hidden_size}); + s_mem_requests.push_back({(void**)&s_gate_input_, config_.hidden_size * ggml_type_size(ggml_internal_get_type_traits(config_.gate_type).vec_dot_type) / ggml_blck_size(ggml_internal_get_type_traits(config_.gate_type).vec_dot_type)}); + s_mem_requests.push_back({(void**)&s_up_input_, config_.hidden_size * ggml_type_size(ggml_internal_get_type_traits(config_.up_type).vec_dot_type) / ggml_blck_size(ggml_internal_get_type_traits(config_.up_type).vec_dot_type)}); + s_gate_output_.resize(config_.routed_expert_num); + s_up_output_.resize(config_.routed_expert_num); + s_intermediate_fp32_.resize(config_.routed_expert_num); + s_down_input_.resize(config_.routed_expert_num); + s_down_output_.resize(config_.routed_expert_num); + for (int i = 0; i < config_.routed_expert_num; i++) { + s_mem_requests.push_back({(void**)&s_gate_output_[i], sizeof(float) * config_.intermediate_size}); + s_mem_requests.push_back({(void**)&s_up_output_[i], sizeof(float) * config_.intermediate_size}); + s_mem_requests.push_back({(void**)&s_intermediate_fp32_[i], sizeof(float) * config_.intermediate_size}); + s_mem_requests.push_back({(void**)&s_down_input_[i], config_.intermediate_size * ggml_type_size(ggml_internal_get_type_traits(config_.down_type).vec_dot_type) / ggml_blck_size(ggml_internal_get_type_traits(config_.down_type).vec_dot_type)}); + s_mem_requests.push_back({(void**)&s_down_output_[i], sizeof(float) * config_.hidden_size}); + } + s_mem_requests.push_back({(void**)&s_output_fp32_, sizeof(float) * config_.hidden_size}); + + s_down_input_grad_.resize(config_.routed_expert_num); + s_gate_output_grad_fp32_.resize(config_.routed_expert_num); + s_up_output_grad_fp32_.resize(config_.routed_expert_num); + s_gate_output_grad_.resize(config_.routed_expert_num); + s_up_output_grad_.resize(config_.routed_expert_num); + s_gate_input_grad_.resize(config_.routed_expert_num); + s_up_input_grad_.resize(config_.routed_expert_num); + for (int i = 0; i < config_.routed_expert_num; i++) { + s_mem_requests.push_back({(void**)&s_down_input_grad_[i], config_.intermediate_size * sizeof(float)}); + s_mem_requests.push_back({(void**)&s_gate_output_grad_fp32_[i], config_.intermediate_size * sizeof(float)}); + s_mem_requests.push_back({(void**)&s_up_output_grad_fp32_[i], config_.intermediate_size * sizeof(float)}); + s_mem_requests.push_back({(void**)&s_gate_output_grad_[i], config_.intermediate_size * ggml_type_size(config_.grad_type)}); + s_mem_requests.push_back({(void**)&s_up_output_grad_[i], config_.intermediate_size * ggml_type_size(config_.grad_type)}); + s_mem_requests.push_back({(void**)&s_gate_input_grad_[i], config_.hidden_size * sizeof(float)}); + s_mem_requests.push_back({(void**)&s_up_input_grad_[i], config_.hidden_size * sizeof(float)}); + } + s_mem_requests.push_back({(void**)&s_input_grad_fp32_, config_.hidden_size * sizeof(float)}); + + shared_mem_buffer.alloc(this, s_mem_requests); + + std::vector> m_mem_requests; + m_input_fp32_.resize(config_.group_max_len); + m_gate_input_.resize(config_.group_max_len); + m_up_input_.resize(config_.group_max_len); + for (int i = 0; i < config_.group_max_len; i++) { + m_mem_requests.push_back({(void**)&m_input_fp32_[i], sizeof(float) * config_.hidden_size}); + m_mem_requests.push_back({(void**)&m_gate_input_[i], config_.hidden_size * ggml_type_size(ggml_internal_get_type_traits(config_.gate_type).vec_dot_type) / ggml_blck_size(ggml_internal_get_type_traits(config_.gate_type).vec_dot_type)}); + m_mem_requests.push_back({(void**)&m_up_input_[i], config_.hidden_size * ggml_type_size(ggml_internal_get_type_traits(config_.up_type).vec_dot_type) / ggml_blck_size(ggml_internal_get_type_traits(config_.up_type).vec_dot_type)}); + } + m_mem_requests.push_back({(void**)&m_local_gate_input_, config_.routed_expert_num * config_.group_max_len * config_.hidden_size * ggml_type_size(ggml_internal_get_type_traits(config_.gate_type).vec_dot_type) / ggml_blck_size(ggml_internal_get_type_traits(config_.gate_type).vec_dot_type)}); + m_mem_requests.push_back({(void**)&m_local_up_input_, config_.routed_expert_num * config_.group_max_len * config_.hidden_size * ggml_type_size(ggml_internal_get_type_traits(config_.up_type).vec_dot_type) / ggml_blck_size(ggml_internal_get_type_traits(config_.up_type).vec_dot_type)}); + m_mem_requests.push_back({(void**)&m_local_gate_output_, sizeof(float) * config_.routed_expert_num * config_.group_max_len * config_.intermediate_size}); + m_mem_requests.push_back({(void**)&m_local_up_output_, sizeof(float) * config_.routed_expert_num * config_.group_max_len * config_.intermediate_size}); + m_mem_requests.push_back({(void**)&m_local_intermediate_fp32_, sizeof(float) * config_.routed_expert_num * config_.group_max_len * config_.intermediate_size}); + m_mem_requests.push_back({(void**)&m_local_down_input_, config_.routed_expert_num * config_.group_max_len * config_.intermediate_size * ggml_type_size(ggml_internal_get_type_traits(config_.down_type).vec_dot_type) / ggml_blck_size(ggml_internal_get_type_traits(config_.down_type).vec_dot_type)}); + m_mem_requests.push_back({(void**)&m_local_down_output_, sizeof(float) * config_.routed_expert_num * config_.group_max_len * config_.hidden_size}); + m_output_fp32_.resize(config_.group_max_len); + for (int i = 0; i < config_.group_max_len; i++) { + m_mem_requests.push_back({(void**)&m_output_fp32_[i], sizeof(float) * config_.hidden_size}); + } + + m_mem_requests.push_back({(void**)&m_local_down_output_grad_, config_.routed_expert_num * config_.group_max_len * config_.hidden_size * ggml_type_size(config_.grad_type)}); + m_mem_requests.push_back({(void**)&m_local_down_input_grad_, sizeof(float) * config_.routed_expert_num * config_.group_max_len * config_.intermediate_size}); + m_mem_requests.push_back({(void**)&m_local_gate_output_grad_fp32_, sizeof(float) * config_.routed_expert_num * config_.group_max_len * config_.intermediate_size}); + m_mem_requests.push_back({(void**)&m_local_up_output_grad_fp32_, sizeof(float) * config_.routed_expert_num * config_.group_max_len * config_.intermediate_size}); + m_mem_requests.push_back({(void**)&m_local_gate_output_grad_, config_.routed_expert_num * config_.group_max_len * config_.intermediate_size * ggml_type_size(config_.grad_type)}); + m_mem_requests.push_back({(void**)&m_local_up_output_grad_, config_.routed_expert_num * config_.group_max_len * config_.intermediate_size * ggml_type_size(config_.grad_type)}); + m_mem_requests.push_back({(void**)&m_local_gate_input_grad_, sizeof(float) * config_.routed_expert_num * config_.group_max_len * config_.hidden_size}); + m_mem_requests.push_back({(void**)&m_local_up_input_grad_, sizeof(float) * config_.routed_expert_num * config_.group_max_len * config_.hidden_size}); + m_mem_requests.push_back({(void**)&m_local_token_indices_, sizeof(int) * config_.routed_expert_num * config_.group_max_len}); + m_mem_requests.push_back({(void**)&m_local_expert_positions_, sizeof(int) * config_.routed_expert_num * config_.group_max_len}); + m_grad_input_fp32_.resize(config_.group_max_len); + for (int i = 0; i < config_.group_max_len; i++) { + m_mem_requests.push_back({(void**)&m_grad_input_fp32_[i], sizeof(float) * config_.hidden_size}); + } + + shared_mem_buffer.alloc(this, m_mem_requests); + + m_local_pos_.resize(config_.group_max_len); + for (int i = 0; i < config_.group_max_len; i++) { + m_local_pos_[i].resize(config_.routed_expert_num); + } + m_local_num_.resize(config_.expert_num); + m_local_gate_input_ptr_.resize(config_.expert_num); + m_local_up_input_ptr_.resize(config_.expert_num); + m_local_gate_output_ptr_.resize(config_.expert_num); + m_local_up_output_ptr_.resize(config_.expert_num); + m_local_intermediate_fp32_ptr_.resize(config_.expert_num); + m_local_down_input_ptr_.resize(config_.expert_num); + m_local_down_output_ptr_.resize(config_.expert_num); + + // backward_many 专用指针数组初始化 + m_local_down_output_grad_ptr_.resize(config_.expert_num); + m_local_down_input_grad_ptr_.resize(config_.expert_num); + m_local_gate_output_grad_fp32_ptr_.resize(config_.expert_num); + m_local_up_output_grad_fp32_ptr_.resize(config_.expert_num); + m_local_gate_output_grad_ptr_.resize(config_.expert_num); + m_local_up_output_grad_ptr_.resize(config_.expert_num); + m_local_gate_input_grad_ptr_.resize(config_.expert_num); + m_local_up_input_grad_ptr_.resize(config_.expert_num); + + // fwd_cache访问映射指针数组初始化 + m_local_token_indices_ptr_.resize(config_.expert_num); + m_local_expert_positions_ptr_.resize(config_.expert_num); +} + +SFT_MOE::~SFT_MOE() { + shared_mem_buffer.dealloc(this); + + #ifdef USE_NUMA + int numa_nodes = numa_num_configured_nodes(); + for (int i = 0; i < numa_nodes; i++) { + numa_free(gate_proj_numa_[i], config_.expert_num * config_.intermediate_size * config_.hidden_size * ggml_type_size(config_.gate_type) / ggml_blck_size(config_.gate_type)); + numa_free(up_proj_numa_[i], config_.expert_num * config_.intermediate_size * config_.hidden_size * ggml_type_size(config_.up_type) / ggml_blck_size(config_.up_type)); + numa_free(down_proj_numa_[i], config_.expert_num * config_.hidden_size * config_.intermediate_size * ggml_type_size(config_.down_type) / ggml_blck_size(config_.down_type)); + } + #endif +} + +void SFT_MOE::warm_up(Backend* backend) { + std::vector input_fp32(config_.hidden_size); + std::vector input(config_.hidden_size * ggml_type_size(config_.hidden_type) / ggml_blck_size(config_.hidden_type)); + std::vector output(config_.hidden_size * ggml_type_size(config_.hidden_type) / ggml_blck_size(config_.hidden_type)); + for (int i = 0; i < config_.hidden_size; i++) { + input_fp32[i] = 0; + } + from_float(input_fp32.data(), input.data(), config_.hidden_size, config_.hidden_type); + /* ---------- 仅用于占位的 ForwardCache ---------- */ + SFT_MoEForwardCache dummy_cache; // 内容无用，只为满足接口 + dummy_cache.init(/*k=*/1, config_.intermediate_size); + for (int i = 0; i < config_.expert_num; i++) { + uint64_t expert_ids = i; + float weights = 0; + forward_one(1, &expert_ids, &weights, input.data(), output.data(), backend, &dummy_cache); + } +} + +static float act_fn(float x) { + return x / (1.0f + expf(-x)); +} + +void SFT_MOE::ensure_fwd_cache(int qlen, int k) +{ + // if ((int)fw_cache_.size() < qlen) + // fw_cache_.resize(qlen); + // /* 只在扩容的那部分做 init，防止重复开辟 */ + // for (int i = 0; i < qlen; ++i) + // fw_cache_[i].init(k, config_.intermediate_size); + + int old_sz = fw_cache_.size(); + if (old_sz < qlen) + { + fw_cache_.resize(qlen); + for (int i = old_sz; i < qlen; ++i) // 仅初始化新增元素 + fw_cache_[i].init(k, config_.intermediate_size); + } + + + // if ((int)fw_cache_.size() < qlen) + // fw_cache_.resize(qlen); + + // for (int i = 0; i < qlen; ++i) // 每轮都 init + // fw_cache_[i].init(k, config_.intermediate_size); // 但 无重 alloc + +} + +SFT_MoEForwardCache* SFT_MOE::fwd_cache_ptr() +{ + return fw_cache_.empty() ? nullptr : fw_cache_.data(); +} + +void SFT_MOE::forward_one(int k, const uint64_t* expert_ids, const float* weights, const void* input, void* output, Backend* backend, SFT_MoEForwardCache* fwd_cache) { + const void* gate_input_ptr; + const void* up_input_ptr; + if (config_.hidden_type == ggml_internal_get_type_traits(config_.gate_type).vec_dot_type && config_.hidden_type == ggml_internal_get_type_traits(config_.up_type).vec_dot_type) { + gate_input_ptr = up_input_ptr = input; + } else { + to_float(input, s_input_fp32_, config_.hidden_size, config_.hidden_type); + if (ggml_internal_get_type_traits(config_.gate_type).vec_dot_type == ggml_internal_get_type_traits(config_.up_type).vec_dot_type) { + from_float(s_input_fp32_, s_gate_input_, config_.hidden_size, ggml_internal_get_type_traits(config_.gate_type).vec_dot_type); + gate_input_ptr = up_input_ptr = s_gate_input_; + } else { + if (config_.hidden_type != ggml_internal_get_type_traits(config_.gate_type).vec_dot_type) { + from_float(s_input_fp32_, s_gate_input_, config_.hidden_size, ggml_internal_get_type_traits(config_.gate_type).vec_dot_type); + gate_input_ptr = s_gate_input_; + } else { + gate_input_ptr = input; + } + if (config_.hidden_type != ggml_internal_get_type_traits(config_.up_type).vec_dot_type) { + from_float(s_input_fp32_, s_up_input_, config_.hidden_size, ggml_internal_get_type_traits(config_.up_type).vec_dot_type); + up_input_ptr = s_up_input_; + } else { + up_input_ptr = input; + } + } + } + int nth = config_.intermediate_size / config_.stride; + backend->do_work_stealing_job(nth * k, nullptr, [&](int task_id) { + int expert_idx = task_id / nth; + uint64_t expert_id = expert_ids[expert_idx]; + int ith = task_id % nth; + + #ifdef USE_NUMA + void* gate_proj_ptr = (uint8_t*)gate_proj_numa_[Backend::numa_node] + (expert_id * config_.intermediate_size + ith * config_.stride) * config_.hidden_size * ggml_type_size(config_.gate_type) / ggml_blck_size(config_.gate_type); + #else + void* gate_proj_ptr = (uint8_t*)gate_proj_ + (expert_id * config_.intermediate_size + ith * config_.stride) * config_.hidden_size * ggml_type_size(config_.gate_type) / ggml_blck_size(config_.gate_type); + #endif + + float* gate_output_ptr = s_gate_output_[expert_idx] + ith * config_.stride; + llamafile_sgemm(config_.stride, 1, config_.hidden_size / ggml_blck_size(config_.gate_type), gate_proj_ptr, config_.hidden_size / ggml_blck_size(config_.gate_type), gate_input_ptr, config_.hidden_size / ggml_blck_size(config_.gate_type), gate_output_ptr, config_.stride, 0, 1, GGML_TASK_TYPE_COMPUTE, config_.gate_type, ggml_internal_get_type_traits(config_.gate_type).vec_dot_type, GGML_TYPE_F32, GGML_PREC_DEFAULT); + + #ifdef USE_NUMA + void* up_proj_ptr = (uint8_t*)up_proj_numa_[Backend::numa_node] + (expert_id * config_.intermediate_size + ith * config_.stride) * config_.hidden_size * ggml_type_size(config_.up_type) / ggml_blck_size(config_.up_type); + #else + void* up_proj_ptr = (uint8_t*)up_proj_ + (expert_id * config_.intermediate_size + ith * config_.stride) * config_.hidden_size * ggml_type_size(config_.up_type) / ggml_blck_size(config_.up_type); + #endif + + float* up_output_ptr = s_up_output_[expert_idx] + ith * config_.stride; + llamafile_sgemm(config_.stride, 1, config_.hidden_size / ggml_blck_size(config_.up_type), up_proj_ptr, config_.hidden_size / ggml_blck_size(config_.up_type), up_input_ptr, config_.hidden_size / ggml_blck_size(config_.up_type), up_output_ptr, config_.stride, 0, 1, GGML_TASK_TYPE_COMPUTE, config_.up_type, ggml_internal_get_type_traits(config_.up_type).vec_dot_type, GGML_TYPE_F32, GGML_PREC_DEFAULT); + for (int i = ith * config_.stride; i < (ith + 1) * config_.stride; i++) { + s_intermediate_fp32_[expert_idx][i] = act_fn(s_gate_output_[expert_idx][i]) * s_up_output_[expert_idx][i]; + } + if (config_.stride % ggml_blck_size(ggml_internal_get_type_traits(config_.down_type).vec_dot_type) == 0) { + float* intermediate_fp32_ptr = s_intermediate_fp32_[expert_idx] + ith * config_.stride; + void* down_input_ptr = s_down_input_[expert_idx] + ith * config_.stride * ggml_type_size(ggml_internal_get_type_traits(config_.down_type).vec_dot_type) / ggml_blck_size(ggml_internal_get_type_traits(config_.down_type).vec_dot_type); + from_float(intermediate_fp32_ptr, down_input_ptr, config_.stride, ggml_internal_get_type_traits(config_.down_type).vec_dot_type); + } + }, nullptr); + if (config_.stride % ggml_blck_size(ggml_internal_get_type_traits(config_.down_type).vec_dot_type) != 0) { + for (int i = 0; i < k; i++) { + from_float(s_intermediate_fp32_[i], s_down_input_[i], config_.intermediate_size, ggml_internal_get_type_traits(config_.down_type).vec_dot_type); + } + } + nth = config_.hidden_size / config_.stride; + backend->do_work_stealing_job(nth, nullptr, [&](int task_id) { + int ith = task_id; + for (int i = ith * config_.stride; i < (ith + 1) * config_.stride; i++) { + s_output_fp32_[i] = 0; + } + for (int expert_idx = 0; expert_idx < k; expert_idx++) { + uint64_t expert_id = expert_ids[expert_idx]; + + #ifdef USE_NUMA + void* down_proj_ptr = (uint8_t*)down_proj_numa_[Backend::numa_node] + (expert_id * config_.hidden_size + ith * config_.stride) * config_.intermediate_size * ggml_type_size(config_.down_type) / ggml_blck_size(config_.down_type); + #else + void* down_proj_ptr = (uint8_t*)down_proj_ + (expert_id * config_.hidden_size + ith * config_.stride) * config_.intermediate_size * ggml_type_size(config_.down_type) / ggml_blck_size(config_.down_type); + #endif + + float* down_output_ptr = s_down_output_[expert_idx] + ith * config_.stride; + llamafile_sgemm(config_.stride, 1, config_.intermediate_size / ggml_blck_size(config_.down_type), down_proj_ptr, config_.intermediate_size / ggml_blck_size(config_.down_type), s_down_input_[expert_idx], config_.intermediate_size / ggml_blck_size(config_.down_type), down_output_ptr, config_.stride, 0, 1, GGML_TASK_TYPE_COMPUTE, config_.down_type, ggml_internal_get_type_traits(config_.down_type).vec_dot_type, GGML_TYPE_F32, GGML_PREC_DEFAULT); + for (int i = ith * config_.stride; i < (ith + 1) * config_.stride; i++) { + s_output_fp32_[i] += s_down_output_[expert_idx][i] * weights[expert_idx]; + } + } + if (config_.stride % ggml_blck_size(config_.hidden_type) == 0) { + float* output_fp32_ptr = s_output_fp32_ + ith * config_.stride; + void* output_ptr = (uint8_t*)output + ith * config_.stride * ggml_type_size(config_.hidden_type) / ggml_blck_size(config_.hidden_type); + from_float(output_fp32_ptr, output_ptr, config_.stride, config_.hidden_type); + } + }, nullptr); + + for (int e = 0; e < k; ++e) { + // gate_output_: float[inter_size] per expert + std::memcpy(fwd_cache->gate_u[e].data(), + s_gate_output_[e], + sizeof(float) * config_.intermediate_size); + + std::memcpy(fwd_cache->up_v[e].data(), + s_up_output_[e], + sizeof(float) * config_.intermediate_size); + + // 可选保存 z + // std::memcpy(fwd_cache->z[e].data(), + // s_intermediate_fp32_[e], + // sizeof(float) * config_.intermediate_size); + } +} + +void SFT_MOE::forward_many(int qlen, int k, const uint64_t* expert_ids, const float* weights, const void* input, void* output, Backend* backend, SFT_MoEForwardCache* fwd_cache) { + for (int i = 0; i < config_.expert_num; i++) { + m_local_num_[i] = 0; + } + for (int i = 0; i < qlen; i++) { + for (int j = 0; j < k; j++) { + m_local_pos_[i][j] = m_local_num_[expert_ids[i * k + j]]++; + } + } + uint64_t offset = 0; + for (int i = 0; i < config_.expert_num; i++) { + m_local_gate_input_ptr_[i] = m_local_gate_input_ + offset * config_.hidden_size * ggml_type_size(ggml_internal_get_type_traits(config_.gate_type).vec_dot_type) / ggml_blck_size(ggml_internal_get_type_traits(config_.gate_type).vec_dot_type); + m_local_up_input_ptr_[i] = m_local_up_input_ + offset * config_.hidden_size * ggml_type_size(ggml_internal_get_type_traits(config_.up_type).vec_dot_type) / ggml_blck_size(ggml_internal_get_type_traits(config_.up_type).vec_dot_type); + m_local_gate_output_ptr_[i] = m_local_gate_output_ + offset * config_.intermediate_size; + m_local_up_output_ptr_[i] = m_local_up_output_ + offset * config_.intermediate_size; + m_local_intermediate_fp32_ptr_[i] = m_local_intermediate_fp32_ + offset * config_.intermediate_size; + m_local_down_input_ptr_[i] = m_local_down_input_ + offset * config_.intermediate_size * ggml_type_size(ggml_internal_get_type_traits(config_.down_type).vec_dot_type) / ggml_blck_size(ggml_internal_get_type_traits(config_.down_type).vec_dot_type); + m_local_down_output_ptr_[i] = m_local_down_output_ + offset * config_.hidden_size; + offset += m_local_num_[i]; + } + backend->do_work_stealing_job(qlen, nullptr, [&](int i) { + const void* gate_input_ptr; + const void* up_input_ptr; + if (config_.hidden_type == ggml_internal_get_type_traits(config_.gate_type).vec_dot_type && config_.hidden_type == ggml_internal_get_type_traits(config_.up_type).vec_dot_type) { + gate_input_ptr = up_input_ptr = (uint8_t*)input + i * config_.hidden_size * ggml_type_size(config_.hidden_type) / ggml_blck_size(config_.hidden_type); + } else { + to_float((uint8_t*)input + i * config_.hidden_size * ggml_type_size(config_.hidden_type) / ggml_blck_size(config_.hidden_type), m_input_fp32_[i], config_.hidden_size, config_.hidden_type); + if (ggml_internal_get_type_traits(config_.gate_type).vec_dot_type == ggml_internal_get_type_traits(config_.up_type).vec_dot_type) { + from_float(m_input_fp32_[i], m_gate_input_[i], config_.hidden_size, ggml_internal_get_type_traits(config_.gate_type).vec_dot_type); + gate_input_ptr = up_input_ptr = m_gate_input_[i]; + } else { + if (config_.hidden_type != ggml_internal_get_type_traits(config_.gate_type).vec_dot_type) { + from_float(m_input_fp32_[i], m_gate_input_[i], config_.hidden_size, ggml_internal_get_type_traits(config_.gate_type).vec_dot_type); + gate_input_ptr = m_gate_input_[i]; + } else { + gate_input_ptr = (uint8_t*)input + i * config_.hidden_size * ggml_type_size(config_.hidden_type) / ggml_blck_size(config_.hidden_type); + } + if (config_.hidden_type != ggml_internal_get_type_traits(config_.up_type).vec_dot_type) { + from_float(m_input_fp32_[i], m_up_input_[i], config_.hidden_size, ggml_internal_get_type_traits(config_.up_type).vec_dot_type); + up_input_ptr = m_up_input_[i]; + } else { + up_input_ptr = (uint8_t*)input + i * config_.hidden_size * ggml_type_size(config_.hidden_type) / ggml_blck_size(config_.hidden_type); + } + } + } + for (int j = 0; j < k; j++) { + memcpy(m_local_gate_input_ptr_[expert_ids[i * k + j]] + m_local_pos_[i][j] * config_.hidden_size * ggml_type_size(ggml_internal_get_type_traits(config_.gate_type).vec_dot_type) / ggml_blck_size(ggml_internal_get_type_traits(config_.gate_type).vec_dot_type), gate_input_ptr, config_.hidden_size * ggml_type_size(ggml_internal_get_type_traits(config_.gate_type).vec_dot_type) / ggml_blck_size(ggml_internal_get_type_traits(config_.gate_type).vec_dot_type)); + memcpy(m_local_up_input_ptr_[expert_ids[i * k + j]] + m_local_pos_[i][j] * config_.hidden_size * ggml_type_size(ggml_internal_get_type_traits(config_.up_type).vec_dot_type) / ggml_blck_size(ggml_internal_get_type_traits(config_.up_type).vec_dot_type), up_input_ptr, config_.hidden_size * ggml_type_size(ggml_internal_get_type_traits(config_.up_type).vec_dot_type) / ggml_blck_size(ggml_internal_get_type_traits(config_.up_type).vec_dot_type)); + } + }, nullptr); + int stride = QK_K; + int nth = config_.intermediate_size / stride; + backend->do_work_stealing_job(nth * config_.expert_num, nullptr, [&](int task_id) { + uint64_t expert_idx = task_id / nth; + int ith = task_id % nth; + void* gate_input_ptr = m_local_gate_input_ptr_[expert_idx]; + + #ifdef USE_NUMA + void* gate_proj_ptr = (uint8_t*)gate_proj_numa_[Backend::numa_node] + (expert_idx * config_.intermediate_size + ith * stride) * config_.hidden_size * ggml_type_size(config_.gate_type) / ggml_blck_size(config_.gate_type); + #else + void* gate_proj_ptr = (uint8_t*)gate_proj_ + (expert_idx * config_.intermediate_size + ith * stride) * config_.hidden_size * ggml_type_size(config_.gate_type) / ggml_blck_size(config_.gate_type); + #endif + + float* gate_output_ptr = m_local_gate_output_ptr_[expert_idx] + ith * stride; + llamafile_sgemm(stride, m_local_num_[expert_idx], config_.hidden_size / ggml_blck_size(config_.gate_type), gate_proj_ptr, config_.hidden_size / ggml_blck_size(config_.gate_type), gate_input_ptr, config_.hidden_size / ggml_blck_size(config_.gate_type), gate_output_ptr, config_.intermediate_size, 0, 1, GGML_TASK_TYPE_COMPUTE, config_.gate_type, ggml_internal_get_type_traits(config_.gate_type).vec_dot_type, GGML_TYPE_F32, GGML_PREC_DEFAULT); + void* up_input_ptr = m_local_up_input_ptr_[expert_idx]; + + #ifdef USE_NUMA + void* up_proj_ptr = (uint8_t*)up_proj_numa_[Backend::numa_node] + (expert_idx * config_.intermediate_size + ith * stride) * config_.hidden_size * ggml_type_size(config_.up_type) / ggml_blck_size(config_.up_type); + #else + void* up_proj_ptr = (uint8_t*)up_proj_ + (expert_idx * config_.intermediate_size + ith * stride) * config_.hidden_size * ggml_type_size(config_.up_type) / ggml_blck_size(config_.up_type); + #endif + + float* up_output_ptr = m_local_up_output_ptr_[expert_idx] + ith * stride; + llamafile_sgemm(stride, m_local_num_[expert_idx], config_.hidden_size / ggml_blck_size(config_.up_type), up_proj_ptr, config_.hidden_size / ggml_blck_size(config_.up_type), up_input_ptr, config_.hidden_size / ggml_blck_size(config_.up_type), up_output_ptr, config_.intermediate_size, 0, 1, GGML_TASK_TYPE_COMPUTE, config_.up_type, ggml_internal_get_type_traits(config_.up_type).vec_dot_type, GGML_TYPE_F32, GGML_PREC_DEFAULT); + for (int i = 0; i < m_local_num_[expert_idx]; i++) { + for (int j = ith * stride; j < (ith + 1) * stride; j++) { + m_local_intermediate_fp32_ptr_[expert_idx][i * config_.intermediate_size + j] = act_fn(m_local_gate_output_ptr_[expert_idx][i * config_.intermediate_size + j]) * m_local_up_output_ptr_[expert_idx][i * config_.intermediate_size + j]; + } + float* intermediate_fp32_ptr = m_local_intermediate_fp32_ptr_[expert_idx] + i * config_.intermediate_size + ith * stride; + void* down_input_ptr = m_local_down_input_ptr_[expert_idx] + i * config_.intermediate_size * ggml_type_size(ggml_internal_get_type_traits(config_.down_type).vec_dot_type) / ggml_blck_size(ggml_internal_get_type_traits(config_.down_type).vec_dot_type) + ith * stride * ggml_type_size(ggml_internal_get_type_traits(config_.down_type).vec_dot_type) / ggml_blck_size(ggml_internal_get_type_traits(config_.down_type).vec_dot_type); + from_float(intermediate_fp32_ptr, down_input_ptr, stride, ggml_internal_get_type_traits(config_.down_type).vec_dot_type); + } + }, nullptr); + stride = QK_K; + nth = config_.hidden_size / stride; + backend->do_work_stealing_job(nth * config_.expert_num, nullptr, [&](int task_id) { + uint64_t expert_idx = task_id / nth; + int ith = task_id % nth; + void* down_input_ptr = m_local_down_input_ptr_[expert_idx]; + + #ifdef USE_NUMA + void* down_proj_ptr = (uint8_t*)down_proj_numa_[Backend::numa_node] + (expert_idx * config_.hidden_size + ith * stride) * config_.intermediate_size * ggml_type_size(config_.down_type) / ggml_blck_size(config_.down_type); + #else + void* down_proj_ptr = (uint8_t*)down_proj_ + (expert_idx * config_.hidden_size + ith * stride) * config_.intermediate_size * ggml_type_size(config_.down_type) / ggml_blck_size(config_.down_type); + #endif + + float* down_output_ptr = m_local_down_output_ptr_[expert_idx] + ith * stride; + llamafile_sgemm(stride, m_local_num_[expert_idx], config_.intermediate_size / ggml_blck_size(config_.down_type), down_proj_ptr, config_.intermediate_size / ggml_blck_size(config_.down_type), down_input_ptr, config_.intermediate_size / ggml_blck_size(config_.down_type), down_output_ptr, config_.hidden_size, 0, 1, GGML_TASK_TYPE_COMPUTE, config_.down_type, ggml_internal_get_type_traits(config_.down_type).vec_dot_type, GGML_TYPE_F32, GGML_PREC_DEFAULT); + }, nullptr); + backend->do_work_stealing_job(qlen, nullptr, [&](int i) { + for (int e = 0; e < config_.hidden_size; e++) { + m_output_fp32_[i][e] = 0; + } + for (int j = 0; j < k; j++) { + for (int e = 0; e < config_.hidden_size; e++) { + m_output_fp32_[i][e] += m_local_down_output_ptr_[expert_ids[i * k + j]][m_local_pos_[i][j] * config_.hidden_size + e] * weights[i * k + j]; + } + } + from_float(m_output_fp32_[i], (uint8_t*)output + i * config_.hidden_size * ggml_type_size(config_.hidden_type) / ggml_blck_size(config_.hidden_type), config_.hidden_size, config_.hidden_type); + }, nullptr); + + /* 把每个 token-expert 的行复制到各自 cache */ + backend->do_work_stealing_job(qlen, nullptr, [&](int token_idx) { + auto& cache = fwd_cache[token_idx]; + // cache 已在上层 init(k, inter_size) + for (int j = 0; j < k; ++j) { + uint64_t eid = expert_ids[token_idx*k + j]; + int row = m_local_pos_[token_idx][j]; + size_t ofs = row * config_.intermediate_size; + /* gate u */ + std::memcpy(cache.gate_u[j].data(), + m_local_gate_output_ptr_[eid] + ofs, + sizeof(float) * config_.intermediate_size); + /* up v */ + std::memcpy(cache.up_v[j].data(), + m_local_up_output_ptr_[eid] + ofs, + sizeof(float) * config_.intermediate_size); + /* 可选 z */ + // std::memcpy(cache.z[j].data(), + // m_local_intermediate_fp32_ptr_[eid] + ofs, + // sizeof(float) * config_.intermediate_size); + } + }, nullptr); +} + +void SFT_MOE::forward(int qlen, int k, const uint64_t* expert_ids, const float* weights, const void* input, void* output, Backend* backend, SFT_MoEForwardCache* fwd_cache) { + if (qlen < config_.group_min_len) { + for (int i = 0; i < qlen; i++) { + // fwd_cache[i].init(k, config_.intermediate_size); // 预分配 + forward_one(k, expert_ids + i * k, weights + i * k, (uint8_t*)input + i * config_.hidden_size * ggml_type_size(config_.hidden_type) / ggml_blck_size(config_.hidden_type), (uint8_t*)output + i * config_.hidden_size * ggml_type_size(config_.hidden_type) / ggml_blck_size(config_.hidden_type), backend, fwd_cache + i); + } + return; + } + int forward_len = std::min(config_.group_max_len, qlen); + // for (int i = 0; i < forward_len; ++i) + // fwd_cache[i].init(k, config_.intermediate_size); + forward_many(forward_len, k, expert_ids, weights, input, output, backend, fwd_cache); + forward(qlen - forward_len, k, expert_ids + forward_len * k, weights + forward_len * k, (uint8_t*)input + forward_len * config_.hidden_size * ggml_type_size(config_.hidden_type) / ggml_blck_size(config_.hidden_type), (uint8_t*)output + forward_len * config_.hidden_size * ggml_type_size(config_.hidden_type) / ggml_blck_size(config_.hidden_type), backend, fwd_cache + forward_len); +} + +static float act_fn_grad(float x) { + float sigmoid_x = 1.0f / (1.0f + expf(-x)); + return sigmoid_x * (1. + x * (1. - sigmoid_x)); +} + +void SFT_MOE::transpose_expert_matrix(const void* src, void* dst, int R, int C, ggml_type src_type, ggml_type dst_type, uint64_t expert_idx) { + to_float(src, transpose_buffer_fp32_ + (R * C * expert_idx), R * C, src_type); + from_float(transpose_buffer_fp32_ + (R * C * expert_idx), transpose_buffer_ + (R * C * expert_idx) * ggml_type_size(dst_type), R * C, dst_type); + for (int r = 0; r < R; ++r) { + for (int c = 0; c < C; ++c) { + memcpy( + (uint8_t*)dst + (c * R + r) * ggml_type_size(dst_type), + (uint8_t*)transpose_buffer_ + (R * C * expert_idx + r * C + c) * ggml_type_size(dst_type), + ggml_type_size(dst_type)); + } + } +} + +void SFT_MOE::get_transpose(Backend* backend) { + // Transpose gate_proj_ + int R_gate = config_.intermediate_size; + int C_gate = config_.hidden_size; + size_t gate_expert_src_stride_bytes = (size_t)R_gate * C_gate * ggml_type_size(config_.gate_type); + size_t gate_expert_dst_t_stride_bytes = (size_t)C_gate * R_gate * ggml_type_size(config_.grad_type); + backend->do_work_stealing_job(config_.expert_num, nullptr, [&](int expert_idx) { + void* src_expert = (uint8_t*)gate_proj_ + expert_idx * gate_expert_src_stride_bytes; + void* dst_expert_t = (uint8_t*)gate_proj_t_ + expert_idx * gate_expert_dst_t_stride_bytes; + transpose_expert_matrix(src_expert, dst_expert_t, R_gate, C_gate, config_.gate_type, config_.grad_type, expert_idx); + }, nullptr); + + // Transpose up_proj_ + int R_up = config_.intermediate_size; + int C_up = config_.hidden_size; + size_t up_expert_src_stride_bytes = (size_t)R_up * C_up * ggml_type_size(config_.up_type); + size_t up_expert_dst_t_stride_bytes = (size_t)C_up * R_up * ggml_type_size(config_.grad_type); + backend->do_work_stealing_job(config_.expert_num, nullptr, [&](int expert_idx) { + void* src_expert = (uint8_t*)up_proj_ + expert_idx * up_expert_src_stride_bytes; + void* dst_expert_t = (uint8_t*)up_proj_t_ + expert_idx * up_expert_dst_t_stride_bytes; + transpose_expert_matrix(src_expert, dst_expert_t, R_up, C_up, config_.up_type, config_.grad_type, expert_idx); + }, nullptr); + + // Transpose down_proj_ + int R_down = config_.hidden_size; + int C_down = config_.intermediate_size; + size_t down_expert_src_stride_bytes = (size_t)R_down * C_down * ggml_type_size(config_.down_type); + size_t down_expert_dst_t_stride_bytes = (size_t)C_down * R_down * ggml_type_size(config_.grad_type); + backend->do_work_stealing_job(config_.expert_num, nullptr, [&](int expert_idx) { + void* src_expert = (uint8_t*)down_proj_ + expert_idx * down_expert_src_stride_bytes; + void* dst_expert_t = (uint8_t*)down_proj_t_ + expert_idx * down_expert_dst_t_stride_bytes; + transpose_expert_matrix(src_expert, dst_expert_t, R_down, C_down, config_.down_type, config_.grad_type, expert_idx); + }, nullptr); +} + +void SFT_MOE::backward_one(int k, const uint64_t* expert_ids, const float* weights, const void* output_grad, void* input_grad, Backend* backend, const SFT_MoEForwardCache* fwd_cache) { + // clock_t clk1, clk2, clk3, clk4; + // clock_t clkz1, clkz2, clkz3, clkz4, clkz5; + // clk1 = clock(); + // clk2 = clock(); + int nth = config_.intermediate_size / config_.stride; + backend->do_work_stealing_job(nth * k, nullptr, [&](int task_id) { + int expert_idx = task_id / nth; + uint64_t expert_id = expert_ids[expert_idx]; + int ith = task_id % nth; + // clkz1 = clock(); + void* down_proj_t_ptr = (uint8_t*)down_proj_t_ + (expert_id * config_.intermediate_size + ith * config_.stride) * config_.hidden_size * ggml_type_size(config_.grad_type); + float* down_input_grad_ptr = s_down_input_grad_[expert_idx] + ith * config_.stride; + // clkz2 = clock(); + llamafile_sgemm(config_.stride, 1, config_.hidden_size, down_proj_t_ptr, config_.hidden_size, output_grad, config_.hidden_size, down_input_grad_ptr, config_.stride, 0, 1, GGML_TASK_TYPE_COMPUTE, config_.grad_type, config_.grad_type, GGML_TYPE_F32, GGML_PREC_DEFAULT); + // clkz3 = clock(); + for (int i = ith * config_.stride; i < (ith + 1) * config_.stride; i++) { + s_down_input_grad_[expert_idx][i] *= weights[expert_idx]; + + s_gate_output_grad_fp32_[expert_idx][i] = s_down_input_grad_[expert_idx][i] * fwd_cache->up_v[expert_idx][i] * act_fn_grad(fwd_cache->gate_u[expert_idx][i]); + s_up_output_grad_fp32_[expert_idx][i] = s_down_input_grad_[expert_idx][i] * act_fn(fwd_cache->gate_u[expert_idx][i]); + } + // clkz4 = clock(); + from_float(s_gate_output_grad_fp32_[expert_idx] + ith * config_.stride, s_gate_output_grad_[expert_idx] + ith * config_.stride * ggml_type_size(config_.grad_type), config_.stride, config_.grad_type); + from_float(s_up_output_grad_fp32_[expert_idx] + ith * config_.stride, s_up_output_grad_[expert_idx] + ith * config_.stride * ggml_type_size(config_.grad_type), config_.stride, config_.grad_type); + // clkz5 = clock(); + }, nullptr); + + // clk3 = clock(); + nth = config_.hidden_size / config_.stride; + backend->do_work_stealing_job(nth, nullptr, [&](int task_id) { + int ith = task_id; + for (int i = ith * config_.stride; i < (ith + 1) * config_.stride; i++) { + s_input_grad_fp32_[i] = 0; + } + for (int expert_idx = 0; expert_idx < k; expert_idx++) { + uint64_t expert_id = expert_ids[expert_idx]; + + void* gate_proj_t_ptr = (uint8_t*)gate_proj_t_ + (expert_id * config_.hidden_size + ith * config_.stride) * config_.intermediate_size * ggml_type_size(config_.grad_type); + float* gate_input_grad_ptr = s_gate_input_grad_[expert_idx] + ith * config_.stride; + llamafile_sgemm(config_.stride, 1, config_.intermediate_size, gate_proj_t_ptr, config_.intermediate_size, s_gate_output_grad_[expert_idx], config_.intermediate_size, gate_input_grad_ptr, config_.stride, 0, 1, GGML_TASK_TYPE_COMPUTE, config_.grad_type, config_.grad_type, GGML_TYPE_F32, GGML_PREC_DEFAULT); + + void* up_proj_t_ptr = (uint8_t*)up_proj_t_ + (expert_id * config_.hidden_size + ith * config_.stride) * config_.intermediate_size * ggml_type_size(config_.grad_type); + float* up_input_grad_ptr = s_up_input_grad_[expert_idx] + ith * config_.stride; + llamafile_sgemm(config_.stride, 1, config_.intermediate_size, up_proj_t_ptr, config_.intermediate_size, s_up_output_grad_[expert_idx], config_.intermediate_size, up_input_grad_ptr, config_.stride, 0, 1, GGML_TASK_TYPE_COMPUTE, config_.grad_type, config_.grad_type, GGML_TYPE_F32, GGML_PREC_DEFAULT); + + for (int i = ith * config_.stride; i < (ith + 1) * config_.stride; i++) { + s_input_grad_fp32_[i] += s_gate_input_grad_[expert_idx][i] + s_up_input_grad_[expert_idx][i]; + } + } + from_float(s_input_grad_fp32_ + ith * config_.stride, (uint8_t*)input_grad + ith * config_.stride * ggml_type_size(config_.grad_type), config_.stride, config_.grad_type); + }, nullptr); + // clk4 = clock(); + // std::cout << "[Δclk12] " << (clk2 - clk1) / static_cast(CLOCKS_PER_SEC) * 1000 + // << " ms [Δclk23] " << (clk3 - clk2) / static_cast(CLOCKS_PER_SEC) * 1000 + // << " ms [Δclk34] " << (clk4 - clk3) / static_cast(CLOCKS_PER_SEC) * 1000 + // << " ms [Δclkz12] " << (clkz2 - clkz1) / static_cast(CLOCKS_PER_SEC) * 1000 + // << " ms [Δclkz23] " << (clkz3 - clkz2) / static_cast(CLOCKS_PER_SEC) * 1000 + // << " ms [Δclkz34] " << (clkz4 - clkz3) / static_cast(CLOCKS_PER_SEC) * 1000 + // << " ms [Δclkz45] " << (clkz5 - clkz4) / static_cast(CLOCKS_PER_SEC) * 1000 + // << " ms\n"; + +} + +void SFT_MOE::backward_many(int qlen, int k, const uint64_t* expert_ids, const float* weights, const void* output_grad, void* input_grad, Backend* backend, const SFT_MoEForwardCache* fwd_cache) { + for (int i = 0; i < config_.expert_num; i++) { + m_local_num_[i] = 0; + } + for (int i = 0; i < qlen; i++) { + for (int j = 0; j < k; j++) { + m_local_pos_[i][j] = m_local_num_[expert_ids[i * k + j]]++; + } + } + uint64_t offset = 0; + for (int i = 0; i < config_.expert_num; i++) { + m_local_down_output_grad_ptr_[i] = m_local_down_output_grad_ + offset * config_.hidden_size * ggml_type_size(config_.grad_type); + m_local_down_input_grad_ptr_[i] = m_local_down_input_grad_ + offset * config_.intermediate_size; + m_local_gate_output_grad_fp32_ptr_[i] = m_local_gate_output_grad_fp32_ + offset * config_.intermediate_size; + m_local_up_output_grad_fp32_ptr_[i] = m_local_up_output_grad_fp32_ + offset * config_.intermediate_size; + m_local_gate_output_grad_ptr_[i] = m_local_gate_output_grad_ + offset * config_.intermediate_size * ggml_type_size(config_.grad_type); + m_local_up_output_grad_ptr_[i] = m_local_up_output_grad_ + offset * config_.intermediate_size * ggml_type_size(config_.grad_type); + m_local_gate_input_grad_ptr_[i] = m_local_gate_input_grad_ + offset * config_.hidden_size; + m_local_up_input_grad_ptr_[i] = m_local_up_input_grad_ + offset * config_.hidden_size; + m_local_token_indices_ptr_[i] = m_local_token_indices_ + offset; + m_local_expert_positions_ptr_[i] = m_local_expert_positions_ + offset; + offset += m_local_num_[i]; + } + + backend->do_work_stealing_job(qlen, nullptr, [&](int i) { + for (int j = 0; j < k; j++) { + uint64_t expert_id = expert_ids[i * k + j]; + int local_row = m_local_pos_[i][j]; + memcpy(m_local_down_output_grad_ptr_[expert_id] + local_row * config_.hidden_size * ggml_type_size(config_.grad_type), (uint8_t*)output_grad + i * config_.hidden_size * ggml_type_size(config_.grad_type), config_.hidden_size * ggml_type_size(config_.grad_type)); + m_local_token_indices_ptr_[expert_id][local_row] = i; + m_local_expert_positions_ptr_[expert_id][local_row] = j; + } + }, nullptr); + + // get_transpose(backend); + + int stride = QK_K; + int nth = config_.intermediate_size / stride; + backend->do_work_stealing_job(nth * config_.expert_num, nullptr, [&](int task_id) { + uint64_t expert_idx = task_id / nth; + int ith = task_id % nth; + + void* down_proj_t_ptr = (uint8_t*)down_proj_t_ + (expert_idx * config_.intermediate_size + ith * stride) * config_.hidden_size * ggml_type_size(config_.grad_type); + void* down_output_grad_ptr = m_local_down_output_grad_ptr_[expert_idx]; + float* down_input_grad_ptr = m_local_down_input_grad_ptr_[expert_idx] + ith * stride; + + llamafile_sgemm(stride, m_local_num_[expert_idx], config_.hidden_size, down_proj_t_ptr, config_.hidden_size, down_output_grad_ptr, config_.hidden_size, down_input_grad_ptr, config_.intermediate_size, 0, 1, GGML_TASK_TYPE_COMPUTE, config_.grad_type, config_.grad_type, GGML_TYPE_F32, GGML_PREC_DEFAULT); + + for (int i = 0; i < m_local_num_[expert_idx]; i++) { + int token_idx = m_local_token_indices_ptr_[expert_idx][i]; + int expert_pos = m_local_expert_positions_ptr_[expert_idx][i]; + float weight = weights[token_idx * k + expert_pos]; + + for (int j = ith * stride; j < (ith + 1) * stride; j++) { + m_local_down_input_grad_ptr_[expert_idx][i * config_.intermediate_size + j] *= weight; + + float down_input_grad = m_local_down_input_grad_ptr_[expert_idx][i * config_.intermediate_size + j]; + m_local_gate_output_grad_fp32_ptr_[expert_idx][i * config_.intermediate_size + j] = down_input_grad * fwd_cache[token_idx].up_v[expert_pos][j] * act_fn_grad(fwd_cache[token_idx].gate_u[expert_pos][j]); + m_local_up_output_grad_fp32_ptr_[expert_idx][i * config_.intermediate_size + j] = down_input_grad * act_fn(fwd_cache[token_idx].gate_u[expert_pos][j]); + } + + float* gate_output_grad_fp32_ptr = m_local_gate_output_grad_fp32_ptr_[expert_idx] + i * config_.intermediate_size + ith * stride; + void* gate_output_grad_ptr = m_local_gate_output_grad_ptr_[expert_idx] + (i * config_.intermediate_size + ith * stride) * ggml_type_size(config_.grad_type); + from_float(gate_output_grad_fp32_ptr, gate_output_grad_ptr, stride, config_.grad_type); + + float* up_output_grad_fp32_ptr = m_local_up_output_grad_fp32_ptr_[expert_idx] + i * config_.intermediate_size + ith * stride; + void* up_output_grad_ptr = m_local_up_output_grad_ptr_[expert_idx] + (i * config_.intermediate_size + ith * stride) * ggml_type_size(config_.grad_type); + from_float(up_output_grad_fp32_ptr, up_output_grad_ptr, stride, config_.grad_type); + } + }, nullptr); + stride = QK_K; + nth = config_.hidden_size / stride; + backend->do_work_stealing_job(nth * config_.expert_num, nullptr, [&](int task_id) { + uint64_t expert_idx = task_id / nth; + int ith = task_id % nth; + + void* gate_proj_t_ptr = (uint8_t*)gate_proj_t_ + (expert_idx * config_.hidden_size + ith * stride) * config_.intermediate_size * ggml_type_size(config_.grad_type); + void* up_proj_t_ptr = (uint8_t*)up_proj_t_ + (expert_idx * config_.hidden_size + ith * stride) * config_.intermediate_size * ggml_type_size(config_.grad_type); + void* gate_output_grad_ptr = m_local_gate_output_grad_ptr_[expert_idx]; + void* up_output_grad_ptr = m_local_up_output_grad_ptr_[expert_idx]; + float* gate_input_grad_ptr = m_local_gate_input_grad_ptr_[expert_idx] + ith * stride; + float* up_input_grad_ptr = m_local_up_input_grad_ptr_[expert_idx] + ith * stride; + + llamafile_sgemm(stride, m_local_num_[expert_idx], config_.intermediate_size, gate_proj_t_ptr, config_.intermediate_size, gate_output_grad_ptr, config_.intermediate_size, gate_input_grad_ptr, config_.hidden_size, 0, 1, GGML_TASK_TYPE_COMPUTE, config_.grad_type, config_.grad_type, GGML_TYPE_F32, GGML_PREC_DEFAULT); + llamafile_sgemm(stride, m_local_num_[expert_idx], config_.intermediate_size, up_proj_t_ptr, config_.intermediate_size, up_output_grad_ptr, config_.intermediate_size, up_input_grad_ptr, config_.hidden_size, 0, 1, GGML_TASK_TYPE_COMPUTE, config_.grad_type, config_.grad_type, GGML_TYPE_F32, GGML_PREC_DEFAULT); + }, nullptr); + backend->do_work_stealing_job(qlen, nullptr, [&](int i) { + for (int e = 0; e < config_.hidden_size; e++) { + m_grad_input_fp32_[i][e] = 0; + } + for (int j = 0; j < k; j++) { + for (int e = 0; e < config_.hidden_size; e++) { + m_grad_input_fp32_[i][e] += m_local_gate_input_grad_ptr_[expert_ids[i * k + j]][m_local_pos_[i][j] * config_.hidden_size + e] + m_local_up_input_grad_ptr_[expert_ids[i * k + j]][m_local_pos_[i][j] * config_.hidden_size + e]; + } + } + from_float(m_grad_input_fp32_[i], (uint8_t*)input_grad + i * config_.hidden_size * ggml_type_size(config_.grad_type), config_.hidden_size, config_.grad_type); + }, nullptr); +} + +// TODO: input和layer_idx参数可以删除 +void SFT_MOE::backward(int layer_idx, int qlen, int k, const uint64_t* expert_ids, const float* weights, + const void* input, const void* grad_output, void* grad_input, Backend* backend, const SFT_MoEForwardCache* fwd_cache) { + + get_transpose(backend); + int remaining_qlen = qlen; + int processed_offset = 0; + + while (remaining_qlen > 0) { + // config_.group_min_len = 10000000; + if (remaining_qlen < config_.group_min_len) { + for (int i = 0; i < remaining_qlen; i++) { + backward_one(k, + expert_ids + (processed_offset + i) * k, + weights + (processed_offset + i) * k, + (uint8_t*)grad_output + (processed_offset + i) * config_.hidden_size * ggml_type_size(config_.grad_type), + (uint8_t*)grad_input + (processed_offset + i) * config_.hidden_size * ggml_type_size(config_.grad_type), + backend, + fwd_cache + processed_offset + i); + } + break; + } else { + int backward_len = std::min(config_.group_max_len, remaining_qlen); + backward_many(backward_len, + k, + expert_ids + processed_offset * k, + weights + processed_offset * k, + (uint8_t*)grad_output + processed_offset * config_.hidden_size * ggml_type_size(config_.grad_type), + (uint8_t*)grad_input + processed_offset * config_.hidden_size * ggml_type_size(config_.grad_type), + backend, + fwd_cache + processed_offset); + + remaining_qlen -= backward_len; + processed_offset += backward_len; + } + } +} \ No newline at end of file diff --git a/archive/kt-sft/csrc/ktransformers_ext/operators/llamafile/sft_moe.h b/archive/kt-sft/csrc/ktransformers_ext/operators/llamafile/sft_moe.h new file mode 100644 index 0000000..ad8efe8 --- /dev/null +++ b/archive/kt-sft/csrc/ktransformers_ext/operators/llamafile/sft_moe.h @@ -0,0 +1,156 @@ +/** + * @Description : + * @Author : chenht2022 + * @Date : 2024-07-22 02:03:22 + * @Version : 1.0.0 + * @LastEditors : chenht2022 + * @LastEditTime : 2024-07-25 10:35:10 + * @Copyright (c) 2024 by KVCache.AI, All Rights Reserved. + **/ +#ifndef CPUINFER_OPERATOR_SFT_MOE_H +#define CPUINFER_OPERATOR_SFT_MOE_H + +#include +#include +#include +#include +#include + +#include "../../cpu_backend/backend.h" +#include "../../cpu_backend/shared_mem_buffer.h" +#include "conversion.h" +#include "llama.cpp/ggml-impl.h" +#include "llama.cpp/ggml-quants.h" +#include "llama.cpp/ggml.h" +#include "llamafile/sgemm.h" +#include "sft_moe_forward_cache.h" + +struct SFT_MOEConfig { + long expert_num; + int routed_expert_num; + long hidden_size; + long intermediate_size; + int stride; + int group_min_len; + int group_max_len; + void* gate_proj; + void* up_proj; + void* down_proj; + ggml_type gate_type; + ggml_type up_type; + ggml_type down_type; + ggml_type hidden_type; + ggml_type grad_type = GGML_TYPE_BF16; + + SFT_MOEConfig() {} + + SFT_MOEConfig(int expert_num, int routed_expert_num, int hidden_size, int intermediate_size, int stride, int group_min_len, int group_max_len, void* gate_proj, void* up_proj, void* down_proj, ggml_type gate_type, ggml_type up_type, ggml_type down_type, ggml_type hidden_type) + : expert_num(expert_num), routed_expert_num(routed_expert_num), hidden_size(hidden_size), intermediate_size(intermediate_size), stride(stride), group_min_len(group_min_len), group_max_len(group_max_len), gate_proj(gate_proj), up_proj(up_proj), down_proj(down_proj), gate_type(gate_type), up_type(up_type), down_type(down_type), hidden_type(hidden_type) {} +}; + +class SFT_MOE { + public: + SFT_MOE(SFT_MOEConfig); + ~SFT_MOE(); + void warm_up(Backend* backend); + void forward_one(int k, const uint64_t* expert_ids, const float* weights, const void* input, void* output, Backend* backend, SFT_MoEForwardCache* fwd_cache); + void forward_many(int qlen, int k, const uint64_t* expert_ids, const float* weights, const void* input, void* output, Backend* backend, SFT_MoEForwardCache* fwd_cache); + void forward(int qlen, int k, const uint64_t* expert_ids, const float* weights, const void* input, void* output, Backend* backend, SFT_MoEForwardCache* fwd_cache); + void backward_one(int k, const uint64_t* expert_ids, const float* weights, const void* output_grad, void* input_grad, Backend* backend, const SFT_MoEForwardCache* fwd_cache); + void backward_many(int qlen, int k, const uint64_t* expert_ids, const float* weights, const void* output_grad, void* input_grad, Backend* backend, const SFT_MoEForwardCache* fwd_cache); + void backward(int layer_idx, int qlen, int k, const uint64_t* expert_ids, const float* weights, + const void* input, const void* grad_output, void* grad_input, Backend* backend, const SFT_MoEForwardCache* fwd_cache); // FIXME: expert backward definition for C++ + + void transpose_expert_matrix(const void* src, void* dst, int R, int C, ggml_type src_type, ggml_type dst_type, uint64_t expert_idx); + void ensure_fwd_cache(int qlen, int k); + void get_transpose(Backend* backend); + SFT_MoEForwardCache* fwd_cache_ptr(); + + private: + SFT_MOEConfig config_; + void* gate_proj_; // [expert_num * intermediate_size * hidden_size ( /32 if quantized)] + void* up_proj_; // [expert_num * intermediate_size * hidden_size ( /32 if quantized)] + void* down_proj_; // [expert_num * hidden_size * intermediate_size ( /32 if quantized)] + + float* transpose_buffer_fp32_; // [expert_num * intermediate_size * hidden_size] + uint8_t* transpose_buffer_; // [expert_num * intermediate_size * hidden_size] + + uint8_t* gate_proj_t_; // [expert_num * hidden_size * intermediate_size] + uint8_t* up_proj_t_; // [expert_num * hidden_size * intermediate_size] + uint8_t* down_proj_t_; // [expert_num * intermediate_size * hidden_size] + + #ifdef USE_NUMA + std::vector gate_proj_numa_; // [numa_num, expert_num * intermediate_size * hidden_size ( /32 if quantized)] + std::vector up_proj_numa_; // [numa_num, expert_num * intermediate_size * hidden_size ( /32 if quantized)] + std::vector down_proj_numa_; // [numa_num, expert_num * hidden_size * intermediate_size ( /32 if quantized)] + #endif + + float* s_input_fp32_; // [hidden_size] + uint8_t* s_gate_input_; // [hidden_size * ggml_type_size(ggml_internal_get_type_traits(gate_type).vec_dot_type) / ggml_blck_size(ggml_internal_get_type_traits(gate_type).vec_dot_type)] + uint8_t* s_up_input_; // [hidden_size * ggml_type_size(ggml_internal_get_type_traits(up_type).vec_dot_type) / ggml_blck_size(ggml_internal_get_type_traits(up_type).vec_dot_type)] + std::vector s_gate_output_; // [routed_expert_num, intermediate_size] + std::vector s_up_output_; // [routed_expert_num, intermediate_size] + std::vector s_intermediate_fp32_; // [routed_expert_num, intermediate_size] + std::vector s_down_input_; // [routed_expert_num, intermediate_size * ggml_type_size(ggml_internal_get_type_traits(down_type).vec_dot_type) / ggml_blck_size(ggml_internal_get_type_traits(down_type).vec_dot_type)] + std::vector s_down_output_; // [routed_expert_num, hidden_size] + float* s_output_fp32_; // [hidden_size] + + std::vector s_down_input_grad_; // [routed_expert_num, intermediate_size] + std::vector s_gate_output_grad_fp32_; // [routed_expert_num, intermediate_size] + std::vector s_up_output_grad_fp32_; // [routed_expert_num, intermediate_size] + std::vector s_gate_output_grad_; // [routed_expert_num, intermediate_size * ggml_type_size(grad_type)] + std::vector s_up_output_grad_; // [routed_expert_num, intermediate_size * ggml_type_size(grad_type)] + std::vector s_gate_input_grad_; // [routed_expert_num, hidden_size] + std::vector s_up_input_grad_; // [routed_expert_num, hidden_size] + float* s_input_grad_fp32_; // [hidden_size] + + std::vector m_input_fp32_; // [group_max_len, hidden_size] + std::vector m_gate_input_; // [group_max_len, hidden_size * ggml_type_size(ggml_internal_get_type_traits(gate_type).vec_dot_type) / ggml_blck_size(ggml_internal_get_type_traits(gate_type).vec_dot_type)] + std::vector m_up_input_; // [group_max_len, hidden_size * ggml_type_size(ggml_internal_get_type_traits(up_type).vec_dot_type) / ggml_blck_size(ggml_internal_get_type_traits(up_type).vec_dot_type)] + uint8_t* m_local_gate_input_; // [routed_expert_num * group_max_len * hidden_size * ggml_type_size(ggml_internal_get_type_traits(gate_type).vec_dot_type) / ggml_blck_size(ggml_internal_get_type_traits(gate_type).vec_dot_type)] + uint8_t* m_local_up_input_; // [routed_expert_num * group_max_len * hidden_size * ggml_type_size(ggml_internal_get_type_traits(up_type).vec_dot_type) / ggml_blck_size(ggml_internal_get_type_traits(up_type).vec_dot_type)] + float* m_local_gate_output_; // [routed_expert_num * group_max_len * intermediate_size] + float* m_local_up_output_; // [routed_expert_num * group_max_len * intermediate_size] + float* m_local_intermediate_fp32_; // [routed_expert_num * group_max_len * intermediate_size] + uint8_t* m_local_down_input_; // [routed_expert_num * group_max_len * intermediate_size * ggml_type_size(ggml_internal_get_type_traits(down_type).vec_dot_type) / ggml_blck_size(ggml_internal_get_type_traits(down_type).vec_dot_type)] + float* m_local_down_output_; // [routed_expert_num * group_max_len * hidden_size] + std::vector m_output_fp32_; // [group_max_len, hidden_size] + + std::vector> m_local_pos_; // [group_max_len, routed_expert_num] + std::vector m_local_num_; // [expert_num] + std::vector m_local_gate_input_ptr_; // [expert_num] + std::vector m_local_up_input_ptr_; // [expert_num] + std::vector m_local_gate_output_ptr_; // [expert_num] + std::vector m_local_up_output_ptr_; // [expert_num] + std::vector m_local_intermediate_fp32_ptr_; // [expert_num] + std::vector m_local_down_input_ptr_; // [expert_num] + std::vector m_local_down_output_ptr_; // [expert_num] + + uint8_t* m_local_down_output_grad_; // [routed_expert_num * group_max_len * hidden_size * ggml_type_size(grad_type)] + float* m_local_down_input_grad_; // [routed_expert_num * group_max_len * intermediate_size] + float* m_local_gate_output_grad_fp32_; // [routed_expert_num * group_max_len * intermediate_size] + float* m_local_up_output_grad_fp32_; // [routed_expert_num * group_max_len * intermediate_size] + uint8_t* m_local_gate_output_grad_; // [routed_expert_num * group_max_len * intermediate_size * ggml_type_size(grad_type)] + uint8_t* m_local_up_output_grad_; // [routed_expert_num * group_max_len * intermediate_size * ggml_type_size(grad_type)] + float* m_local_gate_input_grad_; // [routed_expert_num * group_max_len * hidden_size] + float* m_local_up_input_grad_; // [routed_expert_num * group_max_len * hidden_size] + std::vector m_grad_input_fp32_; // [group_max_len, hidden_size] + + std::vector m_local_down_output_grad_ptr_; // [expert_num] + std::vector m_local_down_input_grad_ptr_; // [expert_num] + std::vector m_local_gate_output_grad_fp32_ptr_; // [expert_num] + std::vector m_local_up_output_grad_fp32_ptr_; // [expert_num] + std::vector m_local_gate_output_grad_ptr_; // [expert_num] + std::vector m_local_up_output_grad_ptr_; // [expert_num] + std::vector m_local_gate_input_grad_ptr_; // [expert_num] + std::vector m_local_up_input_grad_ptr_; // [expert_num] + + int* m_local_token_indices_; // [routed_expert_num * group_max_len] + int* m_local_expert_positions_; // [routed_expert_num * group_max_len] + std::vector m_local_token_indices_ptr_; // [expert_num] + std::vector m_local_expert_positions_ptr_; // [expert_num] + + std::vector fw_cache_; // 持久缓存，便于backward读取到forward_cache +}; + +#endif \ No newline at end of file diff --git a/archive/kt-sft/csrc/ktransformers_ext/operators/llamafile/sft_moe_forward_cache.h b/archive/kt-sft/csrc/ktransformers_ext/operators/llamafile/sft_moe_forward_cache.h new file mode 100644 index 0000000..3a21565 --- /dev/null +++ b/archive/kt-sft/csrc/ktransformers_ext/operators/llamafile/sft_moe_forward_cache.h @@ -0,0 +1,32 @@ +#pragma once +#include + +struct SFT_MoEForwardCache { + // 每个 token 按 expert 分块保存 + std::vector> gate_u; // u = W_gate x + std::vector> up_v; // v = W_up x + // 若希望反向直接用 z = σ(u)⊙v，则再加一份 + // std::vector> z; + void init(int k, int inter_size) { + /* ---- 只增不减：capacity 不够时才增，永不缩小，避免多线程情况下的use-after-free ---- */ + if (k > (int)gate_u.size()) { + gate_u.resize(k); + up_v .resize(k); + // z .resize(k); + } + + for (int i = 0; i < k; ++i) { + if ((int)gate_u[i].capacity() < inter_size) + gate_u[i].reserve(inter_size); // 只增 capacity + if ((int)up_v[i].capacity() < inter_size) + up_v[i].reserve(inter_size); + // if ((int)z[i].capacity() < inter_size) + // z[i].reserve(inter_size); + + // size() 更新为 inter_size 以便直接下标写入 + gate_u[i].resize(inter_size); + up_v[i] .resize(inter_size); + // z[i] .resize(inter_size); + } + } +}; \ No newline at end of file diff --git a/archive/kt-sft/csrc/ktransformers_ext/vendors/cuda.h b/archive/kt-sft/csrc/ktransformers_ext/vendors/cuda.h new file mode 100644 index 0000000..1746b07 --- /dev/null +++ b/archive/kt-sft/csrc/ktransformers_ext/vendors/cuda.h @@ -0,0 +1,15 @@ +#pragma once + +#include +#include +#include +#include +#include + +#if CUDART_VERSION < 11020 +#define CU_DEVICE_ATTRIBUTE_VIRTUAL_MEMORY_MANAGEMENT_SUPPORTED CU_DEVICE_ATTRIBUTE_VIRTUAL_ADDRESS_MANAGEMENT_SUPPORTED +#define CUBLAS_TF32_TENSOR_OP_MATH CUBLAS_TENSOR_OP_MATH +#define CUBLAS_COMPUTE_16F CUDA_R_16F +#define CUBLAS_COMPUTE_32F CUDA_R_32F +#define cublasComputeType_t cudaDataType_t +#endif // CUDART_VERSION < 11020 diff --git a/archive/kt-sft/csrc/ktransformers_ext/vendors/hip.h b/archive/kt-sft/csrc/ktransformers_ext/vendors/hip.h new file mode 100644 index 0000000..abbc1e8 --- /dev/null +++ b/archive/kt-sft/csrc/ktransformers_ext/vendors/hip.h @@ -0,0 +1,172 @@ +#pragma once + +#define HIP_ENABLE_WARP_SYNC_BUILTINS 1 +#include +#include +#include +#include +#ifdef __HIP_PLATFORM_AMD__ +// for rocblas_initialize() +#include "rocblas/rocblas.h" +#endif // __HIP_PLATFORM_AMD__ + +#define CUBLAS_COMPUTE_16F HIPBLAS_R_16F +#define CUBLAS_COMPUTE_32F HIPBLAS_R_32F +#define CUBLAS_COMPUTE_32F_FAST_16F HIPBLAS_R_32F +#define CUBLAS_GEMM_DEFAULT HIPBLAS_GEMM_DEFAULT +#define CUBLAS_GEMM_DEFAULT_TENSOR_OP HIPBLAS_GEMM_DEFAULT +#define CUBLAS_OP_N HIPBLAS_OP_N +#define CUBLAS_OP_T HIPBLAS_OP_T +#define CUBLAS_STATUS_SUCCESS HIPBLAS_STATUS_SUCCESS +#define CUBLAS_TF32_TENSOR_OP_MATH 0 +#define CUDA_R_16F HIPBLAS_R_16F +#define CUDA_R_32F HIPBLAS_R_32F +#define CU_DEVICE_ATTRIBUTE_VIRTUAL_MEMORY_MANAGEMENT_SUPPORTED hipDeviceAttributeVirtualMemoryManagementSupported +#define CU_MEM_ALLOC_GRANULARITY_RECOMMENDED hipMemAllocationGranularityRecommended +#define CU_MEM_ALLOCATION_TYPE_PINNED hipMemAllocationTypePinned +#define CU_MEM_LOCATION_TYPE_DEVICE hipMemLocationTypeDevice +#define CU_MEM_ACCESS_FLAGS_PROT_READWRITE hipMemAccessFlagsProtReadWrite +#define CU_CHECK(fn) {hipError_t err = fn; if(err != hipSuccess) { GGML_ABORT("HipVMM Failure: %s\n", hipGetErrorString(err)); }} +#define __shfl_sync(mask, var, laneMask, width) __shfl(var, laneMask, width) +#define __shfl_xor_sync(mask, var, laneMask, width) __shfl_xor(var, laneMask, width) +#define cublasComputeType_t hipblasDatatype_t //deprecated, new hipblasComputeType_t not in 5.6 +#define cublasCreate hipblasCreate +#define cublasDestroy hipblasDestroy +#define cublasGemmEx hipblasGemmEx +#define cublasGemmBatchedEx hipblasGemmBatchedEx +#define cublasGemmStridedBatchedEx hipblasGemmStridedBatchedEx +#define cublasHandle_t hipblasHandle_t +#define cublasSetMathMode(handle, mode) CUBLAS_STATUS_SUCCESS +#define cublasSetStream hipblasSetStream +#define cublasSgemm hipblasSgemm +#define cublasStatus_t hipblasStatus_t +#define cublasOperation_t hipblasOperation_t +#define cudaDataType_t hipblasDatatype_t //deprecated, new hipblasDatatype not in 5.6 +#define cudaDeviceCanAccessPeer hipDeviceCanAccessPeer +#define cudaDeviceDisablePeerAccess hipDeviceDisablePeerAccess +#define cudaDeviceEnablePeerAccess hipDeviceEnablePeerAccess +#define cudaDeviceProp hipDeviceProp_t +#define cudaDeviceSynchronize hipDeviceSynchronize +#define cudaError_t hipError_t +#define cudaErrorPeerAccessAlreadyEnabled hipErrorPeerAccessAlreadyEnabled +#define cudaErrorPeerAccessNotEnabled hipErrorPeerAccessNotEnabled +#define cudaEventCreateWithFlags hipEventCreateWithFlags +#define cudaEventDisableTiming hipEventDisableTiming +#define cudaEventRecord hipEventRecord +#define cudaEventSynchronize hipEventSynchronize +#define cudaEvent_t hipEvent_t +#define cudaEventDestroy hipEventDestroy +#define cudaFree hipFree +#define cudaFreeHost hipHostFree +#define cudaGetDevice hipGetDevice +#define cudaGetDeviceCount hipGetDeviceCount +#define cudaGetDeviceProperties hipGetDeviceProperties +#define cudaGetErrorString hipGetErrorString +#define cudaGetLastError hipGetLastError +#define cudaHostRegister hipHostRegister +#define cudaHostRegisterPortable hipHostRegisterPortable +#define cudaHostRegisterReadOnly hipHostRegisterReadOnly +#define cudaHostUnregister hipHostUnregister +#define cudaLaunchHostFunc hipLaunchHostFunc +#define cudaMalloc hipMalloc +#define cudaMallocHost(ptr, size) hipHostMalloc(ptr, size, hipHostMallocDefault) +#define cudaMemcpy hipMemcpy +#define cudaMemcpyAsync hipMemcpyAsync +#define cudaMemcpyPeerAsync hipMemcpyPeerAsync +#define cudaMemcpy2DAsync hipMemcpy2DAsync +#define cudaMemcpyDeviceToDevice hipMemcpyDeviceToDevice +#define cudaMemcpyDeviceToHost hipMemcpyDeviceToHost +#define cudaMemcpyHostToDevice hipMemcpyHostToDevice +#define cudaMemcpyKind hipMemcpyKind +#define cudaMemset hipMemset +#define cudaMemsetAsync hipMemsetAsync +#define cudaMemGetInfo hipMemGetInfo +#define cudaOccupancyMaxPotentialBlockSize hipOccupancyMaxPotentialBlockSize +#define cudaSetDevice hipSetDevice +#define cuDeviceGet hipDeviceGet +#define CUdevice hipDevice_t +#define CUdeviceptr hipDeviceptr_t +#define cuMemUnmap hipMemUnmap +#define CUmemAccessDesc hipMemAccessDesc +#define cuMemAddressFree hipMemAddressFree +#define cuMemRelease hipMemRelease +#define CUmemGenericAllocationHandle hipMemGenericAllocationHandle_t +#define cuMemCreate hipMemCreate +#define cuMemAddressReserve hipMemAddressReserve +#define cuMemMap hipMemMap +#define cuMemSetAccess hipMemSetAccess +#define cuMemGetAllocationGranularity hipMemGetAllocationGranularity +#define CUmemAllocationProp hipMemAllocationProp +#define cuDeviceGetAttribute hipDeviceGetAttribute +#define cudaStreamCreateWithFlags hipStreamCreateWithFlags +#define cudaStreamDestroy hipStreamDestroy +#define cudaStreamFireAndForget hipStreamFireAndForget +#define cudaStreamNonBlocking hipStreamNonBlocking +#define cudaStreamPerThread hipStreamPerThread +#define cudaStreamSynchronize hipStreamSynchronize +#define cudaStreamWaitEvent(stream, event, flags) hipStreamWaitEvent(stream, event, flags) +#define cudaGraphExec_t hipGraphExec_t +#define cudaGraphNode_t hipGraphNode_t +#define cudaKernelNodeParams hipKernelNodeParams +#define cudaKernelNodeParams hipKernelNodeParams +#define cudaGraphExecDestroy hipGraphExecDestroy +#define cudaGraphLaunch hipGraphLaunch +#define cudaErrorGraphExecUpdateFailure hipErrorGraphExecUpdateFailure +#define cudaGraphExecUpdateResultInfo hipGraphExecUpdateResult +#define cudaGraphNodeType hipGraphNodeType +#define cudaGraphNodeTypeKernel hipGraphNodeTypeKernel +#define cudaGraphInstantiate hipGraphInstantiate +#define cudaStreamEndCapture hipStreamEndCapture +#define cudaGraphDestroy hipGraphDestroy +#define cudaGraphKernelNodeSetParams hipGraphKernelNodeSetParams +#define cudaErrorInvalidDeviceFunction hipErrorInvalidDeviceFunction +#define cudaGraphKernelNodeGetParams hipGraphKernelNodeGetParams +#define cudaGraphNodeGetType hipGraphNodeGetType +#define cudaGraphGetNodes hipGraphGetNodes +#define cudaGraphExecUpdate hipGraphExecUpdate +#define cudaStreamCaptureModeRelaxed hipStreamCaptureModeRelaxed +#define cudaStreamBeginCapture hipStreamBeginCapture +#define cudaGraph_t hipGraph_t +#define cudaStream_t hipStream_t +#define cudaSuccess hipSuccess +#define cudaHostFn_t hipHostFn_t +#define __trap() do { abort(); __builtin_unreachable(); } while(0) +#define CUBLAS_STATUS_SUCCESS HIPBLAS_STATUS_SUCCESS +#define CUBLAS_STATUS_NOT_INITIALIZED HIPBLAS_STATUS_NOT_INITIALIZED +#define CUBLAS_STATUS_ALLOC_FAILED HIPBLAS_STATUS_ALLOC_FAILED +#define CUBLAS_STATUS_INVALID_VALUE HIPBLAS_STATUS_INVALID_VALUE +#define CUBLAS_STATUS_ARCH_MISMATCH HIPBLAS_STATUS_ARCH_MISMATCH +#define CUBLAS_STATUS_MAPPING_ERROR HIPBLAS_STATUS_MAPPING_ERROR +#define CUBLAS_STATUS_EXECUTION_FAILED HIPBLAS_STATUS_EXECUTION_FAILED +#define CUBLAS_STATUS_INTERNAL_ERROR HIPBLAS_STATUS_INTERNAL_ERROR +#define CUBLAS_STATUS_NOT_SUPPORTED HIPBLAS_STATUS_NOT_SUPPORTED + +#define __CUDA_ARCH__ 1300 + +#if defined(__gfx803__) || defined(__gfx900__) || defined(__gfx906__) +#define GCN +#endif + +#if defined(__gfx908__) || defined(__gfx90a__) || defined(__gfx942__) +#define CDNA +#endif + +#if defined(__gfx1100__) || defined(__gfx1101__) || defined(__gfx1102__) || defined(__gfx1103__) || \ + defined(__gfx1150__) || defined(__gfx1151__) +#define RDNA3 +#endif + +#if defined(__gfx1030__) || defined(__gfx1031__) || defined(__gfx1032__) || defined(__gfx1033__) || \ + defined(__gfx1034__) || defined(__gfx1035__) || defined(__gfx1036__) || defined(__gfx1037__) +#define RDNA2 +#endif + +#if defined(__gfx1010__) || defined(__gfx1012__) +#define RDNA1 +#endif + +#ifndef __has_builtin + #define __has_builtin(x) 0 +#endif + +typedef hip_bfloat16 nv_bfloat16; diff --git a/archive/kt-sft/csrc/ktransformers_ext/vendors/musa.h b/archive/kt-sft/csrc/ktransformers_ext/vendors/musa.h new file mode 100644 index 0000000..6cc1b69 --- /dev/null +++ b/archive/kt-sft/csrc/ktransformers_ext/vendors/musa.h @@ -0,0 +1,137 @@ +#pragma once + +#include +#include +#include +#include +#include +#define CUBLAS_COMPUTE_16F CUDA_R_16F +#define CUBLAS_COMPUTE_32F CUDA_R_32F +#define CUBLAS_COMPUTE_32F_FAST_16F MUBLAS_COMPUTE_32F_FAST_16F +#define CUBLAS_GEMM_DEFAULT MUBLAS_GEMM_DEFAULT +#define CUBLAS_GEMM_DEFAULT_TENSOR_OP MUBLAS_GEMM_DEFAULT +#define CUBLAS_OP_N MUBLAS_OP_N +#define CUBLAS_OP_T MUBLAS_OP_T +#define CUBLAS_STATUS_SUCCESS MUBLAS_STATUS_SUCCESS +#define CUBLAS_TF32_TENSOR_OP_MATH MUBLAS_MATH_MODE_DEFAULT +#define CUDA_R_16F MUSA_R_16F +#define CUDA_R_32F MUSA_R_32F +#define cublasComputeType_t cudaDataType_t +#define cublasCreate mublasCreate +#define cublasDestroy mublasDestroy +#define cublasGemmEx mublasGemmEx +#define cublasGemmBatchedEx mublasGemmBatchedEx +#define cublasGemmStridedBatchedEx mublasGemmStridedBatchedEx +#define cublasHandle_t mublasHandle_t +#define cublasSetMathMode mublasSetMathMode +#define cublasSetStream mublasSetStream +#define cublasSgemm mublasSgemm +#define cublasStatus_t mublasStatus_t +#define cublasOperation_t mublasOperation_t +#define cublasGetStatusString mublasStatus_to_string +#define cudaDataType_t musaDataType_t +#define cudaDeviceCanAccessPeer musaDeviceCanAccessPeer +#define cudaDeviceDisablePeerAccess musaDeviceDisablePeerAccess +#define cudaDeviceEnablePeerAccess musaDeviceEnablePeerAccess +#define cudaDeviceProp musaDeviceProp +#define cudaDeviceSynchronize musaDeviceSynchronize +#define cudaError_t musaError_t +#define cudaErrorPeerAccessAlreadyEnabled musaErrorPeerAccessAlreadyEnabled +#define cudaErrorPeerAccessNotEnabled musaErrorPeerAccessNotEnabled +#define cudaEventCreateWithFlags musaEventCreateWithFlags +#define cudaEventDisableTiming musaEventDisableTiming +#define cudaEventRecord musaEventRecord +#define cudaEventSynchronize musaEventSynchronize +#define cudaEvent_t musaEvent_t +#define cudaEventDestroy musaEventDestroy +#define cudaFree musaFree +#define cudaFreeHost musaFreeHost +#define cudaGetDevice musaGetDevice +#define cudaGetDeviceCount musaGetDeviceCount +#define cudaGetDeviceProperties musaGetDeviceProperties +#define cudaGetErrorString musaGetErrorString +#define cudaGetLastError musaGetLastError +#define cudaHostRegister musaHostRegister +#define cudaHostRegisterPortable musaHostRegisterPortable +#define cudaHostRegisterReadOnly musaHostRegisterReadOnly +#define cudaHostUnregister musaHostUnregister +#define cudaLaunchHostFunc musaLaunchHostFunc +#define cudaMalloc musaMalloc +#define cudaMallocHost musaMallocHost +#define cudaMallocManaged musaMallocManaged +#define cudaMemcpy musaMemcpy +#define cudaMemcpyAsync musaMemcpyAsync +#define cudaMemcpyPeerAsync musaMemcpyPeerAsync +#define cudaMemcpy2DAsync musaMemcpy2DAsync +#define cudaMemcpyDeviceToDevice musaMemcpyDeviceToDevice +#define cudaMemcpyDeviceToHost musaMemcpyDeviceToHost +#define cudaMemcpyHostToDevice musaMemcpyHostToDevice +#define cudaMemcpyKind musaMemcpyKind +#define cudaMemset musaMemset +#define cudaMemsetAsync musaMemsetAsync +#define cudaMemGetInfo musaMemGetInfo +#define cudaOccupancyMaxPotentialBlockSize musaOccupancyMaxPotentialBlockSize +#define cudaSetDevice musaSetDevice +#define cudaStreamCreateWithFlags musaStreamCreateWithFlags +#define cudaStreamDestroy musaStreamDestroy +#define cudaStreamFireAndForget musaStreamFireAndForget +#define cudaStreamNonBlocking musaStreamNonBlocking +#define cudaStreamPerThread musaStreamPerThread +#define cudaStreamSynchronize musaStreamSynchronize +#define cudaStreamWaitEvent musaStreamWaitEvent +#define cudaStream_t musaStream_t +#define cudaSuccess musaSuccess + +// Additional mappings for MUSA virtual memory pool +#define CU_DEVICE_ATTRIBUTE_VIRTUAL_MEMORY_MANAGEMENT_SUPPORTED MU_DEVICE_ATTRIBUTE_VIRTUAL_ADDRESS_MANAGEMENT_SUPPORTED +#define CU_MEM_ACCESS_FLAGS_PROT_READWRITE MU_MEM_ACCESS_FLAGS_PROT_READWRITE +#define CU_MEM_ALLOC_GRANULARITY_RECOMMENDED MU_MEM_ALLOC_GRANULARITY_RECOMMENDED +#define CU_MEM_ALLOCATION_TYPE_PINNED MU_MEM_ALLOCATION_TYPE_PINNED +#define CU_MEM_LOCATION_TYPE_DEVICE MU_MEM_LOCATION_TYPE_DEVICE +#define CUdevice MUdevice +#define CUdeviceptr MUdeviceptr +#define CUmemAccessDesc MUmemAccessDesc +#define CUmemAllocationProp MUmemAllocationProp +#define CUmemGenericAllocationHandle MUmemGenericAllocationHandle +#define cuDeviceGet muDeviceGet +#define cuDeviceGetAttribute muDeviceGetAttribute +#define cuMemAddressFree muMemAddressFree +#define cuMemAddressReserve muMemAddressReserve +#define cuMemCreate muMemCreate +#define cuMemGetAllocationGranularity muMemGetAllocationGranularity +#define cuMemMap muMemMap +#define cuMemRelease muMemRelease +#define cuMemSetAccess muMemSetAccess +#define cuMemUnmap muMemUnmap +#define cudaFuncAttributeMaxDynamicSharedMemorySize musaFuncAttributeMaxDynamicSharedMemorySize +#define cudaFuncSetAttribute musaFuncSetAttribute +#define cudaMemcpy3DPeerParms musaMemcpy3DPeerParms +#define make_cudaExtent make_musaExtent +#define make_cudaPitchedPtr make_musaPitchedPtr + +// Additional mappings for MUSA graphs +#define CUDA_SUCCESS MUSA_SUCCESS +#define CUresult MUresult +#define cuGetErrorString muGetErrorString +#define cudaErrorGraphExecUpdateFailure musaErrorGraphExecUpdateFailure +#define cudaErrorInvalidDeviceFunction musaErrorInvalidDeviceFunction +#define cudaGraphDestroy musaGraphDestroy +#define cudaGraphExecDestroy musaGraphExecDestroy +#define cudaGraphExec_t musaGraphExec_t +#define cudaGraphExecUpdate musaGraphExecUpdate +#define cudaGraphExecUpdateResultInfo musaGraphExecUpdateResult +#define cudaGraphGetNodes musaGraphGetNodes +#define cudaGraphInstantiate musaGraphInstantiate +#define cudaGraphKernelNodeGetParams musaGraphKernelNodeGetParams +#define cudaGraphKernelNodeSetParams musaGraphKernelNodeSetParams +#define cudaGraphLaunch musaGraphLaunch +#define cudaGraphNodeGetType musaGraphNodeGetType +#define cudaGraphNode_t musaGraphNode_t +#define cudaGraphNodeType musaGraphNodeType +#define cudaGraphNodeTypeKernel musaGraphNodeTypeKernel +#define cudaGraph_t musaGraph_t +#define cudaKernelNodeParams musaKernelNodeParams +#define cudaStreamCaptureModeRelaxed musaStreamCaptureModeRelaxed +#define cudaStreamEndCapture musaStreamEndCapture + +typedef mt_bfloat16 nv_bfloat16; diff --git a/archive/kt-sft/csrc/ktransformers_ext/vendors/vendor.h b/archive/kt-sft/csrc/ktransformers_ext/vendors/vendor.h new file mode 100644 index 0000000..8470438 --- /dev/null +++ b/archive/kt-sft/csrc/ktransformers_ext/vendors/vendor.h @@ -0,0 +1,13 @@ +#ifndef CPUINFER_VENDOR_VENDOR_H +#define CPUINFER_VENDOR_VENDOR_H + +#ifdef USE_CUDA +#include "cuda.h" +#elif USE_HIP +#define __HIP_PLATFORM_AMD__ +#include "hip.h" +#elif USE_MUSA +#include "musa.h" +#endif + +#endif // CPUINFER_VENDOR_VENDOR_H \ No newline at end of file diff --git a/archive/kt-sft/install-with-cache.sh b/archive/kt-sft/install-with-cache.sh new file mode 100755 index 0000000..cef4341 --- /dev/null +++ b/archive/kt-sft/install-with-cache.sh @@ -0,0 +1,26 @@ +#!/bin/bash +set -e + +# clear build dirs +# rm -rf build +# rm -rf *.egg-info +# rm -rf csrc/build +# rm -rf csrc/ktransformers_ext/build +# rm -rf csrc/ktransformers_ext/cuda/build +# rm -rf csrc/ktransformers_ext/cuda/dist +# rm -rf csrc/ktransformers_ext/cuda/*.egg-info +rm -rf ~/.ktransformers +echo "Installing python dependencies from requirements.txt" +pip install -r requirements-local_chat.txt +pip install -r ktransformers/server/requirements.txt +echo "Installing ktransformers" +KTRANSFORMERS_FORCE_BUILD=TRUE USE_BALANCE_SERVE=1 pip install -v . --no-build-isolation +pip install third_party/custom_flashinfer/ -v + +# SITE_PACKAGES=$(python -c "import site; print(site.getsitepackages()[0])") +# echo "Copying thirdparty libs to $SITE_PACKAGES" +# cp -a csrc/balance_serve/build/third_party/prometheus-cpp/lib/libprometheus-cpp-*.so* $SITE_PACKAGES/ +# patchelf --set-rpath '$ORIGIN' $SITE_PACKAGES/sched_ext.cpython* + + +echo "Installation completed successfully" diff --git a/archive/kt-sft/install.bat b/archive/kt-sft/install.bat new file mode 100644 index 0000000..dc429e4 --- /dev/null +++ b/archive/kt-sft/install.bat @@ -0,0 +1,16 @@ +@echo off + +REM clear build dirs +rmdir /S /Q ktransformers\ktransformers_ext\build +rmdir /S /Q ktransformers\ktransformers_ext\cuda\build +rmdir /S /Q ktransformers\ktransformers_ext\cuda\dist +rmdir /S /Q ktransformers\ktransformers_ext\out +del /F /Q ktransformers\ktransformers_ext\cuda\*.egg-info + +echo Installing python dependencies from requirements.txt +pip install -r requirements-local_chat.txt + +echo Installing ktransformers +set KTRANSFORMERS_FORCE_BUILD=TRUE +pip install . --no-build-isolation +echo Installation completed successfully \ No newline at end of file diff --git a/archive/kt-sft/install.sh b/archive/kt-sft/install.sh new file mode 100755 index 0000000..e05b5d8 --- /dev/null +++ b/archive/kt-sft/install.sh @@ -0,0 +1,44 @@ +#!/bin/bash +set -e + +CWD=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd) + +# default backend +DEV="cuda" + +# parse --dev argument +while [[ "$#" -gt 0 ]]; do + case $1 in + --dev) DEV="$2"; shift ;; + *) echo "Unknown parameter passed: $1"; exit 1 ;; + esac + shift +done +export DEV_BACKEND="$DEV" +echo "Selected backend: $DEV_BACKEND" + +# clear build dirs +rm -rf build +rm -rf *.egg-info +rm -rf csrc/build +rm -rf csrc/ktransformers_ext/build +rm -rf csrc/ktransformers_ext/cuda/build +rm -rf csrc/ktransformers_ext/cuda/dist +rm -rf csrc/ktransformers_ext/cuda/*.egg-info +rm -rf ~/.ktransformers +echo "Installing python dependencies from requirements.txt" +pip install -r "${CWD}/requirements-sft.txt" + +echo "Installing ktransformers" +KTRANSFORMERS_FORCE_BUILD=TRUE pip install -v "${CWD}" --no-build-isolation + +if [[ "$DEV_BACKEND" == "cuda" ]]; then + echo "Installing custom_flashinfer for CUDA backend" + pip install "${CWD}/../third_party/custom_flashinfer/" +fi +# SITE_PACKAGES=$(python -c "import site; print(site.getsitepackages()[0])") +# echo "Copying thirdparty libs to $SITE_PACKAGES" +# cp -a csrc/balance_serve/build/third_party/prometheus-cpp/lib/libprometheus-cpp-*.so* $SITE_PACKAGES/ +# patchelf --set-rpath '$ORIGIN' $SITE_PACKAGES/sched_ext.cpython* + +echo "Installation completed successfully" \ No newline at end of file diff --git a/archive/kt-sft/ktransformers/__init__.py b/archive/kt-sft/ktransformers/__init__.py new file mode 100644 index 0000000..ab915fa --- /dev/null +++ b/archive/kt-sft/ktransformers/__init__.py @@ -0,0 +1,20 @@ +#!/usr/bin/env python +# coding=utf-8 +''' +Description : +Author : kkk1nak0 +Date : 2024-08-15 07:34:46 +Version : 1.0.0 +LastEditors : chenxl +LastEditTime : 2025-02-15 03:53:02 +''' +import sys +import os + +# Import version from shared version.py at project root +_root_dir = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) +sys.path.insert(0, _root_dir) +try: + from version import __version__ +finally: + sys.path.pop(0) diff --git a/archive/kt-sft/ktransformers/configs/config.yaml b/archive/kt-sft/ktransformers/configs/config.yaml new file mode 100644 index 0000000..ed1713b --- /dev/null +++ b/archive/kt-sft/ktransformers/configs/config.yaml @@ -0,0 +1,72 @@ +log: + dir: "logs" + file: "lexllama.log" + #log level: debug, info, warn, error, crit + level: "debug" + backup_count: -1 + +server: + ip: 0.0.0.0 + port: 10002 + +db: + type: "sqllite" + database: "server.db" + host: "./" + pool_size: 10 + +user: + secret_key: "981f1dd2a44e27d68759d0252a486568ed43480b4e616a26e3af3709c3a7ce73" + algorithm: "HS256" + +model: + # type: transformers + # type: balance_serve + type: ktransformers + + name: DeepSeek-Coder-V2-Instruct + path: deepseek-ai/DeepSeek-V2-Lite-Chat + gguf_path: ./DeepSeek-V2-Lite-Chat-GGUF + + device: cuda:0 + cache_lens: 16384 + max_new_tokens: 500 +web: + mount: False + open_cross_domain: True + +ext: + cpu_infer: 10 + +long_context: + max_seq_len: 32000 + block_size: 128 + local_windows_len: 4096 + second_select_num: 32 + anchor_type: DYNAMIC + kv_type: FP16 + dense_layer_num: 2 + anchor_num: 1 + preselect_block: True + head_select_mode: SHARED + preselect_block_count: 32 + layer_step: 1 + token_step: + +local_chat: + prompt_file: "" + +async_server: + sched_strategy: "FCFS" + sched_port: 56441 + sched_metrics_port: 54321 + kvc2_metrics_port: 54391 + max_batch_size: 4 # decode count + prefill count, in one mini batch + +attn: + page_size: 256 + chunk_size: 256 +kvc2: + gpu_only: true + utilization_percentage: 1.0 + cpu_memory_size_GB: 500 diff --git a/archive/kt-sft/ktransformers/configs/log_config.ini b/archive/kt-sft/ktransformers/configs/log_config.ini new file mode 100644 index 0000000..8b5b2e0 --- /dev/null +++ b/archive/kt-sft/ktransformers/configs/log_config.ini @@ -0,0 +1,46 @@ +[loggers] +keys=root,uvicorn,uvicornError,uvicornAccess + +[handlers] +keys=consoleHandler,fileHandler + +[formatters] +keys=detailedFormatter + +[logger_root] +level=INFO +handlers=consoleHandler + +[logger_uvicorn] +level=INFO +handlers=consoleHandler,fileHandler +qualname=uvicorn +propagate=0 + +[logger_uvicornError] +level=ERROR +handlers=consoleHandler,fileHandler +qualname=uvicorn.error +propagate=0 + +[logger_uvicornAccess] +level=INFO +handlers=consoleHandler,fileHandler +qualname=uvicorn.access +propagate=0 + +[handler_consoleHandler] +class=StreamHandler +level=INFO +formatter=detailedFormatter +args=(sys.stdout,) + +[handler_fileHandler] +class=logging.FileHandler +level=INFO +formatter=detailedFormatter +args=('uvicorn_logs.log', 'a') + +[formatter_detailedFormatter] +format=%(asctime)s - %(name)s - %(levelname)s - [%(filename)s:%(lineno)d] - %(message)s +datefmt=%Y-%m-%d %H:%M:%S diff --git a/archive/kt-sft/ktransformers/configs/model_config/config.json b/archive/kt-sft/ktransformers/configs/model_config/config.json new file mode 100644 index 0000000..d215df1 --- /dev/null +++ b/archive/kt-sft/ktransformers/configs/model_config/config.json @@ -0,0 +1,59 @@ +{ + "architectures": [ + "DeepseekV2ForCausalLM" + ], + "attention_bias": false, + "attention_dropout": 0.0, + "auto_map": { + "AutoConfig": "configuration_deepseek.DeepseekV2Config", + "AutoModel": "modeling_deepseek.DeepseekV2Model", + "AutoModelForCausalLM": "modeling_deepseek.DeepseekV2ForCausalLM" + }, + "aux_loss_alpha": 0.001, + "bos_token_id": 100000, + "eos_token_id": 100001, + "first_k_dense_replace": 1, + "hidden_act": "silu", + "hidden_size": 2048, + "initializer_range": 0.02, + "intermediate_size": 10944, + "kv_lora_rank": 512, + "max_position_embeddings": 163840, + "model_type": "deepseek_v2", + "moe_intermediate_size": 1408, + "moe_layer_freq": 1, + "n_group": 1, + "n_routed_experts": 64, + "n_shared_experts": 2, + "norm_topk_prob": false, + "num_attention_heads": 16, + "num_experts_per_tok": 6, + "num_hidden_layers": 27, + "num_key_value_heads": 16, + "pretraining_tp": 1, + "q_lora_rank": null, + "qk_nope_head_dim": 128, + "qk_rope_head_dim": 64, + "rms_norm_eps": 1e-06, + "rope_scaling": { + "beta_fast": 32, + "beta_slow": 1, + "factor": 40, + "mscale": 0.707, + "mscale_all_dim": 0.707, + "original_max_position_embeddings": 4096, + "type": "yarn" + }, + "rope_theta": 10000, + "routed_scaling_factor": 1.0, + "scoring_func": "softmax", + "seq_aux": true, + "tie_word_embeddings": false, + "topk_group": 1, + "topk_method": "greedy", + "torch_dtype": "bfloat16", + "transformers_version": "4.33.1", + "use_cache": true, + "v_head_dim": 128, + "vocab_size": 102400 +} \ No newline at end of file diff --git a/archive/kt-sft/ktransformers/configs/model_config/configuration_deepseek.py b/archive/kt-sft/ktransformers/configs/model_config/configuration_deepseek.py new file mode 100644 index 0000000..82e0f5d --- /dev/null +++ b/archive/kt-sft/ktransformers/configs/model_config/configuration_deepseek.py @@ -0,0 +1,206 @@ +from transformers.configuration_utils import PretrainedConfig +from transformers.utils import logging + +logger = logging.get_logger(__name__) + +DEEPSEEK_PRETRAINED_CONFIG_ARCHIVE_MAP = {} +class DeepseekV2Config(PretrainedConfig): + r""" + This is the configuration class to store the configuration of a [`DeepseekV2Model`]. It is used to instantiate an DeepSeek + model according to the specified arguments, defining the model architecture. Instantiating a configuration with the + defaults will yield a similar configuration to that of the DeepSeek-V2. + + 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 102400): + Vocabulary size of the Deep model. Defines the number of different tokens that can be represented by the + `inputs_ids` passed when calling [`DeepseekV2Model`] + hidden_size (`int`, *optional*, defaults to 4096): + Dimension of the hidden representations. + intermediate_size (`int`, *optional*, defaults to 11008): + Dimension of the MLP representations. + moe_intermediate_size (`int`, *optional*, defaults to 1407): + Dimension of the MoE representations. + num_hidden_layers (`int`, *optional*, defaults to 32): + Number of hidden layers in the Transformer decoder. + num_attention_heads (`int`, *optional*, defaults to 32): + Number of attention heads for each attention layer in the Transformer decoder. + n_shared_experts (`int`, *optional*, defaults to None): + Number of shared experts, None means dense model. + n_routed_experts (`int`, *optional*, defaults to None): + Number of routed experts, None means dense model. + routed_scaling_factor (`float`, *optional*, defaults to 1.0): + Scaling factor or routed experts. + topk_method (`str`, *optional*, defaults to `gready`): + Topk method used in routed gate. + n_group (`int`, *optional*, defaults to None): + Number of groups for routed experts. + topk_group (`int`, *optional*, defaults to None): + Number of selected groups for each token(for each token, ensuring the selected experts is only within `topk_group` groups). + num_experts_per_tok (`int`, *optional*, defaults to None): + Number of selected experts, None means dense model. + moe_layer_freq (`int`, *optional*, defaults to 1): + The frequency of the MoE layer: one expert layer for every `moe_layer_freq - 1` dense layers. + first_k_dense_replace (`int`, *optional*, defaults to 0): + Number of dense layers in shallow layers(embed->dense->dense->...->dense->moe->moe...->lm_head). + \--k dense layers--/ + norm_topk_prob (`bool`, *optional*, defaults to False): + Whether to normalize the weights of the routed experts. + scoring_func (`str`, *optional*, defaults to 'softmax'): + Method of computing expert weights. + aux_loss_alpha (`float`, *optional*, defaults to 0.001): + Auxiliary loss weight coefficient. + seq_aux = (`bool`, *optional*, defaults to True): + Whether to compute the auxiliary loss for each individual sample. + num_key_value_heads (`int`, *optional*): + This is the number of key_value heads that should be used to implement Grouped Query Attention. If + `num_key_value_heads=num_attention_heads`, the model will use Multi Head Attention (MHA), if + `num_key_value_heads=1 the model will use Multi Query Attention (MQA) otherwise GQA is used. When + converting a multi-head checkpoint to a GQA checkpoint, each group key and value head should be constructed + by meanpooling all the original heads within that group. For more details checkout [this + paper](https://arxiv.org/pdf/2305.13245.pdf). If it is not specified, will default to + `num_attention_heads`. + hidden_act (`str` or `function`, *optional*, defaults to `"silu"`): + The non-linear activation function (function or string) in the decoder. + max_position_embeddings (`int`, *optional*, defaults to 2048): + The maximum sequence length that this model might ever be used with. + 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`. + pad_token_id (`int`, *optional*): + Padding token id. + bos_token_id (`int`, *optional*, defaults to 1): + Beginning of stream token id. + eos_token_id (`int`, *optional*, defaults to 2): + End of stream token id. + pretraining_tp (`int`, *optional*, defaults to 1): + Experimental feature. Tensor parallelism rank used during pretraining. Please refer to [this + document](https://huggingface.co/docs/transformers/parallelism) to understand more about it. This value is + necessary to ensure exact reproducibility of the pretraining results. Please refer to [this + issue](https://github.com/pytorch/pytorch/issues/76232). + tie_word_embeddings (`bool`, *optional*, defaults to `False`): + Whether to tie weight embeddings + 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. Currently supports two scaling + strategies: linear and dynamic. Their scaling factor must be a float greater than 1. The expected format is + `{"type": strategy name, "factor": scaling factor}`. When using this flag, don't update + `max_position_embeddings` to the expected new maximum. + 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. + attention_dropout (`float`, *optional*, defaults to 0.0): + The dropout ratio for the attention probabilities. + + ```python + >>> from transformers import DeepseekV2Model, DeepseekV2Config + + >>> # Initializing a Deepseek-V2 style configuration + >>> configuration = DeepseekV2Config() + + >>> # Accessing the model configuration + >>> configuration = model.config + ```""" + + model_type = "deepseek_v2" + keys_to_ignore_at_inference = ["past_key_values"] + + def __init__( + self, + vocab_size=102400, + hidden_size=4096, + intermediate_size=11008, + moe_intermediate_size = 1407, + num_hidden_layers=30, + num_attention_heads=32, + num_key_value_heads=32, + n_shared_experts = None, + n_routed_experts = None, + ep_size = 1, + routed_scaling_factor = 1.0, + kv_lora_rank = 512, + q_lora_rank = 1536, + qk_rope_head_dim = 64, + v_head_dim = 128, + qk_nope_head_dim = 128, + topk_method = 'gready', + n_group = None, + topk_group = None, + num_experts_per_tok = None, + moe_layer_freq = 1, + first_k_dense_replace = 0, + norm_topk_prob = False, + scoring_func = 'softmax', + aux_loss_alpha = 0.001, + seq_aux = True, + hidden_act="silu", + max_position_embeddings=2048, + initializer_range=0.02, + rms_norm_eps=1e-6, + use_cache=True, + pad_token_id=None, + bos_token_id=100000, + eos_token_id=100001, + pretraining_tp=1, + tie_word_embeddings=False, + rope_theta=10000.0, + rope_scaling=None, + attention_bias=False, + attention_dropout=0.0, + **kwargs, + ): + self.vocab_size = vocab_size + self.max_position_embeddings = max_position_embeddings + self.hidden_size = hidden_size + self.intermediate_size = intermediate_size + self.moe_intermediate_size = moe_intermediate_size + self.num_hidden_layers = num_hidden_layers + self.num_attention_heads = num_attention_heads + self.n_shared_experts = n_shared_experts + self.n_routed_experts = n_routed_experts + self.ep_size = ep_size + self.routed_scaling_factor = routed_scaling_factor + self.kv_lora_rank = kv_lora_rank + self.q_lora_rank = q_lora_rank + self.qk_rope_head_dim = qk_rope_head_dim + self.v_head_dim = v_head_dim + self.qk_nope_head_dim = qk_nope_head_dim + self.topk_method = topk_method + self.n_group = n_group + self.topk_group = topk_group + self.num_experts_per_tok = num_experts_per_tok + self.moe_layer_freq = moe_layer_freq + self.first_k_dense_replace = first_k_dense_replace + self.norm_topk_prob = norm_topk_prob + self.scoring_func = scoring_func + self.aux_loss_alpha = aux_loss_alpha + self.seq_aux = seq_aux + # for backward compatibility + if num_key_value_heads is None: + num_key_value_heads = num_attention_heads + + self.num_key_value_heads = num_key_value_heads + self.hidden_act = hidden_act + self.initializer_range = initializer_range + self.rms_norm_eps = rms_norm_eps + self.pretraining_tp = pretraining_tp + self.use_cache = use_cache + self.rope_theta = rope_theta + self.rope_scaling = rope_scaling + self.attention_bias = attention_bias + self.attention_dropout = attention_dropout + + super().__init__( + pad_token_id=pad_token_id, + bos_token_id=bos_token_id, + eos_token_id=eos_token_id, + tie_word_embeddings=tie_word_embeddings, + **kwargs, + ) \ No newline at end of file diff --git a/archive/kt-sft/ktransformers/ktransformers_ext/operators/custom_marlin/quantize/utils/__init__.py b/archive/kt-sft/ktransformers/ktransformers_ext/operators/custom_marlin/quantize/utils/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/archive/kt-sft/ktransformers/ktransformers_ext/operators/custom_marlin/quantize/utils/format_24.py b/archive/kt-sft/ktransformers/ktransformers_ext/operators/custom_marlin/quantize/utils/format_24.py new file mode 100644 index 0000000..01c8cf7 --- /dev/null +++ b/archive/kt-sft/ktransformers/ktransformers_ext/operators/custom_marlin/quantize/utils/format_24.py @@ -0,0 +1,308 @@ +# +# Modified by Roberto Lopez Castro (roberto.lopez.castro@udc.es). +# + +import torch + + +# This is PyTorch implementation of main part of reorder_meta() +# function, from tools/util/include/cutlass/util/host_reorder.h file +# of CUTLASS source tree. Furthermore, CUTLASS template for sparse +# GEMM decides upon layout of this matrix, and at the moment for the +# sparse GEMM executed on tensor cores, this is layout described by +# ColumnMajorInterleaved<2> data structure, in +# include/cutlass/layout/matrix.h of CUTLASS source tree. The +# reordering of meta matrix into meta_reordered matrix calculated +# according to these segments of CUTLASS code is re-implemented here. +# Note that this calculation produces offsets for scattering metadata +# matrix elements into reordered metadata matrix elements (or, +# equivalently, for gathering reordered metadata matrix element back +# into metadata matrix elements). +def _calculate_meta_reordering_scatter_offsets(m, meta_ncols, meta_dtype, + device): + dst_rows = torch.arange(0, m, device=device)[:, None].repeat(1, meta_ncols) + dst_cols = torch.arange(0, meta_ncols, device=device).repeat(m, 1) + + # Reorder the rows, then swizzle the 2x2 blocks. + group_x = 64 + group_y = 32 if meta_dtype.itemsize == 2 else 16 + + dst_rows = (dst_rows // group_x * group_x + (dst_rows % 2) * 2 + + (dst_rows % 8) // 4 + ((dst_rows % group_y) % 4) // 2 * 32 + + ((dst_rows % group_x) // 8) * 4) + + topright = ((dst_rows % 2 == 0) & (dst_cols % 2 == 1)).to(torch.int8) + bottomleft = ((dst_rows % 2 == 1) & (dst_cols % 2 == 0)).to(torch.int8) + dst_rows += topright - bottomleft + dst_cols -= topright - bottomleft + + # Assumed that meta tensor is to be stored in CUTLASS + # InterleavedColumnMajor layout, and reverse engineered + # corresponding code to store values into this tensor. + interleave = 2 + cols_maj = dst_cols // interleave + cols_min = dst_cols % interleave + return (cols_maj * m * interleave + dst_rows * interleave + + cols_min).view(-1) + + +# This function converts dense matrix into sparse semi-structured +# representation, producing "compressed" matrix, in the layout used by +# CUTLASS backend, and corresponding metadata matrix. +def sparse_semi_structured_from_dense_cutlass(dense): + if dense.dim() != 2: + raise RuntimeError( + f"Expected 2-dimensional dense tensor, got {dense.dim()}-dimensional tensor" # noqa: E501 + ) + + m, k = dense.shape + device = dense.device + + meta_dtype = torch.int8 + if dense.dtype == torch.int8: + meta_dtype = torch.int32 + elif dense.dtype in [torch.half, torch.bfloat16, torch.float, torch.int32]: + meta_dtype = torch.int16 + else: + raise RuntimeError(f"Invalid datatype {dense.dtype} of dense matrix") + quadbits_per_meta_elem = meta_dtype.itemsize * 8 // 4 + if quadbits_per_meta_elem not in (4, 8): + raise RuntimeError( + "Invalid number of elements per meta element calculated") + + if meta_dtype == torch.int32: + if m % 16 != 0: + raise RuntimeError( + f"Number of rows of dense matrix {m} must be divisible by 16") + else: + if m % 32 != 0: + raise RuntimeError( + f"Number of rows of dense matrix {m} must be divisible by 32") + if k % (4 * quadbits_per_meta_elem) != 0: + raise RuntimeError( + f"Number of columns of dense matrix {k} must be divisible by {4 * quadbits_per_meta_elem}" # noqa: E501 + ) + + if dense.dtype != torch.float: + ksparse = 4 + dense_4 = dense.view(-1, k // ksparse, ksparse) + m0, m1, m2, m3 = (dense_4 != 0).unbind(-1) + else: + ksparse = 2 + dense_2 = dense.view(-1, k // ksparse, ksparse) + m0, m2 = m1, m3 = (dense_2 != 0).unbind(-1) + meta_ncols = k // (ksparse * quadbits_per_meta_elem) + + # Encoding quadruples of True/False values as follows: + # [True, True, False, False] -> 0b0100 + # [True, False, True, False] -> 0b1000 + # [False, True, True, False] -> 0b1001 + # [True, False, False, True ] -> 0b1100 + # [False, True, False, True ] -> 0b1101 + # [False, False, True, True ] -> 0b1110 + # Thus, lower two bits in the encoding are index of the True value + # at the lowest index in the quadruple, and the higher two bits in + # the encoding are index of the other True value in the quadruple. + # In case there are less than two True values, than False value or + # values at some index or indices are considered True for the + # encoding. In case there are more than two True values, then the + # excess True value(s) at some indices are considered False for + # the encoding. The exact encodings used for these cases are as + # follows: + # [False, False, False, False] -> 0b1110 + # [False, False, False, True ] -> 0b1110 + # [False, False, True, False] -> 0b1110 + # [False, True, False, False] -> 0b1001 + # [False, True, True, True ] -> 0b1101 + # [True, False, False, False] -> 0b1000 + # [True, False, True, True ] -> 0b1100 + # [True, True, False, True ] -> 0b0100 + # [True, True, True, False] -> 0b0100 + # [True, True, True, True ] -> 0b0100 + # These particular encodings are chosen, with the help of Espresso + # logic minimizer software, for the purpose of minimization of + # corresponding Boolean functions, that translate non-zero flags + # into encoding bits. Note also possible choices for the first + # and last of these encodings were limited only to (0b0100, + # 0b1110), in order to produce valid encodings for 1:2 sparsity + # case. + + expr0 = m0 & m1 + expr1 = ~m0 & m1 + expr2 = ~m0 & ~m1 + bit0 = expr1 + bit1 = expr2 + bit2 = expr0 | expr2 | m3 + bit3 = expr1 | ~m1 + idxs0 = bit0 | (bit1.to(torch.int64) << 1) + idxs1 = bit2 | (bit3.to(torch.int64) << 1) + + if dense.dtype != torch.float: + sparse0 = dense_4.gather( + -1, idxs0.unsqueeze(-1)) # type: ignore[possibly-undefined] + sparse1 = dense_4.gather(-1, idxs1.unsqueeze(-1)) + sparse = torch.stack((sparse0, sparse1), dim=-1).view(m, k // 2) + else: + sparse = dense_2.gather(-1, + idxs0.unsqueeze(-1) // 2).view( + m, + k // 2) # type: ignore[possibly-undefined] + + meta_4 = idxs0 | (idxs1 << 2) + meta_n = meta_4.view( + (-1, meta_ncols, quadbits_per_meta_elem)).to(meta_dtype) + + if quadbits_per_meta_elem == 4: + meta = (meta_n[:, :, 0] + | (meta_n[:, :, 1] << 4) + | (meta_n[:, :, 2] << 8) + | (meta_n[:, :, 3] << 12)) + elif quadbits_per_meta_elem == 8: + meta = (meta_n[:, :, 0] + | (meta_n[:, :, 1] << 4) + | (meta_n[:, :, 2] << 8) + | (meta_n[:, :, 3] << 12) + | (meta_n[:, :, 4] << 16) + | (meta_n[:, :, 5] << 20) + | (meta_n[:, :, 6] << 24) + | (meta_n[:, :, 7] << 28)) + + # Reorder meta tensor elements. + meta_reordered = meta.new_empty( + (m * meta_ncols, )) # type: ignore[possibly-undefined] + meta_offsets = _calculate_meta_reordering_scatter_offsets( + m, meta_ncols, meta_dtype, device) + meta_reordered.scatter_(0, meta_offsets, meta.view(-1)) + + return (sparse, meta_reordered.view(m, meta_ncols)) + + +# This function performs reverse of the function above - it +# reconstructs dense matrix from a pair of "compressed" matrix, given +# in the layout used by CUTLASS backend, and accompanying metadata +# matrix. +def sparse_semi_structured_to_dense_cutlass(sparse, meta_reordered): + if sparse.dim() != 2: + raise RuntimeError( + f"Expected 2-dimensional sparse tensor, got {sparse.dim()}-dimensional tensor" # noqa: E501 + ) + + m, k = sparse.shape + device = sparse.device + + if meta_reordered.dim() != 2: + raise RuntimeError( + f"Expected 2-dimensional meta tensor, got {meta_reordered.dim()}-dimensional tensor" # noqa: E501 + ) + if meta_reordered.device != device: + raise RuntimeError( + f"Expected meta matrix to be on {device} device, got matrix on {meta_reordered.device} device" # noqa: E501 + ) + + meta_dtype = meta_reordered.dtype + if meta_dtype not in (torch.int16, torch.int32): + raise RuntimeError(f"Invalid datatype {meta_dtype} of meta matrix") + quadbits_per_meta_elem = meta_dtype.itemsize * 8 // 4 + + ksparse = 4 if sparse.dtype != torch.float else 2 + + meta_nrows, meta_ncols = meta_reordered.shape + if meta_nrows != m: + raise RuntimeError( + f"Number of rows of meta matrix {meta_nrows} must be equal to number of columns of spase matrix {m}" # noqa: E501 + ) + if meta_ncols * ksparse * quadbits_per_meta_elem != 2 * k: + raise RuntimeError( + f"Number of columns of sparse matrix {k} different from the {meta_ncols * ksparse * quadbits_per_meta_elem // 2}, " # noqa: E501 + "expected according to the number of columns of meta matrix") + + # Undo meta tensor elements reordering. + meta_offsets = _calculate_meta_reordering_scatter_offsets( + m, meta_ncols, meta_dtype, device) + meta = torch.gather(meta_reordered.view(-1), 0, + meta_offsets).view(m, meta_ncols) + + # Unpack sparse tensor back to original dense tensor, using + # information provided by meta tensor. Note that torch.float + # datatype is handled pretty much the same as + # torch.half/torch.bfloat16, as metadata for a pair of torch.float + # value is encoded as if underlying 8 bytes contain four + # torch.half/torch.bfloat16 values, where either first two or last + # two are zeros. + meta_2 = torch.empty( + (m, meta_ncols, 2 * quadbits_per_meta_elem), + dtype=meta_dtype, + device=device, + ) + if quadbits_per_meta_elem == 4: + meta_2[:, :, 0] = meta & 0b11 + meta_2[:, :, 1] = (meta >> 2) & 0b11 + meta_2[:, :, 2] = (meta >> 4) & 0b11 + meta_2[:, :, 3] = (meta >> 6) & 0b11 + meta_2[:, :, 4] = (meta >> 8) & 0b11 + meta_2[:, :, 5] = (meta >> 10) & 0b11 + meta_2[:, :, 6] = (meta >> 12) & 0b11 + meta_2[:, :, 7] = (meta >> 14) & 0b11 + elif quadbits_per_meta_elem == 8: + meta_2[:, :, 0] = meta & 0b11 + meta_2[:, :, 1] = (meta >> 2) & 0b11 + meta_2[:, :, 2] = (meta >> 4) & 0b11 + meta_2[:, :, 3] = (meta >> 6) & 0b11 + meta_2[:, :, 4] = (meta >> 8) & 0b11 + meta_2[:, :, 5] = (meta >> 10) & 0b11 + meta_2[:, :, 6] = (meta >> 12) & 0b11 + meta_2[:, :, 7] = (meta >> 14) & 0b11 + meta_2[:, :, 8] = (meta >> 16) & 0b11 + meta_2[:, :, 9] = (meta >> 18) & 0b11 + meta_2[:, :, 10] = (meta >> 20) & 0b11 + meta_2[:, :, 11] = (meta >> 22) & 0b11 + meta_2[:, :, 12] = (meta >> 24) & 0b11 + meta_2[:, :, 13] = (meta >> 26) & 0b11 + meta_2[:, :, 14] = (meta >> 28) & 0b11 + meta_2[:, :, 15] = (meta >> 30) & 0b11 + + dense_offsets = meta_2.view(-1) + ( + torch.arange(0, 2 * m * k // ksparse, device=device) * 4).view( + -1, 1).repeat(1, 2).view(-1) + + dense = torch.zeros((m * 2 * k, ), dtype=sparse.dtype, device=device) + if sparse.dtype != torch.float: + # dense.scatter_(0, dense_offsets, sparse.view(-1)) + dense.scatter_(0, dense_offsets, sparse.reshape(-1)) + else: + dense.view(torch.half).scatter_(0, dense_offsets, + sparse.view(torch.half).view(-1)) + + return dense.view(m, 2 * k) + + +def mask_creator(tensor): + """ + Class for creating N:M sparsity masks. + Masks will be created using the N:M ratio, where for every block of + M weights, N will be pruned based on ranked weight value. Each mask + will correspond to the given tensor. + + :param N: The number of weights in a group to keep + :param M: The size of a weight group + """ + N = 2 + M = 4 + + mask = None + # for i, tensor in enumerate(tensors): + if tensor.numel() % M != 0: + raise ValueError( + f"Tensor of size {tensor.shape} can't be evenly divided into " + f"{M} groups") + + num_groups = tensor.numel() // M + + # N:M sparsity for linear layers + tensor_temp = tensor.detach().abs().reshape(num_groups, M) + index = torch.argsort(tensor_temp, dim=1)[:, :int(M - N)] + + w_b = torch.ones(tensor_temp.shape, device=tensor_temp.device) + mask = w_b.scatter_(dim=1, index=index, value=0).reshape(tensor.shape) + + return mask diff --git a/archive/kt-sft/ktransformers/ktransformers_ext/operators/custom_marlin/quantize/utils/marlin_24_perms.py b/archive/kt-sft/ktransformers/ktransformers_ext/operators/custom_marlin/quantize/utils/marlin_24_perms.py new file mode 100644 index 0000000..93f65a2 --- /dev/null +++ b/archive/kt-sft/ktransformers/ktransformers_ext/operators/custom_marlin/quantize/utils/marlin_24_perms.py @@ -0,0 +1,60 @@ +"""This file is used for /tests and /benchmarks""" +from typing import Dict, List + +import numpy +import torch + + +# Precompute permutations for Marlin24 weight and scale shuffling # noqa: E501 +# +# Marlin works on [16*2,64] tiles. The goal of the permutations is to reorder the weight data so that it is compatible noqa: # noqa: E501 +# with the tensor-core format that is described here: +# https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#matrix-fragments-for-mma-m16n8k16-with-floating-point-type # noqa: E501 +# +# As a result of this reordering, the vector loads inside the kernel will get the data as it is needed for tensor-core # noqa: E501 +# (without the need to use ldmatrix instructions) # noqa: E501 +def get_perms_24(num_bits: int): + perm_list: List[int] = [] + for i in range(32): + perm1: List[int] = [] + col = i // 4 + col_o = col // 2 + for block in [0, 1]: + for row in [ + 2 * (i % 4), + 2 * (i % 4) + 1, + 2 * (i % 4 + 4), + 2 * (i % 4 + 4) + 1, + ]: + perm1.append(16 * row + col_o * 256 + 8 * (col % 2) + + 4 * block) + for j in range(4): + perm_list.extend([p + 1 * j for p in perm1]) + perm = numpy.array(perm_list) + + if num_bits == 4: + interleave = numpy.array([0, 2, 4, 6, 1, 3, 5, 7]) + elif num_bits == 8: + interleave = numpy.array([0, 2, 1, 3]) + else: + raise ValueError("num_bits must be 4 or 8, got {}".format(num_bits)) + + perm = perm.reshape((-1, len(interleave)))[:, interleave].ravel() + perm = torch.from_numpy(perm) + scale_perm: List[int] = [] + for i in range(8): + scale_perm.extend([i * 8 + j for j in [0, 4, 1, 5, 2, 6, 3, 7]]) + scale_perm_single: List[int] = [] + for i in range(8): + scale_perm_single.extend([8 * i + j for j in [0, 1, 2, 3, 4, 5, 6, 7]]) + return perm, scale_perm, scale_perm_single + + +marlin_24_perm: Dict[int, torch.Tensor] = {} +marlin_24_scale_perm: Dict[int, List[int]] = {} +marlin_24_scale_perm_single: Dict[int, List[int]] = {} +for num_bits in [4, 8]: + perm_24, scale_perm_24, scale_perm_single_24 = get_perms_24(num_bits) + marlin_24_perm[num_bits] = perm_24 + marlin_24_scale_perm[num_bits] = scale_perm_24 + marlin_24_scale_perm_single[num_bits] = scale_perm_single_24 diff --git a/archive/kt-sft/ktransformers/ktransformers_ext/operators/custom_marlin/quantize/utils/marlin_perms.py b/archive/kt-sft/ktransformers/ktransformers_ext/operators/custom_marlin/quantize/utils/marlin_perms.py new file mode 100644 index 0000000..db5e685 --- /dev/null +++ b/archive/kt-sft/ktransformers/ktransformers_ext/operators/custom_marlin/quantize/utils/marlin_perms.py @@ -0,0 +1,60 @@ +"""This file is used for /tests and /benchmarks""" +from typing import Dict, List + +import numpy +import torch + + +# Precompute permutations for Marlin weight and scale shuffling # noqa: E501 +# +# Marlin works on [16,64] tiles. The goal of the permutations is to reorder the weight data so that it is compatible noqa: # noqa: E501 +# with the tensor-core format that is described here: +# https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#matrix-fragments-for-mma-m16n8k16-with-floating-point-type # noqa: E501 +# +# As a result of this reordering, the vector loads inside the kernel will get the data as it is needed for tensor-core # noqa: E501 +# (without the need to use ldmatrix instructions) # noqa: E501 +def get_perms(num_bits: int): + perm_list: List[int] = [] + for i in range(32): + perm1: List[int] = [] + col = i // 4 + for block in [0, 1]: + for row in [ + 2 * (i % 4), + 2 * (i % 4) + 1, + 2 * (i % 4 + 4), + 2 * (i % 4 + 4) + 1, + ]: + perm1.append(16 * row + col + 8 * block) + for j in range(4): + perm_list.extend([p + 256 * j for p in perm1]) + + perm = numpy.array(perm_list) + + if num_bits == 4: + interleave = numpy.array([0, 2, 4, 6, 1, 3, 5, 7]) + elif num_bits == 8: + interleave = numpy.array([0, 2, 1, 3]) + else: + raise Exception("num_bits must be 4 or 8, got {}".format(num_bits)) + + perm = perm.reshape((-1, len(interleave)))[:, interleave].ravel() + perm = torch.from_numpy(perm) + 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 perm, scale_perm, scale_perm_single + + +marlin_perm: Dict[int, torch.Tensor] = {} +marlin_scale_perm: Dict[int, List[int]] = {} +marlin_scale_perm_single: Dict[int, List[int]] = {} +for num_bits in [4, 8]: + perm, scale_perm, scale_perm_single = get_perms(num_bits) + marlin_perm[num_bits] = perm + marlin_scale_perm[num_bits] = scale_perm + marlin_scale_perm_single[num_bits] = scale_perm_single diff --git a/archive/kt-sft/ktransformers/ktransformers_ext/operators/custom_marlin/quantize/utils/marlin_utils.py b/archive/kt-sft/ktransformers/ktransformers_ext/operators/custom_marlin/quantize/utils/marlin_utils.py new file mode 100644 index 0000000..5a56c13 --- /dev/null +++ b/archive/kt-sft/ktransformers/ktransformers_ext/operators/custom_marlin/quantize/utils/marlin_utils.py @@ -0,0 +1,270 @@ +"""This file is used for /tests and /benchmarks""" +import random + +import numpy +import torch + +from ktransformers.ktransformers_ext.operators.custom_marlin.quantize.utils.format_24 import ( + mask_creator, sparse_semi_structured_from_dense_cutlass) +from ktransformers.ktransformers_ext.operators.custom_marlin.quantize.utils.marlin_24_perms import ( + marlin_24_perm, marlin_24_scale_perm, marlin_24_scale_perm_single) +from ktransformers.ktransformers_ext.operators.custom_marlin.quantize.utils.marlin_perms import ( + marlin_perm, marlin_scale_perm, marlin_scale_perm_single) +from ktransformers.ktransformers_ext.operators.custom_marlin.quantize.utils.quant_utils import ( + get_pack_factor, quantize_weights, sort_weights) + +__cuda_arch = torch.cuda.get_device_capability() + +MARLIN_TILE = 16 + +GPTQ_MARLIN_TILE = 16 +GPTQ_MARLIN_MIN_THREAD_N = 64 +GPTQ_MARLIN_MIN_THREAD_K = 128 +GPTQ_MARLIN_MAX_PARALLEL = 16 + +GPTQ_MARLIN_SUPPORTED_NUM_BITS = [4, 8] +GPTQ_MARLIN_SUPPORTED_GROUP_SIZES = [-1, 32, 64, 128] +GPTQ_MARLIN_SUPPORTED_SYM = [True] + +def is_marlin_supported(): + return __cuda_arch[0] >= 8 + + +def marlin_permute_weights(q_w, size_k, size_n, perm, tile=MARLIN_TILE): + assert q_w.shape == (size_k, size_n) + assert size_k % tile == 0, f"size_k = {size_k}, tile = {tile}" + assert size_n % tile == 0, f"size_k = {size_n}, tile = {tile}" + + # Permute weights to 16x64 marlin tiles + q_w = q_w.reshape((size_k // tile, tile, size_n // tile, tile)) + q_w = q_w.permute((0, 2, 1, 3)) + q_w = q_w.reshape((size_k // tile, size_n * tile)) + + q_w = q_w.reshape((-1, perm.numel()))[:, perm].reshape(q_w.shape) + + return q_w + + +def marlin_weights(q_w, size_k, size_n, num_bits, perm): + # Permute + q_w = marlin_permute_weights(q_w, size_k, size_n, perm) + + # Pack + pack_factor = get_pack_factor(num_bits) + orig_device = q_w.device + + q_w = q_w.cpu().numpy().astype(numpy.uint32) + + q_packed = numpy.zeros((q_w.shape[0], q_w.shape[1] // pack_factor), + dtype=numpy.uint32) + for i in range(pack_factor): + q_packed |= q_w[:, i::pack_factor] << num_bits * i + + q_packed = torch.from_numpy(q_packed.astype(numpy.int32)).to(orig_device) + + return q_packed + + +def marlin_permute_scales(s, size_k, size_n, group_size, scale_perm, + scale_perm_single): + 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_quantize( + w: torch.Tensor, + num_bits: int, + group_size: int, + act_order: bool, +): + size_k, size_n = w.shape + + # Normalize group_size + if group_size == -1: + group_size = size_k + assert group_size <= size_k + + # Quantize (and apply act_order if provided) + q_w, s, g_idx, rand_perm = quantize_weights(w, num_bits, group_size, + act_order) + + # For act_order, sort the "weights" and "g_idx" so that group ids are + # increasing + sort_indices = torch.empty(0, dtype=torch.int, device=w.device) + if act_order: + q_w, g_idx, sort_indices = sort_weights(q_w, g_idx) + + # Reformat to marlin + marlin_q_w = marlin_weights(q_w, size_k, size_n, num_bits, + marlin_perm[num_bits]) + marlin_s = marlin_permute_scales(s, size_k, size_n, group_size, + marlin_scale_perm[num_bits], + marlin_scale_perm_single[num_bits]) + + # Create result + res_list = [marlin_q_w, marlin_s, g_idx, sort_indices, rand_perm] + for i in range(len(res_list)): + res_list[i] = res_list[i].to(w.device) + + return res_list + + +def vllm_marlin_quantize( + w: torch.Tensor, + num_bits: int, + group_size: int, + act_order: bool, +): + size_k, size_n = w.shape + + # Normalize group_size + if group_size == -1: + group_size = size_k + assert group_size <= size_k + + # Quantize (and apply act_order if provided) + w_ref, q_w, s, g_idx, rand_perm = quantize_weights(w, num_bits, group_size, + act_order) + + # For act_order, sort the "weights" and "g_idx" so that group ids are + # increasing + sort_indices = torch.empty(0, dtype=torch.int, device=w.device) + if act_order: + q_w, g_idx, sort_indices = sort_weights(q_w, g_idx) + + # Reformat to marlin + marlin_q_w = marlin_weights(q_w, size_k, size_n, num_bits, + marlin_perm[num_bits]) + marlin_s = marlin_permute_scales(s, size_k, size_n, group_size, + marlin_scale_perm[num_bits], + marlin_scale_perm_single[num_bits]) + + # Create result + res_list = [w_ref, marlin_q_w, marlin_s, g_idx, sort_indices, rand_perm] + for i in range(len(res_list)): + res_list[i] = res_list[i].to(w.device) + + return res_list + + +def inject_24(w, size_k, size_n): + assert w.shape == (size_k, size_n) + + mask = mask_creator(w.t()).t().cuda().bool() + + return (mask * w).contiguous(), mask.contiguous() + + +def check_24(w, num_rows_to_sample=50, _verbose=False): + BLOCK_SIZE = 4 + MAX_NON_ZEROS = 2 + + w = w.t().contiguous() + + print("check_24: w.shape = {}".format(w.shape)) + + num_rows, num_cols = w.shape + sampled_row_idxs = random.choices(range(num_rows), k=num_rows_to_sample) + if _verbose: + print(f"Sampled row idxs = {sampled_row_idxs}") + + total_segments = 0 + non_24_segments = 0 + for i in sampled_row_idxs: + for j in range(0, num_cols - BLOCK_SIZE, BLOCK_SIZE): + total_segments += 1 + block = w[i, j:j + BLOCK_SIZE] + num_nonzero = torch.count_nonzero(block) + if num_nonzero > MAX_NON_ZEROS: + print("i = {} j = {} block = {}".format(i, j, block)) + non_24_segments += 1 + + print(f"{non_24_segments} / {total_segments} do not have 2:4 structure.") + + +def compress_quantized_24_weight(q_24, size_k, size_n, num_bits): + assert q_24.shape == (size_k, size_n) + + # Remove zp to normalize over 0 + max_q_val = (1 << num_bits) - 1 + zp = (max_q_val + 1) // 2 + q_24_no_zp = q_24 - zp + + # Compress + q_24_no_zp = q_24_no_zp.t().contiguous() + q_24_no_zp_comp, meta = sparse_semi_structured_from_dense_cutlass( + q_24_no_zp) + q_24_no_zp_comp = q_24_no_zp_comp.t().contiguous() + + # Restore zp + q_24_comp = q_24_no_zp_comp + zp + + # Resize meta to its actual shape (without moving any data) + meta = meta.resize_(meta.shape[1] // 2, meta.shape[0] * 2) + + return q_24_comp, meta + + +def marlin_24_quantize( + w: torch.Tensor, + num_bits: int, + group_size: int, +): + size_k, size_n = w.shape + + # Normalize group_size + if group_size == -1: + group_size = size_k + assert group_size <= size_k + + # Inject 2:4 sparsity + w_24, mask_24 = inject_24(w, size_k, size_n) + + # Quantize + w_24_ref, q_w_24, s, g_idx, rand_perm = quantize_weights(w_24, + num_bits, + group_size, + act_order=False) + + # Compress quantized weight + q_w_24_comp, meta = compress_quantized_24_weight(q_w_24, size_k, size_n, + num_bits) + size_k_comp = size_k // 2 + + # Reformat to marlin + marlin_24_q_w_comp = marlin_weights(q_w_24_comp, size_k_comp, size_n, + num_bits, marlin_24_perm[num_bits]) + marlin_24_s = marlin_permute_scales(s, size_k, size_n, group_size, + marlin_24_scale_perm[num_bits], + marlin_24_scale_perm_single[num_bits]) + + # Create result + res_list = [w_24_ref, marlin_24_q_w_comp, meta, marlin_24_s] + for i in range(len(res_list)): + res_list[i] = res_list[i].to(w.device) + + return res_list + + +def compute_max_diff(output, output_ref): + return torch.mean(torch.abs(output - output_ref)) / torch.mean( + torch.abs(output_ref)) + + +class MarlinWorkspace: + + def __init__(self, out_features, min_thread_n, max_parallel, device): + assert (out_features % min_thread_n == 0), ( + "out_features = {} is undivisible by min_thread_n = {}".format( + out_features, min_thread_n)) + + max_workspace_size = ((out_features // min_thread_n) * max_parallel) + + self.scratch = torch.zeros(max_workspace_size, + dtype=torch.int, + device=device) diff --git a/archive/kt-sft/ktransformers/ktransformers_ext/operators/custom_marlin/quantize/utils/quant_utils.py b/archive/kt-sft/ktransformers/ktransformers_ext/operators/custom_marlin/quantize/utils/quant_utils.py new file mode 100644 index 0000000..de73667 --- /dev/null +++ b/archive/kt-sft/ktransformers/ktransformers_ext/operators/custom_marlin/quantize/utils/quant_utils.py @@ -0,0 +1,138 @@ +"""This file is used for /tests and /benchmarks""" +import numpy +import torch + +SUPPORTED_NUM_BITS = [4, 8] +SUPPORTED_GROUP_SIZES = [-1, 32, 64, 128] + + +def get_pack_factor(num_bits): + assert num_bits in SUPPORTED_NUM_BITS, f"Unsupported num_bits = {num_bits}" + return 32 // num_bits + + +def permute_rows(q_w: torch.Tensor, group_size: int): + + orig_device = q_w.device + k_size, _ = q_w.shape + + g_idx = torch.zeros((k_size, ), dtype=torch.int32) + for i in range(k_size): + g_idx[i] = i // group_size + + # Simulate act_order by doing a random permutation on K + rand_perm = torch.randperm(k_size) + + g_idx = g_idx[rand_perm].contiguous() + q_w = q_w[rand_perm, :].contiguous() + + return ( + q_w.to(device=orig_device), + g_idx.to(device=orig_device), + rand_perm.to(device=orig_device), + ) + + +def quantize_weights(w: torch.Tensor, num_bits: int, group_size: int, + act_order: bool): + orig_device = w.device + size_k, size_n = w.shape + + assert w.is_floating_point(), "w must be float" + assert num_bits in SUPPORTED_NUM_BITS, f"Unsupported num_bits = {num_bits}" + assert group_size in SUPPORTED_GROUP_SIZES + [ + size_k + ], f"Unsupported groupsize = {group_size}" + + if group_size == -1: + group_size = size_k + assert group_size <= size_k + + max_q_val = 2**num_bits - 1 + half_q_val = (max_q_val + 1) // 2 + + # Reshape to [groupsize, -1] + if group_size < size_k: + w = w.view((-1, group_size, size_n)) + w = w.permute(1, 0, 2) + w = w.reshape((group_size, -1)) + + # Compute scale for each group + s = torch.max(torch.abs(w), 0, keepdim=True)[0] + s *= 2 / max_q_val # 2 => symmetric + + # Quantize + q_w = torch.round(w / s).int() + q_w += half_q_val + q_w = torch.clamp(q_w, 0, max_q_val) + + # Restore original shapes + if group_size < size_k: + + def reshape_w(w): + w = w.reshape((group_size, -1, size_n)) + w = w.permute(1, 0, 2) + w = w.reshape((size_k, size_n)).contiguous() + return w + + q_w = reshape_w(q_w) + + s = s.reshape((-1, size_n)).contiguous() + + # Apply act_order + g_idx = torch.empty(0, dtype=torch.int, device=w.device) + rand_perm = torch.empty(0, dtype=torch.int, device=w.device) + if act_order: + assert ( + group_size < size_k + ), "For act_order, groupsize = {} must be less than size_k = {}".format( + group_size, size_k) + + q_w, g_idx, rand_perm = permute_rows(q_w, group_size) + + return ( + q_w.to(device=orig_device), + s.to(device=orig_device), + g_idx.to(device=orig_device), + rand_perm.to(device=orig_device), + ) + + +def sort_weights(q_w: torch.Tensor, g_idx: torch.Tensor): + orig_device = q_w.device + + sort_indices = torch.argsort(g_idx).to( + dtype=torch.int32) # Sort based on g_idx + + g_idx = g_idx[sort_indices].contiguous() + q_w = q_w[sort_indices, :].contiguous() + + return ( + q_w.to(device=orig_device), + g_idx.to(device=orig_device), + sort_indices.to(device=orig_device), + ) + + +def gptq_pack( + q_w: torch.Tensor, + num_bits: int, + size_k: int, + size_n: int, +): + assert q_w.shape == (size_k, size_n) + + pack_factor = get_pack_factor(num_bits) + assert size_k % pack_factor == 0 + + orig_device = q_w.device + + q_w = q_w.cpu().numpy().astype(numpy.uint32) + + q_res = numpy.zeros((size_k // pack_factor, size_n), dtype=numpy.uint32) + + for i in range(pack_factor): + q_res |= q_w[i::pack_factor, :] << num_bits * i + + q_res = torch.from_numpy(q_res.astype(numpy.int32)).to(orig_device) + return q_res diff --git a/archive/kt-sft/ktransformers/ktransformers_ext/triton/fp8gemm.py b/archive/kt-sft/ktransformers/ktransformers_ext/triton/fp8gemm.py new file mode 100644 index 0000000..d5c913d --- /dev/null +++ b/archive/kt-sft/ktransformers/ktransformers_ext/triton/fp8gemm.py @@ -0,0 +1,193 @@ +# Adopted from https://huggingface.co/deepseek-ai/DeepSeek-V3/blob/main/inference/kernel.py +from typing import Tuple + +import torch +import triton +import triton.language as tl +from triton import Config + + +@triton.jit +def act_quant_kernel(x_ptr, y_ptr, s_ptr, BLOCK_SIZE: tl.constexpr): + """ + Quantizes the input tensor `x_ptr` and stores the result in `y_ptr` and the scaling factor in `s_ptr`. + + Args: + x_ptr (triton.Pointer): Pointer to the input tensor. + y_ptr (triton.Pointer): Pointer to the output tensor where quantized values will be stored. + s_ptr (triton.Pointer): Pointer to the output tensor where scaling factors will be stored. + BLOCK_SIZE (tl.constexpr): The size of the block to be processed by each program instance. + + Returns: + None + """ + pid = tl.program_id(axis=0) + offs = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + x = tl.load(x_ptr + offs).to(tl.float32) + s = tl.max(tl.abs(x)) / 448. + y = x / s + y = y.to(y_ptr.dtype.element_ty) + tl.store(y_ptr + offs, y) + tl.store(s_ptr + pid, s) + + +def act_quant(x: torch.Tensor, block_size: int = 128) -> Tuple[torch.Tensor, torch.Tensor]: + """ + Quantizes the input tensor `x` using block-wise quantization. + + Args: + x (torch.Tensor): The input tensor to be quantized. Must be contiguous and its last dimension size must be divisible by `block_size`. + block_size (int, optional): The size of the blocks to be used for quantization. Default is 128. + + Returns: + Tuple[torch.Tensor, torch.Tensor]: A tuple containing: + - The quantized tensor with dtype `torch.float8_e4m3fn`. + - A tensor of scaling factors with dtype `torch.float32`. + """ + assert x.is_contiguous(), 'Input tensor must be contiguous' + assert x.size(-1) % block_size == 0, f'Last dimension size must be divisible by block_size (block_size={block_size})' + y = torch.empty_like(x, dtype=torch.float8_e4m3fn) + s = x.new_empty(*x.size()[:-1], x.size(-1) // block_size, dtype=torch.float32) + grid = lambda meta: (triton.cdiv(x.numel(), meta['BLOCK_SIZE']), ) + act_quant_kernel[grid](x, y, s, BLOCK_SIZE=block_size) + return y, s + + +@triton.jit +def weight_dequant_kernel(x_ptr, s_ptr, y_ptr, M, N, BLOCK_SIZE: tl.constexpr): + """ + Dequantizes weights using the provided scaling factors and stores the result. + + Args: + x_ptr (tl.pointer): Pointer to the quantized weights. + s_ptr (tl.pointer): Pointer to the scaling factors. + y_ptr (tl.pointer): Pointer to the output buffer for dequantized weights. + M (int): Number of rows in the weight matrix. + N (int): Number of columns in the weight matrix. + BLOCK_SIZE (tl.constexpr): Size of the block for tiling. + + Returns: + None + """ + pid_m = tl.program_id(axis=0) + pid_n = tl.program_id(axis=1) + n = tl.cdiv(N, BLOCK_SIZE) + offs_m = pid_m * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + offs_n = pid_n * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + offs = offs_m[:, None] * N + offs_n[None, :] + mask = (offs_m[:, None] < M) & (offs_n[None, :] < N) + x = tl.load(x_ptr + offs, mask=mask).to(tl.float32) + s = tl.load(s_ptr + pid_m * n + pid_n) + y = x * s + tl.store(y_ptr + offs, y, mask=mask) + + +def weight_dequant(x: torch.Tensor, s: torch.Tensor, block_size: int = 128) -> torch.Tensor: + """ + Dequantizes the given weight tensor using the provided scale tensor. + + Args: + x (torch.Tensor): The quantized weight tensor of shape (M, N). + s (torch.Tensor): The scale tensor of shape (M, N). + block_size (int, optional): The block size to use for dequantization. Defaults to 128. + + Returns: + torch.Tensor: The dequantized weight tensor of the same shape as `x`. + + Raises: + AssertionError: If `x` or `s` are not contiguous or if their dimensions are not 2. + """ + assert x.is_contiguous() and s.is_contiguous(), 'Input tensors must be contiguous' + assert x.dim() == 2 and s.dim() == 2, 'Input tensors must have 2 dimensions' + M, N = x.size() + y = torch.empty_like(x, dtype=torch.get_default_dtype()) + grid = lambda meta: (triton.cdiv(M, meta['BLOCK_SIZE']), triton.cdiv(N, meta['BLOCK_SIZE'])) + with torch.cuda.device(x.device): + weight_dequant_kernel[grid](x, s, y, M, N, BLOCK_SIZE=block_size) + return y + + +fp8_gemm_configs = [ + Config({'BLOCK_SIZE_M': block_m, 'BLOCK_SIZE_N': block_n, 'BLOCK_SIZE_K': 128}, num_stages=num_stages, num_warps=8) + for block_m in [16, 32, 64] for block_n in [32, 64, 128] for num_stages in [3, 4, 5, 6] +] + +@triton.autotune(configs=fp8_gemm_configs, key=['N', 'K']) +@triton.jit +def fp8_gemm_kernel(a_ptr, b_ptr, c_ptr, + a_s_ptr, b_s_ptr, + M, N: tl.constexpr, K: tl.constexpr, + BLOCK_SIZE_M: tl.constexpr, + BLOCK_SIZE_N: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr): + """ + Performs a matrix multiplication operation on FP8 matrices with scaling factors. + + Args: + a_ptr (tl.tensor): Pointer to the first input matrix A. + b_ptr (tl.tensor): Pointer to the second input matrix B. + c_ptr (tl.tensor): Pointer to the output matrix C. + a_s_ptr (tl.tensor): Pointer to the scaling factors for matrix A. + b_s_ptr (tl.tensor): Pointer to the scaling factors for matrix B. + M (int): Number of rows in matrix A and C. + N (tl.constexpr): Number of columns in matrix B and C. + K (tl.constexpr): Number of columns in matrix A and rows in matrix B. + BLOCK_SIZE_M (tl.constexpr): Block size for the M dimension. + BLOCK_SIZE_N (tl.constexpr): Block size for the N dimension. + BLOCK_SIZE_K (tl.constexpr): Block size for the K dimension. + + Returns: + None + """ + pid_m = tl.program_id(axis=0) + pid_n = tl.program_id(axis=1) + k = tl.cdiv(K, BLOCK_SIZE_K) + offs_m = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M + offs_n = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)) % N + offs_k = tl.arange(0, BLOCK_SIZE_K) + a_ptrs = a_ptr + offs_m[:, None] * K + offs_k[None, :] + b_ptrs = b_ptr + offs_n[None, :] * K + offs_k[:, None] + a_s_ptrs = a_s_ptr + offs_m * k + b_s_ptrs = b_s_ptr + (offs_n // BLOCK_SIZE_K) * k + + accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32) + for i in range(k): + a = tl.load(a_ptrs, mask=offs_k[None, :] < K - i * BLOCK_SIZE_K, other=0.0) + b = tl.load(b_ptrs, mask=offs_k[:, None] < K - i * BLOCK_SIZE_K, other=0.0) + a_s = tl.load(a_s_ptrs) + b_s = tl.load(b_s_ptrs) + accumulator += tl.dot(a, b) * a_s[:, None] * b_s[None, :] + a_ptrs += BLOCK_SIZE_K + b_ptrs += BLOCK_SIZE_K + a_s_ptrs += 1 + b_s_ptrs += 1 + c = accumulator.to(c_ptr.dtype.element_ty) + offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) + offs_n = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) + c_ptrs = c_ptr + offs_m[:, None] * N + offs_n[None, :] + mask = (offs_m[:, None] < M) & (offs_n[None, :] < N) + tl.store(c_ptrs, c, mask=mask) + + +def fp8_gemm(a: torch.Tensor, a_s: torch.Tensor, b: torch.Tensor, b_s: torch.Tensor): + """ + Perform a matrix multiplication using FP8 precision. + + Args: + a (torch.Tensor): The first input matrix, must be contiguous. + a_s (torch.Tensor): The scaling factor for the first input matrix, must be contiguous. + b (torch.Tensor): The second input matrix, must be contiguous. + b_s (torch.Tensor): The scaling factor for the second input matrix, must be contiguous. + + Returns: + torch.Tensor: The result of the matrix multiplication. + """ + assert a.is_contiguous() and b.is_contiguous(), 'Input tensors must be contiguous' + assert a_s.is_contiguous() and b_s.is_contiguous(), 'Scaling factor tensors must be contiguous' + K = a.size(-1) + M = a.numel() // K + N = b.size(0) + c = a.new_empty(*a.size()[:-1], N, dtype=torch.get_default_dtype()) + grid = lambda META: (triton.cdiv(M, META['BLOCK_SIZE_M']), triton.cdiv(N, META['BLOCK_SIZE_N'])) + fp8_gemm_kernel[grid](a, b, c, a_s, b_s, M, N, K) + return c \ No newline at end of file diff --git a/archive/kt-sft/ktransformers/local_chat.py b/archive/kt-sft/ktransformers/local_chat.py new file mode 100644 index 0000000..432782a --- /dev/null +++ b/archive/kt-sft/ktransformers/local_chat.py @@ -0,0 +1,452 @@ +""" +Description : +Author : Boxin Zhang, Azure-Tang +Version : 0.1.0 +Copyright (c) 2024 by KVCache.AI, All Rights Reserved. +""" + +import os +import platform +import sys + +project_dir = os.path.dirname(os.path.dirname(__file__)) +sys.path.insert(0, project_dir) +import argparse +import torch +import logging +from transformers import ( + AutoTokenizer, + AutoConfig, + AutoModelForCausalLM, + GenerationConfig, + TextStreamer, + EvalPrediction, +) +import json +from pathlib import Path +from tqdm import tqdm +from torchviz import make_dot +import fire +from ktransformers.optimize.optimize import optimize_and_load_gguf +from ktransformers.models.modeling_deepseek import DeepseekV2ForCausalLM +from ktransformers.models.modeling_qwen2_moe import Qwen2MoeForCausalLM +from ktransformers.models.modeling_deepseek_v3 import DeepseekV3ForCausalLM +from ktransformers.models.modeling_llama import LlamaForCausalLM +from ktransformers.models.modeling_mixtral import MixtralForCausalLM +from ktransformers.util.utils import load_weights, prefill_and_generate, prefill_and_generate_capture, get_compute_capability, xpu_fp16_model +from ktransformers.server.config.config import Config +from ktransformers.operators.flashinfer_wrapper import flashinfer_enabled +from ktransformers.util.vendors import device_manager, get_device, to_device, GPUVendor +from ktransformers.sft.lora import inject_lora_layer, lora_and_load_adapter +from ktransformers.util.custom_loader import GGUFLoader, SafeTensorLoader +from ktransformers.util.globals import GLOBAL_CONFIG +from ktransformers.sft.metrics import ComputeSimilarity +from ktransformers.sft.monkey_patch_torch_module import install_patch, restore_patch + +os.environ['CUDA_VISIBLE_DEVICES'] = '1' + +# for debug +def print_module_tree(module, indent=0): + print(" " + f"{module.__class__.__name__}(training={module.training})") + for name, child in module.named_children(): + print(" " + f"└─{name}: ", end="") + print_module_tree(child, indent + 4) + +# for debug +def write_to_file(content, file_path: str = 'ktransformers/mark_content.txt', mode: str = 'a', encoding: str = 'utf-8') -> None: + """ + 将字符串写入指定文件 + :param content: 要写入的字符串内容 + :param file_path: 目标文件路径 + :param mode: 文件打开模式（默认'w'为覆盖写入，可选'a'追加写入） + :param encoding: 文件编码（默认utf-8） + """ + with open(file_path, mode, encoding=encoding) as f: + f.write(content) + +custom_models = { + "DeepseekV2ForCausalLM": DeepseekV2ForCausalLM, + "DeepseekV3ForCausalLM": DeepseekV3ForCausalLM, + "Qwen2MoeForCausalLM": Qwen2MoeForCausalLM, + "LlamaForCausalLM": LlamaForCausalLM, + "MixtralForCausalLM": MixtralForCausalLM, +} + +ktransformer_rules_dir = ( + os.path.dirname(os.path.abspath(__file__)) + "/optimize/optimize_rules/" +) +default_optimize_rules = { + "DeepseekV2ForCausalLM": ktransformer_rules_dir + "DeepSeek-V2-Chat.yaml", + "DeepseekV3ForCausalLM": ktransformer_rules_dir + "DeepSeek-V3-Chat.yaml", + "Qwen2MoeForCausalLM": ktransformer_rules_dir + "Qwen2-57B-A14B-Instruct.yaml", + "LlamaForCausalLM": ktransformer_rules_dir + "Internlm2_5-7b-Chat-1m.yaml", + "MixtralForCausalLM": ktransformer_rules_dir + "Mixtral.yaml", +} + + +def local_chat( + model_path: str | None = None, + model_config_path: str | None = None, + optimize_config_path: str = None, + gguf_path: str | None = None, + max_new_tokens: int = 1000, + cpu_infer: int = Config().cpu_infer, + use_cuda_graph: bool = True, # modify to false if using KExpertsTorch + prompt_file : str | None = None, + mode: str = "normal", + force_think: bool = False, + chunk_size: int = 8192, + device: str = "cuda", + is_sft: bool = False, + sft_data_path: str | None = None, + save_adapter_path: str | None = None, + use_adapter: bool = False, + use_adapter_path: str | None = None, + is_test_data: bool = False, + test_data_path: str | None = None, + output_dir: str | None = None, +): + + if not is_sft: + torch.set_grad_enabled(False) + + if is_sft == True or use_adapter == True: + GLOBAL_CONFIG._config["mod"] = "sft" + else: + GLOBAL_CONFIG._config["mod"] = "infer" + + Config().cpu_infer = cpu_infer + Config().chunk_size = chunk_size + if torch.xpu.is_available(): + use_cuda_graph = False + + tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True) + if model_config_path == None: + config = AutoConfig.from_pretrained(model_path, trust_remote_code=True) + else: + config = AutoConfig.from_pretrained(model_config_path, trust_remote_code=True) + if mode == 'long_context': + assert config.architectures[0] == "LlamaForCausalLM", "only LlamaForCausalLM support long_context mode" + torch.set_default_dtype(torch.float16) + elif xpu_fp16_model(config): + torch.set_default_dtype(torch.float16) + else: + torch.set_default_dtype(config.torch_dtype) + + with torch.device("meta"): + if config.architectures[0] in custom_models: + print("using custom modeling_xxx.py.") + if ( + "Qwen2Moe" in config.architectures[0] + ): # Qwen2Moe must use flash_attention_2 to avoid overflow. + config._attn_implementation = "flash_attention_2" + if "Llama" in config.architectures[0]: + config._attn_implementation = "eager" + if "Mixtral" in config.architectures[0]: + config._attn_implementation = "flash_attention_2" + if torch.xpu.is_available(): + config._attn_implementation = "eager" + model = custom_models[config.architectures[0]](config) + else: + if torch.xpu.is_available(): + attn_implementation = "eager" + else: + attn_implementation = "flash_attention_2" + model = AutoModelForCausalLM.from_config( + config, trust_remote_code=True, attn_implementation=attn_implementation + ) + + if optimize_config_path is None: + if config.architectures[0] in default_optimize_rules: + print("using default_optimize_rule for", config.architectures[0]) + optimize_config_path = default_optimize_rules[config.architectures[0]] + else: + optimize_config_path = input( + "please input the path of your rule file(yaml file containing optimize rules):" + ) + + if gguf_path is None: + gguf_path = input( + "please input the path of your gguf file(gguf file in the dir containing input gguf file must all belong to current model):" + ) + + GLOBAL_CONFIG._config["mod"] = "infer" + optimize_and_load_gguf(model, optimize_config_path, gguf_path, config) + + model.train() + + if is_sft == True: + if use_adapter == True or is_test_data == True: + raise AttributeError("We do not support to run sft and inference at the same time.") + GLOBAL_CONFIG._config["mod"] = "sft" + print(f"sft with lora in dataset: {sft_data_path} ...") + print(f"use_cuda_graph:{use_cuda_graph}") + lora_and_load_adapter(model, tokenizer, sft_data_path, save_adapter_path) + + if use_adapter == True: + GLOBAL_CONFIG._config["mod"] = "sft" + if is_sft == True: + raise AttributeError("We do not support more than one adapter up to now...") + + if use_adapter_path.endswith('.gguf'): + inject_lora_layer(model, use_adapter_path) + adapter_gguf_loader = GGUFLoader(use_adapter_path) + load_weights(model, adapter_gguf_loader, adapter_gguf=True) + model.train() + else: + inject_lora_layer(model, use_adapter_path) + + adapter_loader = SafeTensorLoader(use_adapter_path) + device = next(model.parameters()).device + + # for name, param in model.named_parameters(): + # print(name, param.shape) + + for key in adapter_loader.tensor_file_map.keys(): + try: + tensor = adapter_loader.load_tensor(key, device=device) + + model_key = key.replace("base_model.model.", "") + model_key = model_key.replace(".weight", ".default.weight") + + param = model.get_parameter(model_key) + param.data.copy_(tensor.data) + + print(f"Loaded adapter weight: {key} -> {model_key}") + except AttributeError as e: + print(f"Skipping {key}: not a model parameter") + except KeyError as e: + print(f"Key not found in model: {model_key} (original: {key})") + + + try: + model.generation_config = GenerationConfig.from_pretrained(model_path) + except Exception as e: + print(f"generation config can't auto create, make default. Message: {e}") + gen_config = GenerationConfig( + temperature=0.6, + top_p=0.95, + do_sample=True + ) + model.generation_config = gen_config + # model.generation_config = GenerationConfig.from_pretrained(model_path) + if model.generation_config.pad_token_id is None: + model.generation_config.pad_token_id = model.generation_config.eos_token_id + model.eval() + logging.basicConfig(level=logging.INFO) + + # @torch.no_grad() + # def first_token_argmax_baseline(model, tokenizer, prompt_text, device): + # model.eval() + # enc = tokenizer.apply_chat_template([{"role":"user","content":prompt_text}], + # add_generation_prompt=True, return_tensors="pt") + # x = enc.to(device) + # logits = model(input_ids=x, use_cache=False, return_dict=False)[0] + # return int(torch.argmax(logits[:, -1, :], dim=-1)[0]) + + # try: + # device_map = model.gguf_loader.tensor_device_map + # from ktransformers.util.utils import get_device, torch_device_mapping + # torch_device = get_device('model.layers.0.self_attn', device_map) + # torch_device = torch_device_mapping.get(torch_device, torch_device) + # print(f"[FIRST-TOKEN PROBE] argmax id = {probe_id} ({tokenizer.decode([probe_id])!r})") + # except Exception as e: + # print("[FIRST-TOKEN PROBE] failed:", e) + # return + + system = platform.system() + # for debug + # if system == "Windows": + # os.system("cls") + # else: + # os.system("clear") + + if GLOBAL_CONFIG._config["mod"] == "sft" : + model.model.embed_tokens.to("cpu") + + if is_test_data: + data_path = Path(test_data_path) + with data_path.open("r", encoding="utf-8") as f: + dataset = json.load(f) + preds, refs = [], [] + + for sample in tqdm(dataset, desc="Processing samples"): + inst = sample.get("instruction", "") + prompt = sample.get("input", "") + prompt = prompt+inst + # print(f"prompt: {prompt}") + label = sample.get("output", "") + + messages = [{"role": "user", "content": prompt}] + input_tensor = tokenizer.apply_chat_template( + messages, add_generation_prompt=True, return_tensors="pt" + ) + if force_think: + token_thinks = torch.tensor([tokenizer.encode("\\n",add_special_tokens=False)],device=input_tensor.device) + input_tensor = torch.cat( + [input_tensor, token_thinks], dim=1 + ) + if mode == 'long_context': + assert Config().long_context_config['max_seq_len'] > input_tensor.shape[1] + max_new_tokens, \ + "please change max_seq_len in ~/.ktransformers/config.yaml" + + if system != "Windows" and (config.architectures[0] == "DeepseekV2ForCausalLM" or config.architectures[0] == "DeepseekV3ForCausalLM") and flashinfer_enabled and get_compute_capability() >= 8 and device_manager.gpu_vendor == GPUVendor.NVIDIA: + prediction = prefill_and_generate_capture( + model, tokenizer, input_tensor.to(device), max_new_tokens, use_cuda_graph, mode = mode, force_think = force_think, chunk_size = chunk_size, + use_flashinfer_mla = True, num_heads = config.num_attention_heads, head_dim_ckv = config.kv_lora_rank, head_dim_kpe = config.qk_rope_head_dim, q_head_dim = config.qk_rope_head_dim + config.qk_nope_head_dim, echo_stream=False + ) + else: + prediction = prefill_and_generate_capture( + model, tokenizer, input_tensor.to(device), max_new_tokens, use_cuda_graph, mode = mode, force_think = force_think, chunk_size = chunk_size,echo_stream=False, + ) + # print(f"prediction:{prediction}") + sample["label"] = label + sample["prediction"] = prediction + sample.pop("output", None) + + preds.append(prediction) + refs.append(label) + + pred_file = Path(output_dir) / 'predictions.json' + pred_file.parent.mkdir(parents=True, exist_ok=True) + + with pred_file.open("w", encoding="utf-8") as f: + json.dump(dataset, f, ensure_ascii=False, indent=2) + + compute_metrics = ComputeSimilarity(tokenizer) + # print(f"metrics:{metrics}") + + enc_pred = tokenizer(preds, add_special_tokens=False, padding=True, return_tensors="np") + enc_ref = tokenizer(refs, add_special_tokens=False, padding=True, return_tensors="np") + + ep = EvalPrediction( + predictions=enc_pred["input_ids"], + label_ids=enc_ref["input_ids"] + ) + + metrics = compute_metrics(ep, compute_result=True) + + metric_file = Path(output_dir) / 'metrics.json' + with metric_file.open("w", encoding="utf-8") as f: + json.dump(metrics, f, ensure_ascii=False, indent=2) + + print(f"Results of predictions saved in {pred_file}") + print(f"Results of metrics saved in {metric_file}") + + while not is_test_data: + GLOBAL_CONFIG._config["mod"] = "infer" + content = input("Chat: ") + if content.startswith('"""'): # prefix """ + # multi lines input + content = content[3:] + "\n" + while True: + line = input("") + if line.endswith('"""'): + # end multi lines input + line = line[:-3] # suffix """ + if line: + content += line + "\n" + break + else: + content += line + "\n" + + if content == "": + if prompt_file != None: + content = open(prompt_file, "r").read() + else: + content = "Please write a piece of quicksort code in C++." + elif os.path.isfile(content): + content = open(content, "r").read() + + messages = [{"role": "user", "content": content}] + input_tensor = tokenizer.apply_chat_template( + messages, add_generation_prompt=True, return_tensors="pt" + ) + if force_think: + token_thinks = torch.tensor([tokenizer.encode("\\n",add_special_tokens=False)],device=input_tensor.device) + input_tensor = torch.cat( + [input_tensor, token_thinks], dim=1 + ) + if mode == 'long_context': + assert Config().long_context_config['max_seq_len'] > input_tensor.shape[1] + max_new_tokens, \ + "please change max_seq_len in ~/.ktransformers/config.yaml" + + if system != "Windows" and (config.architectures[0] == "DeepseekV2ForCausalLM" or config.architectures[0] == "DeepseekV3ForCausalLM") and flashinfer_enabled and get_compute_capability() >= 8 and device_manager.gpu_vendor == GPUVendor.NVIDIA: + generated = prefill_and_generate( + model, tokenizer, input_tensor.to(device), max_new_tokens, use_cuda_graph, mode = mode, force_think = force_think, chunk_size = chunk_size, + use_flashinfer_mla = True, num_heads = config.num_attention_heads, head_dim_ckv = config.kv_lora_rank, head_dim_kpe = config.qk_rope_head_dim, q_head_dim = config.qk_rope_head_dim + config.qk_nope_head_dim + ) + else: + generated = prefill_and_generate( + model, tokenizer, input_tensor.to(device), max_new_tokens, use_cuda_graph, mode = mode, force_think = force_think, chunk_size = chunk_size, + ) + + +if __name__ == "__main__": + install_patch() + IS_DEBUG = True + + if IS_DEBUG == False: + parser = argparse.ArgumentParser() + + parser.add_argument("--model_path", required=True) + parser.add_argument("--model_config_path", default=None) + parser.add_argument("--gguf_path", required=True) + parser.add_argument("--cpu_infer", type=int, default=32) + parser.add_argument("--max_new_tokens", type=int, default=1000) + parser.add_argument("--force_think", action="store_true") + parser.add_argument("--optimize_config_path", required=True) + parser.add_argument("--is_sft", type=lambda x: x.lower() == "true", default=False) + parser.add_argument("--sft_data_path", default=None) + parser.add_argument("--save_adapter_path", default=None) + parser.add_argument("--use_adapter", type=lambda x: x.lower() == "true", default=False) + parser.add_argument("--use_adapter_path", default=None) + parser.add_argument("--is_test_data", type=lambda x: x.lower() == "true", default=False) + parser.add_argument("--test_data_path", default=None) + parser.add_argument("--output_dir", default=None) + + args = parser.parse_args() + + local_chat( + model_path=args.model_path, + model_config_path=args.model_config_path, + gguf_path=args.gguf_path, + cpu_infer=args.cpu_infer, + max_new_tokens=args.max_new_tokens, + force_think=args.force_think, + optimize_config_path=args.optimize_config_path, + is_sft=args.is_sft, + sft_data_path=args.sft_data_path, + save_adapter_path=args.save_adapter_path, + use_adapter=args.use_adapter, + use_adapter_path=args.use_adapter_path, + is_test_data=args.is_test_data, + test_data_path=args.test_data_path, + output_dir= args.output_dir + ) + + else: + local_chat( + # model_path="/mnt/data/data/DeepSeek-V3-671B-BF16", + # model_config_path="/mnt/data/data/DeepSeek-V3-671B-BF16", + # gguf_path="/mnt/data/data/DeepSeek-V3-671B-BF16", + model_path="/mnt/data/models/DeepSeek-V2-Lite-Chat", + model_config_path="/mnt/data/models/DeepSeek-V2-Lite-Chat", + gguf_path="/mnt/data/models/DeepSeek-V2-Lite-Chat", + cpu_infer=32, + max_new_tokens=1000, + force_think=False, + # optimize_config_path="ktransformers/optimize/optimize_rules/DeepSeek-V3-Chat-multi-gpu.yaml", + optimize_config_path="ktransformers/optimize/optimize_rules/DeepSeek-V2-Lite-Chat-sft-amx.yaml", + is_sft=True, + sft_data_path="test_adapter/western_train.json", + # sft_data_path="test_adapter/western_train.json", + # sft_data_path="test_adapter/500token_test.json", + save_adapter_path="/mnt/data/lpl/test_adapter/Kwhl_test_py312_torch28_DeepSeekV2_WEST", + use_adapter=False, + use_adapter_path="/mnt/data/lpl/test_adapter/Kllama_deepseekV2_AfriMed_mcq", + is_test_data=False, + test_data_path="/home/lpl/LLaMA-Factory-KT/data/mcq_test.json", + output_dir="/mnt/data/lpl/test_adapter/Kllama_deepseekV2_AfriMed_mcq/baselines", + ) + \ No newline at end of file diff --git a/archive/kt-sft/ktransformers/local_chat.sh b/archive/kt-sft/ktransformers/local_chat.sh new file mode 100755 index 0000000..2654a6f --- /dev/null +++ b/archive/kt-sft/ktransformers/local_chat.sh @@ -0,0 +1,17 @@ +#!/bin/bash + +python3 ktransformers/local_chat.py \ + --model_path "/mnt/data/models/DeepSeek-V2-Lite-Chat" \ + --model_config_path "/mnt/data/models/DeepSeek-V2-Lite-Chat" \ + --gguf_path "/mnt/data/models/DeepSeek-V2-Lite-Chat" \ + --cpu_infer 32 \ + --max_new_tokens 1000 \ + --optimize_config_path "ktransformers/optimize/optimize_rules/DeepSeek-V2-Lite-Chat-sft-amx.yaml" \ + --is_sft False \ + --sft_data_path "test_adapter/sft_translation.json" \ + --save_adapter_path "test_adapter/demo_adapter_KT_target_kv" \ + --use_adapter True \ + --use_adapter_path "/mnt/data/lpl/test_adapter/KT_newLoader_singleGPU_deepseekV2_Neko_AFS/checkpoint-566" \ + --is_test_data False \ + --test_data_path "test_adapter/demo_adapter_origin_target_kv" \ + --output_dir "test_adapter/demo_adapter_origin_target_kv" \ \ No newline at end of file diff --git a/archive/kt-sft/ktransformers/lora_test_module.py b/archive/kt-sft/ktransformers/lora_test_module.py new file mode 100644 index 0000000..0ec5855 --- /dev/null +++ b/archive/kt-sft/ktransformers/lora_test_module.py @@ -0,0 +1,158 @@ +import os +import platform +import sys + +project_dir = os.path.dirname(os.path.dirname(__file__)) +sys.path.insert(0, project_dir) + +from torchviz import make_dot +from torch import nn +import torch +from transformers import ( + AutoTokenizer, + AutoConfig, + AutoModelForCausalLM, + GenerationConfig, + TextStreamer, +) + +from ktransformers.operators.linear import KLinearTorch, KTransformersLinear +from ktransformers.sft.peft_utils.lora_layer import KTransformersLinearLora +from ktransformers.util.custom_loader import GGUFLoader +from ktransformers.util.inference_state import InferenceState + +import hiddenlayer as hl + +gguf_loader = GGUFLoader(gguf_path="/home/yj/ktransformers/GGUF-DeepSeek-V2-Lite-Chat") +config = AutoConfig.from_pretrained("/home/yj/ktransformers/DeepSeek-V2-Lite-Chat", trust_remote_code=True) +torch.set_default_dtype(config.torch_dtype) + +class TestModelLora(nn.Module): + def __init__(self): + super().__init__() + + random_linear_layer = nn.Linear(in_features=3072, out_features=2048, bias=False) + + orig_linear = KTransformersLinear( + key='blk.0.attn_q', + gguf_loader=gguf_loader, + config=config, + orig_module=random_linear_layer, + generate_op="KLinearTorch" + ) + self.layer = KTransformersLinearLora( + orig_module=orig_linear, + adapter_name="lora_test", + r=8, + lora_alpha=16 + ) + self.layer.generate_linear.weight = torch.randn(3072, 2048).to("cuda") + + def forward(self, x): + return self.layer(x) + +class TestModelBase(nn.Module): + def __init__(self): + super().__init__() + self.layer = KTransformersLinear( + key="linear", + gguf_loader=gguf_loader, + config=config, + orig_module=nn.Linear(in_features=3072, out_features=2048, bias=False), + generate_op="KLinearTorch" + ) + # self.layer.generate_linear.weight = torch.randn(3072, 2048).to("cuda") + weight = torch.randn(3072, 2048, device="cuda") + self.layer.load(w=nn.Parameter(weight), mode = InferenceState.GENERATE) + # self.layer.generate_linear.weight = nn.Parameter(torch.randn(3072, 2048).to("cuda")) + self.fc1 = nn.Linear(3072, 2048, bias=False) + self.relu = nn.ReLU() + self.fc2 = nn.Linear(2048, 3072, bias=False) + # self.layer.load(mode=InferenceState.GENERATE) + + def forward(self, x): + x = self.layer(x) + # x = self.fc1(x) + x = self.relu(x) + x = self.fc2(x) + return x + +class TestModelTorch(nn.Module): + def __init__(self): + super().__init__() + self.layer = KLinearTorch( + key="linear", + gguf_loader=gguf_loader, + config=config, + orig_module=nn.Linear(in_features=3072, out_features=2048, bias=False) + ) + # self.layer.weight = nn.Parameter(torch.randn(3072, 2048).to("cuda")) + # self.layer.weight = torch.randn(3072, 2048).to("cuda") + weight = torch.randn(3072, 2048, device="cuda") + self.layer.load(w=nn.Parameter(weight), device="cuda") + self.fc1 = nn.Linear(3072, 2048, bias=False) + self.relu = nn.ReLU() + self.fc2 = nn.Linear(2048, 3072, bias=False) + # self.layer.load(mode=InferenceState.GENERATE) + + def forward(self, x): + x = self.layer(x) + # x = self.fc1(x) + x = self.relu(x) + x = self.fc2(x) + return x + + +# # KLinearTorch Well DONE for test! +# model = TestModelTorch() +# x = torch.randn(2048, 3072, requires_grad=True) +# out = model(x) +# make_dot(out, params=dict(model.named_parameters())).render("KTLinear_graph", format="svg") + + +# model = TestModelBase() +# x = torch.randn(2048, 3072, requires_grad=True) +# out = model(x) +# make_dot(out, params=dict(model.named_parameters())).render("base_graph", format="svg") + +# MyConvNet_graph=hl.build_graph(model,torch.zeros(size=[2048, 3072])) +# MyConvNet_graph.theme=hl.graph.THEMES['blue'].copy() +# MyConvNet_graph.save(path='./base_graph.png',format='png') + +# model = TestModelLora() +# x = torch.randn(2048, 3072, requires_grad=True) +# out = model(x) +# make_dot(out, params=dict(model.named_parameters())).render("lora_graph", format="svg") + + +from peft import LoraConfig, get_peft_model + +class BaseModel(nn.Module): + def __init__(self): + super().__init__() + self.linear = nn.Linear(3072, 2048, bias=False) + + def forward(self, x): + return self.linear(x) + +model = BaseModel().to("cuda") + +lora_config = LoraConfig( + r=8, + lora_alpha=16, + target_modules=["linear"], + lora_dropout=0.0, + bias="none", +) + +peft_model = get_peft_model(model, lora_config) +print(peft_model) + +x = torch.randn(2048, 3072, requires_grad=True).to("cuda") + +out = peft_model(x) + +dot = make_dot(out, + params=dict(peft_model.named_parameters())) + +dot.render("origin_lora_graph", format="svg") \ No newline at end of file diff --git a/archive/kt-sft/ktransformers/models/__init__.py b/archive/kt-sft/ktransformers/models/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/archive/kt-sft/ktransformers/models/configuration_deepseek.py b/archive/kt-sft/ktransformers/models/configuration_deepseek.py new file mode 100644 index 0000000..5341fd9 --- /dev/null +++ b/archive/kt-sft/ktransformers/models/configuration_deepseek.py @@ -0,0 +1,207 @@ +# Adapted from +# https://huggingface.co/deepseek-ai/DeepSeek-V2-Chat-0628/blob/main/configuration_deepseek.py +# Copyright 2023 DeepSeek-AI and The HuggingFace Inc. team. All rights reserved. +# Copyright (c) 2024 by KVCache.AI, All Rights Reserved. +from transformers.configuration_utils import PretrainedConfig +from transformers.utils import logging + +logger = logging.get_logger(__name__) + +DEEPSEEK_PRETRAINED_CONFIG_ARCHIVE_MAP = {} +class DeepseekV2Config(PretrainedConfig): + r""" + This is the configuration class to store the configuration of a [`DeepseekV2Model`]. It is used to instantiate an DeepSeek + model according to the specified arguments, defining the model architecture. Instantiating a configuration with the + defaults will yield a similar configuration to that of the DeepSeek-V2. + 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 102400): + Vocabulary size of the Deep model. Defines the number of different tokens that can be represented by the + `inputs_ids` passed when calling [`DeepseekV2Model`] + hidden_size (`int`, *optional*, defaults to 4096): + Dimension of the hidden representations. + intermediate_size (`int`, *optional*, defaults to 11008): + Dimension of the MLP representations. + moe_intermediate_size (`int`, *optional*, defaults to 1407): + Dimension of the MoE representations. + num_hidden_layers (`int`, *optional*, defaults to 32): + Number of hidden layers in the Transformer decoder. + num_attention_heads (`int`, *optional*, defaults to 32): + Number of attention heads for each attention layer in the Transformer decoder. + n_shared_experts (`int`, *optional*, defaults to None): + Number of shared experts, None means dense model. + n_routed_experts (`int`, *optional*, defaults to None): + Number of routed experts, None means dense model. + routed_scaling_factor (`float`, *optional*, defaults to 1.0): + Scaling factor or routed experts. + topk_method (`str`, *optional*, defaults to `gready`): + Topk method used in routed gate. + n_group (`int`, *optional*, defaults to None): + Number of groups for routed experts. + topk_group (`int`, *optional*, defaults to None): + Number of selected groups for each token(for each token, ensuring the selected experts is only within `topk_group` groups). + num_experts_per_tok (`int`, *optional*, defaults to None): + Number of selected experts, None means dense model. + moe_layer_freq (`int`, *optional*, defaults to 1): + The frequency of the MoE layer: one expert layer for every `moe_layer_freq - 1` dense layers. + first_k_dense_replace (`int`, *optional*, defaults to 0): + Number of dense layers in shallow layers(embed->dense->dense->...->dense->moe->moe...->lm_head). + \--k dense layers--/ + norm_topk_prob (`bool`, *optional*, defaults to False): + Whether to normalize the weights of the routed experts. + scoring_func (`str`, *optional*, defaults to 'softmax'): + Method of computing expert weights. + aux_loss_alpha (`float`, *optional*, defaults to 0.001): + Auxiliary loss weight coefficient. + seq_aux = (`bool`, *optional*, defaults to True): + Whether to compute the auxiliary loss for each individual sample. + num_key_value_heads (`int`, *optional*): + This is the number of key_value heads that should be used to implement Grouped Query Attention. If + `num_key_value_heads=num_attention_heads`, the model will use Multi Head Attention (MHA), if + `num_key_value_heads=1 the model will use Multi Query Attention (MQA) otherwise GQA is used. When + converting a multi-head checkpoint to a GQA checkpoint, each group key and value head should be constructed + by meanpooling all the original heads within that group. For more details checkout [this + paper](https://arxiv.org/pdf/2305.13245.pdf). If it is not specified, will default to + `num_attention_heads`. + hidden_act (`str` or `function`, *optional*, defaults to `"silu"`): + The non-linear activation function (function or string) in the decoder. + max_position_embeddings (`int`, *optional*, defaults to 2048): + The maximum sequence length that this model might ever be used with. + 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`. + pad_token_id (`int`, *optional*): + Padding token id. + bos_token_id (`int`, *optional*, defaults to 1): + Beginning of stream token id. + eos_token_id (`int`, *optional*, defaults to 2): + End of stream token id. + pretraining_tp (`int`, *optional*, defaults to 1): + Experimental feature. Tensor parallelism rank used during pretraining. Please refer to [this + document](https://huggingface.co/docs/transformers/parallelism) to understand more about it. This value is + necessary to ensure exact reproducibility of the pretraining results. Please refer to [this + issue](https://github.com/pytorch/pytorch/issues/76232). + tie_word_embeddings (`bool`, *optional*, defaults to `False`): + Whether to tie weight embeddings + 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. Currently supports two scaling + strategies: linear and dynamic. Their scaling factor must be a float greater than 1. The expected format is + `{"type": strategy name, "factor": scaling factor}`. When using this flag, don't update + `max_position_embeddings` to the expected new maximum. + 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. + attention_dropout (`float`, *optional*, defaults to 0.0): + The dropout ratio for the attention probabilities. + ```python + >>> from transformers import DeepseekV2Model, DeepseekV2Config + >>> # Initializing a Deepseek-V2 style configuration + >>> configuration = DeepseekV2Config() + >>> # Accessing the model configuration + >>> configuration = model.config + ```""" + + model_type = "deepseek_v2" + keys_to_ignore_at_inference = ["past_key_values"] + + def __init__( + self, + vocab_size=102400, + hidden_size=4096, + intermediate_size=11008, + moe_intermediate_size = 1407, + num_hidden_layers=30, + num_attention_heads=32, + num_key_value_heads=32, + n_shared_experts = None, + n_routed_experts = None, + ep_size = 1, + routed_scaling_factor = 1.0, + kv_lora_rank = 512, + q_lora_rank = 1536, + qk_rope_head_dim = 64, + v_head_dim = 128, + qk_nope_head_dim = 128, + topk_method = 'gready', + n_group = None, + topk_group = None, + num_experts_per_tok = None, + moe_layer_freq = 1, + first_k_dense_replace = 0, + norm_topk_prob = False, + scoring_func = 'softmax', + aux_loss_alpha = 0.001, + seq_aux = True, + hidden_act="silu", + max_position_embeddings=2048, + initializer_range=0.02, + rms_norm_eps=1e-6, + use_cache=True, + pad_token_id=None, + bos_token_id=100000, + eos_token_id=100001, + pretraining_tp=1, + tie_word_embeddings=False, + rope_theta=10000.0, + rope_scaling=None, + attention_bias=False, + attention_dropout=0.0, + cpu_quant=None, + **kwargs, + ): + self.vocab_size = vocab_size + self.max_position_embeddings = max_position_embeddings + self.hidden_size = hidden_size + self.intermediate_size = intermediate_size + self.moe_intermediate_size = moe_intermediate_size + self.num_hidden_layers = num_hidden_layers + self.num_attention_heads = num_attention_heads + self.n_shared_experts = n_shared_experts + self.n_routed_experts = n_routed_experts + self.ep_size = ep_size + self.routed_scaling_factor = routed_scaling_factor + self.kv_lora_rank = kv_lora_rank + self.q_lora_rank = q_lora_rank + self.qk_rope_head_dim = qk_rope_head_dim + self.v_head_dim = v_head_dim + self.qk_nope_head_dim = qk_nope_head_dim + self.topk_method = topk_method + self.n_group = n_group + self.topk_group = topk_group + self.num_experts_per_tok = num_experts_per_tok + self.moe_layer_freq = moe_layer_freq + self.first_k_dense_replace = first_k_dense_replace + self.norm_topk_prob = norm_topk_prob + self.scoring_func = scoring_func + self.aux_loss_alpha = aux_loss_alpha + self.seq_aux = seq_aux + # for backward compatibility + if num_key_value_heads is None: + num_key_value_heads = num_attention_heads + + self.num_key_value_heads = num_key_value_heads + self.hidden_act = hidden_act + self.initializer_range = initializer_range + self.rms_norm_eps = rms_norm_eps + self.pretraining_tp = pretraining_tp + self.use_cache = use_cache + self.rope_theta = rope_theta + self.rope_scaling = rope_scaling + self.attention_bias = attention_bias + self.attention_dropout = attention_dropout + + self.cpu_quant = cpu_quant + + super().__init__( + pad_token_id=pad_token_id, + bos_token_id=bos_token_id, + eos_token_id=eos_token_id, + tie_word_embeddings=tie_word_embeddings, + **kwargs, + ) diff --git a/archive/kt-sft/ktransformers/models/configuration_deepseek_v3.py b/archive/kt-sft/ktransformers/models/configuration_deepseek_v3.py new file mode 100644 index 0000000..235b7b0 --- /dev/null +++ b/archive/kt-sft/ktransformers/models/configuration_deepseek_v3.py @@ -0,0 +1,193 @@ +from transformers.configuration_utils import PretrainedConfig +from transformers.utils import logging + +logger = logging.get_logger(__name__) + +DEEPSEEK_PRETRAINED_CONFIG_ARCHIVE_MAP = {} +class DeepseekV3Config(PretrainedConfig): + r""" + This is the configuration class to store the configuration of a [`DeepseekV3Model`]. It is used to instantiate an DeepSeek + model according to the specified arguments, defining the model architecture. Instantiating a configuration with the + defaults will yield a similar configuration to that of the DeepSeek-V3. + 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 129280): + Vocabulary size of the Deep model. Defines the number of different tokens that can be represented by the + `inputs_ids` passed when calling [`DeepseekV3Model`] + hidden_size (`int`, *optional*, defaults to 4096): + Dimension of the hidden representations. + intermediate_size (`int`, *optional*, defaults to 11008): + Dimension of the MLP representations. + moe_intermediate_size (`int`, *optional*, defaults to 1407): + Dimension of the MoE representations. + num_hidden_layers (`int`, *optional*, defaults to 32): + Number of hidden layers in the Transformer decoder. + num_nextn_predict_layers (`int`, *optional*, defaults to 1): + Number of nextn predict layers in the DeepSeekV3 Model. + num_attention_heads (`int`, *optional*, defaults to 32): + Number of attention heads for each attention layer in the Transformer decoder. + n_shared_experts (`int`, *optional*, defaults to None): + Number of shared experts, None means dense model. + n_routed_experts (`int`, *optional*, defaults to None): + Number of routed experts, None means dense model. + routed_scaling_factor (`float`, *optional*, defaults to 1.0): + Scaling factor or routed experts. + topk_method (`str`, *optional*, defaults to `gready`): + Topk method used in routed gate. + n_group (`int`, *optional*, defaults to None): + Number of groups for routed experts. + topk_group (`int`, *optional*, defaults to None): + Number of selected groups for each token(for each token, ensuring the selected experts is only within `topk_group` groups). + num_experts_per_tok (`int`, *optional*, defaults to None): + Number of selected experts, None means dense model. + moe_layer_freq (`int`, *optional*, defaults to 1): + The frequency of the MoE layer: one expert layer for every `moe_layer_freq - 1` dense layers. + first_k_dense_replace (`int`, *optional*, defaults to 0): + Number of dense layers in shallow layers(embed->dense->dense->...->dense->moe->moe...->lm_head). + \--k dense layers--/ + norm_topk_prob (`bool`, *optional*, defaults to False): + Whether to normalize the weights of the routed experts. + scoring_func (`str`, *optional*, defaults to 'softmax'): + Method of computing expert weights. + aux_loss_alpha (`float`, *optional*, defaults to 0.001): + Auxiliary loss weight coefficient. + seq_aux = (`bool`, *optional*, defaults to True): + Whether to compute the auxiliary loss for each individual sample. + num_key_value_heads (`int`, *optional*): + This is the number of key_value heads that should be used to implement Grouped Query Attention. If + `num_key_value_heads=num_attention_heads`, the model will use Multi Head Attention (MHA), if + `num_key_value_heads=1 the model will use Multi Query Attention (MQA) otherwise GQA is used. When + converting a multi-head checkpoint to a GQA checkpoint, each group key and value head should be constructed + by meanpooling all the original heads within that group. For more details checkout [this + paper](https://arxiv.org/pdf/2305.13245.pdf). If it is not specified, will default to + `num_attention_heads`. + hidden_act (`str` or `function`, *optional*, defaults to `"silu"`): + The non-linear activation function (function or string) in the decoder. + max_position_embeddings (`int`, *optional*, defaults to 2048): + The maximum sequence length that this model might ever be used with. + 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`. + pad_token_id (`int`, *optional*): + Padding token id. + bos_token_id (`int`, *optional*, defaults to 1): + Beginning of stream token id. + eos_token_id (`int`, *optional*, defaults to 2): + End of stream token id. + tie_word_embeddings (`bool`, *optional*, defaults to `False`): + Whether to tie weight embeddings + 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. Currently supports two scaling + strategies: linear and dynamic. Their scaling factor must be a float greater than 1. The expected format is + `{"type": strategy name, "factor": scaling factor}`. When using this flag, don't update + `max_position_embeddings` to the expected new maximum. + 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. + attention_dropout (`float`, *optional*, defaults to 0.0): + The dropout ratio for the attention probabilities. + ```python + >>> from transformers import DeepseekV3Model, DeepseekV3Config + >>> # Initializing a Deepseek-V3 style configuration + >>> configuration = DeepseekV3Config() + >>> # Accessing the model configuration + >>> configuration = model.config + ```""" + + model_type = "deepseek_v3" + keys_to_ignore_at_inference = ["past_key_values"] + + def __init__( + self, + vocab_size=129280, + hidden_size=7168, + intermediate_size=18432, + moe_intermediate_size = 2048, + num_hidden_layers=61, + num_nextn_predict_layers=1, + num_attention_heads=128, + num_key_value_heads=128, + n_shared_experts = 1, + n_routed_experts = 256, + ep_size = 1, + routed_scaling_factor = 2.5, + kv_lora_rank = 512, + q_lora_rank = 1536, + qk_rope_head_dim = 64, + v_head_dim = 128, + qk_nope_head_dim = 128, + topk_method = 'noaux_tc', + n_group = 8, + topk_group = 4, + num_experts_per_tok = 8, + moe_layer_freq = 1, + first_k_dense_replace = 3, + norm_topk_prob = True, + scoring_func = 'sigmoid', + hidden_act="silu", + max_position_embeddings=4096, + initializer_range=0.02, + rms_norm_eps=1e-6, + use_cache=True, + pad_token_id=None, + bos_token_id=0, + eos_token_id=1, + tie_word_embeddings=False, + rope_theta=10000.0, + rope_scaling=None, + attention_bias=False, + 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.moe_intermediate_size = moe_intermediate_size + self.num_hidden_layers = num_hidden_layers + self.num_nextn_predict_layers = num_nextn_predict_layers + self.num_attention_heads = num_attention_heads + self.n_shared_experts = n_shared_experts + self.n_routed_experts = n_routed_experts + self.ep_size = ep_size + self.routed_scaling_factor = routed_scaling_factor + self.kv_lora_rank = kv_lora_rank + self.q_lora_rank = q_lora_rank + self.qk_rope_head_dim = qk_rope_head_dim + self.v_head_dim = v_head_dim + self.qk_nope_head_dim = qk_nope_head_dim + self.topk_method = topk_method + self.n_group = n_group + self.topk_group = topk_group + self.num_experts_per_tok = num_experts_per_tok + self.moe_layer_freq = moe_layer_freq + self.first_k_dense_replace = first_k_dense_replace + self.norm_topk_prob = norm_topk_prob + self.scoring_func = scoring_func + # for backward compatibility + if num_key_value_heads is None: + num_key_value_heads = num_attention_heads + + self.num_key_value_heads = num_key_value_heads + self.hidden_act = hidden_act + self.initializer_range = initializer_range + self.rms_norm_eps = 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 + + super().__init__( + pad_token_id=pad_token_id, + bos_token_id=bos_token_id, + eos_token_id=eos_token_id, + tie_word_embeddings=tie_word_embeddings, + **kwargs, + ) \ No newline at end of file diff --git a/archive/kt-sft/ktransformers/models/configuration_llama.py b/archive/kt-sft/ktransformers/models/configuration_llama.py new file mode 100644 index 0000000..2b4f4db --- /dev/null +++ b/archive/kt-sft/ktransformers/models/configuration_llama.py @@ -0,0 +1,203 @@ +# coding=utf-8 +# Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved. +# +# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX +# and OPT implementations in this library. It has been modified from its +# original forms to accommodate minor architectural differences compared +# to GPT-NeoX and OPT used by the Meta AI team that trained the model. +# +# 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. +"""LLaMA model configuration""" + +from transformers.configuration_utils import PretrainedConfig +from transformers.modeling_rope_utils import rope_config_validation + + +class LlamaConfig(PretrainedConfig): + r""" + This is the configuration class to store the configuration of a [`LlamaModel`]. It is used to instantiate an LLaMA + model according to the specified arguments, defining the model architecture. Instantiating a configuration with the + defaults will yield a similar configuration to that of the LLaMA-7B. + + Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the + documentation from [`PretrainedConfig`] for more information. + + + Args: + vocab_size (`int`, *optional*, defaults to 32000): + Vocabulary size of the LLaMA model. Defines the number of different tokens that can be represented by the + `inputs_ids` passed when calling [`LlamaModel`] + hidden_size (`int`, *optional*, defaults to 4096): + Dimension of the hidden representations. + intermediate_size (`int`, *optional*, defaults to 11008): + Dimension of the MLP representations. + num_hidden_layers (`int`, *optional*, defaults to 32): + Number of hidden layers in the Transformer decoder. + num_attention_heads (`int`, *optional*, defaults to 32): + Number of attention heads for each attention layer in the Transformer decoder. + num_key_value_heads (`int`, *optional*): + This is the number of key_value heads that should be used to implement Grouped Query Attention. If + `num_key_value_heads=num_attention_heads`, the model will use Multi Head Attention (MHA), if + `num_key_value_heads=1` the model will use Multi Query Attention (MQA) otherwise GQA is used. When + converting a multi-head checkpoint to a GQA checkpoint, each group key and value head should be constructed + by meanpooling all the original heads within that group. For more details checkout [this + paper](https://arxiv.org/pdf/2305.13245.pdf). If it is not specified, will default to + `num_attention_heads`. + hidden_act (`str` or `function`, *optional*, defaults to `"silu"`): + The non-linear activation function (function or string) in the decoder. + max_position_embeddings (`int`, *optional*, defaults to 2048): + The maximum sequence length that this model might ever be used with. Llama 1 supports up to 2048 tokens, + Llama 2 up to 4096, CodeLlama up to 16384. + 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`. + pad_token_id (`int`, *optional*): + Padding token id. + bos_token_id (`int`, *optional*, defaults to 1): + Beginning of stream token id. + eos_token_id (`int`, *optional*, defaults to 2): + End of stream token id. + pretraining_tp (`int`, *optional*, defaults to 1): + Experimental feature. Tensor parallelism rank used during pretraining. Please refer to [this + document](https://huggingface.co/docs/transformers/main/perf_train_gpu_many#tensor-parallelism) to + understand more about it. This value is necessary to ensure exact reproducibility of the pretraining + results. Please refer to [this issue](https://github.com/pytorch/pytorch/issues/76232). + tie_word_embeddings (`bool`, *optional*, defaults to `False`): + Whether to tie weight embeddings + rope_theta (`float`, *optional*, defaults to 10000.0): + The base period of the RoPE embeddings. + rope_scaling (`Dict`, *optional*): + Dictionary containing the scaling configuration for the RoPE embeddings. NOTE: if you apply new rope type + and you expect the model to work on longer `max_position_embeddings`, we recommend you to update this value + accordingly. + Expected contents: + `rope_type` (`str`): + The sub-variant of RoPE to use. Can be one of ['default', 'linear', 'dynamic', 'yarn', 'longrope', + 'llama3'], with 'default' being the original RoPE implementation. + `factor` (`float`, *optional*): + Used with all rope types except 'default'. The scaling factor to apply to the RoPE embeddings. In + most scaling types, a `factor` of x will enable the model to handle sequences of length x * + original maximum pre-trained length. + `original_max_position_embeddings` (`int`, *optional*): + Used with 'dynamic', 'longrope' and 'llama3'. The original max position embeddings used during + pretraining. + `attention_factor` (`float`, *optional*): + Used with 'yarn' and 'longrope'. The scaling factor to be applied on the attention + computation. If unspecified, it defaults to value recommended by the implementation, using the + `factor` field to infer the suggested value. + `beta_fast` (`float`, *optional*): + Only used with 'yarn'. Parameter to set the boundary for extrapolation (only) in the linear + ramp function. If unspecified, it defaults to 32. + `beta_slow` (`float`, *optional*): + Only used with 'yarn'. Parameter to set the boundary for interpolation (only) in the linear + ramp function. If unspecified, it defaults to 1. + `short_factor` (`List[float]`, *optional*): + Only used with 'longrope'. The scaling factor to be applied to short contexts (< + `original_max_position_embeddings`). Must be a list of numbers with the same length as the hidden + size divided by the number of attention heads divided by 2 + `long_factor` (`List[float]`, *optional*): + Only used with 'longrope'. The scaling factor to be applied to long contexts (< + `original_max_position_embeddings`). Must be a list of numbers with the same length as the hidden + size divided by the number of attention heads divided by 2 + `low_freq_factor` (`float`, *optional*): + Only used with 'llama3'. Scaling factor applied to low frequency components of the RoPE + `high_freq_factor` (`float`, *optional*): + Only used with 'llama3'. Scaling factor applied to high frequency components of the RoPE + 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. + mlp_bias (`bool`, *optional*, defaults to `False`): + Whether to use a bias in up_proj, down_proj and gate_proj layers in the MLP layers. + + ```python + >>> from transformers import LlamaModel, LlamaConfig + + >>> # Initializing a LLaMA llama-7b style configuration + >>> configuration = LlamaConfig() + + >>> # Initializing a model from the llama-7b style configuration + >>> model = LlamaModel(configuration) + + >>> # Accessing the model configuration + >>> configuration = model.config + ```""" + + model_type = "llama" + keys_to_ignore_at_inference = ["past_key_values"] + + def __init__( + self, + vocab_size=32000, + hidden_size=4096, + intermediate_size=11008, + num_hidden_layers=32, + num_attention_heads=32, + num_key_value_heads=None, + hidden_act="silu", + max_position_embeddings=2048, + initializer_range=0.02, + rms_norm_eps=1e-6, + use_cache=True, + pad_token_id=None, + bos_token_id=1, + eos_token_id=2, + pretraining_tp=1, + tie_word_embeddings=False, + rope_theta=10000.0, + rope_scaling=None, + attention_bias=False, + attention_dropout=0.0, + mlp_bias=False, + **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 + + # for backward compatibility + if num_key_value_heads is None: + num_key_value_heads = num_attention_heads + + self.num_key_value_heads = num_key_value_heads + self.hidden_act = hidden_act + self.initializer_range = initializer_range + self.rms_norm_eps = rms_norm_eps + self.pretraining_tp = pretraining_tp + self.use_cache = use_cache + self.rope_theta = rope_theta + self.rope_scaling = rope_scaling + self.attention_bias = attention_bias + self.attention_dropout = attention_dropout + self.mlp_bias = mlp_bias + + # 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) + + super().__init__( + pad_token_id=pad_token_id, + bos_token_id=bos_token_id, + eos_token_id=eos_token_id, + tie_word_embeddings=tie_word_embeddings, + **kwargs, + ) diff --git a/archive/kt-sft/ktransformers/models/configuration_qwen2_moe.py b/archive/kt-sft/ktransformers/models/configuration_qwen2_moe.py new file mode 100644 index 0000000..345af76 --- /dev/null +++ b/archive/kt-sft/ktransformers/models/configuration_qwen2_moe.py @@ -0,0 +1,177 @@ +# coding=utf-8 +# Copyright 2024 The Qwen team, Alibaba Group and the HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Qwen2MoE model configuration""" + +from transformers.configuration_utils import PretrainedConfig +from transformers.utils import logging + + +logger = logging.get_logger(__name__) + + +class Qwen2MoeConfig(PretrainedConfig): + r""" + This is the configuration class to store the configuration of a [`Qwen2MoeModel`]. It is used to instantiate a + Qwen2MoE model according to the specified arguments, defining the model architecture. Instantiating a configuration + with the defaults will yield a similar configuration to that of + Qwen1.5-MoE-A2.7B" [Qwen/Qwen1.5-MoE-A2.7B"](https://huggingface.co/Qwen/Qwen1.5-MoE-A2.7B"). + + Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the + documentation from [`PretrainedConfig`] for more information. + + + Args: + vocab_size (`int`, *optional*, defaults to 151936): + Vocabulary size of the Qwen2MoE model. Defines the number of different tokens that can be represented by the + `inputs_ids` passed when calling [`Qwen2MoeModel`] + 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 24): + 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 16): + 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` 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. + 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. + attention_dropout (`float`, *optional*, defaults to 0.0): + The dropout ratio for the attention probabilities. + decoder_sparse_step (`int`, *optional*, defaults to 1): + The frequency of the MoE layer. + moe_intermediate_size (`int`, *optional*, defaults to 1408): + Intermediate size of the routed expert. + shared_expert_intermediate_size (`int`, *optional*, defaults to 5632): + Intermediate size of the shared expert. + num_experts_per_tok (`int`, *optional*, defaults to 4): + Number of selected experts. + num_experts (`int`, *optional*, defaults to 60): + Number of routed experts. + norm_topk_prob (`bool`, *optional*, defaults to `False`): + 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. Enabeling 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 Qwen2MoeMLP rather than Qwen2MoeSparseMoeBlock + 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. + + ```python + >>> from transformers import Qwen2MoeModel, Qwen2MoeConfig + + >>> # Initializing a Qwen2MoE style configuration + >>> configuration = Qwen2MoeConfig() + + >>> # Initializing a model from the Qwen1.5-MoE-A2.7B" style configuration + >>> model = Qwen2MoeModel(configuration) + + >>> # Accessing the model configuration + >>> configuration = model.config + ```""" + + model_type = "qwen2_moe" + keys_to_ignore_at_inference = ["past_key_values"] + + def __init__( + self, + vocab_size=151936, + hidden_size=2048, + intermediate_size=5632, + num_hidden_layers=24, + num_attention_heads=16, + num_key_value_heads=16, + 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, + use_sliding_window=False, + sliding_window=4096, + max_window_layers=28, + attention_dropout=0.0, + decoder_sparse_step=1, + moe_intermediate_size=1408, + shared_expert_intermediate_size=5632, + num_experts_per_tok=4, + num_experts=60, + norm_topk_prob=False, + output_router_logits=False, + router_aux_loss_coef=0.001, + mlp_only_layers=None, + **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 use_sliding_window else None + self.max_window_layers = max_window_layers + + 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.attention_dropout = attention_dropout + + # 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 = [] if mlp_only_layers is None else mlp_only_layers + + super().__init__( + tie_word_embeddings=tie_word_embeddings, + **kwargs, + ) \ No newline at end of file diff --git a/archive/kt-sft/ktransformers/models/configuration_qwen3_moe.py b/archive/kt-sft/ktransformers/models/configuration_qwen3_moe.py new file mode 100644 index 0000000..ebbf87f --- /dev/null +++ b/archive/kt-sft/ktransformers/models/configuration_qwen3_moe.py @@ -0,0 +1,233 @@ +# coding=utf-8 +# Copyright 2024 The Qwen team, Alibaba Group and the HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Qwen3MoE model configuration""" + +from transformers.configuration_utils import PretrainedConfig +from transformers.modeling_rope_utils import rope_config_validation +from transformers.utils import logging + + +logger = logging.get_logger(__name__) + + +class Qwen3MoeConfig(PretrainedConfig): + r""" + This is the configuration class to store the configuration of a [`Qwen3MoeModel`]. It is used to instantiate a + Qwen3MoE model according to the specified arguments, defining the model architecture. Instantiating a configuration + with the defaults will yield a similar configuration to that of [Qwen/Qwen3-MoE-15B-A2B](https://huggingface.co/Qwen/Qwen3-15B-A2B). + 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 Qwen3MoE model. Defines the number of different tokens that can be represented by the + `inputs_ids` passed when calling [`Qwen3MoeModel`] + hidden_size (`int`, *optional*, defaults to 2048): + Dimension of the hidden representations. + intermediate_size (`int`, *optional*, defaults to 6144): + Dimension of the MLP representations. + num_hidden_layers (`int`, *optional*, defaults to 24): + 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 4): + 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` 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. NOTE: if you apply new rope type + and you expect the model to work on longer `max_position_embeddings`, we recommend you to update this value + accordingly. + Expected contents: + `rope_type` (`str`): + The sub-variant of RoPE to use. Can be one of ['default', 'linear', 'dynamic', 'yarn', 'longrope', + 'llama3'], with 'default' being the original RoPE implementation. + `factor` (`float`, *optional*): + Used with all rope types except 'default'. The scaling factor to apply to the RoPE embeddings. In + most scaling types, a `factor` of x will enable the model to handle sequences of length x * + original maximum pre-trained length. + `original_max_position_embeddings` (`int`, *optional*): + Used with 'dynamic', 'longrope' and 'llama3'. The original max position embeddings used during + pretraining. + `attention_factor` (`float`, *optional*): + Used with 'yarn' and 'longrope'. The scaling factor to be applied on the attention + computation. If unspecified, it defaults to value recommended by the implementation, using the + `factor` field to infer the suggested value. + `beta_fast` (`float`, *optional*): + Only used with 'yarn'. Parameter to set the boundary for extrapolation (only) in the linear + ramp function. If unspecified, it defaults to 32. + `beta_slow` (`float`, *optional*): + Only used with 'yarn'. Parameter to set the boundary for interpolation (only) in the linear + ramp function. If unspecified, it defaults to 1. + `short_factor` (`List[float]`, *optional*): + Only used with 'longrope'. The scaling factor to be applied to short contexts (< + `original_max_position_embeddings`). Must be a list of numbers with the same length as the hidden + size divided by the number of attention heads divided by 2 + `long_factor` (`List[float]`, *optional*): + Only used with 'longrope'. The scaling factor to be applied to long contexts (< + `original_max_position_embeddings`). Must be a list of numbers with the same length as the hidden + size divided by the number of attention heads divided by 2 + `low_freq_factor` (`float`, *optional*): + Only used with 'llama3'. Scaling factor applied to low frequency components of the RoPE + `high_freq_factor` (`float`, *optional*): + Only used with 'llama3'. Scaling factor applied to high frequency components of the RoPE + 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. + attention_dropout (`float`, *optional*, defaults to 0.0): + The dropout ratio for the attention probabilities. + decoder_sparse_step (`int`, *optional*, defaults to 1): + The frequency of the MoE layer. + moe_intermediate_size (`int`, *optional*, defaults to 768): + Intermediate size of the routed expert. + num_experts_per_tok (`int`, *optional*, defaults to 8): + Number of selected experts. + num_experts (`int`, *optional*, defaults to 128): + Number of routed experts. + norm_topk_prob (`bool`, *optional*, defaults to `False`): + 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. Enabeling 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 Qwen3MoeMLP rather than Qwen3MoeSparseMoeBlock + 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. + ```python + >>> from transformers import Qwen3MoeModel, Qwen3MoeConfig + >>> # Initializing a Qwen3MoE style configuration + >>> configuration = Qwen3MoeConfig() + >>> # Initializing a model from the Qwen3-15B-A2B" style configuration + >>> model = Qwen3MoeModel(configuration) + >>> # Accessing the model configuration + >>> configuration = model.config + ```""" + + model_type = "qwen3_moe" + keys_to_ignore_at_inference = ["past_key_values"] + + # Default tensor parallel plan for base model `Qwen3Moe` + 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=2048, + intermediate_size=6144, + num_hidden_layers=24, + num_attention_heads=32, + num_key_value_heads=4, + 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, + attention_dropout=0.0, + decoder_sparse_step=1, + moe_intermediate_size=768, + num_experts_per_tok=8, + num_experts=128, + norm_topk_prob=False, + output_router_logits=False, + router_aux_loss_coef=0.001, + mlp_only_layers=None, + **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 use_sliding_window else None + self.max_window_layers = max_window_layers + + 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_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) + + # MoE arguments + self.decoder_sparse_step = decoder_sparse_step + self.moe_intermediate_size = moe_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 + + super().__init__( + tie_word_embeddings=tie_word_embeddings, + **kwargs, + ) + + +__all__ = ["Qwen3MoeConfig"] \ No newline at end of file diff --git a/archive/kt-sft/ktransformers/models/custom_cache.py b/archive/kt-sft/ktransformers/models/custom_cache.py new file mode 100644 index 0000000..350af73 --- /dev/null +++ b/archive/kt-sft/ktransformers/models/custom_cache.py @@ -0,0 +1,333 @@ +''' +Description : +Author : Boxin Zhang +Version : 0.1.0 +''' +# Adapted from +# https://github.com/huggingface/transformers/blob/v4.41.2/src/transformers/cache_utils.py +# Copyright 2018- The Hugging Face team. All rights reserved. +# Copyright (c) 2024 by KVCache.AI, All Rights Reserved. +import torch +import torch.nn as nn +import transformers +from transformers import Cache, PretrainedConfig +from typing import List, Optional, Dict, Any, Tuple +try: + from ktransformers.server.balance_serve.settings import sched_ext +except: + print("no balance_serve") +class StaticCache(transformers.StaticCache): + """ + Static Cache class to be used with `torch.compile(model)`. + + Parameters: + config (`PretrainedConfig): + The configuration file defining the shape-related attributes required to initialize the static cache. + max_batch_size (`int`): + The maximum batch size with which the model will be used. + max_cache_len (`int`): + The maximum sequence length with which the model will be used. + device (`torch.device` or `dict`): + The device on which the cache should be initialized. Should be the same as the layer. + If a `dict`, it should contain the `device` key with the device name as the value. + dtype (*optional*, defaults to `torch.float32`): + The default `dtype` to use when initializing the layer. + """ + + def __init__(self, config: PretrainedConfig, max_batch_size: int, max_cache_len: int, device: torch.device| dict, dtype=None) -> None: + Cache.__init__(self) + self.max_batch_size = max_batch_size + self.max_cache_len = config.max_position_embeddings if max_cache_len is None else max_cache_len + # Some model define a custom `head_dim` != config.hidden_size // config.num_attention_heads + if config.architectures[0] == "DeepseekV3ForCausalLM": + self.head_dim = config.qk_rope_head_dim + else: + self.head_dim = ( + config.head_dim if hasattr(config, "head_dim") else config.hidden_size // config.num_attention_heads + ) + + self.dtype = dtype if dtype is not None else torch.float32 + self.num_key_value_heads = ( + config.num_attention_heads if config.num_key_value_heads is None else config.num_key_value_heads + ) + + self.key_cache: List[torch.Tensor] = [] + self.value_cache: List[torch.Tensor] = [] + cache_shape = (max_batch_size, self.num_key_value_heads, self.max_cache_len, self.head_dim) + if config.architectures[0] == "DeepseekV2ForCausalLM" or config.architectures[0] == "DeepseekV3ForCausalLM": + # TODO: for deepseek, cache_shape is different whether using Absorbed MLA, check it automatically + self.page_size = 64 + self.max_pages = (self.max_cache_len + self.page_size - 1) // self.page_size + latent_shape = (self.max_pages, self.page_size, 1, config.kv_lora_rank + config.qk_rope_head_dim) + self.kv_lora_rank = config.kv_lora_rank + self.qk_rope_head_dim = config.qk_rope_head_dim + # TODO: support real page table + self.page_table_map = dict() + self.page_table_list = [] + for idx in range(config.num_hidden_layers): + if isinstance(device, dict): + target_device = device[f"model.layers.{idx}.self_attn"]["generate_device"] + else: + target_device = device + + if target_device not in self.page_table_map: + page_table = torch.zeros((max_batch_size, self.max_pages), dtype=torch.int32, device=target_device) + for seq_id in range(max_batch_size): + page_table[seq_id, :] = torch.arange(seq_id * self.max_pages, seq_id * self.max_pages + self.max_pages, dtype=torch.int32, device=target_device) + self.page_table_map[target_device] = page_table + + self.page_table_list.append(self.page_table_map[target_device]) + + self.is_MLA = True + self.is_page = True + else: + key_shape = cache_shape + value_shape = cache_shape + self.is_MLA = False + + self.past_tokens = [] + self.num_hidden_layers = config.num_hidden_layers + for idx in range(self.num_hidden_layers): + # Note: `mark_static_address` is used to tag the cache as an fixed data pointer, preventing cuda graph + # breaks when updating the cache. + if isinstance(device, dict): + target_device = device[f"model.layers.{idx}.self_attn"]["generate_device"] + else: + target_device = device + + if self.is_MLA: + new_layer_key_cache = torch.zeros(latent_shape, dtype=self.dtype, device=target_device) + new_layer_value_cache = None + torch._dynamo.mark_static_address(new_layer_key_cache) + else: + new_layer_key_cache = torch.zeros(key_shape, dtype=self.dtype, device=target_device) + new_layer_value_cache = torch.zeros(value_shape, dtype=self.dtype, device=target_device) + torch._dynamo.mark_static_address(new_layer_key_cache) + torch._dynamo.mark_static_address(new_layer_value_cache) + + self.key_cache.append(new_layer_key_cache) + self.value_cache.append(new_layer_value_cache) + self.past_tokens.append(0) + + def update( + self, + key_states: torch.Tensor, + value_states: torch.Tensor, + layer_idx: int, + cache_kwargs: Optional[Dict[str, Any]] = None, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """ + Updates the cache with the new `key_states` and `value_states` for the layer `layer_idx`. + It is VERY important to index using a tensor, otherwise you introduce a copy to the device. + + Parameters: + key_states (`torch.Tensor`): + The new key states to cache. + value_states (`torch.Tensor`): + The new value states to cache. + layer_idx (`int`): + The index of the layer to cache the states for. + cache_kwargs (`Dict[str, Any]`, `optional`): + Additional arguments for the cache subclass. The `StaticCache` needs the `cache_position` input + to know how where to write in the cache. + + Return: + A tuple containing the updated key and value states. + """ + cache_position = cache_kwargs.get("cache_position") + k_out = self.key_cache[layer_idx] + v_out = self.value_cache[layer_idx] + self.past_tokens[layer_idx] += cache_position.size(0) + #print(cache_position) + if self.is_MLA: + page_idx = cache_position // self.page_size + page_offset = cache_position % self.page_size + # key shape (self.max_pages, self.page_size, 1, config.kv_lora_rank + config.qk_rope_head_dim) + k_out[page_idx, page_offset, :, :self.kv_lora_rank] = key_states + k_out[page_idx, page_offset, :, self.kv_lora_rank:] = value_states + return k_out, self.page_table_list[layer_idx] + else: + k_out[:, :, cache_position] = key_states + v_out[:, :, cache_position] = value_states + return k_out, v_out + + def get_seq_length(self, layer_idx: Optional[int] = 0) -> int: + """Returns the sequence length of the cached states that were seen by the model.""" + # Occupied cache == any slot in the 3rd dim (sequence length) holds a non-zero value. To save on compute, let's + # limit the check to the first batch member and head dimension. + # TODO: deprecate this function in favor of `cache_position` + return self.past_tokens[layer_idx] + + def change_seq_length(self, bias: Optional[int] = 0) -> int: + """Returns the sequence length of the cached states that were seen by the model.""" + # Occupied cache == any slot in the 3rd dim (sequence length) holds a non-zero value. To save on compute, let's + # limit the check to the first batch member and head dimension. + # TODO: deprecate this function in favor of `cache_position` + for layer_idx in range(self.num_hidden_layers): + self.past_tokens[layer_idx] += bias + + def get_max_length(self) -> Optional[int]: + """Returns the maximum sequence length of the cached states.""" + return self.max_cache_len + + def reset(self): + """Resets the cache values while preserving the objects""" + for layer_idx in range(len(self.key_cache)): + # In-place ops prevent breaking the static address + self.key_cache[layer_idx].zero_() + if self.value_cache[layer_idx] is not None: + self.value_cache[layer_idx].zero_() + self.past_tokens[layer_idx] = 0 + + def remove_suffix(self, start_pos): + for layer_idx in range(len(self.key_cache)): + # In-place ops prevent breaking the static address + if self.is_MLA: + k_cache = self.key_cache[layer_idx] + k_cache.view(-1, k_cache.shape[-1])[start_pos:].zero_() + else: + self.key_cache[layer_idx][..., start_pos:, :].zero_() + self.value_cache[layer_idx][..., start_pos:, :].zero_() + self.past_tokens[layer_idx] = start_pos + + def get_max_cache_shape(self) -> Tuple[int, int, int, int]: + """Returns the maximum shape of the cache.""" + return self.max_cache_len + +class KDeepSeekV3Cache(nn.Module): + def __init__( + self, + config: PretrainedConfig, + page_size: int = 256, + dtype=torch.bfloat16, + device=torch.device("cuda:0"), + + ): + super().__init__() + self.config = config + self.dtype = dtype + self.device = device + self.kv_lora_rank = config.kv_lora_rank + self.page_size = page_size + self.k_caches = [] + self.v_caches = [] + + + def load(self, inference_context: "sched_ext.InferenceContext"): + + for i in range(self.config.num_hidden_layers): + self.k_caches.append( + inference_context.k_cache[0][i] + ) + self.max_cache_len = self.k_caches[0].shape[0]*self.k_caches[0].shape[1] + + def update( + self, + key_states: torch.Tensor, + value_states: torch.Tensor, + layer_idx: int, + + page_idx: torch.Tensor, + page_offset: torch.Tensor, + + cache_kwargs: Optional[Dict[str, Any]] = None, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """ + Updates the cache with the new `key_states` and `value_states` for the layer `layer_idx`. + It is VERY important to index using a tensor, otherwise you introduce a copy to the device. + + Parameters: + key_states (`torch.Tensor`): + The new key states to cache. + value_states (`torch.Tensor`): + The new value states to cache. + layer_idx (`int`): + The index of the layer to cache the states for. + cache_kwargs (`Dict[str, Any]`, `optional`): + Additional arguments for the cache subclass. The `StaticCache` needs the `cache_position` input + to know how where to write in the cache. + + Return: + A tuple containing the updated key and value states. + """ + k_out = self.k_caches[layer_idx] + + k_out[page_idx, page_offset, :, :self.kv_lora_rank] = key_states.reshape(-1, *key_states.shape[2:]) + k_out[page_idx, page_offset, :, self.kv_lora_rank:] = value_states.reshape(-1, *value_states.shape[2:]) + return k_out + + + def get_page_table(self, cache_position: torch.Tensor, q_indptr: torch.Tensor, kv_indptr: torch.Tensor, kv_indices: torch.Tensor, bsz_tensors: torch.tensor): + page_offset = cache_position % self.page_size + page_idx_local = cache_position // self.page_size + query_ids = torch.zeros_like(cache_position) + for i in range(len(q_indptr) - 1): + start_idx = q_indptr[i] + end_idx = q_indptr[i + 1] + query_ids[start_idx:end_idx] = i + page_idx = torch.zeros_like(page_idx_local) + for i in range(bsz_tensors[0]): + query_id = query_ids[i] + local_block = page_idx_local[i] + start_block = kv_indptr[query_id] + if local_block < kv_indptr[query_id + 1] - kv_indptr[query_id]: + page_idx[i] = kv_indices[start_block + local_block] + + return page_idx, page_offset + +class KGQACache(nn.Module): + def __init__( + self, + config: PretrainedConfig, + page_size: int = 256, + dtype=torch.bfloat16, + device=torch.device("cuda:0"), + + ): + super().__init__() + self.config = config + self.dtype = dtype + self.device = device + self.page_size = page_size + self.k_caches = [] + self.v_caches = [] + + + def load(self, inference_context: "sched_ext.InferenceContext"): + print(self.config.num_hidden_layers) + for i in range(self.config.num_hidden_layers): + self.k_caches.append( + inference_context.k_cache[0][i] + ) + self.v_caches.append( + inference_context.v_cache[0][i] + ) + + + self.max_cache_len = self.k_caches[0].shape[0]*self.k_caches[0].shape[1] + + + + def get_page_table(self, cache_position: torch.Tensor, q_indptr: torch.Tensor, kv_indptr: torch.Tensor, kv_indices: torch.Tensor, bsz_tensors: torch.tensor): + page_offset = cache_position % self.page_size + page_idx_local = cache_position // self.page_size + query_ids = torch.zeros_like(cache_position) + for i in range(len(q_indptr) - 1): + start_idx = q_indptr[i] + end_idx = q_indptr[i + 1] + query_ids[start_idx:end_idx] = i + page_idx = torch.zeros_like(page_idx_local) + for i in range(bsz_tensors[0]): + query_id = query_ids[i] + local_block = page_idx_local[i] + start_block = kv_indptr[query_id] + if local_block < kv_indptr[query_id + 1] - kv_indptr[query_id]: + page_idx[i] = kv_indices[start_block + local_block] + + return page_idx, page_offset + + def get_k_cache(self, layer_idx): + return self.k_caches[layer_idx] + + def get_v_cache(self, layer_idx): + return self.v_caches[layer_idx] \ No newline at end of file diff --git a/archive/kt-sft/ktransformers/models/custom_modeling_deepseek_v2.py b/archive/kt-sft/ktransformers/models/custom_modeling_deepseek_v2.py new file mode 100644 index 0000000..c8f63bf --- /dev/null +++ b/archive/kt-sft/ktransformers/models/custom_modeling_deepseek_v2.py @@ -0,0 +1,153 @@ +import math +from dataclasses import dataclass +import torch +import torch.nn as nn +from torch.nn import functional as F +import math +from typing import List, Optional, Tuple, Union +import torch +import torch.utils.checkpoint +from torch import nn +from ktransformers.server.balance_serve.inference.forward_batch import ForwardBatchInput, ForwardBatchOutput +from ktransformers.models.custom_cache import KDeepSeekV3Cache +from ktransformers.models.modeling_deepseek import DeepseekV2Model, DeepseekV2PreTrainedModel +from ktransformers.models.configuration_deepseek import DeepseekV2Config + + +torch.set_grad_enabled(False) +torch.set_default_dtype(torch.bfloat16) +import flashinfer + +class KDeepseekV2ForCausalLM(DeepseekV2PreTrainedModel): + + kv_cache: KDeepSeekV3Cache + use_cuda_graph = False + def __init__( + self, + config, + kv_cache, + + ): + super().__init__(config) + self.model = DeepseekV2Model(config) + self.config = config + self.kv_cache = kv_cache + + self.vocab_size = config.vocab_size + self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) + + + def init_wrapper(self, use_cuda_graph, device, max_batch_size, max_pages): + self.use_cuda_graph = use_cuda_graph + self.workspace_buffer = torch.empty(128 * 1024 * 1024, dtype=torch.int8).to(0) + self.qo_indptr_buf = torch.empty((max_batch_size+1,), dtype=torch.int32, device=device) + self.paged_kv_indptr_buf = torch.empty((max_batch_size+1,), dtype=torch.int32, device=device) + self.paged_kv_indices_buf = torch.empty((max_pages,), dtype=torch.int32, device=device) + self.paged_kv_len_buf = torch.empty((max_batch_size,), dtype=torch.int32, device=device) + + + + self.wrapper = flashinfer.mla.BatchMLAPagedAttentionWrapper( + self.workspace_buffer, use_cuda_graph=use_cuda_graph, + qo_indptr=self.qo_indptr_buf,kv_indptr=self.paged_kv_indptr_buf, + kv_indices=self.paged_kv_indices_buf,kv_len_arr=self.paged_kv_len_buf, + backend = "fa2", + ) + + def batch_embeddings(self, batch: ForwardBatchInput, device="cuda:0"): + features = [] + for i in range(batch.batch_size): + tokens = batch.minibatch.tokens.contiguous() + feature = ( + self.model.embed_tokens(tokens.to(torch.device('cpu'))) + .to(torch.bfloat16) + .to(device=device) + ) + features.append(feature) + + return features + + + def forward( + self, + batch: ForwardBatchInput | None = None, + features: List[torch.Tensor] | None = None, + bsz_tensors: torch.Tensor | None = None, + num_tokens_tensors: torch.Tensor | None = None, + page_idx: torch.Tensor | None = None, + page_offset: torch.Tensor | None = None, + ) -> ForwardBatchOutput: + current_stream = torch.cuda.current_stream() + + forward_batch_output = ForwardBatchOutput() + + + hidden_states = features[0] + + + with torch.cuda.stream(current_stream): + residual = torch.zeros_like(hidden_states) + for i, decode_layer in enumerate(self.model.layers): + if self.model.transfer_map is not None and i in self.model.transfer_map: + prev_stream = torch.cuda.current_stream() + cur_device = self.model.transfer_map[i] + if cur_device not in self.model.stream_device_map: + self.model.stream_device_map[cur_device] = torch.cuda.Stream(cur_device) + torch.cuda.set_device(cur_device) + self.model.stream_device_map[cur_device].wait_stream(prev_stream) + torch.cuda.set_stream(self.model.stream_device_map[cur_device]) + hidden_states = hidden_states.to( + self.model.transfer_map[i], non_blocking=True + ) + + batch.minibatch.position_ids = ( + batch.minibatch.position_ids.to(self.model.transfer_map[i], non_blocking=True) + if batch.minibatch.position_ids is not None + else None + ) + hidden_states, residual = decode_layer.input_layernorm(hidden_states, num_tokens_tensors, residual) + hidden_states = decode_layer.self_attn(hidden_states, self.kv_cache, + position_ids=batch.minibatch.position_ids, + wrapper=self.wrapper, bsz_tensors=num_tokens_tensors, + cache_position=batch.minibatch.positions, + batch_indices=batch.minibatch.batch_indices, + kv_indices=batch.minibatch.kv_indices, + kv_indptr=batch.minibatch.kv_indptr, + kv_last_page_len=batch.minibatch.kv_last_page_len, + q_indptr=batch.minibatch.q_indptr, + page_idx=page_idx, + page_offset=page_offset + ) + + hidden_states, residual = decode_layer.post_attention_layernorm(hidden_states, num_tokens_tensors, residual) + if i < 3: + hidden_states = decode_layer.mlp(hidden_states, num_tokens_tensors) + else: + hidden_states = decode_layer.mlp(hidden_states.unsqueeze(0), num_tokens_tensors) + hidden_states = hidden_states.squeeze(0) + forward_batch_output = ForwardBatchOutput() + assert batch.batch_size == 1 + with torch.cuda.stream(current_stream): + + local_logit = self.lm_head(self.model.norm(hidden_states[batch.minibatch.logits_start], num_tokens_tensors, residual[batch.minibatch.logits_start])[0]) + # local_logit = local_logit[batch.minibatch.logits_start] + forward_batch_output.logits.append(local_logit) + + return forward_batch_output + + + + def flash_infer_attn_plan(self, batch: ForwardBatchInput, bsz_tensors, num_tokens_tensors, + num_heads: int, + head_dim_ckv: int, + head_dim_kpe: int, + page_size: int, + causal: bool, + sm_scale: float, + q_data_type: torch.dtype, + kv_data_type: torch.dtype,): + minibatch = batch.minibatch + + self.wrapper.plan(minibatch.q_indptr, minibatch.kv_indptr, minibatch.kv_indices, + minibatch.kv_len, num_heads, head_dim_ckv, head_dim_kpe, page_size, causal, sm_scale, q_data_type, kv_data_type) + \ No newline at end of file diff --git a/archive/kt-sft/ktransformers/models/custom_modeling_deepseek_v3.py b/archive/kt-sft/ktransformers/models/custom_modeling_deepseek_v3.py new file mode 100644 index 0000000..589f6c3 --- /dev/null +++ b/archive/kt-sft/ktransformers/models/custom_modeling_deepseek_v3.py @@ -0,0 +1,148 @@ +""" +Date: 2024-11-06 10:05:11 +LastEditors: djw +LastEditTime: 2024-11-13 07:50:51 +""" + +import math +from dataclasses import dataclass +import torch +import torch.nn as nn +from torch.nn import functional as F +import math +from typing import List, Optional, Tuple, Union +import torch +import torch.utils.checkpoint +from torch import nn +from ktransformers.server.balance_serve.inference.forward_batch import ForwardBatchInput, ForwardBatchOutput +from ktransformers.models.custom_cache import KDeepSeekV3Cache +from ktransformers.models.modeling_deepseek_v3 import DeepseekV3Model, DeepseekV3PreTrainedModel +from ktransformers.models.configuration_deepseek_v3 import DeepseekV3Config + + +torch.set_grad_enabled(False) +torch.set_default_dtype(torch.bfloat16) +import flashinfer + +class KDeepseekV3ForCausalLM(DeepseekV3PreTrainedModel): + + cache: KDeepSeekV3Cache + use_cuda_graph = False + def __init__( + self, + config: DeepseekV3Config, + cache, + ): + super().__init__(config) + self.model = DeepseekV3Model(config) + self.config = config + self.cache = cache + self.vocab_size = config.vocab_size + self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) + + def init_wrapper(self, use_cuda_graph, device, max_batch_size, max_pages): + self.use_cuda_graph = use_cuda_graph + self.workspace_buffer = torch.empty(128 * 1024 * 1024, dtype=torch.int8).to(0) + self.qo_indptr_buf = torch.empty((max_batch_size+2,), dtype=torch.int32, device=device) + self.paged_kv_indptr_buf = torch.empty((max_batch_size+2,), dtype=torch.int32, device=device) + self.paged_kv_indices_buf = torch.empty((max_pages,), dtype=torch.int32, device=device) + self.paged_kv_len_buf = torch.empty((max_batch_size+1,), dtype=torch.int32, device=device) + self.bsz_tensor_buf = torch.empty((1, ), dtype=torch.int32, device=device) + + + self.wrapper = flashinfer.mla.BatchMLAPagedAttentionWrapper( + self.workspace_buffer, use_cuda_graph=use_cuda_graph, + qo_indptr=self.qo_indptr_buf,kv_indptr=self.paged_kv_indptr_buf, + kv_indices=self.paged_kv_indices_buf,kv_len_arr=self.paged_kv_len_buf, + bsz_tensor=self.bsz_tensor_buf, + backend = "fa2", + ) + + def batch_embeddings(self, batch: ForwardBatchInput, device="cuda:0"): + features = [] + for i in range(batch.batch_size): + tokens = batch.minibatch.tokens.contiguous() + feature = ( + self.model.embed_tokens(tokens.to(torch.device('cpu'))) + .to(torch.bfloat16) + .to(device=device) + ) + features.append(feature) + + return features + + + def forward( + self, + batch: ForwardBatchInput | None = None, + features: List[torch.Tensor] | None = None, + bsz_tensors: torch.Tensor | None = None, + num_tokens_tensors: torch.Tensor | None = None, + page_idx: torch.Tensor | None = None, + page_offset: torch.Tensor | None = None, + cuda_graph_idx: int | None = -1 + ) -> ForwardBatchOutput: + current_stream = torch.cuda.current_stream() + + forward_batch_output = ForwardBatchOutput() + + + hidden_states = features[0] + + with torch.cuda.stream(current_stream): + residual = torch.zeros_like(hidden_states) + for i, decode_layer in enumerate(self.model.layers): + # can't use now, only one flashinfer wrapper + if self.model.transfer_map is not None and i in self.model.transfer_map: + prev_stream = torch.cuda.current_stream() + cur_device = self.model.transfer_map[i] + if cur_device not in self.model.stream_device_map: + self.model.stream_device_map[cur_device] = torch.cuda.Stream(cur_device) + torch.cuda.set_device(cur_device) + self.model.stream_device_map[cur_device].wait_stream(prev_stream) + torch.cuda.set_stream(self.model.stream_device_map[cur_device]) + hidden_states = hidden_states.to( + self.model.transfer_map[i], non_blocking=True + ) + + batch.minibatch.position_ids = ( + batch.minibatch.position_ids.to(self.model.transfer_map[i], non_blocking=True) + if batch.minibatch.position_ids is not None + else None + ) + hidden_states, residual = decode_layer.input_layernorm(hidden_states, num_tokens_tensors, residual) + hidden_states = decode_layer.self_attn(hidden_states, self.cache, + position_ids=batch.minibatch.position_ids, + wrapper=self.wrapper, num_tokens_tensors=num_tokens_tensors, + page_idx=page_idx, + page_offset=page_offset + ) + + hidden_states, residual = decode_layer.post_attention_layernorm(hidden_states, num_tokens_tensors, residual) + if i < self.config.first_k_dense_replace: + hidden_states = decode_layer.mlp(hidden_states, num_tokens_tensors) + else: + hidden_states = decode_layer.mlp(hidden_states.unsqueeze(0), num_tokens_tensors, cuda_graph_idx) + hidden_states = hidden_states.squeeze(0) + forward_batch_output = ForwardBatchOutput() + with torch.cuda.stream(current_stream): + local_logit = self.lm_head(self.model.norm(hidden_states, num_tokens_tensors, residual)[0], num_tokens_tensors) + forward_batch_output.logits.append(local_logit) + + return forward_batch_output + + + + def flash_infer_attn_plan(self, batch: ForwardBatchInput, bsz_tensors, num_tokens_tensors, + num_heads: int, + head_dim_ckv: int, + head_dim_kpe: int, + page_size: int, + causal: bool, + sm_scale: float, + q_data_type: torch.dtype, + kv_data_type: torch.dtype,): + minibatch = batch.minibatch + self.wrapper.plan(minibatch.q_indptr, minibatch.kv_indptr, minibatch.kv_indices, + minibatch.kv_len, num_heads, head_dim_ckv, head_dim_kpe, page_size, causal, sm_scale, q_data_type, kv_data_type, bsz_tensors) + \ No newline at end of file diff --git a/archive/kt-sft/ktransformers/models/custom_modeling_qwen2_moe.py b/archive/kt-sft/ktransformers/models/custom_modeling_qwen2_moe.py new file mode 100644 index 0000000..1c84cbf --- /dev/null +++ b/archive/kt-sft/ktransformers/models/custom_modeling_qwen2_moe.py @@ -0,0 +1,133 @@ +""" +Date: 2024-11-06 10:05:11 +LastEditors: djw +LastEditTime: 2024-11-13 07:50:51 +""" + +import math +from dataclasses import dataclass +import torch +import torch.nn as nn +from torch.nn import functional as F +import math +from typing import List, Optional, Tuple, Union +import torch +import torch.utils.checkpoint +from torch import nn +from ktransformers.server.balance_serve.inference.forward_batch import ForwardBatchInput, ForwardBatchOutput +from ktransformers.models.custom_cache import KGQACache +from ktransformers.models.modeling_qwen2_moe import Qwen2MoeModel, Qwen2MoePreTrainedModel +from ktransformers.models.configuration_qwen2_moe import Qwen2MoeConfig +from ktransformers.operators.flashinfer_batch_prefill_wrapper import flashInferAttn + +torch.set_grad_enabled(False) +torch.set_default_dtype(torch.bfloat16) +import flashinfer + +class KQwen2MoeForCausalLM(Qwen2MoePreTrainedModel): + + cache: KGQACache + use_cuda_graph = False + def __init__( + self, + config: Qwen2MoeConfig, + cache, + ): + super().__init__(config) + self.model = Qwen2MoeModel(config) + self.config = config + self.cache = cache + self.vocab_size = config.vocab_size + self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) + self.attn = [None] * 100 + + def init_wrapper(self, use_cuda_graph, device, max_batch_token, max_batch_size, max_pages, cuda_graph_idx = 0): + self.attn[cuda_graph_idx] = flashInferAttn(use_cuda_graph=use_cuda_graph, max_batch_token=max_batch_token, max_batch_size=max_batch_size, max_pages=max_pages, device=device) + + + def batch_embeddings(self, batch: ForwardBatchInput, device="cuda:0"): + features = [] + for i in range(batch.batch_size): + tokens = batch.minibatch.tokens.contiguous() + feature = ( + self.model.embed_tokens(tokens.to(torch.device('cpu'))) + .to(torch.bfloat16) + .to(device=device) + ) + features.append(feature) + + return features + + + def forward( + self, + batch: ForwardBatchInput | None = None, + features: List[torch.Tensor] | None = None, + bsz_tensors: torch.Tensor | None = None, + num_tokens_tensors: torch.Tensor | None = None, + page_idx: torch.Tensor | None = None, + page_offset: torch.Tensor | None = None, + cuda_graph_idx: int | None = 0 + ) -> ForwardBatchOutput: + current_stream = torch.cuda.current_stream() + + forward_batch_output = ForwardBatchOutput() + + + hidden_states = features[0] + self.attn[cuda_graph_idx].calc_batch_indices(hidden_states.shape[0]) + + with torch.cuda.stream(current_stream): + residual = torch.zeros_like(hidden_states) + for i, decode_layer in enumerate(self.model.layers): + if self.model.transfer_map is not None and i in self.model.transfer_map: + prev_stream = torch.cuda.current_stream() + cur_device = self.model.transfer_map[i] + if cur_device not in self.model.stream_device_map: + self.model.stream_device_map[cur_device] = torch.cuda.Stream(cur_device) + torch.cuda.set_device(cur_device) + self.model.stream_device_map[cur_device].wait_stream(prev_stream) + torch.cuda.set_stream(self.model.stream_device_map[cur_device]) + hidden_states = hidden_states.to( + self.model.transfer_map[i], non_blocking=True + ) + + batch.minibatch.position_ids = ( + batch.minibatch.position_ids.to(self.model.transfer_map[i], non_blocking=True) + if batch.minibatch.position_ids is not None + else None + ) + hidden_states, residual = decode_layer.input_layernorm(hidden_states, num_tokens_tensors, residual) + hidden_states = decode_layer.self_attn(hidden_states, self.cache, + position_ids=batch.minibatch.position_ids, + wrapper=self.attn[cuda_graph_idx], bsz_tensors=num_tokens_tensors, + page_idx=page_idx, + page_offset=page_offset + ) + + hidden_states, residual = decode_layer.post_attention_layernorm(hidden_states, num_tokens_tensors, residual) + hidden_states = decode_layer.mlp(hidden_states.unsqueeze(0), num_tokens_tensors, cuda_graph_idx) + hidden_states = hidden_states.squeeze(0) + forward_batch_output = ForwardBatchOutput() + with torch.cuda.stream(current_stream): + local_logit = self.lm_head(self.model.norm(hidden_states, num_tokens_tensors, residual)[0], num_tokens_tensors) + forward_batch_output.logits.append(local_logit) + + return forward_batch_output + + + + def flash_infer_attn_plan(self, batch: ForwardBatchInput, bsz_tensors, num_tokens_tensors, + num_q_heads: int, + num_kv_heads: int, + head_dim: int, + page_size: int, + causal: bool, + q_data_type: torch.dtype, + kv_data_type: torch.dtype, + cuda_graph_idx: int = 0 + ): + minibatch = batch.minibatch + self.attn[cuda_graph_idx].plan(minibatch.q_indptr, minibatch.kv_indptr, minibatch.kv_indices, + minibatch.kv_last_page_len, bsz_tensors, num_tokens_tensors,num_q_heads, num_kv_heads, head_dim, page_size, causal=causal, q_data_type=q_data_type, kv_data_type=kv_data_type) + \ No newline at end of file diff --git a/archive/kt-sft/ktransformers/models/custom_modeling_qwen3_moe.py b/archive/kt-sft/ktransformers/models/custom_modeling_qwen3_moe.py new file mode 100644 index 0000000..32b9797 --- /dev/null +++ b/archive/kt-sft/ktransformers/models/custom_modeling_qwen3_moe.py @@ -0,0 +1,133 @@ +""" +Date: 2024-11-06 10:05:11 +LastEditors: djw +LastEditTime: 2024-11-13 07:50:51 +""" + +import math +from dataclasses import dataclass +import torch +import torch.nn as nn +from torch.nn import functional as F +import math +from typing import List, Optional, Tuple, Union +import torch +import torch.utils.checkpoint +from torch import nn +from ktransformers.server.balance_serve.inference.forward_batch import ForwardBatchInput, ForwardBatchOutput +from ktransformers.models.custom_cache import KGQACache +from ktransformers.models.modeling_qwen3_moe import Qwen3MoeModel, Qwen3MoePreTrainedModel +from ktransformers.models.configuration_qwen3_moe import Qwen3MoeConfig +from ktransformers.operators.flashinfer_batch_prefill_wrapper import flashInferAttn + +torch.set_grad_enabled(False) +torch.set_default_dtype(torch.bfloat16) +import flashinfer + +class KQwen3MoeForCausalLM(Qwen3MoePreTrainedModel): + + cache: KGQACache + use_cuda_graph = False + def __init__( + self, + config: Qwen3MoeConfig, + cache = None, + ): + super().__init__(config) + self.model = Qwen3MoeModel(config) + self.config = config + self.cache = cache + self.vocab_size = config.vocab_size + self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) + self.attn = [None] * 100 + + def init_wrapper(self, use_cuda_graph, device, max_batch_token, max_batch_size, max_pages, cuda_graph_idx = 0): + self.attn[cuda_graph_idx] = flashInferAttn(use_cuda_graph=use_cuda_graph, max_batch_token=max_batch_token, max_batch_size=max_batch_size, max_pages=max_pages, device=device) + + + def batch_embeddings(self, batch: ForwardBatchInput, device="cuda:0"): + features = [] + for i in range(batch.batch_size): + tokens = batch.minibatch.tokens.contiguous() + feature = ( + self.model.embed_tokens(tokens.to(torch.device('cpu'))) + .to(torch.bfloat16) + .to(device=device) + ) + features.append(feature) + + return features + + + def forward( + self, + batch: ForwardBatchInput | None = None, + features: List[torch.Tensor] | None = None, + bsz_tensors: torch.Tensor | None = None, + num_tokens_tensors: torch.Tensor | None = None, + page_idx: torch.Tensor | None = None, + page_offset: torch.Tensor | None = None, + cuda_graph_idx: int | None = 0 + ) -> ForwardBatchOutput: + current_stream = torch.cuda.current_stream() + + forward_batch_output = ForwardBatchOutput() + + + hidden_states = features[0] + self.attn[cuda_graph_idx].calc_batch_indices(hidden_states.shape[0]) + + with torch.cuda.stream(current_stream): + residual = torch.zeros_like(hidden_states) + for i, decode_layer in enumerate(self.model.layers): + if self.model.transfer_map is not None and i in self.model.transfer_map: + prev_stream = torch.cuda.current_stream() + cur_device = self.model.transfer_map[i] + if cur_device not in self.model.stream_device_map: + self.model.stream_device_map[cur_device] = torch.cuda.Stream(cur_device) + torch.cuda.set_device(cur_device) + self.model.stream_device_map[cur_device].wait_stream(prev_stream) + torch.cuda.set_stream(self.model.stream_device_map[cur_device]) + hidden_states = hidden_states.to( + self.model.transfer_map[i], non_blocking=True + ) + + batch.minibatch.position_ids = ( + batch.minibatch.position_ids.to(self.model.transfer_map[i], non_blocking=True) + if batch.minibatch.position_ids is not None + else None + ) + hidden_states, residual = decode_layer.input_layernorm(hidden_states, num_tokens_tensors, residual) + hidden_states = decode_layer.self_attn(hidden_states, self.cache, + position_ids=batch.minibatch.position_ids, + wrapper=self.attn[cuda_graph_idx], bsz_tensors=num_tokens_tensors, + page_idx=page_idx, + page_offset=page_offset + ) + + hidden_states, residual = decode_layer.post_attention_layernorm(hidden_states, num_tokens_tensors, residual) + hidden_states = decode_layer.mlp(hidden_states.unsqueeze(0), num_tokens_tensors, cuda_graph_idx) + hidden_states = hidden_states.squeeze(0) + forward_batch_output = ForwardBatchOutput() + with torch.cuda.stream(current_stream): + local_logit = self.lm_head(self.model.norm(hidden_states, num_tokens_tensors, residual)[0], num_tokens_tensors) + forward_batch_output.logits.append(local_logit) + + return forward_batch_output + + + + def flash_infer_attn_plan(self, batch: ForwardBatchInput, bsz_tensors, num_tokens_tensors, + num_q_heads: int, + num_kv_heads: int, + head_dim: int, + page_size: int, + causal: bool, + q_data_type: torch.dtype, + kv_data_type: torch.dtype, + cuda_graph_idx: int = 0 + ): + minibatch = batch.minibatch + self.attn[cuda_graph_idx].plan(minibatch.q_indptr, minibatch.kv_indptr, minibatch.kv_indices, + minibatch.kv_last_page_len, bsz_tensors, num_tokens_tensors, num_q_heads, num_kv_heads, head_dim, page_size, causal=causal, q_data_type=q_data_type, kv_data_type=kv_data_type) + \ No newline at end of file diff --git a/archive/kt-sft/ktransformers/models/modeling_deepseek.py b/archive/kt-sft/ktransformers/models/modeling_deepseek.py new file mode 100644 index 0000000..0d630d9 --- /dev/null +++ b/archive/kt-sft/ktransformers/models/modeling_deepseek.py @@ -0,0 +1,1996 @@ +# coding=utf-8 +''' +Description : +Author : Boxin Zhang +Version : 0.1.0 +''' +# Adapted from +# https://huggingface.co/deepseek-ai/DeepSeek-V2-Chat-0628/blob/main/modeling_deepseek.py +# Copyright 2023 DeepSeek-AI and The HuggingFace Inc. team. All rights reserved. +# Copyright (c) 2024 by KVCache.AI, All Rights Reserved. +# +# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX +# and OPT implementations in this library. It has been modified from its +# original forms to accommodate minor architectural differences compared +# to GPT-NeoX and OPT used by the Meta AI team that trained the model. +# +# 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. +""" PyTorch DeepSeek model.""" +import math +import warnings +from typing import List, Optional, Tuple, Union + +import torch +import torch.nn.functional as F +import torch.utils.checkpoint +from torch import nn +from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss + +from transformers.activations import ACT2FN +from transformers.cache_utils import Cache, DynamicCache, StaticCache +from transformers.modeling_attn_mask_utils import ( + AttentionMaskConverter, + _prepare_4d_attention_mask, + _prepare_4d_causal_attention_mask, +) +from transformers.modeling_outputs import ( + BaseModelOutputWithPast, + CausalLMOutputWithPast, + SequenceClassifierOutputWithPast, +) +from transformers.modeling_utils import PreTrainedModel +from transformers.pytorch_utils import ( + ALL_LAYERNORM_LAYERS, + is_torch_greater_or_equal_than_1_13, +) +from transformers.utils import ( + add_start_docstrings, + add_start_docstrings_to_model_forward, + is_flash_attn_2_available, + is_flash_attn_greater_or_equal_2_10, + logging, + replace_return_docstrings, +) +from transformers.utils.import_utils import is_torch_fx_available +from .configuration_deepseek import DeepseekV2Config +import torch.distributed as dist +import numpy as np + +if is_flash_attn_2_available(): + from flash_attn import flash_attn_func, flash_attn_varlen_func, flash_attn_with_kvcache + from flash_attn.bert_padding import index_first_axis, pad_input, unpad_input # noqa + +from ktransformers.util.grad_wrapper import maybe_no_grad + +# This makes `_prepare_4d_causal_attention_mask` a leaf function in the FX graph. +# It means that the function will not be traced through and simply appear as a node in the graph. +if is_torch_fx_available(): + if not is_torch_greater_or_equal_than_1_13: + import torch.fx + + _prepare_4d_causal_attention_mask = torch.fx.wrap(_prepare_4d_causal_attention_mask) + + +logger = logging.get_logger(__name__) + +_CONFIG_FOR_DOC = "DeepseekV2Config" + + +def _get_unpad_data(attention_mask): + seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32) + indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten() + max_seqlen_in_batch = seqlens_in_batch.max().item() + cu_seqlens = F.pad( + torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.torch.int32), (1, 0) + ) + return ( + indices, + cu_seqlens, + max_seqlen_in_batch, + ) + + +class DeepseekV2RMSNorm(nn.Module): + def __init__(self, hidden_size, eps=1e-6): + """ + DeepseekV2RMSNorm is equivalent to T5LayerNorm + """ + super().__init__() + self.weight = nn.Parameter(torch.ones(hidden_size)) + self.variance_epsilon = eps + self.hidden_size = hidden_size + + def forward(self, hidden_states): + input_dtype = hidden_states.dtype + hidden_states = hidden_states.to(torch.float32) + variance = hidden_states.pow(2).mean(-1, keepdim=True) + hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon) + return (self.weight * hidden_states).to(input_dtype) + + +ALL_LAYERNORM_LAYERS.append(DeepseekV2RMSNorm) + +# Copied from transformers.models.mistral.modeling_mistral.MistralRotaryEmbedding with Mistral->DeepseekV2 +class DeepseekV2RotaryEmbedding(nn.Module): + def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None, scaling_factor=1.0): + super().__init__() + self.scaling_factor = scaling_factor + self.dim = dim + self.max_position_embeddings = max_position_embeddings + self.base = base + inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2, dtype=torch.int64).float().to(device) / self.dim)) + self.register_buffer("inv_freq", inv_freq, persistent=False) + # For BC we register cos and sin cached + self.max_seq_len_cached = max_position_embeddings + + @maybe_no_grad() + def forward(self, x, position_ids): + # x: [bs, num_attention_heads, seq_len, head_size] + inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1) + position_ids_expanded = position_ids[:, None, :].float() + # Force float32 since bfloat16 loses precision on long contexts + # See https://github.com/huggingface/transformers/pull/29285 + device_type = x.device.type + device_type = device_type if isinstance(device_type, str) and device_type != "mps" else "cpu" + with torch.autocast(device_type=device_type, enabled=False): + freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2) + emb = torch.cat((freqs, freqs), dim=-1) + cos = emb.cos() + sin = emb.sin() + return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype) + +# Copied from transformers.models.llama.modeling_llama.LlamaLinearScalingRotaryEmbedding with Llama->DeepseekV2 +class DeepseekV2LinearScalingRotaryEmbedding(DeepseekV2RotaryEmbedding): + """DeepseekV2RotaryEmbedding extended with linear scaling. Credits to the Reddit user /u/kaiokendev""" + + def __init__( + self, + dim, + max_position_embeddings=2048, + base=10000, + device=None, + scaling_factor=1.0, + ): + raise NotImplementedError("LinearScalingRotaryEmbedding is not supported now.") + self.scaling_factor = scaling_factor + super().__init__(dim, max_position_embeddings, base, device) + + def _set_cos_sin_cache(self, seq_len, device, dtype): + self.max_seq_len_cached = seq_len + t = torch.arange( + self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype + ) + t = t / self.scaling_factor + + freqs = torch.outer(t, self.inv_freq) + # Different from paper, but it uses a different permutation in order to obtain the same calculation + emb = torch.cat((freqs, freqs), dim=-1) + self.register_buffer("cos_cached", emb.cos().to(dtype), persistent=False) + self.register_buffer("sin_cached", emb.sin().to(dtype), persistent=False) + + +# Copied from transformers.models.llama.modeling_llama.LlamaDynamicNTKScalingRotaryEmbedding with Llama->DeepseekV2 +class DeepseekV2DynamicNTKScalingRotaryEmbedding(DeepseekV2RotaryEmbedding): + """DeepseekV2RotaryEmbedding extended with Dynamic NTK scaling. Credits to the Reddit users /u/bloc97 and /u/emozilla""" + + def __init__( + self, + dim, + max_position_embeddings=2048, + base=10000, + device=None, + scaling_factor=1.0, + ): + raise NotImplementedError("DynamicNTKScalingRotaryEmbedding is not supported now.") + self.scaling_factor = scaling_factor + super().__init__(dim, max_position_embeddings, base, device) + + def _set_cos_sin_cache(self, seq_len, device, dtype): + self.max_seq_len_cached = seq_len + + if seq_len > self.max_position_embeddings: + base = self.base * ( + (self.scaling_factor * seq_len / self.max_position_embeddings) + - (self.scaling_factor - 1) + ) ** (self.dim / (self.dim - 2)) + inv_freq = 1.0 / ( + base ** (torch.arange(0, self.dim, 2).float().to(device) / self.dim) + ) + self.register_buffer("inv_freq", inv_freq, persistent=False) + + t = torch.arange( + self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype + ) + + freqs = torch.outer(t, self.inv_freq) + # Different from paper, but it uses a different permutation in order to obtain the same calculation + emb = torch.cat((freqs, freqs), dim=-1) + self.register_buffer("cos_cached", emb.cos().to(dtype), persistent=False) + self.register_buffer("sin_cached", emb.sin().to(dtype), persistent=False) + + +# Inverse dim formula to find dim based on number of rotations +def yarn_find_correction_dim( + num_rotations, dim, base=10000, max_position_embeddings=2048 +): + 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, high_rot, dim, base=10000, max_position_embeddings=2048 +): + 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_get_mscale(scale=1, mscale=1): + if scale <= 1: + return 1.0 + return 0.1 * mscale * math.log(scale) + 1.0 + + +def yarn_linear_ramp_mask(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 + +class DeepseekV2YarnRotaryEmbedding(DeepseekV2RotaryEmbedding): + def __init__( + self, + dim, + max_position_embeddings=2048, + base=10000, + device=None, + scaling_factor=1.0, + original_max_position_embeddings=4096, + beta_fast=32, + beta_slow=1, + mscale=1, + mscale_all_dim=0, + ): + nn.Module.__init__(self) + self.scaling_factor = scaling_factor + self.original_max_position_embeddings = original_max_position_embeddings + self.beta_fast = beta_fast + self.beta_slow = beta_slow + self.mscale = mscale + self.mscale_all_dim = mscale_all_dim + self.scaling_factor = scaling_factor + self.dim = dim + self.max_position_embeddings = max_position_embeddings + self.base = base + + freq_extra = 1.0 / ( + self.base + ** (torch.arange(0, dim, 2, dtype=torch.float32, device=device) / dim) + ) + freq_inter = 1.0 / ( + self.scaling_factor + * self.base + ** (torch.arange(0, dim, 2, dtype=torch.float32, device=device) / dim) + ) + + low, high = yarn_find_correction_range( + self.beta_fast, + self.beta_slow, + dim, + self.base, + self.original_max_position_embeddings, + ) + inv_freq_mask = 1.0 - yarn_linear_ramp_mask(low, high, dim // 2).to( + device=device, dtype=torch.float32 + ) + inv_freq = freq_inter * (1 - inv_freq_mask) + freq_extra * inv_freq_mask + self.register_buffer("inv_freq", inv_freq, persistent=False) + self._mscale = float( + yarn_get_mscale(self.scaling_factor, self.mscale) + / yarn_get_mscale(self.scaling_factor, self.mscale_all_dim) + ) + # For BC we register cos and sin cached + self.max_seq_len_cached = max_position_embeddings + + @maybe_no_grad() + def forward(self, x, position_ids): + # x: [bs, num_attention_heads, seq_len, head_size] + inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1) + position_ids_expanded = position_ids[:, None, :].float() + # Force float32 since bfloat16 loses precision on long contexts + # See https://github.com/huggingface/transformers/pull/29285 + device_type = x.device.type + device_type = device_type if isinstance(device_type, str) and device_type != "mps" else "cpu" + with torch.autocast(device_type=device_type, enabled=False): + freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2) + emb = torch.cat((freqs, freqs), dim=-1) + cos = emb.cos()* self._mscale + sin = emb.sin()* self._mscale + return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype) + +# Copied from transformers.models.llama.modeling_llama.rotate_half +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) + + +# Copied from transformers.models.mistral.modeling_mistral.apply_rotary_pos_emb +def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1): + """Applies Rotary Position Embedding to the query and key tensors. + + Args: + q (`torch.Tensor`): The query tensor. + k (`torch.Tensor`): The key tensor. + cos (`torch.Tensor`): The cosine part of the rotary embedding. + sin (`torch.Tensor`): The sine part of the rotary embedding. + position_ids (`torch.Tensor`): + Deprecated and unused. + unsqueeze_dim (`int`, *optional*, defaults to 1): + The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and + sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note + that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and + k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes + cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have + the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2. + Returns: + `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding. + """ + cos = cos.unsqueeze(unsqueeze_dim) + sin = sin.unsqueeze(unsqueeze_dim) + b, h, s, d = q.shape + q = q.view(b, h, s, d // 2, 2).transpose(4, 3).reshape(b, h, s, d) + b, h, s, d = k.shape + k = k.view(b, h, s, d // 2, 2).transpose(4, 3).reshape(b, h, s, d) + q_embed = (q * cos) + (rotate_half(q) * sin) + k_embed = (k * cos) + (rotate_half(k) * sin) + return q_embed, k_embed + +class DeepseekV2MLP(nn.Module): + def __init__(self, config, hidden_size=None, intermediate_size=None): + super().__init__() + self.config = config + self.hidden_size = config.hidden_size if hidden_size is None else hidden_size + self.intermediate_size = ( + config.intermediate_size if intermediate_size is None else intermediate_size + ) + + self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) + self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) + self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False) + self.act_fn = ACT2FN[config.hidden_act] + + def forward(self, x): + act = self.act_fn(self.gate_proj(x)) * self.up_proj(x) + down_proj = self.down_proj(act) + return down_proj + +class MoEGate(nn.Module): + def __init__(self, config): + super().__init__() + self.config = config + self.top_k = config.num_experts_per_tok + self.n_routed_experts = config.n_routed_experts + self.routed_scaling_factor = config.routed_scaling_factor + self.scoring_func = config.scoring_func + self.alpha = config.aux_loss_alpha + self.seq_aux = config.seq_aux + self.topk_method = config.topk_method + self.n_group = config.n_group + self.topk_group = config.topk_group + + # topk selection algorithm + self.norm_topk_prob = config.norm_topk_prob + self.gating_dim = config.hidden_size + self.weight = nn.Parameter( + torch.empty((self.n_routed_experts, self.gating_dim)) + ) + self.reset_parameters() + + def reset_parameters(self) -> None: + import torch.nn.init as init + + init.kaiming_uniform_(self.weight, a=math.sqrt(5)) + + def forward(self, hidden_states): + bsz, seq_len, h = hidden_states.shape + ### compute gating score + hidden_states = hidden_states.view(-1, h) + logits = F.linear( + hidden_states.type(torch.float32), self.weight.type(torch.float32), None + ) + if self.scoring_func == "softmax": + scores = logits.softmax(dim=-1, dtype=torch.float32) + else: + raise NotImplementedError( + f"insupportable scoring function for MoE gating: {self.scoring_func}" + ) + + ### select top-k experts + if self.topk_method == "greedy": + topk_weight, topk_idx = torch.topk( + scores, k=self.top_k, dim=-1, sorted=False + ) + elif self.topk_method == "group_limited_greedy": + group_scores = ( + scores.view(bsz * seq_len, self.n_group, -1).max(dim=-1).values + ) # [n, n_group] + group_idx = torch.topk( + group_scores, k=self.topk_group, dim=-1, sorted=False + )[ + 1 + ] # [n, top_k_group] + group_mask = torch.zeros_like(group_scores) # [n, n_group] + group_mask.scatter_(1, group_idx, 1) # [n, n_group] + score_mask = ( + group_mask.unsqueeze(-1) + .expand( + bsz * seq_len, self.n_group, self.n_routed_experts // self.n_group + ) + .reshape(bsz * seq_len, -1) + ) # [n, e] + tmp_scores = scores.masked_fill(~score_mask.bool(), 0.0) # [n, e] + topk_weight, topk_idx = torch.topk( + tmp_scores, k=self.top_k, dim=-1, sorted=False + ) + + ### norm gate to sum 1 + if self.top_k > 1 and self.norm_topk_prob: + denominator = topk_weight.sum(dim=-1, keepdim=True) + 1e-20 + topk_weight = topk_weight / denominator + else: + topk_weight = topk_weight * self.routed_scaling_factor + ### expert-level computation auxiliary loss + if self.training and self.alpha > 0.0: + scores_for_aux = scores + aux_topk = self.top_k + # always compute aux loss based on the naive greedy topk method + topk_idx_for_aux_loss = topk_idx.view(bsz, -1) + if self.seq_aux: + scores_for_seq_aux = scores_for_aux.view(bsz, seq_len, -1) + ce = torch.zeros( + bsz, self.n_routed_experts, device=hidden_states.device + ) + ce.scatter_add_( + 1, + topk_idx_for_aux_loss, + torch.ones(bsz, seq_len * aux_topk, device=hidden_states.device), + ).div_(seq_len * aux_topk / self.n_routed_experts) + aux_loss = (ce * scores_for_seq_aux.mean(dim=1)).sum( + dim=1 + ).mean() * self.alpha + else: + mask_ce = F.one_hot( + topk_idx_for_aux_loss.view(-1), num_classes=self.n_routed_experts + ) + ce = mask_ce.float().mean(0) + Pi = scores_for_aux.mean(0) + fi = ce * self.n_routed_experts + aux_loss = (Pi * fi).sum() * self.alpha + else: + aux_loss = None + return topk_idx, topk_weight, aux_loss + + +class AddAuxiliaryLoss(torch.autograd.Function): + """ + The trick function of adding auxiliary (aux) loss, + which includes the gradient of the aux loss during backpropagation. + """ + + @staticmethod + def forward(ctx, x, loss): + assert loss.numel() == 1 + ctx.dtype = loss.dtype + ctx.required_aux_loss = loss.requires_grad + return x + + @staticmethod + def backward(ctx, grad_output): + grad_loss = None + if ctx.required_aux_loss: + grad_loss = torch.ones(1, dtype=ctx.dtype, device=grad_output.device) + return grad_output, grad_loss + +class DeepseekV2MoE(nn.Module): + """ + A mixed expert module containing shared experts. + """ + + def __init__(self, config): + super().__init__() + self.config = config + self.num_experts_per_tok = config.num_experts_per_tok + + if hasattr(config, "ep_size") and config.ep_size > 1: + assert config.ep_size == dist.get_world_size() + self.ep_size = config.ep_size + self.experts_per_rank = config.n_routed_experts // config.ep_size + self.ep_rank = dist.get_rank() + self.experts = nn.ModuleList( + [ + ( + DeepseekV2MLP( + config, intermediate_size=config.moe_intermediate_size + ) + if i >= self.ep_rank * self.experts_per_rank + and i < (self.ep_rank + 1) * self.experts_per_rank + else None + ) + for i in range(config.n_routed_experts) + ] + ) + else: + self.ep_size = 1 + self.experts_per_rank = config.n_routed_experts + self.ep_rank = 0 + self.experts = nn.ModuleList( + [ + DeepseekV2MLP(config, intermediate_size=config.moe_intermediate_size) + for i in range(config.n_routed_experts) + ] + ) + self.gate = MoEGate(config) + if config.n_shared_experts is not None: + intermediate_size = config.moe_intermediate_size * config.n_shared_experts + self.shared_experts = DeepseekV2MLP( + config=config, intermediate_size=intermediate_size + ) + + def forward(self, hidden_states): + identity = hidden_states + orig_shape = hidden_states.shape + topk_idx, topk_weight, aux_loss = self.gate(hidden_states) + hidden_states = hidden_states.view(-1, hidden_states.shape[-1]) + flat_topk_idx = topk_idx.view(-1) + if self.training: + hidden_states = hidden_states.repeat_interleave( + self.num_experts_per_tok, dim=0 + ) + y = torch.empty_like(hidden_states) + for i, expert in enumerate(self.experts): + y[flat_topk_idx == i] = expert(hidden_states[flat_topk_idx == i]) + y = (y.view(*topk_weight.shape, -1) * topk_weight.unsqueeze(-1)).sum(dim=1) + y = y.view(*orig_shape) + y = AddAuxiliaryLoss.apply(y, aux_loss) + else: + y = self.moe_infer(hidden_states, topk_idx, topk_weight).view(*orig_shape) + if self.config.n_shared_experts is not None: + y = y + self.shared_experts(identity) + return y + + @maybe_no_grad() + def moe_infer(self, x, topk_ids, topk_weight): + cnts = topk_ids.new_zeros((topk_ids.shape[0], len(self.experts))) + cnts.scatter_(1, topk_ids, 1) + tokens_per_expert = cnts.sum(dim=0) + idxs = topk_ids.view(-1).argsort() + sorted_tokens = x[idxs // topk_ids.shape[1]] + sorted_tokens_shape = sorted_tokens.shape + if self.ep_size > 1: + tokens_per_ep_rank = tokens_per_expert.view(self.ep_size, -1).sum(dim=1) + tokens_per_expert_group = tokens_per_expert.new_empty( + tokens_per_expert.shape[0] + ) + dist.all_to_all_single(tokens_per_expert_group, tokens_per_expert) + output_splits = ( + tokens_per_expert_group.view(self.ep_size, -1) + .sum(1) + .cpu() + .numpy() + .tolist() + ) + gathered_tokens = sorted_tokens.new_empty( + tokens_per_expert_group.sum(dim=0).cpu().item(), sorted_tokens.shape[1] + ) + input_split_sizes = tokens_per_ep_rank.cpu().numpy().tolist() + dist.all_to_all( + list(gathered_tokens.split(output_splits)), + list(sorted_tokens.split(input_split_sizes)), + ) + tokens_per_expert_post_gather = tokens_per_expert_group.view( + self.ep_size, self.experts_per_rank + ).sum(dim=0) + gatherd_idxs = np.zeros(shape=(gathered_tokens.shape[0],), dtype=np.int32) + s = 0 + for i, k in enumerate(tokens_per_expert_group.cpu().numpy()): + gatherd_idxs[s : s + k] = i % self.experts_per_rank + s += k + gatherd_idxs = gatherd_idxs.argsort() + sorted_tokens = gathered_tokens[gatherd_idxs] + tokens_per_expert = tokens_per_expert_post_gather + tokens_per_expert = tokens_per_expert.cpu().numpy() + + outputs = [] + start_idx = 0 + for i, num_tokens in enumerate(tokens_per_expert): + end_idx = start_idx + num_tokens + if num_tokens == 0: + continue + expert = self.experts[i + self.ep_rank * self.experts_per_rank] + tokens_for_this_expert = sorted_tokens[start_idx:end_idx] + expert_out = expert(tokens_for_this_expert) + outputs.append(expert_out) + start_idx = end_idx + + outs = torch.cat(outputs, dim=0) if len(outputs) else sorted_tokens.new_empty(0) + if self.ep_size > 1: + new_x = torch.empty_like(outs) + new_x[gatherd_idxs] = outs + gathered_tokens = new_x.new_empty(*sorted_tokens_shape) + dist.all_to_all( + list(gathered_tokens.split(input_split_sizes)), + list(new_x.split(output_splits)), + ) + outs = gathered_tokens + + new_x = torch.empty_like(outs) + new_x[idxs] = outs + final_out = ( + new_x.view(*topk_ids.shape, -1) + .type(topk_weight.dtype) + .mul_(topk_weight.unsqueeze(dim=-1)) + .sum(dim=1) + .type(new_x.dtype) + ) + return final_out + +# Copied from transformers.models.llama.modeling_llama.repeat_kv +def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: + """ + This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch, + num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim) + """ + batch, num_key_value_heads, slen, head_dim = hidden_states.shape + if n_rep == 1: + return hidden_states + hidden_states = hidden_states[:, :, None, :, :].expand( + batch, num_key_value_heads, n_rep, slen, head_dim + ) + return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim) + +# Copied from transformers.models.llama.modeling_llama.LlamaAttention with Llama->DeepseekV2 +class DeepseekV2Attention(nn.Module): + """Multi-headed attention from 'Attention Is All You Need' paper""" + + def __init__(self, config: DeepseekV2Config, layer_idx: Optional[int] = None): + super().__init__() + self.config = config + self.layer_idx = layer_idx + if layer_idx is None: + logger.warning_once( + f"Instantiating {self.__class__.__name__} without passing `layer_idx` is not recommended and will " + "to errors during the forward call, if caching is used. Please make sure to provide a `layer_idx` " + "when creating this class." + ) + + self.attention_dropout = config.attention_dropout + self.hidden_size = config.hidden_size + self.num_heads = config.num_attention_heads + + self.max_position_embeddings = config.max_position_embeddings + self.rope_theta = config.rope_theta + self.q_lora_rank = config.q_lora_rank + self.qk_rope_head_dim = config.qk_rope_head_dim + self.kv_lora_rank = config.kv_lora_rank + self.v_head_dim = config.v_head_dim + self.qk_nope_head_dim = config.qk_nope_head_dim + self.q_head_dim = config.qk_nope_head_dim + config.qk_rope_head_dim + + self.is_causal = True + + if self.q_lora_rank is None: + self.q_proj = nn.Linear( + self.hidden_size, self.num_heads * self.q_head_dim, bias=False + ) + else: + self.q_a_proj = nn.Linear( + self.hidden_size, config.q_lora_rank, bias=config.attention_bias + ) + self.q_a_layernorm = DeepseekV2RMSNorm(config.q_lora_rank) + self.q_b_proj = nn.Linear( + config.q_lora_rank, self.num_heads * self.q_head_dim, bias=False + ) + + self.kv_a_proj_with_mqa = nn.Linear( + self.hidden_size, + config.kv_lora_rank + config.qk_rope_head_dim, + bias=config.attention_bias, + ) + self.kv_a_layernorm = DeepseekV2RMSNorm(config.kv_lora_rank) + self.kv_b_proj = nn.Linear( + config.kv_lora_rank, + self.num_heads + * (self.q_head_dim - self.qk_rope_head_dim + self.v_head_dim), + bias=False, + ) + + self.o_proj = nn.Linear( + self.num_heads * self.v_head_dim, + self.hidden_size, + bias=config.attention_bias, + ) + self._init_rope() + + self.softmax_scale = self.q_head_dim ** (-0.5) + if self.config.rope_scaling is not None: + mscale_all_dim = self.config.rope_scaling.get("mscale_all_dim", 0) + scaling_factor = self.config.rope_scaling["factor"] + if mscale_all_dim: + mscale = yarn_get_mscale(scaling_factor, mscale_all_dim) + self.softmax_scale = self.softmax_scale * mscale * mscale + + def _init_rope(self): + if self.config.rope_scaling is None: + self.rotary_emb = DeepseekV2RotaryEmbedding( + self.qk_rope_head_dim, + max_position_embeddings=self.max_position_embeddings, + base=self.rope_theta, + ) + else: + scaling_type = self.config.rope_scaling["type"] + scaling_factor = self.config.rope_scaling["factor"] + if scaling_type == "linear": + self.rotary_emb = DeepseekV2LinearScalingRotaryEmbedding( + self.qk_rope_head_dim, + max_position_embeddings=self.max_position_embeddings, + scaling_factor=scaling_factor, + base=self.rope_theta, + ) + elif scaling_type == "dynamic": + self.rotary_emb = DeepseekV2DynamicNTKScalingRotaryEmbedding( + self.qk_rope_head_dim, + max_position_embeddings=self.max_position_embeddings, + scaling_factor=scaling_factor, + base=self.rope_theta, + ) + elif scaling_type == "yarn": + kwargs = { + key: self.config.rope_scaling[key] + for key in [ + "original_max_position_embeddings", + "beta_fast", + "beta_slow", + "mscale", + "mscale_all_dim", + ] + if key in self.config.rope_scaling + } + self.rotary_emb = DeepseekV2YarnRotaryEmbedding( + self.qk_rope_head_dim, + max_position_embeddings=self.max_position_embeddings, + scaling_factor=scaling_factor, + base=self.rope_theta, + **kwargs, + ) + else: + raise ValueError(f"Unknown RoPE scaling type {scaling_type}") + + def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int): + return ( + tensor.view(bsz, seq_len, self.num_heads, self.v_head_dim) + .transpose(1, 2) + .contiguous() + ) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_value: Optional[Cache] = None, + output_attentions: bool = False, + use_cache: bool = False, + cache_position: Optional[torch.LongTensor] = None, + **kwargs, + ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: + if "padding_mask" in kwargs: + warnings.warn( + "Passing `padding_mask` is deprecated and will be removed in v4.37. Please make sure use `attention_mask` instead.`" + ) + bsz, q_len, _ = hidden_states.size() + + if self.q_lora_rank is None: + q = self.q_proj(hidden_states) + else: + q = self.q_b_proj(self.q_a_layernorm(self.q_a_proj(hidden_states))) + q = q.view(bsz, q_len, self.num_heads, self.q_head_dim).transpose(1, 2) + q_nope, q_pe = torch.split( + q, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1 + ) + + compressed_kv = self.kv_a_proj_with_mqa(hidden_states) + compressed_kv, k_pe = torch.split( + compressed_kv, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1 + ) + k_pe = k_pe.view(bsz, q_len, 1, self.qk_rope_head_dim).transpose(1, 2) + kv = ( + self.kv_b_proj(self.kv_a_layernorm(compressed_kv)) + .view(bsz, q_len, self.num_heads, self.qk_nope_head_dim + self.v_head_dim) + .transpose(1, 2) + ) + + k_nope, value_states = torch.split( + kv, [self.qk_nope_head_dim, self.v_head_dim], dim=-1 + ) + kv_seq_len = value_states.shape[-2] + if past_key_value is not None: + if self.layer_idx is None: + raise ValueError( + f"The cache structure has changed since version v4.36. If you are using {self.__class__.__name__} " + "for auto-regressive decoding with k/v caching, please make sure to initialize the attention class " + "with a layer index." + ) + kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx) + cos, sin = self.rotary_emb(q_pe, position_ids) + q_pe, k_pe = apply_rotary_pos_emb(q_pe, k_pe, cos, sin) + + query_states = k_pe.new_empty(bsz, self.num_heads, q_len, self.q_head_dim) + query_states[:, :, :, : self.qk_nope_head_dim] = q_nope + query_states[:, :, :, self.qk_nope_head_dim :] = q_pe + + key_states = k_pe.new_empty(bsz, self.num_heads, q_len, self.q_head_dim) + key_states[:, :, :, : self.qk_nope_head_dim] = k_nope + key_states[:, :, :, self.qk_nope_head_dim :] = k_pe + if past_key_value is not None: + cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position} # Specific to RoPE models + key_states, value_states = past_key_value.update( + key_states, value_states, self.layer_idx, cache_kwargs + ) + + attn_weights = ( + torch.matmul(query_states, key_states.transpose(2, 3)) * self.softmax_scale + ) + + if attention_mask is not None: + attn_weights = attn_weights + attention_mask + + # upcast attention to fp32 + attn_weights = nn.functional.softmax( + attn_weights, dim=-1, dtype=torch.float32 + ).to(query_states.dtype) + attn_weights = nn.functional.dropout( + attn_weights, p=self.attention_dropout, training=self.training + ) + attn_output = torch.matmul(attn_weights, value_states) + + if attn_output.size() != (bsz, self.num_heads, q_len, self.v_head_dim): + raise ValueError( + f"`attn_output` should be of size {(bsz, self.num_heads, q_len, self.v_head_dim)}, but is" + f" {attn_output.size()}" + ) + + attn_output = attn_output.transpose(1, 2).contiguous() + + attn_output = attn_output.reshape(bsz, q_len, self.num_heads * self.v_head_dim) + + attn_output = self.o_proj(attn_output) + + if not output_attentions: + attn_weights = None + + return attn_output, attn_weights, past_key_value + +# Copied from transformers.models.llama.modeling_llama.LlamaFlashAttention2 with Llama->DeepseekV2 +class DeepseekV2FlashAttention2(DeepseekV2Attention): + """ + DeepseekV2 flash attention module. This module inherits from `DeepseekV2Attention` as the weights of the module stays + untouched. The only required change would be on the forward pass where it needs to correctly call the public API of + flash attention and deal with padding tokens in case the input contains any of them. + """ + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + + # TODO: Should be removed once Flash Attention for RoCm is bumped to 2.1. + # flash_attn<2.1 generates top-left aligned causal mask, while what is needed here is bottom-right alignement, that was made default for flash_attn>=2.1. This attribute is used to handle this difference. Reference: https://github.com/Dao-AILab/flash-attention/releases/tag/v2.1.0. + # Beware that with flash_attn<2.1, using q_seqlen != k_seqlen (except for the case q_seqlen == 1) produces a wrong mask (top-left). + self._flash_attn_uses_top_left_mask = not is_flash_attn_greater_or_equal_2_10() + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.LongTensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_value: Optional[Cache] = None, + output_attentions: bool = False, + use_cache: bool = False, + cache_position: Optional[torch.LongTensor] = None, + **kwargs, + ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: + # DeepseekV2FlashAttention2 attention does not support output_attentions + if "padding_mask" in kwargs: + warnings.warn( + "Passing `padding_mask` is deprecated and will be removed in v4.37. Please make sure use `attention_mask` instead.`" + ) + + # overwrite attention_mask with padding_mask + attention_mask = kwargs.pop("padding_mask") + + output_attentions = False + + bsz, q_len, _ = hidden_states.size() + + q = self.q_b_proj(self.q_a_layernorm(self.q_a_proj(hidden_states))) + q = q.view(bsz, q_len, self.num_heads, self.q_head_dim).transpose(1, 2) + q_nope, q_pe = torch.split( + q, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1 + ) + + # Flash attention requires the input to have the shape + # batch_size x seq_length x head_dim x hidden_dim + # therefore we just need to keep the original shape + compressed_kv = self.kv_a_proj_with_mqa(hidden_states) + compressed_kv, k_pe = torch.split( + compressed_kv, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1 + ) + k_pe = k_pe.view(bsz, q_len, 1, self.qk_rope_head_dim).transpose(1, 2) + kv = ( + self.kv_b_proj(self.kv_a_layernorm(compressed_kv)) + .view(bsz, q_len, self.num_heads, self.qk_nope_head_dim + self.v_head_dim) + .transpose(1, 2) + ) + + k_nope, value_states = torch.split( + kv, [self.qk_nope_head_dim, self.v_head_dim], dim=-1 + ) + kv_seq_len = value_states.shape[-2] + + kv_seq_len = value_states.shape[-2] + if past_key_value is not None: + kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx) + + cos, sin = self.rotary_emb(q_pe, position_ids) + q_pe, k_pe = apply_rotary_pos_emb(q_pe, k_pe, cos, sin) + + query_states = k_pe.new_empty(bsz, self.num_heads, q_len, self.q_head_dim) + query_states[:, :, :, : self.qk_nope_head_dim] = q_nope + query_states[:, :, :, self.qk_nope_head_dim :] = q_pe + + key_states = k_pe.new_empty(bsz, self.num_heads, q_len, self.q_head_dim) + key_states[:, :, :, : self.qk_nope_head_dim] = k_nope + key_states[:, :, :, self.qk_nope_head_dim :] = k_pe + + if self.q_head_dim != self.v_head_dim: + value_states = F.pad(value_states, [0, self.q_head_dim - self.v_head_dim]) + + if past_key_value is not None: + cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position} # Specific to RoPE models + key_states, value_states = past_key_value.update( + key_states, value_states, self.layer_idx, cache_kwargs + ) + + # TODO: These transpose are quite inefficient but Flash Attention requires the layout [batch_size, sequence_length, num_heads, head_dim]. We would need to refactor the KV cache + # to be able to avoid many of these transpose/reshape/view. + query_states = query_states.transpose(1, 2) + key_states = key_states.transpose(1, 2) + value_states = value_states.transpose(1, 2) + + dropout_rate = self.attention_dropout if self.training else 0.0 + + # In PEFT, usually we cast the layer norms in float32 for training stability reasons + # therefore the input hidden states gets silently casted in float32. Hence, we need + # cast them back in the correct dtype just to be sure everything works as expected. + # This might slowdown training & inference so it is recommended to not cast the LayerNorms + # in fp32. (DeepseekV2RMSNorm handles it correctly) + + input_dtype = query_states.dtype + if input_dtype == torch.float32: + # Handle the case where the model is quantized + if hasattr(self.config, "_pre_quantization_dtype"): + target_dtype = self.config._pre_quantization_dtype + elif torch.is_autocast_enabled(): + target_dtype = torch.get_autocast_gpu_dtype() + else: + target_dtype = self.q_a_proj.weight.dtype + + logger.warning_once( + f"The input hidden states seems to be silently casted in float32, this might be related to" + f" the fact you have upcasted embedding or layer norm layers in float32. We will cast back the input in" + f" {target_dtype}." + ) + + query_states = query_states.to(target_dtype) + key_states = key_states.to(target_dtype) + value_states = value_states.to(target_dtype) + + attn_output = self._flash_attention_forward( + query_states, + key_states, + value_states, + attention_mask, + q_len, + position_ids=position_ids, + dropout=dropout_rate, + softmax_scale=self.softmax_scale, + ) + if self.q_head_dim != self.v_head_dim: + attn_output = attn_output[:, :, :, : self.v_head_dim] + + attn_output = attn_output.reshape( + bsz, q_len, self.num_heads * self.v_head_dim + ).contiguous() + attn_output = self.o_proj(attn_output) + + if not output_attentions: + attn_weights = None + + return attn_output, attn_weights, past_key_value + + def _flash_attention_forward( + self, + query_states, + key_states, + value_states, + attention_mask, + query_length, + position_ids, + dropout=0.0, + softmax_scale=None, + ): + """ + Calls the forward method of Flash Attention - if the input hidden states contain at least one padding token + first unpad the input, then computes the attention scores and pad the final attention scores. + # Args: + query_states (`torch.Tensor`): + Input query states to be passed to Flash Attention API + key_states (`torch.Tensor`): + Input key states to be passed to Flash Attention API + value_states (`torch.Tensor`): + Input value states to be passed to Flash Attention API + attention_mask (`torch.Tensor`): + The padding mask - corresponds to a tensor of size `(batch_size, seq_len)` where 0 stands for the + position of padding tokens and 1 for the position of non-padding tokens. + dropout (`int`, *optional*): + Attention dropout + softmax_scale (`float`, *optional*): + The scaling of QK^T before applying softmax. Default to 1 / sqrt(head_dim) + """ + if not self._flash_attn_uses_top_left_mask: + causal = self.is_causal + else: + # TODO: Remove the `query_length != 1` check once Flash Attention for RoCm is bumped to 2.1. For details, please see the comment in DeepseekV2FlashAttention2 __init__. + causal = self.is_causal and query_length != 1 + + # Contains at least one padding token in the sequence + if attention_mask is not None: + batch_size = query_states.shape[0] + ( + query_states, + key_states, + value_states, + indices_q, + cu_seq_lens, + max_seq_lens, + ) = self._upad_input( + query_states, key_states, value_states, attention_mask, query_length + ) + + cu_seqlens_q, cu_seqlens_k = cu_seq_lens + max_seqlen_in_batch_q, max_seqlen_in_batch_k = max_seq_lens + + attn_output_unpad = flash_attn_varlen_func( + query_states, + key_states, + value_states, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k=cu_seqlens_k, + max_seqlen_q=max_seqlen_in_batch_q, + max_seqlen_k=max_seqlen_in_batch_k, + dropout_p=dropout, + softmax_scale=softmax_scale, + causal=causal, + ) + + attn_output = pad_input( + attn_output_unpad, indices_q, batch_size, query_length + ) + else: + if query_length == 1: + position_ids = position_ids.to(dtype=torch.int32).squeeze(1) + attn_output = flash_attn_with_kvcache( + query_states, + key_states, + value_states, + cache_seqlens=position_ids, + softmax_scale=softmax_scale, + causal=causal, + ) + else: + attn_output = flash_attn_func( + query_states, + key_states, + value_states, + dropout, + softmax_scale=softmax_scale, + causal=causal, + ) + + return attn_output + + def _upad_input( + self, query_layer, key_layer, value_layer, attention_mask, query_length + ): + indices_k, cu_seqlens_k, max_seqlen_in_batch_k = _get_unpad_data(attention_mask) + batch_size, kv_seq_len, num_key_value_heads, head_dim = key_layer.shape + + key_layer = index_first_axis( + key_layer.reshape(batch_size * kv_seq_len, num_key_value_heads, head_dim), + indices_k, + ) + value_layer = index_first_axis( + value_layer.reshape(batch_size * kv_seq_len, num_key_value_heads, head_dim), + indices_k, + ) + if query_length == kv_seq_len: + query_layer = index_first_axis( + query_layer.reshape(batch_size * kv_seq_len, self.num_heads, head_dim), + indices_k, + ) + cu_seqlens_q = cu_seqlens_k + max_seqlen_in_batch_q = max_seqlen_in_batch_k + indices_q = indices_k + elif query_length == 1: + max_seqlen_in_batch_q = 1 + cu_seqlens_q = torch.arange( + batch_size + 1, dtype=torch.int32, device=query_layer.device + ) # There is a memcpy here, that is very bad. + indices_q = cu_seqlens_q[:-1] + query_layer = query_layer.squeeze(1) + else: + # The -q_len: slice assumes left padding. + attention_mask = attention_mask[:, -query_length:] + query_layer, indices_q, cu_seqlens_q, max_seqlen_in_batch_q = unpad_input( + query_layer, attention_mask + ) + + return ( + query_layer, + key_layer, + value_layer, + indices_q, + (cu_seqlens_q, cu_seqlens_k), + (max_seqlen_in_batch_q, max_seqlen_in_batch_k), + ) + + +ATTENTION_CLASSES = { + "eager": DeepseekV2Attention, + "flash_attention_2": DeepseekV2FlashAttention2, +} + +class DeepseekV2DecoderLayer(nn.Module): + def __init__(self, config: DeepseekV2Config, layer_idx: int): + super().__init__() + self.hidden_size = config.hidden_size + + self.self_attn = ATTENTION_CLASSES[config._attn_implementation]( + config=config, layer_idx=layer_idx + ) + + self.mlp = ( + DeepseekV2MoE(config) + if ( + config.n_routed_experts is not None + and layer_idx >= config.first_k_dense_replace + and layer_idx % config.moe_layer_freq == 0 + ) + else DeepseekV2MLP(config) + ) + self.input_layernorm = DeepseekV2RMSNorm( + config.hidden_size, eps=config.rms_norm_eps + ) + self.post_attention_layernorm = DeepseekV2RMSNorm( + config.hidden_size, eps=config.rms_norm_eps + ) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_value: Optional[Tuple[torch.Tensor]] = None, + output_attentions: Optional[bool] = False, + use_cache: Optional[bool] = False, + cache_position: Optional[torch.LongTensor] = None, + **kwargs, + ) -> Tuple[ + torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]] + ]: + """ + Args: + hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)` + attention_mask (`torch.FloatTensor`, *optional*): + attention mask of size `(batch_size, sequence_length)` if flash attention is used or `(batch_size, 1, + query_sequence_length, key_sequence_length)` if default attention is used. + output_attentions (`bool`, *optional*): + Whether or not to return the attentions tensors of all attention layers. See `attentions` under + returned tensors for more detail. + use_cache (`bool`, *optional*): + If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding + (see `past_key_values`). + past_key_value (`Tuple(torch.FloatTensor)`, *optional*): cached past key and value projection states + """ + if "padding_mask" in kwargs: + warnings.warn( + "Passing `padding_mask` is deprecated and will be removed in v4.37. Please make sure use `attention_mask` instead.`" + ) + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + # Self Attention + hidden_states, self_attn_weights, present_key_value = self.self_attn( + hidden_states=hidden_states, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_value=past_key_value, + output_attentions=output_attentions, + use_cache=use_cache, + cache_position=cache_position, + **kwargs, + ) + + hidden_states = residual + hidden_states + + # Fully Connected + residual = hidden_states + hidden_states = self.post_attention_layernorm(hidden_states) + hidden_states = self.mlp(hidden_states) + + hidden_states = residual + hidden_states + + outputs = (hidden_states,) + + if output_attentions: + outputs += (self_attn_weights,) + + if use_cache: + outputs += (present_key_value,) + + return outputs + + +DeepseekV2_START_DOCSTRING = r""" + This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the + library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads + etc.) + + This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass. + Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage + and behavior. + + Parameters: + config ([`DeepseekV2Config`]): + Model configuration class with all the parameters of the model. Initializing with a config file does not + load the weights associated with the model, only the configuration. Check out the + [`~PreTrainedModel.from_pretrained`] method to load the model weights. +""" + + +@add_start_docstrings( + "The bare DeepseekV2 Model outputting raw hidden-states without any specific head on top.", + DeepseekV2_START_DOCSTRING, +) +class DeepseekV2PreTrainedModel(PreTrainedModel): + config_class = DeepseekV2Config + base_model_prefix = "model" + supports_gradient_checkpointing = True + _no_split_modules = ["DeepseekV2DecoderLayer"] + _skip_keys_device_placement = "past_key_values" + _supports_flash_attn_2 = True + _supports_cache_class = True + _supports_static_cache = True + + def _init_weights(self, module): + std = self.config.initializer_range + if isinstance(module, nn.Linear): + module.weight.data.normal_(mean=0.0, std=std) + if module.bias is not None: + module.bias.data.zero_() + elif isinstance(module, nn.Embedding): + module.weight.data.normal_(mean=0.0, std=std) + if module.padding_idx is not None: + module.weight.data[module.padding_idx].zero_() + + +DeepseekV2_INPUTS_DOCSTRING = r""" + Args: + input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): + Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide + it. + + Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and + [`PreTrainedTokenizer.__call__`] for details. + + [What are input IDs?](../glossary#input-ids) + attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*): + Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`: + + - 1 for tokens that are **not masked**, + - 0 for tokens that are **masked**. + + [What are attention masks?](../glossary#attention-mask) + + Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and + [`PreTrainedTokenizer.__call__`] for details. + + If `past_key_values` is used, optionally only the last `input_ids` have to be input (see + `past_key_values`). + + If you want to change padding behavior, you should read [`modeling_opt._prepare_decoder_attention_mask`] + and modify to your needs. See diagram 1 in [the paper](https://arxiv.org/abs/1910.13461) for more + information on the default strategy. + + - 1 indicates the head is **not masked**, + - 0 indicates the head is **masked**. + position_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0, + config.n_positions - 1]`. + + [What are position IDs?](../glossary#position-ids) + past_key_values (`Cache` or `tuple(tuple(torch.FloatTensor))`, *optional*): + Pre-computed hidden-states (key and values in the self-attention blocks and in the cross-attention + blocks) that can be used to speed up sequential decoding. This typically consists in the `past_key_values` + returned by the model at a previous stage of decoding, when `use_cache=True` or `config.use_cache=True`. + + Two formats are allowed: + - a [`~cache_utils.Cache`] instance; + - Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of + shape `(batch_size, num_heads, sequence_length, embed_size_per_head)`). This is also known as the legacy + cache format. + + The model will output the same cache format that is fed as input. If no `past_key_values` are passed, the + legacy cache format will be returned. + + If `past_key_values` are used, the user can optionally input only the last `input_ids` (those that don't + have their past key value states given to this model) of shape `(batch_size, 1)` instead of all `input_ids` + of shape `(batch_size, sequence_length)`. + inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*): + Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This + is useful if you want more control over how to convert `input_ids` indices into associated vectors than the + model's internal embedding lookup matrix. + use_cache (`bool`, *optional*): + If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see + `past_key_values`). + output_attentions (`bool`, *optional*): + Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned + tensors for more detail. + output_hidden_states (`bool`, *optional*): + Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for + more detail. + return_dict (`bool`, *optional*): + Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple. +""" + + +@add_start_docstrings( + "The bare DeepseekV2 Model outputting raw hidden-states without any specific head on top.", + DeepseekV2_START_DOCSTRING, +) +class DeepseekV2Model(DeepseekV2PreTrainedModel): + """ + Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`DeepseekV2DecoderLayer`] + + Args: + config: DeepseekV2Config + """ + + def __init__(self, config: DeepseekV2Config): + super().__init__(config) + self.padding_idx = config.pad_token_id + self.vocab_size = config.vocab_size + + self.embed_tokens = nn.Embedding( + config.vocab_size, config.hidden_size, self.padding_idx + ) + self.layers = nn.ModuleList( + [ + DeepseekV2DecoderLayer(config, layer_idx) + for layer_idx in range(config.num_hidden_layers) + ] + ) + self._use_flash_attention_2 = config._attn_implementation == "flash_attention_2" + self.norm = DeepseekV2RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + self.gradient_checkpointing = False + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self): + return self.embed_tokens + + def set_input_embeddings(self, value): + self.embed_tokens = value + + @add_start_docstrings_to_model_forward(DeepseekV2_INPUTS_DOCSTRING) + def forward( + self, + input_ids: torch.LongTensor = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[List[torch.FloatTensor]] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + cache_position: Optional[torch.LongTensor] = None, + ) -> Union[Tuple, BaseModelOutputWithPast]: + output_attentions = ( + output_attentions + if output_attentions is not None + else self.config.output_attentions + ) + output_hidden_states = ( + output_hidden_states + if output_hidden_states is not None + else self.config.output_hidden_states + ) + use_cache = use_cache if use_cache is not None else self.config.use_cache + + return_dict = ( + return_dict if return_dict is not None else self.config.use_return_dict + ) + + # retrieve input_ids and inputs_embeds + if input_ids is not None and inputs_embeds is not None: + raise ValueError( + "You cannot specify both input_ids and inputs_embeds at the same time" + ) + elif input_ids is not None: + batch_size, seq_length = input_ids.shape[:2] + elif inputs_embeds is not None: + batch_size, seq_length = inputs_embeds.shape[:2] + else: + raise ValueError("You have to specify either input_ids or inputs_embeds") + + if self.gradient_checkpointing and self.training: + if use_cache: + logger.warning_once( + "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`transformers." + ) + use_cache = False + + past_key_values_length = 0 + if use_cache: + use_legacy_cache = not isinstance(past_key_values, Cache) + if use_legacy_cache: + past_key_values = DynamicCache.from_legacy_cache(past_key_values) + past_key_values_length = past_key_values.get_usable_length(seq_length) + + if cache_position is None: + past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 + cache_position = torch.arange( + past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device + ) + + if position_ids is None: + position_ids = cache_position.unsqueeze(0) + + if inputs_embeds is None: + inputs_embeds = self.embed_tokens(input_ids) + + causal_mask = self._update_causal_mask( + attention_mask, inputs_embeds, cache_position, past_key_values, output_attentions + ) + + # embed positions + hidden_states = inputs_embeds + + # decoder layers + all_hidden_states = () if output_hidden_states else None + all_self_attns = () if output_attentions else None + next_decoder_cache = None + + for decoder_layer in self.layers: + if output_hidden_states: + all_hidden_states += (hidden_states,) + + if self.gradient_checkpointing and self.training: + layer_outputs = self._gradient_checkpointing_func( + decoder_layer.__call__, + hidden_states, + causal_mask, + position_ids, + past_key_values, + output_attentions, + use_cache, + cache_position, + ) + else: + layer_outputs = decoder_layer( + hidden_states, + attention_mask=causal_mask, + position_ids=position_ids, + past_key_value=past_key_values, + output_attentions=output_attentions, + use_cache=use_cache, + cache_position=cache_position, + ) + + hidden_states = layer_outputs[0] + + if use_cache: + next_decoder_cache = layer_outputs[2 if output_attentions else 1] + + if output_attentions: + all_self_attns += (layer_outputs[1],) + + hidden_states = self.norm(hidden_states) + + # add hidden states from the last decoder layer + if output_hidden_states: + all_hidden_states += (hidden_states,) + + next_cache = None + if use_cache: + next_cache = ( + next_decoder_cache.to_legacy_cache() + if use_legacy_cache + else next_decoder_cache + ) + if not return_dict: + return tuple( + v + for v in [hidden_states, next_cache, all_hidden_states, all_self_attns] + if v is not None + ) + return BaseModelOutputWithPast( + last_hidden_state=hidden_states, + past_key_values=next_cache, + hidden_states=all_hidden_states, + attentions=all_self_attns, + ) + + # Copied from transformers.models.llama.modeling_llama.LlamaModel._update_causal_mask + def _update_causal_mask( + self, + attention_mask: torch.Tensor, + input_tensor: torch.Tensor, + cache_position: torch.Tensor, + past_key_values: Cache, + output_attentions: bool, + ): + # TODO: As of torch==2.2.0, the `attention_mask` passed to the model in `generate` is 2D and of dynamic length even when the static + # KV cache is used. This is an issue for torch.compile which then recaptures cudagraphs at each decode steps due to the dynamic shapes. + # (`recording cudagraph tree for symint key 13`, etc.), which is VERY slow. A workaround is `@torch.compiler.disable`, but this prevents using + # `fullgraph=True`. See more context in https://github.com/huggingface/transformers/pull/29114 + + if self.config._attn_implementation == "flash_attention_2": + if attention_mask is not None and 0.0 in attention_mask: + return attention_mask + return None + + # For SDPA, when possible, we will rely on its `is_causal` argument instead of its `attn_mask` argument, in + # order to dispatch on Flash Attention 2. This feature is not compatible with static cache, as SDPA will fail + # to infer the attention mask. + past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 + using_static_cache = isinstance(past_key_values, StaticCache) + + # When output attentions is True, sdpa implementation's forward method calls the eager implementation's forward + if self.config._attn_implementation == "sdpa" and not using_static_cache and not output_attentions: + if AttentionMaskConverter._ignore_causal_mask_sdpa( + attention_mask, + inputs_embeds=input_tensor, + past_key_values_length=past_seen_tokens, + is_training=self.training, + ): + return None + + dtype, device = input_tensor.dtype, input_tensor.device + min_dtype = torch.finfo(dtype).min + sequence_length = input_tensor.shape[1] + if using_static_cache: + target_length = past_key_values.get_max_length() + else: + target_length = ( + attention_mask.shape[-1] + if isinstance(attention_mask, torch.Tensor) + else past_seen_tokens + sequence_length + 1 + ) + + if attention_mask is not None and attention_mask.dim() == 4: + # in this case we assume that the mask comes already in inverted form and requires no inversion or slicing + if attention_mask.max() != 0: + raise ValueError("Custom 4D attention mask should be passed in inverted form with max==0`") + causal_mask = attention_mask + else: + causal_mask = torch.full( + (sequence_length, target_length), fill_value=min_dtype, dtype=dtype, device=device + ) + if sequence_length != 1: + causal_mask = torch.triu(causal_mask, diagonal=1) + causal_mask *= torch.arange(target_length, device=device) > cache_position.reshape(-1, 1) + causal_mask = causal_mask[None, None, :, :].expand(input_tensor.shape[0], 1, -1, -1) + if attention_mask is not None: + causal_mask = causal_mask.clone() # copy to contiguous memory for in-place edit + mask_length = attention_mask.shape[-1] + padding_mask = causal_mask[:, :, :, :mask_length] + attention_mask[:, None, None, :] + padding_mask = padding_mask == 0 + causal_mask[:, :, :, :mask_length] = causal_mask[:, :, :, :mask_length].masked_fill( + padding_mask, min_dtype + ) + if ( + self.config._attn_implementation == "sdpa" + and attention_mask is not None + and attention_mask.device.type == "cuda" + and not output_attentions + ): + # Attend to all tokens in fully masked rows in the causal_mask, for example the relevant first rows when + # using left padding. This is required by F.scaled_dot_product_attention memory-efficient attention path. + # Details: https://github.com/pytorch/pytorch/issues/110213 + causal_mask = AttentionMaskConverter._unmask_unattended(causal_mask, min_dtype) + + return causal_mask + + +class DeepseekV2ForCausalLM(DeepseekV2PreTrainedModel): + _tied_weights_keys = ["lm_head.weight"] + + def __init__(self, config): + super().__init__(config) + self.model = DeepseekV2Model(config) + self.vocab_size = config.vocab_size + self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) + + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self): + return self.model.embed_tokens + + def set_input_embeddings(self, value): + self.model.embed_tokens = value + + def get_output_embeddings(self): + return self.lm_head + + def set_output_embeddings(self, new_embeddings): + self.lm_head = new_embeddings + + def set_decoder(self, decoder): + self.model = decoder + + def get_decoder(self): + return self.model + + @add_start_docstrings_to_model_forward(DeepseekV2_INPUTS_DOCSTRING) + @replace_return_docstrings( + output_type=CausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC + ) + def forward( + self, + input_ids: torch.LongTensor = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[List[torch.FloatTensor]] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + labels: Optional[torch.LongTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + cache_position: Optional[torch.LongTensor] = None, + **kwargs, + ) -> Union[Tuple, CausalLMOutputWithPast]: + r""" + Args: + labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Labels for computing the masked language modeling loss. Indices should either be in `[0, transformers., + config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored + (masked), the loss is only computed for the tokens with labels in `[0, transformers., config.vocab_size]`. + + Returns: + + Example: + + ```python + >>> from transformers import AutoTokenizer, DeepseekV2ForCausalLM + + >>> model = DeepseekV2ForCausalLM.from_pretrained(PATH_TO_CONVERTED_WEIGHTS) + >>> tokenizer = AutoTokenizer.from_pretrained(PATH_TO_CONVERTED_TOKENIZER) + + >>> prompt = "Hey, are you conscious? Can you talk to me?" + >>> inputs = tokenizer(prompt, return_tensors="pt") + + >>> # Generate + >>> generate_ids = model.generate(inputs.input_ids, max_length=30) + >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0] + "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you." + ```""" + output_attentions = ( + output_attentions + if output_attentions is not None + else self.config.output_attentions + ) + output_hidden_states = ( + output_hidden_states + if output_hidden_states is not None + else self.config.output_hidden_states + ) + return_dict = ( + return_dict if return_dict is not None else self.config.use_return_dict + ) + + # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn) + outputs = self.model( + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + cache_position=cache_position, + ) + + hidden_states = outputs[0] + # logits = self.lm_head(hidden_states[:,-1:,:]).float() + + logits = self.lm_head(hidden_states).float() + + loss = None + if labels is not None: + # Shift so that tokens < n predict n + shift_logits = logits[..., :-1, :].contiguous() + shift_labels = labels[..., 1:].contiguous() + # Flatten the tokens + loss_fct = CrossEntropyLoss() + shift_logits = shift_logits.view(-1, self.config.vocab_size) + shift_labels = shift_labels.view(-1) + # Enable model parallelism + shift_labels = shift_labels.to(shift_logits.device) + loss = loss_fct(shift_logits, shift_labels) + + if not return_dict: + output = (logits,) + outputs[1:] + return (loss,) + output if loss is not None else output + + return CausalLMOutputWithPast( + loss=loss, + logits=logits, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) + + def prepare_inputs_for_generation( + self, + input_ids, + past_key_values=None, + attention_mask=None, + inputs_embeds=None, + cache_position=None, + use_cache=True, + **kwargs, + ): + past_length = 0 + # Omit tokens covered by past_key_values + if past_key_values is not None: + if isinstance(past_key_values, Cache): + past_length = cache_position[0] if cache_position is not None else past_key_values.get_seq_length() + max_cache_length = ( + torch.tensor(past_key_values.get_max_length(), device=input_ids.device) + if past_key_values.get_max_length() is not None + else None + ) + cache_length = past_length if max_cache_length is None else torch.min(max_cache_length, past_length) + # TODO joao: remove this `else` after `generate` prioritizes `Cache` objects + else: + cache_length = past_length = past_key_values[0][0].shape[2] + max_cache_length = None + + # Keep only the unprocessed tokens: + # 1 - If the length of the attention_mask exceeds the length of input_ids, then we are in a setting where + # some of the inputs are exclusively passed as part of the cache (e.g. when passing input_embeds as + # input) + if attention_mask is not None and attention_mask.shape[1] > input_ids.shape[1]: + input_ids = input_ids[:, -(attention_mask.shape[1] - past_length) :] + # 2 - If the past_length is smaller than input_ids', then input_ids holds all input tokens. We can discard + # input_ids based on the past_length. + elif past_length < input_ids.shape[1]: + input_ids = input_ids[:, past_length:] + # 3 - Otherwise (past_length >= input_ids.shape[1]), let's assume input_ids only has unprocessed tokens. + + # If we are about to go beyond the maximum cache length, we need to crop the input attention mask. + if ( + max_cache_length is not None + and attention_mask is not None + and cache_length + input_ids.shape[1] > max_cache_length + ): + attention_mask = attention_mask[:, -max_cache_length:] + + position_ids = kwargs.get("position_ids", None) + if attention_mask is not None and position_ids is None: + # create position_ids on the fly for batch generation + position_ids = attention_mask.long().cumsum(-1) - 1 + position_ids.masked_fill_(attention_mask == 0, 1) + if past_key_values: + position_ids = position_ids[:, -input_ids.shape[1] :] + + # if `inputs_embeds` are passed, we only want to use them in the 1st generation step + if inputs_embeds is not None and past_length == 0: + model_inputs = {"inputs_embeds": inputs_embeds} + else: + model_inputs = {"input_ids": input_ids} + + input_length = position_ids.shape[-1] if position_ids is not None else input_ids.shape[-1] + if cache_position is None: + cache_position = torch.arange(past_length, past_length + input_length, device=input_ids.device) + elif use_cache: + cache_position = cache_position[-input_length:] + + model_inputs.update( + { + "position_ids": position_ids, + "past_key_values": past_key_values, + "use_cache": use_cache, + "attention_mask": attention_mask, + "cache_position": cache_position, + } + ) + return model_inputs + + @staticmethod + def _reorder_cache(past_key_values, beam_idx): + reordered_past = () + for layer_past in past_key_values: + reordered_past += ( + tuple( + past_state.index_select(0, beam_idx.to(past_state.device)) + for past_state in layer_past + ), + ) + return reordered_past + + +@add_start_docstrings( + """ + The DeepseekV2 Model transformer with a sequence classification head on top (linear layer). + + [`DeepseekV2ForSequenceClassification`] uses the last token in order to do the classification, as other causal models + (e.g. GPT-2) do. + + Since it does classification on the last token, it requires to know the position of the last token. If a + `pad_token_id` is defined in the configuration, it finds the last token that is not a padding token in each row. If + no `pad_token_id` is defined, it simply takes the last value in each row of the batch. Since it cannot guess the + padding tokens when `inputs_embeds` are passed instead of `input_ids`, it does the same (take the last value in + each row of the batch). + """, + DeepseekV2_START_DOCSTRING, +) +class DeepseekV2ForSequenceClassification(DeepseekV2PreTrainedModel): + def __init__(self, config): + super().__init__(config) + self.num_labels = config.num_labels + self.model = DeepseekV2Model(config) + self.score = nn.Linear(config.hidden_size, self.num_labels, bias=False) + + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self): + return self.model.embed_tokens + + def set_input_embeddings(self, value): + self.model.embed_tokens = value + + @add_start_docstrings_to_model_forward(DeepseekV2_INPUTS_DOCSTRING) + def forward( + self, + input_ids: torch.LongTensor = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[List[torch.FloatTensor]] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + labels: Optional[torch.LongTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + ) -> Union[Tuple, SequenceClassifierOutputWithPast]: + r""" + labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*): + Labels for computing the sequence classification/regression loss. Indices should be in `[0, transformers., + config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If + `config.num_labels > 1` a classification loss is computed (Cross-Entropy). + """ + return_dict = ( + return_dict if return_dict is not None else self.config.use_return_dict + ) + + transformer_outputs = self.model( + input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + hidden_states = transformer_outputs[0] + logits = self.score(hidden_states) + + if input_ids is not None: + batch_size = input_ids.shape[0] + else: + batch_size = inputs_embeds.shape[0] + + if self.config.pad_token_id is None and batch_size != 1: + raise ValueError( + "Cannot handle batch sizes > 1 if no padding token is defined." + ) + if self.config.pad_token_id is None: + sequence_lengths = -1 + else: + if input_ids is not None: + sequence_lengths = ( + torch.eq(input_ids, self.config.pad_token_id).int().argmax(-1) - 1 + ).to(logits.device) + else: + sequence_lengths = -1 + + pooled_logits = logits[ + torch.arange(batch_size, device=logits.device), sequence_lengths + ] + + loss = None + if labels is not None: + labels = labels.to(logits.device) + if self.config.problem_type is None: + if self.num_labels == 1: + self.config.problem_type = "regression" + elif self.num_labels > 1 and ( + labels.dtype == torch.long or labels.dtype == torch.int + ): + self.config.problem_type = "single_label_classification" + else: + self.config.problem_type = "multi_label_classification" + + if self.config.problem_type == "regression": + loss_fct = MSELoss() + if self.num_labels == 1: + loss = loss_fct(pooled_logits.squeeze(), labels.squeeze()) + else: + loss = loss_fct(pooled_logits, labels) + elif self.config.problem_type == "single_label_classification": + loss_fct = CrossEntropyLoss() + loss = loss_fct( + pooled_logits.view(-1, self.num_labels), labels.view(-1) + ) + elif self.config.problem_type == "multi_label_classification": + loss_fct = BCEWithLogitsLoss() + loss = loss_fct(pooled_logits, labels) + if not return_dict: + output = (pooled_logits,) + transformer_outputs[1:] + return ((loss,) + output) if loss is not None else output + + return SequenceClassifierOutputWithPast( + loss=loss, + logits=pooled_logits, + past_key_values=transformer_outputs.past_key_values, + hidden_states=transformer_outputs.hidden_states, + attentions=transformer_outputs.attentions, + ) diff --git a/archive/kt-sft/ktransformers/models/modeling_deepseek_v3.py b/archive/kt-sft/ktransformers/models/modeling_deepseek_v3.py new file mode 100644 index 0000000..51df57d --- /dev/null +++ b/archive/kt-sft/ktransformers/models/modeling_deepseek_v3.py @@ -0,0 +1,1941 @@ +# coding=utf-8 +# Copyright 2023 DeepSeek-AI and The HuggingFace Inc. team. All rights reserved. +# +# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX +# and OPT implementations in this library. It has been modified from its +# original forms to accommodate minor architectural differences compared +# to GPT-NeoX and OPT used by the Meta AI team that trained the model. +# +# 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. +""" PyTorch DeepSeek model.""" +import math +import warnings +from typing import List, Optional, Tuple, Union + +import torch +import torch.nn.functional as F +import torch.utils.checkpoint +from torch import nn +from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss + +from transformers.activations import ACT2FN +from transformers.cache_utils import Cache, DynamicCache, StaticCache +from transformers.generation import GenerationMixin +from transformers.modeling_attn_mask_utils import ( + AttentionMaskConverter, + _prepare_4d_attention_mask, + _prepare_4d_causal_attention_mask, +) +from transformers.modeling_outputs import ( + BaseModelOutputWithPast, + CausalLMOutputWithPast, + SequenceClassifierOutputWithPast, +) +from transformers.modeling_utils import PreTrainedModel +from transformers.pytorch_utils import ( + ALL_LAYERNORM_LAYERS, + is_torch_greater_or_equal_than_1_13, +) +from transformers.utils import ( + add_start_docstrings, + add_start_docstrings_to_model_forward, + is_flash_attn_2_available, + is_flash_attn_greater_or_equal_2_10, + logging, + replace_return_docstrings, +) +from transformers.utils.import_utils import is_torch_fx_available +from .configuration_deepseek_v3 import DeepseekV3Config +import torch.distributed as dist +import numpy as np + +from ktransformers.util.grad_wrapper import maybe_no_grad + +if is_flash_attn_2_available(): + from flash_attn import flash_attn_func, flash_attn_varlen_func + from flash_attn.bert_padding import index_first_axis, pad_input, unpad_input # noqa + + +# This makes `_prepare_4d_causal_attention_mask` a leaf function in the FX graph. +# It means that the function will not be traced through and simply appear as a node in the graph. +if is_torch_fx_available(): + if not is_torch_greater_or_equal_than_1_13: + import torch.fx + + _prepare_4d_causal_attention_mask = torch.fx.wrap(_prepare_4d_causal_attention_mask) + + +logger = logging.get_logger(__name__) + +_CONFIG_FOR_DOC = "DeepseekV3Config" + + +def _get_unpad_data(attention_mask): + seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32) + indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten() + max_seqlen_in_batch = seqlens_in_batch.max().item() + cu_seqlens = F.pad( + torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.torch.int32), (1, 0) + ) + return ( + indices, + cu_seqlens, + max_seqlen_in_batch, + ) + + +class DeepseekV3RMSNorm(nn.Module): + def __init__(self, hidden_size, eps=1e-6): + """ + DeepseekV3RMSNorm is equivalent to T5LayerNorm + """ + super().__init__() + self.weight = nn.Parameter(torch.ones(hidden_size)) + self.variance_epsilon = eps + self.hidden_size = hidden_size + + def forward(self, hidden_states): + input_dtype = hidden_states.dtype + hidden_states = hidden_states.to(torch.float32) + variance = hidden_states.pow(2).mean(-1, keepdim=True) + hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon) + return self.weight * hidden_states.to(input_dtype) + + +ALL_LAYERNORM_LAYERS.append(DeepseekV3RMSNorm) + + +class DeepseekV3RotaryEmbedding(nn.Module): + def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None): + super().__init__() + + self.dim = dim + self.max_position_embeddings = max_position_embeddings + self.base = base + inv_freq = 1.0 / ( + self.base ** (torch.arange(0, self.dim, 2).float().to(device) / self.dim) + ) + self.register_buffer("inv_freq", inv_freq, persistent=False) + + # Build here to make `torch.jit.trace` work. + self._set_cos_sin_cache( + seq_len=max_position_embeddings, + device=self.inv_freq.device, + dtype=torch.get_default_dtype(), + ) + self.max_seq_len_cached = None + + def _set_cos_sin_cache(self, seq_len, device, dtype): + self.max_seq_len_cached = seq_len + t = torch.arange( + self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype + ) + + freqs = torch.outer(t, self.inv_freq.to(t.device)) + # Different from paper, but it uses a different permutation in order to obtain the same calculation + emb = torch.cat((freqs, freqs), dim=-1) + self.register_buffer("cos_cached", emb.cos().to(dtype), persistent=False) + self.register_buffer("sin_cached", emb.sin().to(dtype), persistent=False) + + def forward(self, x, seq_len=None): + # x: [bs, num_attention_heads, seq_len, head_size] + if self.max_seq_len_cached is None or seq_len > self.max_seq_len_cached: + self._set_cos_sin_cache(seq_len=seq_len, device=x.device, dtype=x.dtype) + + return ( + self.cos_cached[:seq_len].to(dtype=x.dtype), + self.sin_cached[:seq_len].to(dtype=x.dtype), + ) + + +# Copied from transformers.models.llama.modeling_llama.LlamaLinearScalingRotaryEmbedding with Llama->DeepseekV3 +class DeepseekV3LinearScalingRotaryEmbedding(DeepseekV3RotaryEmbedding): + """DeepseekV3RotaryEmbedding extended with linear scaling. Credits to the Reddit user /u/kaiokendev""" + + def __init__( + self, + dim, + max_position_embeddings=2048, + base=10000, + device=None, + scaling_factor=1.0, + ): + self.scaling_factor = scaling_factor + super().__init__(dim, max_position_embeddings, base, device) + + def _set_cos_sin_cache(self, seq_len, device, dtype): + self.max_seq_len_cached = seq_len + t = torch.arange( + self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype + ) + t = t / self.scaling_factor + + freqs = torch.outer(t, self.inv_freq) + # Different from paper, but it uses a different permutation in order to obtain the same calculation + emb = torch.cat((freqs, freqs), dim=-1) + self.register_buffer("cos_cached", emb.cos().to(dtype), persistent=False) + self.register_buffer("sin_cached", emb.sin().to(dtype), persistent=False) + + +# Copied from transformers.models.llama.modeling_llama.LlamaDynamicNTKScalingRotaryEmbedding with Llama->DeepseekV3 +class DeepseekV3DynamicNTKScalingRotaryEmbedding(DeepseekV3RotaryEmbedding): + """DeepseekV3RotaryEmbedding extended with Dynamic NTK scaling. Credits to the Reddit users /u/bloc97 and /u/emozilla""" + + def __init__( + self, + dim, + max_position_embeddings=2048, + base=10000, + device=None, + scaling_factor=1.0, + ): + self.scaling_factor = scaling_factor + super().__init__(dim, max_position_embeddings, base, device) + + def _set_cos_sin_cache(self, seq_len, device, dtype): + self.max_seq_len_cached = seq_len + + if seq_len > self.max_position_embeddings: + base = self.base * ( + (self.scaling_factor * seq_len / self.max_position_embeddings) + - (self.scaling_factor - 1) + ) ** (self.dim / (self.dim - 2)) + inv_freq = 1.0 / ( + base ** (torch.arange(0, self.dim, 2).float().to(device) / self.dim) + ) + self.register_buffer("inv_freq", inv_freq, persistent=False) + + t = torch.arange( + self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype + ) + + freqs = torch.outer(t, self.inv_freq) + # Different from paper, but it uses a different permutation in order to obtain the same calculation + emb = torch.cat((freqs, freqs), dim=-1) + self.register_buffer("cos_cached", emb.cos().to(dtype), persistent=False) + self.register_buffer("sin_cached", emb.sin().to(dtype), persistent=False) + + +# Inverse dim formula to find dim based on number of rotations +def yarn_find_correction_dim( + num_rotations, dim, base=10000, max_position_embeddings=2048 +): + 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, high_rot, dim, base=10000, max_position_embeddings=2048 +): + 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_get_mscale(scale=1, mscale=1): + if scale <= 1: + return 1.0 + return 0.1 * mscale * math.log(scale) + 1.0 + + +def yarn_linear_ramp_mask(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 + + +class DeepseekV3YarnRotaryEmbedding(DeepseekV3RotaryEmbedding): + + def __init__( + self, + dim, + max_position_embeddings=2048, + base=10000, + device=None, + scaling_factor=1.0, + original_max_position_embeddings=4096, + beta_fast=32, + beta_slow=1, + mscale=1, + mscale_all_dim=0, + ): + self.scaling_factor = scaling_factor + self.original_max_position_embeddings = original_max_position_embeddings + self.beta_fast = beta_fast + self.beta_slow = beta_slow + self.mscale = mscale + self.mscale_all_dim = mscale_all_dim + super().__init__(dim, max_position_embeddings, base, device) + + def _set_cos_sin_cache(self, seq_len, device, dtype): + self.max_seq_len_cached = seq_len + dim = self.dim + + freq_extra = 1.0 / ( + self.base + ** (torch.arange(0, dim, 2, dtype=torch.float32, device=device) / dim) + ) + freq_inter = 1.0 / ( + self.scaling_factor + * self.base + ** (torch.arange(0, dim, 2, dtype=torch.float32, device=device) / dim) + ) + + low, high = yarn_find_correction_range( + self.beta_fast, + self.beta_slow, + dim, + self.base, + self.original_max_position_embeddings, + ) + inv_freq_mask = 1.0 - yarn_linear_ramp_mask(low, high, dim // 2).to( + device=device, dtype=torch.float32 + ) + inv_freq = freq_inter * (1 - inv_freq_mask) + freq_extra * inv_freq_mask + self.register_buffer("inv_freq", inv_freq, persistent=False) + + t = torch.arange(seq_len, device=device, dtype=torch.float32) + + freqs = torch.outer(t, inv_freq) + + _mscale = float( + yarn_get_mscale(self.scaling_factor, self.mscale) + / yarn_get_mscale(self.scaling_factor, self.mscale_all_dim) + ) + + emb = torch.cat((freqs, freqs), dim=-1) + self.register_buffer( + "cos_cached", (emb.cos() * _mscale).to(dtype), persistent=False + ) + self.register_buffer( + "sin_cached", (emb.sin() * _mscale).to(dtype), persistent=False + ) + + +# Copied from transformers.models.llama.modeling_llama.rotate_half +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) + + +# Copied from transformers.models.llama.modeling_llama.apply_rotary_pos_emb +def apply_rotary_pos_emb(q, k, cos, sin, position_ids, unsqueeze_dim=1): + """Applies Rotary Position Embedding to the query and key tensors. + + Args: + q (`torch.Tensor`): The query tensor. + k (`torch.Tensor`): The key tensor. + cos (`torch.Tensor`): The cosine part of the rotary embedding. + sin (`torch.Tensor`): The sine part of the rotary embedding. + position_ids (`torch.Tensor`): + The position indices of the tokens corresponding to the query and key tensors. For example, this can be + used to pass offsetted position ids when working with a KV-cache. + unsqueeze_dim (`int`, *optional*, defaults to 1): + The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and + sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note + that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and + k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes + cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have + the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2. + Returns: + `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding. + """ + cos = cos[position_ids].unsqueeze(unsqueeze_dim) + sin = sin[position_ids].unsqueeze(unsqueeze_dim) + + b, h, s, d = q.shape + q = q.view(b, h, s, d // 2, 2).transpose(4, 3).reshape(b, h, s, d) + + b, h, s, d = k.shape + k = k.view(b, h, s, d // 2, 2).transpose(4, 3).reshape(b, h, s, d) + + q_embed = (q * cos) + (rotate_half(q) * sin) + k_embed = (k * cos) + (rotate_half(k) * sin) + return q_embed, k_embed + + +class DeepseekV3MLP(nn.Module): + def __init__(self, config, hidden_size=None, intermediate_size=None): + super().__init__() + self.config = config + self.hidden_size = config.hidden_size if hidden_size is None else hidden_size + self.intermediate_size = ( + config.intermediate_size if intermediate_size is None else intermediate_size + ) + + self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) + self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) + self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False) + self.act_fn = ACT2FN[config.hidden_act] + + def forward(self, x): + down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x)) + return down_proj + + +class MoEGate(nn.Module): + def __init__(self, config): + super().__init__() + self.config = config + self.top_k = config.num_experts_per_tok + self.n_routed_experts = config.n_routed_experts + self.routed_scaling_factor = config.routed_scaling_factor + self.scoring_func = config.scoring_func + self.topk_method = config.topk_method + self.n_group = config.n_group + self.topk_group = config.topk_group + + # topk selection algorithm + self.norm_topk_prob = config.norm_topk_prob + self.gating_dim = config.hidden_size + self.weight = nn.Parameter( + torch.empty((self.n_routed_experts, self.gating_dim)) + ) + if self.topk_method == "noaux_tc": + self.e_score_correction_bias = nn.Parameter( + torch.empty((self.n_routed_experts)) + ) + self.reset_parameters() + + def reset_parameters(self) -> None: + import torch.nn.init as init + + init.kaiming_uniform_(self.weight, a=math.sqrt(5)) + + def forward(self, hidden_states): + bsz, seq_len, h = hidden_states.shape + ### compute gating score + hidden_states = hidden_states.view(-1, h) + logits = F.linear( + hidden_states.type(torch.float32), self.weight.type(torch.float32), None + ) + if self.scoring_func == "sigmoid": + scores = logits.sigmoid() + else: + raise NotImplementedError( + f"insupportable scoring function for MoE gating: {self.scoring_func}" + ) + + ### select top-k experts + if self.topk_method == "noaux_tc": + #assert not self.training + scores_for_choice = scores.view(bsz * seq_len, -1) + self.e_score_correction_bias.unsqueeze(0) + group_scores = ( + scores_for_choice.view(bsz * seq_len, self.n_group, -1).topk(2, dim=-1)[0].sum(dim = -1) + ) # [n, n_group] + group_idx = torch.topk( + group_scores, k=self.topk_group, dim=-1, sorted=False + )[ + 1 + ] # [n, top_k_group] + group_mask = torch.zeros_like(group_scores) # [n, n_group] + group_mask.scatter_(1, group_idx, 1) # [n, n_group] + score_mask = ( + group_mask.unsqueeze(-1) + .expand( + bsz * seq_len, self.n_group, self.n_routed_experts // self.n_group + ) + .reshape(bsz * seq_len, -1) + ) # [n, e] + tmp_scores = scores_for_choice.masked_fill(~score_mask.bool(), float("-inf")) # [n, e] + _, topk_idx = torch.topk( + tmp_scores, k=self.top_k, dim=-1, sorted=False + ) + topk_weight = scores.gather(1, topk_idx) + else: + raise NotImplementedError( + f"insupportable TopK function for MoE gating: {self.topk_method}" + ) + + ### norm gate to sum 1 + if self.top_k > 1 and self.norm_topk_prob: + denominator = topk_weight.sum(dim=-1, keepdim=True) + 1e-20 + topk_weight = topk_weight / denominator + topk_weight = topk_weight * self.routed_scaling_factor # must multiply the scaling factor + + return topk_idx, topk_weight + +class DeepseekV3MoE(nn.Module): + """ + A mixed expert module containing shared experts. + """ + + def __init__(self, config): + super().__init__() + self.config = config + self.num_experts_per_tok = config.num_experts_per_tok + + if hasattr(config, "ep_size") and config.ep_size > 1: + assert config.ep_size == dist.get_world_size() + self.ep_size = config.ep_size + self.experts_per_rank = config.n_routed_experts // config.ep_size + self.ep_rank = dist.get_rank() + self.experts = nn.ModuleList( + [ + ( + DeepseekV3MLP( + config, intermediate_size=config.moe_intermediate_size + ) + if i >= self.ep_rank * self.experts_per_rank + and i < (self.ep_rank + 1) * self.experts_per_rank + else None + ) + for i in range(config.n_routed_experts) + ] + ) + else: + self.ep_size = 1 + self.experts_per_rank = config.n_routed_experts + self.ep_rank = 0 + self.experts = nn.ModuleList( + [ + DeepseekV3MLP( + config, intermediate_size=config.moe_intermediate_size + ) + for i in range(config.n_routed_experts) + ] + ) + self.gate = MoEGate(config) + if config.n_shared_experts is not None: + intermediate_size = config.moe_intermediate_size * config.n_shared_experts + self.shared_experts = DeepseekV3MLP( + config=config, intermediate_size=intermediate_size + ) + + def forward(self, hidden_states): + identity = hidden_states + orig_shape = hidden_states.shape + topk_idx, topk_weight = self.gate(hidden_states) + hidden_states = hidden_states.view(-1, hidden_states.shape[-1]) + flat_topk_idx = topk_idx.view(-1) + if not self.training: + y = self.moe_infer(hidden_states, topk_idx, topk_weight).view(*orig_shape) + if self.config.n_shared_experts is not None: + y = y + self.shared_experts(identity) + return y + + @maybe_no_grad() + def moe_infer(self, x, topk_ids, topk_weight): + cnts = topk_ids.new_zeros((topk_ids.shape[0], len(self.experts))) + cnts.scatter_(1, topk_ids, 1) + tokens_per_expert = cnts.sum(dim=0) + idxs = topk_ids.view(-1).argsort() + sorted_tokens = x[idxs // topk_ids.shape[1]] + sorted_tokens_shape = sorted_tokens.shape + if self.ep_size > 1: + tokens_per_ep_rank = tokens_per_expert.view(self.ep_size, -1).sum(dim=1) + tokens_per_expert_group = tokens_per_expert.new_empty( + tokens_per_expert.shape[0] + ) + dist.all_to_all_single(tokens_per_expert_group, tokens_per_expert) + output_splits = ( + tokens_per_expert_group.view(self.ep_size, -1) + .sum(1) + .cpu() + .numpy() + .tolist() + ) + gathered_tokens = sorted_tokens.new_empty( + tokens_per_expert_group.sum(dim=0).cpu().item(), sorted_tokens.shape[1] + ) + input_split_sizes = tokens_per_ep_rank.cpu().numpy().tolist() + dist.all_to_all( + list(gathered_tokens.split(output_splits)), + list(sorted_tokens.split(input_split_sizes)), + ) + tokens_per_expert_post_gather = tokens_per_expert_group.view( + self.ep_size, self.experts_per_rank + ).sum(dim=0) + gatherd_idxs = np.zeros(shape=(gathered_tokens.shape[0],), dtype=np.int32) + s = 0 + for i, k in enumerate(tokens_per_expert_group.cpu().numpy()): + gatherd_idxs[s : s + k] = i % self.experts_per_rank + s += k + gatherd_idxs = gatherd_idxs.argsort() + sorted_tokens = gathered_tokens[gatherd_idxs] + tokens_per_expert = tokens_per_expert_post_gather + tokens_per_expert = tokens_per_expert.cpu().numpy() + + outputs = [] + start_idx = 0 + for i, num_tokens in enumerate(tokens_per_expert): + end_idx = start_idx + num_tokens + if num_tokens == 0: + continue + expert = self.experts[i + self.ep_rank * self.experts_per_rank] + tokens_for_this_expert = sorted_tokens[start_idx:end_idx] + expert_out = expert(tokens_for_this_expert) + outputs.append(expert_out) + start_idx = end_idx + + outs = torch.cat(outputs, dim=0) if len(outputs) else sorted_tokens.new_empty(0) + if self.ep_size > 1: + new_x = torch.empty_like(outs) + new_x[gatherd_idxs] = outs + gathered_tokens = new_x.new_empty(*sorted_tokens_shape) + dist.all_to_all( + list(gathered_tokens.split(input_split_sizes)), + list(new_x.split(output_splits)), + ) + outs = gathered_tokens + + new_x = torch.empty_like(outs) + new_x[idxs] = outs + final_out = ( + new_x.view(*topk_ids.shape, -1) + .type(topk_weight.dtype) + .mul_(topk_weight.unsqueeze(dim=-1)) + .sum(dim=1) + .type(new_x.dtype) + ) + return final_out + + +# Copied from transformers.models.llama.modeling_llama.repeat_kv +def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: + """ + This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch, + num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim) + """ + batch, num_key_value_heads, slen, head_dim = hidden_states.shape + if n_rep == 1: + return hidden_states + hidden_states = hidden_states[:, :, None, :, :].expand( + batch, num_key_value_heads, n_rep, slen, head_dim + ) + return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim) + + +# Copied from transformers.models.llama.modeling_llama.LlamaAttention with Llama->DeepseekV3 +class DeepseekV3Attention(nn.Module): + """Multi-headed attention from 'Attention Is All You Need' paper""" + + def __init__(self, config: DeepseekV3Config, layer_idx: Optional[int] = None): + super().__init__() + self.config = config + self.layer_idx = layer_idx + if layer_idx is None: + logger.warning_once( + f"Instantiating {self.__class__.__name__} without passing `layer_idx` is not recommended and will " + "to errors during the forward call, if caching is used. Please make sure to provide a `layer_idx` " + "when creating this class." + ) + + self.attention_dropout = config.attention_dropout + self.hidden_size = config.hidden_size + self.num_heads = config.num_attention_heads + + self.max_position_embeddings = config.max_position_embeddings + self.rope_theta = config.rope_theta + self.q_lora_rank = config.q_lora_rank + self.qk_rope_head_dim = config.qk_rope_head_dim + self.kv_lora_rank = config.kv_lora_rank + self.v_head_dim = config.v_head_dim + self.qk_nope_head_dim = config.qk_nope_head_dim + self.q_head_dim = config.qk_nope_head_dim + config.qk_rope_head_dim + + self.is_causal = True + + if self.q_lora_rank is None: + self.q_proj = nn.Linear( + self.hidden_size, self.num_heads * self.q_head_dim, bias=False + ) + else: + self.q_a_proj = nn.Linear( + self.hidden_size, config.q_lora_rank, bias=config.attention_bias + ) + self.q_a_layernorm = DeepseekV3RMSNorm(config.q_lora_rank) + self.q_b_proj = nn.Linear( + config.q_lora_rank, self.num_heads * self.q_head_dim, bias=False + ) + + self.kv_a_proj_with_mqa = nn.Linear( + self.hidden_size, + config.kv_lora_rank + config.qk_rope_head_dim, + bias=config.attention_bias, + ) + self.kv_a_layernorm = DeepseekV3RMSNorm(config.kv_lora_rank) + self.kv_b_proj = nn.Linear( + config.kv_lora_rank, + self.num_heads + * (self.q_head_dim - self.qk_rope_head_dim + self.v_head_dim), + bias=False, + ) + + self.o_proj = nn.Linear( + self.num_heads * self.v_head_dim, + self.hidden_size, + bias=config.attention_bias, + ) + self._init_rope() + + self.softmax_scale = self.q_head_dim ** (-0.5) + if self.config.rope_scaling is not None: + mscale_all_dim = self.config.rope_scaling.get("mscale_all_dim", 0) + scaling_factor = self.config.rope_scaling["factor"] + if mscale_all_dim: + mscale = yarn_get_mscale(scaling_factor, mscale_all_dim) + self.softmax_scale = self.softmax_scale * mscale * mscale + + def _init_rope(self): + if self.config.rope_scaling is None: + self.rotary_emb = DeepseekV3RotaryEmbedding( + self.qk_rope_head_dim, + max_position_embeddings=self.max_position_embeddings, + base=self.rope_theta, + ) + else: + scaling_type = self.config.rope_scaling["type"] + scaling_factor = self.config.rope_scaling["factor"] + if scaling_type == "linear": + self.rotary_emb = DeepseekV3LinearScalingRotaryEmbedding( + self.qk_rope_head_dim, + max_position_embeddings=self.max_position_embeddings, + scaling_factor=scaling_factor, + base=self.rope_theta, + ) + elif scaling_type == "dynamic": + self.rotary_emb = DeepseekV3DynamicNTKScalingRotaryEmbedding( + self.qk_rope_head_dim, + max_position_embeddings=self.max_position_embeddings, + scaling_factor=scaling_factor, + base=self.rope_theta, + ) + elif scaling_type == "yarn": + kwargs = { + key: self.config.rope_scaling[key] + for key in [ + "original_max_position_embeddings", + "beta_fast", + "beta_slow", + "mscale", + "mscale_all_dim", + ] + if key in self.config.rope_scaling + } + self.rotary_emb = DeepseekV3YarnRotaryEmbedding( + self.qk_rope_head_dim, + max_position_embeddings=self.max_position_embeddings, + scaling_factor=scaling_factor, + base=self.rope_theta, + **kwargs, + ) + else: + raise ValueError(f"Unknown RoPE scaling type {scaling_type}") + + def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int): + return ( + tensor.view(bsz, seq_len, self.num_heads, self.v_head_dim) + .transpose(1, 2) + .contiguous() + ) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_value: Optional[Cache] = None, + output_attentions: bool = False, + use_cache: bool = False, + cache_position: Optional[torch.LongTensor] = None, + **kwargs, + ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: + if "padding_mask" in kwargs: + warnings.warn( + "Passing `padding_mask` is deprecated and will be removed in v4.37. Please make sure use `attention_mask` instead.`" + ) + bsz, q_len, _ = hidden_states.size() + + if self.q_lora_rank is None: + q = self.q_proj(hidden_states) + else: + q = self.q_b_proj(self.q_a_layernorm(self.q_a_proj(hidden_states))) + q = q.view(bsz, q_len, self.num_heads, self.q_head_dim).transpose(1, 2) + q_nope, q_pe = torch.split( + q, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1 + ) + + compressed_kv = self.kv_a_proj_with_mqa(hidden_states) + compressed_kv, k_pe = torch.split( + compressed_kv, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1 + ) + k_pe = k_pe.view(bsz, q_len, 1, self.qk_rope_head_dim).transpose(1, 2) + kv = ( + self.kv_b_proj(self.kv_a_layernorm(compressed_kv)) + .view(bsz, q_len, self.num_heads, self.qk_nope_head_dim + self.v_head_dim) + .transpose(1, 2) + ) + + k_nope, value_states = torch.split( + kv, [self.qk_nope_head_dim, self.v_head_dim], dim=-1 + ) + kv_seq_len = value_states.shape[-2] + if past_key_value is not None: + if self.layer_idx is None: + raise ValueError( + f"The cache structure has changed since version v4.36. If you are using {self.__class__.__name__} " + "for auto-regressive decoding with k/v caching, please make sure to initialize the attention class " + "with a layer index." + ) + kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx) + cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len) + + q_pe, k_pe = apply_rotary_pos_emb(q_pe, k_pe, cos, sin, position_ids) + + query_states = k_pe.new_empty(bsz, self.num_heads, q_len, self.q_head_dim) + query_states[:, :, :, : self.qk_nope_head_dim] = q_nope + query_states[:, :, :, self.qk_nope_head_dim :] = q_pe + + key_states = k_pe.new_empty(bsz, self.num_heads, q_len, self.q_head_dim) + key_states[:, :, :, : self.qk_nope_head_dim] = k_nope + key_states[:, :, :, self.qk_nope_head_dim :] = k_pe + if past_key_value is not None: + cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position} # Specific to RoPE models + key_states, value_states = past_key_value.update( + key_states, value_states, self.layer_idx, cache_kwargs + ) + + attn_weights = ( + torch.matmul(query_states, key_states.transpose(2, 3)) * self.softmax_scale + ) + + if attn_weights.size() != (bsz, self.num_heads, q_len, kv_seq_len): + raise ValueError( + f"Attention weights should be of size {(bsz, self.num_heads, q_len, kv_seq_len)}, but is" + f" {attn_weights.size()}" + ) + assert attention_mask is not None + if attention_mask is not None: + if attention_mask.size() != (bsz, 1, q_len, kv_seq_len): + raise ValueError( + f"Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}, but is {attention_mask.size()}" + ) + attn_weights = attn_weights + attention_mask + + # upcast attention to fp32 + attn_weights = nn.functional.softmax( + attn_weights, dim=-1, dtype=torch.float32 + ).to(query_states.dtype) + attn_weights = nn.functional.dropout( + attn_weights, p=self.attention_dropout, training=self.training + ) + attn_output = torch.matmul(attn_weights, value_states) + + if attn_output.size() != (bsz, self.num_heads, q_len, self.v_head_dim): + raise ValueError( + f"`attn_output` should be of size {(bsz, self.num_heads, q_len, self.v_head_dim)}, but is" + f" {attn_output.size()}" + ) + + attn_output = attn_output.transpose(1, 2).contiguous() + + attn_output = attn_output.reshape(bsz, q_len, self.num_heads * self.v_head_dim) + + attn_output = self.o_proj(attn_output) + + if not output_attentions: + attn_weights = None + + return attn_output, attn_weights, past_key_value + + +# Copied from transformers.models.llama.modeling_llama.LlamaFlashAttention2 with Llama->DeepseekV3 +class DeepseekV3FlashAttention2(DeepseekV3Attention): + """ + DeepseekV3 flash attention module. This module inherits from `DeepseekV3Attention` as the weights of the module stays + untouched. The only required change would be on the forward pass where it needs to correctly call the public API of + flash attention and deal with padding tokens in case the input contains any of them. + """ + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + + # TODO: Should be removed once Flash Attention for RoCm is bumped to 2.1. + # flash_attn<2.1 generates top-left aligned causal mask, while what is needed here is bottom-right alignement, that was made default for flash_attn>=2.1. This attribute is used to handle this difference. Reference: https://github.com/Dao-AILab/flash-attention/releases/tag/v2.1.0. + # Beware that with flash_attn<2.1, using q_seqlen != k_seqlen (except for the case q_seqlen == 1) produces a wrong mask (top-left). + self._flash_attn_uses_top_left_mask = not is_flash_attn_greater_or_equal_2_10() + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.LongTensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_value: Optional[Cache] = None, + output_attentions: bool = False, + use_cache: bool = False, + cache_position: Optional[torch.LongTensor] = None, + **kwargs, + ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: + # DeepseekV3FlashAttention2 attention does not support output_attentions + if "padding_mask" in kwargs: + warnings.warn( + "Passing `padding_mask` is deprecated and will be removed in v4.37. Please make sure use `attention_mask` instead.`" + ) + + # overwrite attention_mask with padding_mask + attention_mask = kwargs.pop("padding_mask") + + output_attentions = False + + bsz, q_len, _ = hidden_states.size() + + if self.q_lora_rank is None: + q = self.q_proj(hidden_states) + else: + q = self.q_b_proj(self.q_a_layernorm(self.q_a_proj(hidden_states))) + q = q.view(bsz, q_len, self.num_heads, self.q_head_dim).transpose(1, 2) + q_nope, q_pe = torch.split( + q, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1 + ) + + # Flash attention requires the input to have the shape + # batch_size x seq_length x head_dim x hidden_dim + # therefore we just need to keep the original shape + compressed_kv = self.kv_a_proj_with_mqa(hidden_states) + compressed_kv, k_pe = torch.split( + compressed_kv, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1 + ) + k_pe = k_pe.view(bsz, q_len, 1, self.qk_rope_head_dim).transpose(1, 2) + kv = ( + self.kv_b_proj(self.kv_a_layernorm(compressed_kv)) + .view(bsz, q_len, self.num_heads, self.qk_nope_head_dim + self.v_head_dim) + .transpose(1, 2) + ) + + k_nope, value_states = torch.split( + kv, [self.qk_nope_head_dim, self.v_head_dim], dim=-1 + ) + kv_seq_len = value_states.shape[-2] + + kv_seq_len = value_states.shape[-2] + if past_key_value is not None: + kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx) + + cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len) + q_pe, k_pe = apply_rotary_pos_emb(q_pe, k_pe, cos, sin, position_ids) + + query_states = k_pe.new_empty(bsz, self.num_heads, q_len, self.q_head_dim) + query_states[:, :, :, : self.qk_nope_head_dim] = q_nope + query_states[:, :, :, self.qk_nope_head_dim :] = q_pe + + key_states = k_pe.new_empty(bsz, self.num_heads, q_len, self.q_head_dim) + key_states[:, :, :, : self.qk_nope_head_dim] = k_nope + key_states[:, :, :, self.qk_nope_head_dim :] = k_pe + + if self.q_head_dim != self.v_head_dim: + value_states = F.pad(value_states, [0, self.q_head_dim - self.v_head_dim]) + + if past_key_value is not None: + cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position} # Specific to RoPE models + key_states, value_states = past_key_value.update( + key_states, value_states, self.layer_idx, cache_kwargs + ) + + # TODO: These transpose are quite inefficient but Flash Attention requires the layout [batch_size, sequence_length, num_heads, head_dim]. We would need to refactor the KV cache + # to be able to avoid many of these transpose/reshape/view. + query_states = query_states.transpose(1, 2) + key_states = key_states.transpose(1, 2) + value_states = value_states.transpose(1, 2) + + dropout_rate = self.attention_dropout if self.training else 0.0 + + # In PEFT, usually we cast the layer norms in float32 for training stability reasons + # therefore the input hidden states gets silently casted in float32. Hence, we need + # cast them back in the correct dtype just to be sure everything works as expected. + # This might slowdown training & inference so it is recommended to not cast the LayerNorms + # in fp32. (DeepseekV3RMSNorm handles it correctly) + + input_dtype = query_states.dtype + if input_dtype == torch.float32: + # Handle the case where the model is quantized + if hasattr(self.config, "_pre_quantization_dtype"): + target_dtype = self.config._pre_quantization_dtype + elif torch.is_autocast_enabled(): + target_dtype = torch.get_autocast_gpu_dtype() + else: + target_dtype = ( + self.q_proj.weight.dtype + if self.q_lora_rank is None + else self.q_a_proj.weight.dtype + ) + + logger.warning_once( + f"The input hidden states seems to be silently casted in float32, this might be related to" + f" the fact you have upcasted embedding or layer norm layers in float32. We will cast back the input in" + f" {target_dtype}." + ) + + query_states = query_states.to(target_dtype) + key_states = key_states.to(target_dtype) + value_states = value_states.to(target_dtype) + + attn_output = self._flash_attention_forward( + query_states, + key_states, + value_states, + attention_mask, + q_len, + dropout=dropout_rate, + softmax_scale=self.softmax_scale, + ) + if self.q_head_dim != self.v_head_dim: + attn_output = attn_output[:, :, :, : self.v_head_dim] + + attn_output = attn_output.reshape( + bsz, q_len, self.num_heads * self.v_head_dim + ).contiguous() + attn_output = self.o_proj(attn_output) + + if not output_attentions: + attn_weights = None + + return attn_output, attn_weights, past_key_value + + def _flash_attention_forward( + self, + query_states, + key_states, + value_states, + attention_mask, + query_length, + dropout=0.0, + softmax_scale=None, + ): + """ + Calls the forward method of Flash Attention - if the input hidden states contain at least one padding token + first unpad the input, then computes the attention scores and pad the final attention scores. + + Args: + query_states (`torch.Tensor`): + Input query states to be passed to Flash Attention API + key_states (`torch.Tensor`): + Input key states to be passed to Flash Attention API + value_states (`torch.Tensor`): + Input value states to be passed to Flash Attention API + attention_mask (`torch.Tensor`): + The padding mask - corresponds to a tensor of size `(batch_size, seq_len)` where 0 stands for the + position of padding tokens and 1 for the position of non-padding tokens. + dropout (`int`, *optional*): + Attention dropout + softmax_scale (`float`, *optional*): + The scaling of QK^T before applying softmax. Default to 1 / sqrt(head_dim) + """ + if not self._flash_attn_uses_top_left_mask: + causal = self.is_causal + else: + # TODO: Remove the `query_length != 1` check once Flash Attention for RoCm is bumped to 2.1. For details, please see the comment in DeepseekV3FlashAttention2 __init__. + causal = self.is_causal and query_length != 1 + + # Contains at least one padding token in the sequence + if attention_mask is not None: + batch_size = query_states.shape[0] + ( + query_states, + key_states, + value_states, + indices_q, + cu_seq_lens, + max_seq_lens, + ) = self._upad_input( + query_states, key_states, value_states, attention_mask, query_length + ) + + cu_seqlens_q, cu_seqlens_k = cu_seq_lens + max_seqlen_in_batch_q, max_seqlen_in_batch_k = max_seq_lens + + attn_output_unpad = flash_attn_varlen_func( + query_states, + key_states, + value_states, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k=cu_seqlens_k, + max_seqlen_q=max_seqlen_in_batch_q, + max_seqlen_k=max_seqlen_in_batch_k, + dropout_p=dropout, + softmax_scale=softmax_scale, + causal=causal, + ) + + attn_output = pad_input( + attn_output_unpad, indices_q, batch_size, query_length + ) + else: + attn_output = flash_attn_func( + query_states, + key_states, + value_states, + dropout, + softmax_scale=softmax_scale, + causal=causal, + ) + + return attn_output + + def _upad_input( + self, query_layer, key_layer, value_layer, attention_mask, query_length + ): + indices_k, cu_seqlens_k, max_seqlen_in_batch_k = _get_unpad_data(attention_mask) + batch_size, kv_seq_len, num_key_value_heads, head_dim = key_layer.shape + + key_layer = index_first_axis( + key_layer.reshape(batch_size * kv_seq_len, num_key_value_heads, head_dim), + indices_k, + ) + value_layer = index_first_axis( + value_layer.reshape(batch_size * kv_seq_len, num_key_value_heads, head_dim), + indices_k, + ) + if query_length == kv_seq_len: + query_layer = index_first_axis( + query_layer.reshape(batch_size * kv_seq_len, self.num_heads, head_dim), + indices_k, + ) + cu_seqlens_q = cu_seqlens_k + max_seqlen_in_batch_q = max_seqlen_in_batch_k + indices_q = indices_k + elif query_length == 1: + max_seqlen_in_batch_q = 1 + cu_seqlens_q = torch.arange( + batch_size + 1, dtype=torch.int32, device=query_layer.device + ) # There is a memcpy here, that is very bad. + indices_q = cu_seqlens_q[:-1] + query_layer = query_layer.squeeze(1) + else: + # The -q_len: slice assumes left padding. + attention_mask = attention_mask[:, -query_length:] + query_layer, indices_q, cu_seqlens_q, max_seqlen_in_batch_q = unpad_input( + query_layer, attention_mask + ) + + return ( + query_layer, + key_layer, + value_layer, + indices_q, + (cu_seqlens_q, cu_seqlens_k), + (max_seqlen_in_batch_q, max_seqlen_in_batch_k), + ) + + +ATTENTION_CLASSES = { + "eager": DeepseekV3Attention, + "flash_attention_2": DeepseekV3FlashAttention2, +} + + +class DeepseekV3DecoderLayer(nn.Module): + def __init__(self, config: DeepseekV3Config, layer_idx: int): + super().__init__() + self.hidden_size = config.hidden_size + + self.self_attn = ATTENTION_CLASSES[config._attn_implementation]( + config=config, layer_idx=layer_idx + ) + + self.mlp = ( + DeepseekV3MoE(config) + if ( + config.n_routed_experts is not None + and layer_idx >= config.first_k_dense_replace + and layer_idx % config.moe_layer_freq == 0 + ) + else DeepseekV3MLP(config) + ) + self.input_layernorm = DeepseekV3RMSNorm( + config.hidden_size, eps=config.rms_norm_eps + ) + self.post_attention_layernorm = DeepseekV3RMSNorm( + config.hidden_size, eps=config.rms_norm_eps + ) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_value: Optional[Tuple[torch.Tensor]] = None, + output_attentions: Optional[bool] = False, + use_cache: Optional[bool] = False, + cache_position: Optional[torch.LongTensor] = None, + **kwargs, + ) -> Tuple[ + torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]] + ]: + """ + Args: + hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)` + attention_mask (`torch.FloatTensor`, *optional*): + attention mask of size `(batch_size, sequence_length)` if flash attention is used or `(batch_size, 1, + query_sequence_length, key_sequence_length)` if default attention is used. + output_attentions (`bool`, *optional*): + Whether or not to return the attentions tensors of all attention layers. See `attentions` under + returned tensors for more detail. + use_cache (`bool`, *optional*): + If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding + (see `past_key_values`). + past_key_value (`Tuple(torch.FloatTensor)`, *optional*): cached past key and value projection states + """ + if "padding_mask" in kwargs: + warnings.warn( + "Passing `padding_mask` is deprecated and will be removed in v4.37. Please make sure use `attention_mask` instead.`" + ) + residual = hidden_states + + hidden_states = self.input_layernorm(hidden_states) + + # Self Attention + hidden_states, self_attn_weights, present_key_value = self.self_attn( + hidden_states=hidden_states, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_value=past_key_value, + output_attentions=output_attentions, + use_cache=use_cache, + cache_position=cache_position, + **kwargs, + ) + hidden_states = residual + hidden_states + + # Fully Connected + residual = hidden_states + hidden_states = self.post_attention_layernorm(hidden_states) + hidden_states = self.mlp(hidden_states) + hidden_states = residual + hidden_states + + outputs = (hidden_states,) + + if output_attentions: + outputs += (self_attn_weights,) + + if use_cache: + outputs += (present_key_value,) + + return outputs + + +DeepseekV3_START_DOCSTRING = r""" + This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the + library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads + etc.) + + This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass. + Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage + and behavior. + + Parameters: + config ([`DeepseekV3Config`]): + Model configuration class with all the parameters of the model. Initializing with a config file does not + load the weights associated with the model, only the configuration. Check out the + [`~PreTrainedModel.from_pretrained`] method to load the model weights. +""" + + +@add_start_docstrings( + "The bare DeepseekV3 Model outputting raw hidden-states without any specific head on top.", + DeepseekV3_START_DOCSTRING, +) +class DeepseekV3PreTrainedModel(PreTrainedModel): + config_class = DeepseekV3Config + base_model_prefix = "model" + supports_gradient_checkpointing = True + _no_split_modules = ["DeepseekV3DecoderLayer"] + _skip_keys_device_placement = "past_key_values" + _supports_flash_attn_2 = True + _supports_cache_class = True + + def _init_weights(self, module): + std = self.config.initializer_range + if isinstance(module, nn.Linear): + module.weight.data.normal_(mean=0.0, std=std) + if module.bias is not None: + module.bias.data.zero_() + elif isinstance(module, nn.Embedding): + module.weight.data.normal_(mean=0.0, std=std) + if module.padding_idx is not None: + module.weight.data[module.padding_idx].zero_() + + +DeepseekV3_INPUTS_DOCSTRING = r""" + Args: + input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): + Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide + it. + + Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and + [`PreTrainedTokenizer.__call__`] for details. + + [What are input IDs?](../glossary#input-ids) + attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*): + Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`: + + - 1 for tokens that are **not masked**, + - 0 for tokens that are **masked**. + + [What are attention masks?](../glossary#attention-mask) + + Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and + [`PreTrainedTokenizer.__call__`] for details. + + If `past_key_values` is used, optionally only the last `input_ids` have to be input (see + `past_key_values`). + + If you want to change padding behavior, you should read [`modeling_opt._prepare_decoder_attention_mask`] + and modify to your needs. See diagram 1 in [the paper](https://arxiv.org/abs/1910.13461) for more + information on the default strategy. + + - 1 indicates the head is **not masked**, + - 0 indicates the head is **masked**. + position_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0, + config.n_positions - 1]`. + + [What are position IDs?](../glossary#position-ids) + past_key_values (`Cache` or `tuple(tuple(torch.FloatTensor))`, *optional*): + Pre-computed hidden-states (key and values in the self-attention blocks and in the cross-attention + blocks) that can be used to speed up sequential decoding. This typically consists in the `past_key_values` + returned by the model at a previous stage of decoding, when `use_cache=True` or `config.use_cache=True`. + + Two formats are allowed: + - a [`~cache_utils.Cache`] instance; + - Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of + shape `(batch_size, num_heads, sequence_length, embed_size_per_head)`). This is also known as the legacy + cache format. + + The model will output the same cache format that is fed as input. If no `past_key_values` are passed, the + legacy cache format will be returned. + + If `past_key_values` are used, the user can optionally input only the last `input_ids` (those that don't + have their past key value states given to this model) of shape `(batch_size, 1)` instead of all `input_ids` + of shape `(batch_size, sequence_length)`. + inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*): + Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This + is useful if you want more control over how to convert `input_ids` indices into associated vectors than the + model's internal embedding lookup matrix. + use_cache (`bool`, *optional*): + If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see + `past_key_values`). + output_attentions (`bool`, *optional*): + Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned + tensors for more detail. + output_hidden_states (`bool`, *optional*): + Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for + more detail. + return_dict (`bool`, *optional*): + Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple. +""" + + +@add_start_docstrings( + "The bare DeepseekV3 Model outputting raw hidden-states without any specific head on top.", + DeepseekV3_START_DOCSTRING, +) +class DeepseekV3Model(DeepseekV3PreTrainedModel): + """ + Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`DeepseekV3DecoderLayer`] + + Args: + config: DeepseekV3Config + """ + + def __init__(self, config: DeepseekV3Config): + super().__init__(config) + self.padding_idx = config.pad_token_id + self.vocab_size = config.vocab_size + + self.embed_tokens = nn.Embedding( + config.vocab_size, config.hidden_size, self.padding_idx + ) + self.layers = nn.ModuleList( + [ + DeepseekV3DecoderLayer(config, layer_idx) + for layer_idx in range(config.num_hidden_layers) + ] + ) + self._use_flash_attention_2 = config._attn_implementation == "flash_attention_2" + self.norm = DeepseekV3RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + self.gradient_checkpointing = False + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self): + return self.embed_tokens + + def set_input_embeddings(self, value): + self.embed_tokens = value + + @add_start_docstrings_to_model_forward(DeepseekV3_INPUTS_DOCSTRING) + def forward( + self, + input_ids: torch.LongTensor = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[List[torch.FloatTensor]] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + cache_position: Optional[torch.LongTensor] = None, + ) -> Union[Tuple, BaseModelOutputWithPast]: + output_attentions = ( + output_attentions + if output_attentions is not None + else self.config.output_attentions + ) + output_hidden_states = ( + output_hidden_states + if output_hidden_states is not None + else self.config.output_hidden_states + ) + use_cache = use_cache if use_cache is not None else self.config.use_cache + + return_dict = ( + return_dict if return_dict is not None else self.config.use_return_dict + ) + + # retrieve input_ids and inputs_embeds + if input_ids is not None and inputs_embeds is not None: + raise ValueError( + "You cannot specify both input_ids and inputs_embeds at the same time" + ) + elif input_ids is not None: + batch_size, seq_length = input_ids.shape[:2] + elif inputs_embeds is not None: + batch_size, seq_length = inputs_embeds.shape[:2] + else: + raise ValueError("You have to specify either input_ids or inputs_embeds") + + past_key_values_length = 0 + if use_cache: + use_legacy_cache = not isinstance(past_key_values, Cache) + if use_legacy_cache: + past_key_values = DynamicCache.from_legacy_cache(past_key_values) + past_key_values_length = past_key_values.get_usable_length(seq_length) + + if position_ids is None: + device = input_ids.device if input_ids is not None else inputs_embeds.device + position_ids = torch.arange( + past_key_values_length, + seq_length + past_key_values_length, + dtype=torch.long, + device=device, + ) + position_ids = position_ids.unsqueeze(0) + + if inputs_embeds is None: + inputs_embeds = self.embed_tokens(input_ids) + + if self._use_flash_attention_2: + # 2d mask is passed through the layers + attention_mask = ( + attention_mask + if (attention_mask is not None and 0 in attention_mask) + else None + ) + else: + # 4d mask is passed through the layers + attention_mask = _prepare_4d_causal_attention_mask( + attention_mask, + (batch_size, seq_length), + inputs_embeds, + past_key_values_length, + ) + + # embed positions + hidden_states = inputs_embeds + + # decoder layers + all_hidden_states = () if output_hidden_states else None + all_self_attns = () if output_attentions else None + next_decoder_cache = None + + for decoder_layer in self.layers: + if output_hidden_states: + all_hidden_states += (hidden_states,) + + layer_outputs = decoder_layer( + hidden_states, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_value=past_key_values, + output_attentions=output_attentions, + use_cache=use_cache, + cache_position=cache_position, + ) + + hidden_states = layer_outputs[0] + + if use_cache: + next_decoder_cache = layer_outputs[2 if output_attentions else 1] + + if output_attentions: + all_self_attns += (layer_outputs[1],) + + hidden_states = self.norm(hidden_states) + + # add hidden states from the last decoder layer + if output_hidden_states: + all_hidden_states += (hidden_states,) + + next_cache = None + if use_cache: + next_cache = ( + next_decoder_cache.to_legacy_cache() + if use_legacy_cache + else next_decoder_cache + ) + if not return_dict: + return tuple( + v + for v in [hidden_states, next_cache, all_hidden_states, all_self_attns] + if v is not None + ) + return BaseModelOutputWithPast( + last_hidden_state=hidden_states, + past_key_values=next_cache, + hidden_states=all_hidden_states, + attentions=all_self_attns, + ) + + # Copied from transformers.models.llama.modeling_llama.LlamaModel._update_causal_mask + def _update_causal_mask( + self, + attention_mask: torch.Tensor, + input_tensor: torch.Tensor, + cache_position: torch.Tensor, + past_key_values: Cache, + output_attentions: bool, + ): + # TODO: As of torch==2.2.0, the `attention_mask` passed to the model in `generate` is 2D and of dynamic length even when the static + # KV cache is used. This is an issue for torch.compile which then recaptures cudagraphs at each decode steps due to the dynamic shapes. + # (`recording cudagraph tree for symint key 13`, etc.), which is VERY slow. A workaround is `@torch.compiler.disable`, but this prevents using + # `fullgraph=True`. See more context in https://github.com/huggingface/transformers/pull/29114 + + if self.config._attn_implementation == "flash_attention_2": + if attention_mask is not None and 0.0 in attention_mask: + return attention_mask + return None + + # For SDPA, when possible, we will rely on its `is_causal` argument instead of its `attn_mask` argument, in + # order to dispatch on Flash Attention 2. This feature is not compatible with static cache, as SDPA will fail + # to infer the attention mask. + past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 + using_static_cache = isinstance(past_key_values, StaticCache) + + # When output attentions is True, sdpa implementation's forward method calls the eager implementation's forward + if self.config._attn_implementation == "sdpa" and not using_static_cache and not output_attentions: + if AttentionMaskConverter._ignore_causal_mask_sdpa( + attention_mask, + inputs_embeds=input_tensor, + past_key_values_length=past_seen_tokens, + is_training=self.training, + ): + return None + + dtype, device = input_tensor.dtype, input_tensor.device + min_dtype = torch.finfo(dtype).min + sequence_length = input_tensor.shape[1] + if using_static_cache: + target_length = past_key_values.get_max_length() + else: + target_length = ( + attention_mask.shape[-1] + if isinstance(attention_mask, torch.Tensor) + else past_seen_tokens + sequence_length + 1 + ) + + if attention_mask is not None and attention_mask.dim() == 4: + # in this case we assume that the mask comes already in inverted form and requires no inversion or slicing + if attention_mask.max() != 0: + raise ValueError("Custom 4D attention mask should be passed in inverted form with max==0`") + causal_mask = attention_mask + else: + causal_mask = torch.full( + (sequence_length, target_length), fill_value=min_dtype, dtype=dtype, device=device + ) + if sequence_length != 1: + causal_mask = torch.triu(causal_mask, diagonal=1) + causal_mask *= torch.arange(target_length, device=device) > cache_position.reshape(-1, 1) + causal_mask = causal_mask[None, None, :, :].expand(input_tensor.shape[0], 1, -1, -1) + if attention_mask is not None: + causal_mask = causal_mask.clone() # copy to contiguous memory for in-place edit + mask_length = attention_mask.shape[-1] + padding_mask = causal_mask[:, :, :, :mask_length] + attention_mask[:, None, None, :] + padding_mask = padding_mask == 0 + causal_mask[:, :, :, :mask_length] = causal_mask[:, :, :, :mask_length].masked_fill( + padding_mask, min_dtype + ) + if ( + self.config._attn_implementation == "sdpa" + and attention_mask is not None + and attention_mask.device.type == "cuda" + and not output_attentions + ): + # Attend to all tokens in fully masked rows in the causal_mask, for example the relevant first rows when + # using left padding. This is required by F.scaled_dot_product_attention memory-efficient attention path. + # Details: https://github.com/pytorch/pytorch/issues/110213 + causal_mask = AttentionMaskConverter._unmask_unattended(causal_mask, min_dtype) + + return causal_mask + +class DeepseekV3ForCausalLM(DeepseekV3PreTrainedModel, GenerationMixin): + _tied_weights_keys = ["lm_head.weight"] + + def __init__(self, config): + super().__init__(config) + self.model = DeepseekV3Model(config) + self.vocab_size = config.vocab_size + self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) + + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self): + return self.model.embed_tokens + + def set_input_embeddings(self, value): + self.model.embed_tokens = value + + def get_output_embeddings(self): + return self.lm_head + + def set_output_embeddings(self, new_embeddings): + self.lm_head = new_embeddings + + def set_decoder(self, decoder): + self.model = decoder + + def get_decoder(self): + return self.model + + @add_start_docstrings_to_model_forward(DeepseekV3_INPUTS_DOCSTRING) + @replace_return_docstrings( + output_type=CausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC + ) + def forward( + self, + input_ids: torch.LongTensor = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[List[torch.FloatTensor]] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + labels: Optional[torch.LongTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + cache_position: Optional[torch.LongTensor] = None, + ) -> Union[Tuple, CausalLMOutputWithPast]: + r""" + Args: + labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Labels for computing the masked language modeling loss. Indices should either be in `[0, transformers., + config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored + (masked), the loss is only computed for the tokens with labels in `[0, transformers., config.vocab_size]`. + + Returns: + + Example: + + ```python + >>> from transformers import AutoTokenizer, DeepseekV3ForCausalLM + + >>> model = DeepseekV3ForCausalLM.from_pretrained(PATH_TO_CONVERTED_WEIGHTS) + >>> tokenizer = AutoTokenizer.from_pretrained(PATH_TO_CONVERTED_TOKENIZER) + + >>> prompt = "Hey, are you conscious? Can you talk to me?" + >>> inputs = tokenizer(prompt, return_tensors="pt") + + >>> # Generate + >>> generate_ids = model.generate(inputs.input_ids, max_length=30) + >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0] + "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you." + ```""" + output_attentions = ( + output_attentions + if output_attentions is not None + else self.config.output_attentions + ) + output_hidden_states = ( + output_hidden_states + if output_hidden_states is not None + else self.config.output_hidden_states + ) + return_dict = ( + return_dict if return_dict is not None else self.config.use_return_dict + ) + + # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn) + outputs = self.model( + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + cache_position=cache_position, + ) + + hidden_states = outputs[0] + # logits = self.lm_head(hidden_states[:,-1:,:]) + logits = self.lm_head(hidden_states) + logits = logits.float() + + loss = None + if labels is not None: + # Shift so that tokens < n predict n + shift_logits = logits[..., :-1, :].contiguous() + shift_labels = labels[..., 1:].contiguous() + # Flatten the tokens + loss_fct = CrossEntropyLoss() + shift_logits = shift_logits.view(-1, self.config.vocab_size) + shift_labels = shift_labels.view(-1) + # Enable model parallelism + shift_labels = shift_labels.to(shift_logits.device) + loss = loss_fct(shift_logits, shift_labels) + + if not return_dict: + output = (logits,) + outputs[1:] + return (loss,) + output if loss is not None else output + + return CausalLMOutputWithPast( + loss=loss, + logits=logits, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) + + def prepare_inputs_for_generation( + self, + input_ids, + past_key_values=None, + attention_mask=None, + inputs_embeds=None, + **kwargs, + ): + if past_key_values is not None: + if isinstance(past_key_values, Cache): + cache_length = past_key_values.get_seq_length() + past_length = past_key_values.seen_tokens + max_cache_length = past_key_values.get_max_length() + else: + cache_length = past_length = past_key_values[0][0].shape[2] + max_cache_length = None + + # Keep only the unprocessed tokens: + # 1 - If the length of the attention_mask exceeds the length of input_ids, then we are in a setting where + # some of the inputs are exclusivelly passed as part of the cache (e.g. when passing input_embeds as + # input) + if ( + attention_mask is not None + and attention_mask.shape[1] > input_ids.shape[1] + ): + input_ids = input_ids[:, -(attention_mask.shape[1] - past_length) :] + # 2 - If the past_length is smaller than input_ids', then input_ids holds all input tokens. We can discard + # input_ids based on the past_length. + elif past_length < input_ids.shape[1]: + input_ids = input_ids[:, past_length:] + # 3 - Otherwise (past_length >= input_ids.shape[1]), let's assume input_ids only has unprocessed tokens. + + # If we are about to go beyond the maximum cache length, we need to crop the input attention mask. + if ( + max_cache_length is not None + and attention_mask is not None + and cache_length + input_ids.shape[1] > max_cache_length + ): + attention_mask = attention_mask[:, -max_cache_length:] + + position_ids = kwargs.get("position_ids", None) + if attention_mask is not None and position_ids is None: + # create position_ids on the fly for batch generation + position_ids = attention_mask.long().cumsum(-1) - 1 + position_ids.masked_fill_(attention_mask == 0, 1) + if past_key_values: + position_ids = position_ids[:, -input_ids.shape[1] :] + + # if `inputs_embeds` are passed, we only want to use them in the 1st generation step + if inputs_embeds is not None and past_key_values is None: + model_inputs = {"inputs_embeds": inputs_embeds} + else: + model_inputs = {"input_ids": input_ids} + + model_inputs.update( + { + "position_ids": position_ids, + "past_key_values": past_key_values, + "use_cache": kwargs.get("use_cache"), + "attention_mask": attention_mask, + } + ) + return model_inputs + + @staticmethod + def _reorder_cache(past_key_values, beam_idx): + reordered_past = () + for layer_past in past_key_values: + reordered_past += ( + tuple( + past_state.index_select(0, beam_idx.to(past_state.device)) + for past_state in layer_past + ), + ) + return reordered_past + + +@add_start_docstrings( + """ + The DeepseekV3 Model transformer with a sequence classification head on top (linear layer). + + [`DeepseekV3ForSequenceClassification`] uses the last token in order to do the classification, as other causal models + (e.g. GPT-2) do. + + Since it does classification on the last token, it requires to know the position of the last token. If a + `pad_token_id` is defined in the configuration, it finds the last token that is not a padding token in each row. If + no `pad_token_id` is defined, it simply takes the last value in each row of the batch. Since it cannot guess the + padding tokens when `inputs_embeds` are passed instead of `input_ids`, it does the same (take the last value in + each row of the batch). + """, + DeepseekV3_START_DOCSTRING, +) +class DeepseekV3ForSequenceClassification(DeepseekV3PreTrainedModel): + def __init__(self, config): + super().__init__(config) + self.num_labels = config.num_labels + self.model = DeepseekV3Model(config) + self.score = nn.Linear(config.hidden_size, self.num_labels, bias=False) + + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self): + return self.model.embed_tokens + + def set_input_embeddings(self, value): + self.model.embed_tokens = value + + @add_start_docstrings_to_model_forward(DeepseekV3_INPUTS_DOCSTRING) + def forward( + self, + input_ids: torch.LongTensor = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[List[torch.FloatTensor]] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + labels: Optional[torch.LongTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + ) -> Union[Tuple, SequenceClassifierOutputWithPast]: + r""" + labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*): + Labels for computing the sequence classification/regression loss. Indices should be in `[0, transformers., + config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If + `config.num_labels > 1` a classification loss is computed (Cross-Entropy). + """ + return_dict = ( + return_dict if return_dict is not None else self.config.use_return_dict + ) + + transformer_outputs = self.model( + input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + hidden_states = transformer_outputs[0] + logits = self.score(hidden_states) + + if input_ids is not None: + batch_size = input_ids.shape[0] + else: + batch_size = inputs_embeds.shape[0] + + if self.config.pad_token_id is None and batch_size != 1: + raise ValueError( + "Cannot handle batch sizes > 1 if no padding token is defined." + ) + if self.config.pad_token_id is None: + sequence_lengths = -1 + else: + if input_ids is not None: + sequence_lengths = ( + torch.eq(input_ids, self.config.pad_token_id).int().argmax(-1) - 1 + ).to(logits.device) + else: + sequence_lengths = -1 + + pooled_logits = logits[ + torch.arange(batch_size, device=logits.device), sequence_lengths + ] + + loss = None + if labels is not None: + labels = labels.to(logits.device) + if self.config.problem_type is None: + if self.num_labels == 1: + self.config.problem_type = "regression" + elif self.num_labels > 1 and ( + labels.dtype == torch.long or labels.dtype == torch.int + ): + self.config.problem_type = "single_label_classification" + else: + self.config.problem_type = "multi_label_classification" + + if self.config.problem_type == "regression": + loss_fct = MSELoss() + if self.num_labels == 1: + loss = loss_fct(pooled_logits.squeeze(), labels.squeeze()) + else: + loss = loss_fct(pooled_logits, labels) + elif self.config.problem_type == "single_label_classification": + loss_fct = CrossEntropyLoss() + loss = loss_fct( + pooled_logits.view(-1, self.num_labels), labels.view(-1) + ) + elif self.config.problem_type == "multi_label_classification": + loss_fct = BCEWithLogitsLoss() + loss = loss_fct(pooled_logits, labels) + if not return_dict: + output = (pooled_logits,) + transformer_outputs[1:] + return ((loss,) + output) if loss is not None else output + + return SequenceClassifierOutputWithPast( + loss=loss, + logits=pooled_logits, + past_key_values=transformer_outputs.past_key_values, + hidden_states=transformer_outputs.hidden_states, + attentions=transformer_outputs.attentions, + ) diff --git a/archive/kt-sft/ktransformers/models/modeling_llama.py b/archive/kt-sft/ktransformers/models/modeling_llama.py new file mode 100644 index 0000000..9c1b4dc --- /dev/null +++ b/archive/kt-sft/ktransformers/models/modeling_llama.py @@ -0,0 +1,1745 @@ +# coding=utf-8 +# Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved. +# +# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX +# and OPT implementations in this library. It has been modified from its +# original forms to accommodate minor architectural differences compared +# to GPT-NeoX and OPT used by the Meta AI team that trained the model. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import math +from typing import List, Optional, Tuple, Union + +import torch +import torch.nn.functional as F +import torch.utils.checkpoint +from torch import nn +from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss + +from transformers.activations import ACT2FN +from transformers.cache_utils import Cache, DynamicCache, StaticCache +from transformers.modeling_attn_mask_utils import AttentionMaskConverter +from transformers.modeling_flash_attention_utils import _flash_attention_forward +from transformers.modeling_outputs import ( + BaseModelOutputWithPast, + CausalLMOutputWithPast, + QuestionAnsweringModelOutput, + SequenceClassifierOutputWithPast, + TokenClassifierOutput, +) +from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS +from transformers.modeling_utils import PreTrainedModel +from transformers.pytorch_utils import ALL_LAYERNORM_LAYERS +from transformers.utils import ( + add_start_docstrings, + add_start_docstrings_to_model_forward, + is_flash_attn_greater_or_equal_2_10, + logging, + replace_return_docstrings, +) +from .configuration_llama import LlamaConfig + +from ktransformers.util.grad_wrapper import maybe_no_grad + +logger = logging.get_logger(__name__) + +_CONFIG_FOR_DOC = "LlamaConfig" + + +class LlamaRMSNorm(nn.Module): + def __init__(self, hidden_size, eps=1e-6): + """ + LlamaRMSNorm is equivalent to T5LayerNorm + """ + super().__init__() + self.weight = nn.Parameter(torch.ones(hidden_size)) + self.variance_epsilon = eps + + def forward(self, hidden_states): + input_dtype = hidden_states.dtype + hidden_states = hidden_states.to(torch.float32) + variance = hidden_states.pow(2).mean(-1, keepdim=True) + hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon) + return self.weight * hidden_states.to(input_dtype) + + +ALL_LAYERNORM_LAYERS.append(LlamaRMSNorm) + + +class LlamaRotaryEmbedding(nn.Module): + def __init__( + self, + dim=None, + max_position_embeddings=2048, + base=10000, + device=None, + scaling_factor=1.0, + rope_type="default", + config: Optional[LlamaConfig] = None, + ): + super().__init__() + self.dim = dim + self.max_position_embeddings = max_position_embeddings + self.base = base + self.device = device + self.scaling_factor = scaling_factor + self.rope_type = rope_type + self.config = config + # TODO (joao): remove the `if` below, only used for BC + self.rope_kwargs = {} + if config is None: + logger.warning_once( + "`LlamaRotaryEmbedding` can now be fully parameterized by passing the model config through the " + "`config` argument. All other arguments will be removed in v4.45" + ) + self.rope_kwargs = { + "rope_type": rope_type, + "factor": scaling_factor, + "dim": dim, + "base": base, + "max_position_embeddings": max_position_embeddings, + } + self.rope_type = rope_type + self.max_seq_len_cached = max_position_embeddings + self.original_max_seq_len = max_position_embeddings + else: + # BC: "rope_type" was originally "type" + if config.rope_scaling is not None: + self.rope_type = config.rope_scaling.get( + "rope_type", config.rope_scaling.get("type") + ) + else: + self.rope_type = "default" + self.max_seq_len_cached = config.max_position_embeddings + self.original_max_seq_len = config.max_position_embeddings + + self.config = config + self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type] + + inv_freq, self.attention_scaling = self.rope_init_fn( + self.config, device, **self.rope_kwargs + ) + self.register_buffer("inv_freq", inv_freq, persistent=False) + self.original_inv_freq = self.inv_freq + + def _dynamic_frequency_update(self, position_ids, device): + """ + dynamic RoPE layers should recompute `inv_freq` in the following situations: + 1 - growing beyond the cached sequence length (allow scaling) + 2 - the current sequence length is in the original scale (avoid losing precision with small sequences) + """ + seq_len = torch.max(position_ids) + 1 + # seq_len = position_ids[0, -1] + 1 + if seq_len > self.max_seq_len_cached: # growth + inv_freq, self.attention_scaling = self.rope_init_fn( + self.config, device, seq_len=seq_len, **self.rope_kwargs + ) + self.register_buffer( + "inv_freq", inv_freq, persistent=False + ) # TODO joao: may break with compilation + self.max_seq_len_cached = seq_len + + if ( + seq_len < self.original_max_seq_len + and self.max_seq_len_cached > self.original_max_seq_len + ): # reset + self.register_buffer("inv_freq", self.original_inv_freq, persistent=False) + self.max_seq_len_cached = self.original_max_seq_len + + @maybe_no_grad() + def forward(self, x, position_ids): + # if "dynamic" in self.rope_type: + # self._dynamic_frequency_update(position_ids, device=x.device) + + # Core RoPE block + inv_freq_expanded = ( + self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1) + ) + position_ids_expanded = position_ids[:, None, :].float() + # Force float32 (see https://github.com/huggingface/transformers/pull/29285) + device_type = x.device.type + device_type = ( + device_type + if isinstance(device_type, str) and device_type != "mps" + else "cpu" + ) + with torch.autocast(device_type=device_type, enabled=False): + freqs = ( + inv_freq_expanded.float() @ position_ids_expanded.float() + ).transpose(1, 2) + emb = torch.cat((freqs, freqs), dim=-1) + cos = emb.cos() + sin = emb.sin() + + # Advanced RoPE types (e.g. yarn) apply a post-processing scaling factor, equivalent to scaling attention + cos = cos * self.attention_scaling + sin = sin * self.attention_scaling + + return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype) + + +class LlamaLinearScalingRotaryEmbedding(LlamaRotaryEmbedding): + """LlamaRotaryEmbedding extended with linear scaling. Credits to the Reddit user /u/kaiokendev""" + + def __init__(self, *args, **kwargs): + logger.warning_once( + "`LlamaLinearScalingRotaryEmbedding` is deprecated an will be removed in v4.45. Please use " + "`LlamaRotaryEmbedding`, which now also does linear scaling (simply pass the model config to __init__)." + ) + kwargs["rope_type"] = "linear" + super().__init__(*args, **kwargs) + + +class LlamaDynamicNTKScalingRotaryEmbedding(LlamaRotaryEmbedding): + """LlamaRotaryEmbedding extended with Dynamic NTK scaling. Credits to the Reddit users /u/bloc97 and /u/emozilla""" + + def __init__(self, *args, **kwargs): + logger.warning_once( + "`LlamaDynamicNTKScalingRotaryEmbedding` is deprecated an will be removed in v4.45. Please use " + "`LlamaRotaryEmbedding`, which now also does dynamic ntk scaling (simply pass the model config to " + "__init__)." + ) + kwargs["rope_type"] = "dynamic" + super().__init__(*args, **kwargs) + + +def rotate_half(x): + """Rotates half the hidden dims of the input.""" + x1 = x[..., : x.shape[-1] // 2] + x2 = x[..., x.shape[-1] // 2 :] + return torch.cat((-x2, x1), dim=-1) + + +def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1): + """Applies Rotary Position Embedding to the query and key tensors. + + Args: + q (`torch.Tensor`): The query tensor. + k (`torch.Tensor`): The key tensor. + cos (`torch.Tensor`): The cosine part of the rotary embedding. + sin (`torch.Tensor`): The sine part of the rotary embedding. + position_ids (`torch.Tensor`, *optional*): + Deprecated and unused. + unsqueeze_dim (`int`, *optional*, defaults to 1): + The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and + sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note + that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and + k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes + cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have + the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2. + Returns: + `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding. + """ + cos = cos.unsqueeze(unsqueeze_dim) + sin = sin.unsqueeze(unsqueeze_dim) + q_embed = (q * cos) + (rotate_half(q) * sin) + k_embed = (k * cos) + (rotate_half(k) * sin) + return q_embed, k_embed + + +class LlamaMLP(nn.Module): + def __init__(self, config): + super().__init__() + self.config = config + self.hidden_size = config.hidden_size + self.intermediate_size = config.intermediate_size + self.gate_proj = nn.Linear( + self.hidden_size, self.intermediate_size, bias=config.mlp_bias + ) + self.up_proj = nn.Linear( + self.hidden_size, self.intermediate_size, bias=config.mlp_bias + ) + self.down_proj = nn.Linear( + self.intermediate_size, self.hidden_size, bias=config.mlp_bias + ) + self.act_fn = ACT2FN[config.hidden_act] + + def forward(self, x): + if self.config.pretraining_tp > 1: + slice = self.intermediate_size // self.config.pretraining_tp + gate_proj_slices = self.gate_proj.weight.split(slice, dim=0) + up_proj_slices = self.up_proj.weight.split(slice, dim=0) + down_proj_slices = self.down_proj.weight.split(slice, dim=1) + + gate_proj = torch.cat( + [ + F.linear(x, gate_proj_slices[i]) + for i in range(self.config.pretraining_tp) + ], + dim=-1, + ) + up_proj = torch.cat( + [ + F.linear(x, up_proj_slices[i]) + for i in range(self.config.pretraining_tp) + ], + dim=-1, + ) + + intermediate_states = (self.act_fn(gate_proj) * up_proj).split(slice, dim=2) + down_proj = [ + F.linear(intermediate_states[i], down_proj_slices[i]) + for i in range(self.config.pretraining_tp) + ] + down_proj = sum(down_proj) + else: + down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x)) + + return down_proj + + +def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: + """ + This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch, + num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim) + """ + batch, num_key_value_heads, slen, head_dim = hidden_states.shape + if n_rep == 1: + return hidden_states + hidden_states = hidden_states[:, :, None, :, :].expand( + batch, num_key_value_heads, n_rep, slen, head_dim + ) + return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim) + + +class LlamaAttention(nn.Module): + """Multi-headed attention from 'Attention Is All You Need' paper""" + + def __init__(self, config: LlamaConfig, layer_idx: Optional[int] = None): + super().__init__() + self.config = config + self.layer_idx = layer_idx + if layer_idx is None: + logger.warning_once( + f"Instantiating {self.__class__.__name__} without passing a `layer_idx` is not recommended and will " + "lead to errors during the forward call if caching is used. Please make sure to provide a `layer_idx` " + "when creating this class." + ) + + self.attention_dropout = config.attention_dropout + self.hidden_size = config.hidden_size + self.num_heads = config.num_attention_heads + self.head_dim = self.hidden_size // self.num_heads + self.num_key_value_heads = config.num_key_value_heads + self.num_key_value_groups = self.num_heads // self.num_key_value_heads + self.max_position_embeddings = config.max_position_embeddings + self.rope_theta = config.rope_theta + self.is_causal = True + + if (self.head_dim * self.num_heads) != self.hidden_size: + raise ValueError( + f"hidden_size must be divisible by num_heads (got `hidden_size`: {self.hidden_size}" + f" and `num_heads`: {self.num_heads})." + ) + + self.q_proj = nn.Linear( + self.hidden_size, self.num_heads * self.head_dim, bias=config.attention_bias + ) + self.k_proj = nn.Linear( + self.hidden_size, + self.num_key_value_heads * self.head_dim, + bias=config.attention_bias, + ) + self.v_proj = nn.Linear( + self.hidden_size, + self.num_key_value_heads * self.head_dim, + bias=config.attention_bias, + ) + self.o_proj = nn.Linear( + self.hidden_size, self.hidden_size, bias=config.attention_bias + ) + + # TODO (joao): remove in v4.45 (RoPE is computed in the model, not in the decoder layers) + self.rotary_emb = LlamaRotaryEmbedding(config=self.config) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_value: Optional[Cache] = None, + output_attentions: bool = False, + use_cache: bool = False, + cache_position: Optional[torch.LongTensor] = None, + position_embeddings: Optional[ + Tuple[torch.Tensor, torch.Tensor] + ] = None, # will become mandatory in v4.45 + **kwargs, + ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: + bsz, q_len, _ = hidden_states.size() + + if self.config.pretraining_tp > 1: + key_value_slicing = ( + self.num_key_value_heads * self.head_dim + ) // self.config.pretraining_tp + query_slices = self.q_proj.weight.split( + (self.num_heads * self.head_dim) // self.config.pretraining_tp, dim=0 + ) + key_slices = self.k_proj.weight.split(key_value_slicing, dim=0) + value_slices = self.v_proj.weight.split(key_value_slicing, dim=0) + + query_states = [ + F.linear(hidden_states, query_slices[i]) + for i in range(self.config.pretraining_tp) + ] + query_states = torch.cat(query_states, dim=-1) + + key_states = [ + F.linear(hidden_states, key_slices[i]) + for i in range(self.config.pretraining_tp) + ] + key_states = torch.cat(key_states, dim=-1) + + value_states = [ + F.linear(hidden_states, value_slices[i]) + for i in range(self.config.pretraining_tp) + ] + value_states = torch.cat(value_states, dim=-1) + + else: + query_states = self.q_proj(hidden_states) + key_states = self.k_proj(hidden_states) + value_states = self.v_proj(hidden_states) + + query_states = query_states.view( + bsz, q_len, self.num_heads, self.head_dim + ).transpose(1, 2) + key_states = key_states.view( + bsz, q_len, self.num_key_value_heads, self.head_dim + ).transpose(1, 2) + value_states = value_states.view( + bsz, q_len, self.num_key_value_heads, self.head_dim + ).transpose(1, 2) + + if position_embeddings is None: + logger.warning_once( + "The attention layers in this model are transitioning from computing the RoPE embeddings internally " + "through `position_ids` (2D tensor with the indexes of the tokens), to using externally computed " + "`position_embeddings` (Tuple of tensors, containing cos and sin). In v4.45 `position_ids` will be " + "removed and `position_embeddings` will be mandatory." + ) + cos, sin = self.rotary_emb(value_states, position_ids) + else: + cos, sin = position_embeddings + query_states, key_states = apply_rotary_pos_emb( + query_states, key_states, cos, sin + ) + + if past_key_value is not None: + # sin and cos are specific to RoPE models; cache_position needed for the static cache + cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position} + key_states, value_states = past_key_value.update( + key_states, value_states, self.layer_idx, cache_kwargs + ) + + key_states = repeat_kv(key_states, self.num_key_value_groups) + value_states = repeat_kv(value_states, self.num_key_value_groups) + + attn_weights = torch.matmul( + query_states, key_states.transpose(2, 3) + ) / math.sqrt(self.head_dim) + + if attention_mask is not None: # no matter the length, we just slice it + causal_mask = attention_mask[:, :, :, : key_states.shape[-2]] + attn_weights = attn_weights + causal_mask + + # upcast attention to fp32 + attn_weights = nn.functional.softmax( + attn_weights, dim=-1, dtype=torch.float32 + ).to(query_states.dtype) + attn_weights = nn.functional.dropout( + attn_weights, p=self.attention_dropout, training=self.training + ) + attn_output = torch.matmul(attn_weights, value_states) + + if attn_output.size() != (bsz, self.num_heads, q_len, self.head_dim): + raise ValueError( + f"`attn_output` should be of size {(bsz, self.num_heads, q_len, self.head_dim)}, but is" + f" {attn_output.size()}" + ) + + attn_output = attn_output.transpose(1, 2).contiguous() + + attn_output = attn_output.reshape(bsz, q_len, -1) + + if self.config.pretraining_tp > 1: + attn_output = attn_output.split( + self.hidden_size // self.config.pretraining_tp, dim=2 + ) + o_proj_slices = self.o_proj.weight.split( + self.hidden_size // self.config.pretraining_tp, dim=1 + ) + attn_output = sum( + [ + F.linear(attn_output[i], o_proj_slices[i]) + for i in range(self.config.pretraining_tp) + ] + ) + else: + attn_output = self.o_proj(attn_output) + + if not output_attentions: + attn_weights = None + + return attn_output, attn_weights, past_key_value + + +class LlamaFlashAttention2(LlamaAttention): + """ + Llama flash attention module. This module inherits from `LlamaAttention` as the weights of the module stays + untouched. The only required change would be on the forward pass where it needs to correctly call the public API of + flash attention and deal with padding tokens in case the input contains any of them. + """ + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + + # TODO: Should be removed once Flash Attention for RoCm is bumped to 2.1. + # flash_attn<2.1 generates top-left aligned causal mask, while what is needed here is bottom-right alignement, that was made default for flash_attn>=2.1. This attribute is used to handle this difference. Reference: https://github.com/Dao-AILab/flash-attention/releases/tag/v2.1.0. + # Beware that with flash_attn<2.1, using q_seqlen != k_seqlen (except for the case q_seqlen == 1) produces a wrong mask (top-left). + self._flash_attn_uses_top_left_mask = not is_flash_attn_greater_or_equal_2_10() + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.LongTensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_value: Optional[Cache] = None, + output_attentions: bool = False, + use_cache: bool = False, + cache_position: Optional[torch.LongTensor] = None, + position_embeddings: Optional[ + Tuple[torch.Tensor, torch.Tensor] + ] = None, # will become mandatory in v4.45 + ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: + if isinstance(past_key_value, StaticCache): + raise ValueError( + "`static` cache implementation is not compatible with `attn_implementation==flash_attention_2` " + "make sure to use `sdpa` in the mean time, and open an issue at https://github.com/huggingface/transformers" + ) + + output_attentions = False + + bsz, q_len, _ = hidden_states.size() + + query_states = self.q_proj(hidden_states) + key_states = self.k_proj(hidden_states) + value_states = self.v_proj(hidden_states) + + # Flash attention requires the input to have the shape + # batch_size x seq_length x head_dim x hidden_dim + # therefore we just need to keep the original shape + query_states = query_states.view( + bsz, q_len, self.num_heads, self.head_dim + ).transpose(1, 2) + key_states = key_states.view( + bsz, q_len, self.num_key_value_heads, self.head_dim + ).transpose(1, 2) + value_states = value_states.view( + bsz, q_len, self.num_key_value_heads, self.head_dim + ).transpose(1, 2) + + if position_embeddings is None: + logger.warning_once( + "The attention layers in this model are transitioning from computing the RoPE embeddings internally " + "through `position_ids` (2D tensor with the indexes of the tokens), to using externally computed " + "`position_embeddings` (Tuple of tensors, containing cos and sin). In v4.45 `position_ids` will be " + "removed and `position_embeddings` will be mandatory." + ) + cos, sin = self.rotary_emb(value_states, position_ids) + else: + cos, sin = position_embeddings + query_states, key_states = apply_rotary_pos_emb( + query_states, key_states, cos, sin + ) + + if past_key_value is not None: + # sin and cos are specific to RoPE models; cache_position needed for the static cache + cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position} + key_states, value_states = past_key_value.update( + key_states, value_states, self.layer_idx, cache_kwargs + ) + + # TODO: These transpose are quite inefficient but Flash Attention requires the layout [batch_size, sequence_length, num_heads, head_dim]. We would need to refactor the KV cache + # to be able to avoid many of these transpose/reshape/view. + query_states = query_states.transpose(1, 2) + key_states = key_states.transpose(1, 2) + value_states = value_states.transpose(1, 2) + + dropout_rate = self.attention_dropout if self.training else 0.0 + + # In PEFT, usually we cast the layer norms in float32 for training stability reasons + # therefore the input hidden states gets silently casted in float32. Hence, we need + # cast them back in the correct dtype just to be sure everything works as expected. + # This might slowdown training & inference so it is recommended to not cast the LayerNorms + # in fp32. (LlamaRMSNorm handles it correctly) + + input_dtype = query_states.dtype + if input_dtype == torch.float32: + if torch.is_autocast_enabled(): + target_dtype = torch.get_autocast_gpu_dtype() + # Handle the case where the model is quantized + elif hasattr(self.config, "_pre_quantization_dtype"): + target_dtype = self.config._pre_quantization_dtype + else: + target_dtype = self.q_proj.weight.dtype + + logger.warning_once( + f"The input hidden states seems to be silently casted in float32, this might be related to" + f" the fact you have upcasted embedding or layer norm layers in float32. We will cast back the input in" + f" {target_dtype}." + ) + + query_states = query_states.to(target_dtype) + key_states = key_states.to(target_dtype) + value_states = value_states.to(target_dtype) + + attn_output = _flash_attention_forward( + query_states, + key_states, + value_states, + attention_mask, + q_len, + dropout=dropout_rate, + sliding_window=getattr(self, "sliding_window", None), + use_top_left_mask=self._flash_attn_uses_top_left_mask, + is_causal=self.is_causal, + ) + + attn_output = attn_output.reshape(bsz, q_len, -1).contiguous() + attn_output = self.o_proj(attn_output) + + if not output_attentions: + attn_weights = None + + return attn_output, attn_weights, past_key_value + + +class LlamaSdpaAttention(LlamaAttention): + """ + Llama attention module using torch.nn.functional.scaled_dot_product_attention. This module inherits from + `LlamaAttention` as the weights of the module stays untouched. The only changes are on the forward pass to adapt to + SDPA API. + """ + + # Adapted from LlamaAttention.forward + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_value: Optional[Cache] = None, + output_attentions: bool = False, + use_cache: bool = False, + cache_position: Optional[torch.LongTensor] = None, + position_embeddings: Optional[ + Tuple[torch.Tensor, torch.Tensor] + ] = None, # will become mandatory in v4.45 + **kwargs, + ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: + if output_attentions: + # TODO: Improve this warning with e.g. `model.config.attn_implementation = "manual"` once this is implemented. + logger.warning_once( + "LlamaModel is using LlamaSdpaAttention, but `torch.nn.functional.scaled_dot_product_attention` does not support `output_attentions=True`. Falling back to the manual attention implementation, " + 'but specifying the manual implementation will be required from Transformers version v5.0.0 onwards. This warning can be removed using the argument `attn_implementation="eager"` when loading the model.' + ) + return super().forward( + hidden_states=hidden_states, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_value=past_key_value, + output_attentions=output_attentions, + use_cache=use_cache, + cache_position=cache_position, + position_embeddings=position_embeddings, + ) + + bsz, q_len, _ = hidden_states.size() + + query_states = self.q_proj(hidden_states) + key_states = self.k_proj(hidden_states) + value_states = self.v_proj(hidden_states) + + query_states = query_states.view( + bsz, q_len, self.num_heads, self.head_dim + ).transpose(1, 2) + key_states = key_states.view( + bsz, q_len, self.num_key_value_heads, self.head_dim + ).transpose(1, 2) + value_states = value_states.view( + bsz, q_len, self.num_key_value_heads, self.head_dim + ).transpose(1, 2) + + if position_embeddings is None: + logger.warning_once( + "The attention layers in this model are transitioning from computing the RoPE embeddings internally " + "through `position_ids` (2D tensor with the indexes of the tokens), to using externally computed " + "`position_embeddings` (Tuple of tensors, containing cos and sin). In v4.45 `position_ids` will be " + "removed and `position_embeddings` will be mandatory." + ) + cos, sin = self.rotary_emb(value_states, position_ids) + else: + cos, sin = position_embeddings + query_states, key_states = apply_rotary_pos_emb( + query_states, key_states, cos, sin + ) + + if past_key_value is not None: + # sin and cos are specific to RoPE models; cache_position needed for the static cache + cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position} + key_states, value_states = past_key_value.update( + key_states, value_states, self.layer_idx, cache_kwargs + ) + + key_states = repeat_kv(key_states, self.num_key_value_groups) + value_states = repeat_kv(value_states, self.num_key_value_groups) + + causal_mask = attention_mask + if attention_mask is not None: + causal_mask = causal_mask[:, :, :, : key_states.shape[-2]] + + # SDPA with memory-efficient backend is currently (torch==2.1.2) bugged with non-contiguous inputs with custom attn_mask, + # Reference: https://github.com/pytorch/pytorch/issues/112577. + if query_states.device.type == "cuda" and causal_mask is not None: + query_states = query_states.contiguous() + key_states = key_states.contiguous() + value_states = value_states.contiguous() + + # We dispatch to SDPA's Flash Attention or Efficient kernels via this `is_causal` if statement instead of an inline conditional assignment + # in SDPA to support both torch.compile's dynamic shapes and full graph options. An inline conditional prevents dynamic shapes from compiling. + is_causal = True if causal_mask is None and q_len > 1 else False + + attn_output = torch.nn.functional.scaled_dot_product_attention( + query_states, + key_states, + value_states, + attn_mask=causal_mask, + dropout_p=self.attention_dropout if self.training else 0.0, + is_causal=is_causal, + ) + + attn_output = attn_output.transpose(1, 2).contiguous() + attn_output = attn_output.view(bsz, q_len, -1) + + attn_output = self.o_proj(attn_output) + + return attn_output, None, past_key_value + + +LLAMA_ATTENTION_CLASSES = { + "eager": LlamaAttention, + "flash_attention_2": LlamaFlashAttention2, + "sdpa": LlamaSdpaAttention, +} + + +class LlamaDecoderLayer(nn.Module): + def __init__(self, config: LlamaConfig, layer_idx: int): + super().__init__() + self.hidden_size = config.hidden_size + + self.self_attn = LLAMA_ATTENTION_CLASSES[config._attn_implementation]( + config=config, layer_idx=layer_idx + ) + + self.mlp = LlamaMLP(config) + self.input_layernorm = LlamaRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = LlamaRMSNorm( + config.hidden_size, eps=config.rms_norm_eps + ) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_value: Optional[Cache] = None, + output_attentions: Optional[bool] = False, + use_cache: Optional[bool] = False, + cache_position: Optional[torch.LongTensor] = None, + position_embeddings: Optional[ + Tuple[torch.Tensor, torch.Tensor] + ] = None, # will become mandatory in v4.45 + **kwargs, + ) -> Tuple[ + torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]] + ]: + """ + Args: + hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)` + attention_mask (`torch.FloatTensor`, *optional*): + attention mask of size `(batch_size, sequence_length)` if flash attention is used or `(batch_size, 1, + query_sequence_length, key_sequence_length)` if default attention is used. + output_attentions (`bool`, *optional*): + Whether or not to return the attentions tensors of all attention layers. See `attentions` under + returned tensors for more detail. + use_cache (`bool`, *optional*): + If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding + (see `past_key_values`). + past_key_value (`Tuple(torch.FloatTensor)`, *optional*): cached past key and value projection states + cache_position (`torch.LongTensor` of shape `(sequence_length)`, *optional*): + Indices depicting the position of the input sequence tokens in the sequence + position_embeddings (`Tuple[torch.FloatTensor, torch.FloatTensor]`, *optional*): + Tuple containing the cosine and sine positional embeddings of shape `(batch_size, seq_len, head_dim)`, + with `head_dim` being the embedding dimension of each attention head. + kwargs (`dict`, *optional*): + Arbitrary kwargs to be ignored, used for FSDP and other methods that injects code + into the model + """ + residual = hidden_states + + hidden_states = self.input_layernorm(hidden_states) + + # Self Attention + hidden_states, self_attn_weights, present_key_value = self.self_attn( + hidden_states=hidden_states, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_value=past_key_value, + output_attentions=output_attentions, + use_cache=use_cache, + cache_position=cache_position, + position_embeddings=position_embeddings, + **kwargs, + ) + hidden_states = residual + hidden_states + + # Fully Connected + residual = hidden_states + hidden_states = self.post_attention_layernorm(hidden_states) + hidden_states = self.mlp(hidden_states) + hidden_states = residual + hidden_states + + outputs = (hidden_states,) + + if output_attentions: + outputs += (self_attn_weights,) + + if use_cache: + outputs += (present_key_value,) + + return outputs + + +LLAMA_START_DOCSTRING = r""" + This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the + library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads + etc.) + + This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass. + Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage + and behavior. + + Parameters: + config ([`LlamaConfig`]): + Model configuration class with all the parameters of the model. Initializing with a config file does not + load the weights associated with the model, only the configuration. Check out the + [`~PreTrainedModel.from_pretrained`] method to load the model weights. +""" + + +@add_start_docstrings( + "The bare LLaMA Model outputting raw hidden-states without any specific head on top.", + LLAMA_START_DOCSTRING, +) +class LlamaPreTrainedModel(PreTrainedModel): + config_class = LlamaConfig + base_model_prefix = "model" + supports_gradient_checkpointing = True + _no_split_modules = ["LlamaDecoderLayer"] + _skip_keys_device_placement = ["past_key_values"] + _supports_flash_attn_2 = True + _supports_sdpa = True + _supports_cache_class = True + _supports_quantized_cache = True + _supports_static_cache = True + + def _init_weights(self, module): + std = self.config.initializer_range + if isinstance(module, nn.Linear): + module.weight.data.normal_(mean=0.0, std=std) + if module.bias is not None: + module.bias.data.zero_() + elif isinstance(module, nn.Embedding): + module.weight.data.normal_(mean=0.0, std=std) + if module.padding_idx is not None: + module.weight.data[module.padding_idx].zero_() + + +LLAMA_INPUTS_DOCSTRING = r""" + Args: + input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): + Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide + it. + + Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and + [`PreTrainedTokenizer.__call__`] for details. + + [What are input IDs?](../glossary#input-ids) + attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*): + Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`: + + - 1 for tokens that are **not masked**, + - 0 for tokens that are **masked**. + + [What are attention masks?](../glossary#attention-mask) + + Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and + [`PreTrainedTokenizer.__call__`] for details. + + If `past_key_values` is used, optionally only the last `input_ids` have to be input (see + `past_key_values`). + + If you want to change padding behavior, you should read [`modeling_opt._prepare_decoder_attention_mask`] + and modify to your needs. See diagram 1 in [the paper](https://arxiv.org/abs/1910.13461) for more + information on the default strategy. + + - 1 indicates the head is **not masked**, + - 0 indicates the head is **masked**. + position_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0, + config.n_positions - 1]`. + + [What are position IDs?](../glossary#position-ids) + past_key_values (`Cache` or `tuple(tuple(torch.FloatTensor))`, *optional*): + Pre-computed hidden-states (key and values in the self-attention blocks and in the cross-attention + blocks) that can be used to speed up sequential decoding. This typically consists in the `past_key_values` + returned by the model at a previous stage of decoding, when `use_cache=True` or `config.use_cache=True`. + + Two formats are allowed: + - a [`~cache_utils.Cache`] instance; + - Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of + shape `(batch_size, num_heads, sequence_length, embed_size_per_head)`). This is also known as the legacy + cache format. + + The model will output the same cache format that is fed as input. If no `past_key_values` are passed, the + legacy cache format will be returned. + + If `past_key_values` are used, the user can optionally input only the last `input_ids` (those that don't + have their past key value states given to this model) of shape `(batch_size, 1)` instead of all `input_ids` + of shape `(batch_size, sequence_length)`. + inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*): + Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This + is useful if you want more control over how to convert `input_ids` indices into associated vectors than the + model's internal embedding lookup matrix. + use_cache (`bool`, *optional*): + If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see + `past_key_values`). + output_attentions (`bool`, *optional*): + Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned + tensors for more detail. + output_hidden_states (`bool`, *optional*): + Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for + more detail. + return_dict (`bool`, *optional*): + Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple. + cache_position (`torch.LongTensor` of shape `(sequence_length)`, *optional*): + Indices depicting the position of the input sequence tokens in the sequence. Contrarily to `position_ids`, + this tensor is not affected by padding. It is used to update the cache in the correct position and to infer + the complete sequence length. +""" + + +@add_start_docstrings( + "The bare LLaMA Model outputting raw hidden-states without any specific head on top.", + LLAMA_START_DOCSTRING, +) +class LlamaModel(LlamaPreTrainedModel): + """ + Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`LlamaDecoderLayer`] + + Args: + config: LlamaConfig + """ + + def __init__(self, config: LlamaConfig): + super().__init__(config) + self.padding_idx = config.pad_token_id + self.vocab_size = config.vocab_size + + self.embed_tokens = nn.Embedding( + config.vocab_size, config.hidden_size, self.padding_idx + ) + self.layers = nn.ModuleList( + [ + LlamaDecoderLayer(config, layer_idx) + for layer_idx in range(config.num_hidden_layers) + ] + ) + self.norm = LlamaRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.rotary_emb = LlamaRotaryEmbedding(config=config) + self.gradient_checkpointing = False + + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self): + return self.embed_tokens + + def set_input_embeddings(self, value): + self.embed_tokens = value + + @add_start_docstrings_to_model_forward(LLAMA_INPUTS_DOCSTRING) + def forward( + self, + input_ids: torch.LongTensor = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[Union[Cache, List[torch.FloatTensor]]] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + cache_position: Optional[torch.LongTensor] = None, + ) -> Union[Tuple, BaseModelOutputWithPast]: + output_attentions = ( + output_attentions + if output_attentions is not None + else self.config.output_attentions + ) + output_hidden_states = ( + output_hidden_states + if output_hidden_states is not None + else self.config.output_hidden_states + ) + use_cache = use_cache if use_cache is not None else self.config.use_cache + return_dict = ( + return_dict if return_dict is not None else self.config.use_return_dict + ) + + if (input_ids is None) ^ (inputs_embeds is not None): + raise ValueError( + "You cannot specify both input_ids and inputs_embeds at the same time, and must specify either one" + ) + + if self.gradient_checkpointing and self.training and use_cache: + logger.warning_once( + "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`." + ) + use_cache = False + + if inputs_embeds is None: + inputs_embeds = self.embed_tokens(input_ids) + + return_legacy_cache = False + if ( + use_cache and not isinstance(past_key_values, Cache) and not self.training + ): # kept for BC (non `Cache` `past_key_values` inputs) + return_legacy_cache = True + past_key_values = DynamicCache.from_legacy_cache(past_key_values) + logger.warning_once( + "We detected that you are passing `past_key_values` as a tuple and this is deprecated and will be removed in v4.43. " + "Please use an appropriate `Cache` class (https://huggingface.co/docs/transformers/v4.41.3/en/internal/generation_utils#transformers.Cache)" + ) + + if cache_position is None: + past_seen_tokens = ( + past_key_values.get_seq_length() if past_key_values is not None else 0 + ) + cache_position = torch.arange( + past_seen_tokens, + past_seen_tokens + inputs_embeds.shape[1], + device=inputs_embeds.device, + ) + if position_ids is None: + position_ids = cache_position.unsqueeze(0) + + causal_mask = self._update_causal_mask( + attention_mask, + inputs_embeds, + cache_position, + past_key_values, + output_attentions, + ) + hidden_states = inputs_embeds + + # create position embeddings to be shared across the decoder layers + position_embeddings = self.rotary_emb(hidden_states, position_ids) + + # decoder layers + all_hidden_states = () if output_hidden_states else None + all_self_attns = () if output_attentions else None + next_decoder_cache = None + + for decoder_layer in self.layers: + if output_hidden_states: + all_hidden_states += (hidden_states,) + + if self.gradient_checkpointing and self.training: + layer_outputs = self._gradient_checkpointing_func( + decoder_layer.__call__, + hidden_states, + causal_mask, + position_ids, + past_key_values, + output_attentions, + use_cache, + cache_position, + position_embeddings, + ) + else: + layer_outputs = decoder_layer( + hidden_states, + attention_mask=causal_mask, + position_ids=position_ids, + past_key_value=past_key_values, + output_attentions=output_attentions, + use_cache=use_cache, + cache_position=cache_position, + position_embeddings=position_embeddings, + ) + + hidden_states = layer_outputs[0] + + if use_cache: + next_decoder_cache = layer_outputs[2 if output_attentions else 1] + + if output_attentions: + all_self_attns += (layer_outputs[1],) + + hidden_states = self.norm(hidden_states) + + # add hidden states from the last decoder layer + if output_hidden_states: + all_hidden_states += (hidden_states,) + + next_cache = next_decoder_cache if use_cache else None + if return_legacy_cache: + next_cache = next_cache.to_legacy_cache() + + if not return_dict: + return tuple( + v + for v in [hidden_states, next_cache, all_hidden_states, all_self_attns] + if v is not None + ) + return BaseModelOutputWithPast( + last_hidden_state=hidden_states, + past_key_values=next_cache, + hidden_states=all_hidden_states, + attentions=all_self_attns, + ) + + def _update_causal_mask( + self, + attention_mask: torch.Tensor, + input_tensor: torch.Tensor, + cache_position: torch.Tensor, + past_key_values: Cache, + output_attentions: bool, + ): + # TODO: As of torch==2.2.0, the `attention_mask` passed to the model in `generate` is 2D and of dynamic length even when the static + # KV cache is used. This is an issue for torch.compile which then recaptures cudagraphs at each decode steps due to the dynamic shapes. + # (`recording cudagraph tree for symint key 13`, etc.), which is VERY slow. A workaround is `@torch.compiler.disable`, but this prevents using + # `fullgraph=True`. See more context in https://github.com/huggingface/transformers/pull/29114 + + if self.config._attn_implementation == "flash_attention_2": + if attention_mask is not None and 0.0 in attention_mask: + return attention_mask + return None + + # For SDPA, when possible, we will rely on its `is_causal` argument instead of its `attn_mask` argument, in + # order to dispatch on Flash Attention 2. This feature is not compatible with static cache, as SDPA will fail + # to infer the attention mask. + past_seen_tokens = ( + past_key_values.get_seq_length() if past_key_values is not None else 0 + ) + using_static_cache = isinstance(past_key_values, StaticCache) + + # When output attentions is True, sdpa implementation's forward method calls the eager implementation's forward + if ( + self.config._attn_implementation == "sdpa" + and not using_static_cache + and not output_attentions + ): + if AttentionMaskConverter._ignore_causal_mask_sdpa( + attention_mask, + inputs_embeds=input_tensor, + past_key_values_length=past_seen_tokens, + is_training=self.training, + ): + return None + + dtype, device = input_tensor.dtype, input_tensor.device + min_dtype = torch.finfo(dtype).min + sequence_length = input_tensor.shape[1] + if using_static_cache: + target_length = past_key_values.get_max_length() + else: + target_length = ( + attention_mask.shape[-1] + if isinstance(attention_mask, torch.Tensor) + else past_seen_tokens + sequence_length + 1 + ) + + if attention_mask is not None and attention_mask.dim() == 4: + # in this case we assume that the mask comes already in inverted form and requires no inversion or slicing + if attention_mask.max() != 0: + raise ValueError( + "Custom 4D attention mask should be passed in inverted form with max==0`" + ) + causal_mask = attention_mask + else: + causal_mask = torch.full( + (sequence_length, target_length), + fill_value=min_dtype, + dtype=dtype, + device=device, + ) + if sequence_length != 1: + causal_mask = torch.triu(causal_mask, diagonal=1) + causal_mask *= torch.arange( + target_length, device=device + ) > cache_position.reshape(-1, 1) + causal_mask = causal_mask[None, None, :, :].expand( + input_tensor.shape[0], 1, -1, -1 + ) + if attention_mask is not None: + causal_mask = ( + causal_mask.clone() + ) # copy to contiguous memory for in-place edit + mask_length = attention_mask.shape[-1] + padding_mask = ( + causal_mask[:, :, :, :mask_length] + + attention_mask[:, None, None, :] + ) + padding_mask = padding_mask == 0 + causal_mask[:, :, :, :mask_length] = causal_mask[ + :, :, :, :mask_length + ].masked_fill(padding_mask, min_dtype) + if ( + self.config._attn_implementation == "sdpa" + and attention_mask is not None + and attention_mask.device.type == "cuda" + and not output_attentions + ): + # Attend to all tokens in fully masked rows in the causal_mask, for example the relevant first rows when + # using left padding. This is required by F.scaled_dot_product_attention memory-efficient attention path. + # Details: https://github.com/pytorch/pytorch/issues/110213 + causal_mask = AttentionMaskConverter._unmask_unattended( + causal_mask, min_dtype + ) + + return causal_mask + + +class LlamaForCausalLM(LlamaPreTrainedModel): + _tied_weights_keys = ["lm_head.weight"] + + def __init__(self, config): + super().__init__(config) + self.model = LlamaModel(config) + self.vocab_size = config.vocab_size + self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) + + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self): + return self.model.embed_tokens + + def set_input_embeddings(self, value): + self.model.embed_tokens = value + + def get_output_embeddings(self): + return self.lm_head + + def set_output_embeddings(self, new_embeddings): + self.lm_head = new_embeddings + + def set_decoder(self, decoder): + self.model = decoder + + def get_decoder(self): + return self.model + + @add_start_docstrings_to_model_forward(LLAMA_INPUTS_DOCSTRING) + @replace_return_docstrings( + output_type=CausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC + ) + def forward( + self, + input_ids: torch.LongTensor = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[Union[Cache, List[torch.FloatTensor]]] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + labels: Optional[torch.LongTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + cache_position: Optional[torch.LongTensor] = None, + ) -> Union[Tuple, CausalLMOutputWithPast]: + r""" + Args: + labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Labels for computing the masked language modeling loss. Indices should either be in `[0, ..., + config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored + (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`. + + Returns: + + Example: + + ```python + >>> from transformers import AutoTokenizer, LlamaForCausalLM + + >>> model = LlamaForCausalLM.from_pretrained("meta-llama/Llama-2-7b-hf") + >>> tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-hf") + + >>> prompt = "Hey, are you conscious? Can you talk to me?" + >>> inputs = tokenizer(prompt, return_tensors="pt") + + >>> # Generate + >>> generate_ids = model.generate(inputs.input_ids, max_length=30) + >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0] + "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you." + ```""" + output_attentions = ( + output_attentions + if output_attentions is not None + else self.config.output_attentions + ) + output_hidden_states = ( + output_hidden_states + if output_hidden_states is not None + else self.config.output_hidden_states + ) + return_dict = ( + return_dict if return_dict is not None else self.config.use_return_dict + ) + + # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn) + outputs = self.model( + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + cache_position=cache_position, + ) + + hidden_states = outputs[0] + if self.config.pretraining_tp > 1: + lm_head_slices = self.lm_head.weight.split( + self.vocab_size // self.config.pretraining_tp, dim=0 + ) + logits = [ + F.linear(hidden_states, lm_head_slices[i]) + for i in range(self.config.pretraining_tp) + ] + logits = torch.cat(logits, dim=-1) + else: + logits = self.lm_head(hidden_states) + # logits = logits.float() + + loss = None + if labels is not None: + # Shift so that tokens < n predict n + shift_logits = logits[..., :-1, :].contiguous() + shift_labels = labels[..., 1:].contiguous() + # Flatten the tokens + loss_fct = CrossEntropyLoss() + shift_logits = shift_logits.view(-1, self.config.vocab_size) + shift_labels = shift_labels.view(-1) + # Enable model parallelism + shift_labels = shift_labels.to(shift_logits.device) + loss = loss_fct(shift_logits, shift_labels) + + if not return_dict: + output = (logits,) + outputs[1:] + return (loss,) + output if loss is not None else output + + return CausalLMOutputWithPast( + loss=loss, + logits=logits, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) + + def prepare_inputs_for_generation( + self, + input_ids, + past_key_values=None, + attention_mask=None, + inputs_embeds=None, + cache_position=None, + position_ids=None, + use_cache=True, + **kwargs, + ): + # If we have cache: let's slice `input_ids` through `cache_position`, to keep only the unprocessed tokens + # Exception 1: when passing input_embeds, input_ids may be missing entries + # Exception 2: some generation methods do special slicing of input_ids, so we don't need to do it here + if past_key_values is not None: + if inputs_embeds is not None: # Exception 1 + input_ids = input_ids[:, -cache_position.shape[0] :] + elif ( + input_ids.shape[1] != cache_position.shape[0] + ): # Default case (the "else", a no op, is Exception 2) + input_ids = input_ids[:, cache_position] + + if attention_mask is not None and position_ids is None: + # create position_ids on the fly for batch generation + position_ids = attention_mask.long().cumsum(-1) - 1 + position_ids.masked_fill_(attention_mask == 0, 1) + if past_key_values: + position_ids = position_ids[:, -input_ids.shape[1] :] + + # if `inputs_embeds` are passed, we only want to use them in the 1st generation step + if inputs_embeds is not None and cache_position[0] == 0: + model_inputs = {"inputs_embeds": inputs_embeds} + else: + model_inputs = { + "input_ids": input_ids.contiguous() + } # `contiguous()` needed for compilation use cases + + model_inputs.update( + { + "position_ids": position_ids, + "cache_position": cache_position, + "past_key_values": past_key_values, + "use_cache": use_cache, + "attention_mask": attention_mask, + } + ) + return model_inputs + + +@add_start_docstrings( + """ + The LLaMa Model transformer with a sequence classification head on top (linear layer). + + [`LlamaForSequenceClassification`] uses the last token in order to do the classification, as other causal models + (e.g. GPT-2) do. + + Since it does classification on the last token, it requires to know the position of the last token. If a + `pad_token_id` is defined in the configuration, it finds the last token that is not a padding token in each row. If + no `pad_token_id` is defined, it simply takes the last value in each row of the batch. Since it cannot guess the + padding tokens when `inputs_embeds` are passed instead of `input_ids`, it does the same (take the last value in + each row of the batch). + """, + LLAMA_START_DOCSTRING, +) +class LlamaForSequenceClassification(LlamaPreTrainedModel): + def __init__(self, config): + super().__init__(config) + self.num_labels = config.num_labels + self.model = LlamaModel(config) + self.score = nn.Linear(config.hidden_size, self.num_labels, bias=False) + + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self): + return self.model.embed_tokens + + def set_input_embeddings(self, value): + self.model.embed_tokens = value + + @add_start_docstrings_to_model_forward(LLAMA_INPUTS_DOCSTRING) + def forward( + self, + input_ids: Optional[torch.LongTensor] = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[Union[Cache, List[torch.FloatTensor]]] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + labels: Optional[torch.LongTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + ) -> Union[Tuple, SequenceClassifierOutputWithPast]: + r""" + labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*): + Labels for computing the sequence classification/regression loss. Indices should be in `[0, ..., + config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If + `config.num_labels > 1` a classification loss is computed (Cross-Entropy). + """ + return_dict = ( + return_dict if return_dict is not None else self.config.use_return_dict + ) + + transformer_outputs = self.model( + input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + hidden_states = transformer_outputs[0] + logits = self.score(hidden_states) + + if input_ids is not None: + batch_size = input_ids.shape[0] + else: + batch_size = inputs_embeds.shape[0] + + if self.config.pad_token_id is None and batch_size != 1: + raise ValueError( + "Cannot handle batch sizes > 1 if no padding token is defined." + ) + if self.config.pad_token_id is None: + sequence_lengths = -1 + else: + if input_ids is not None: + # if no pad token found, use modulo instead of reverse indexing for ONNX compatibility + sequence_lengths = ( + torch.eq(input_ids, self.config.pad_token_id).int().argmax(-1) - 1 + ) + sequence_lengths = sequence_lengths % input_ids.shape[-1] + sequence_lengths = sequence_lengths.to(logits.device) + else: + sequence_lengths = -1 + + pooled_logits = logits[ + torch.arange(batch_size, device=logits.device), sequence_lengths + ] + + loss = None + if labels is not None: + labels = labels.to(logits.device) + if self.config.problem_type is None: + if self.num_labels == 1: + self.config.problem_type = "regression" + elif self.num_labels > 1 and ( + labels.dtype == torch.long or labels.dtype == torch.int + ): + self.config.problem_type = "single_label_classification" + else: + self.config.problem_type = "multi_label_classification" + + if self.config.problem_type == "regression": + loss_fct = MSELoss() + if self.num_labels == 1: + loss = loss_fct(pooled_logits.squeeze(), labels.squeeze()) + else: + loss = loss_fct(pooled_logits, labels) + elif self.config.problem_type == "single_label_classification": + loss_fct = CrossEntropyLoss() + loss = loss_fct( + pooled_logits.view(-1, self.num_labels), labels.view(-1) + ) + elif self.config.problem_type == "multi_label_classification": + loss_fct = BCEWithLogitsLoss() + loss = loss_fct(pooled_logits, labels) + if not return_dict: + output = (pooled_logits,) + transformer_outputs[1:] + return ((loss,) + output) if loss is not None else output + + return SequenceClassifierOutputWithPast( + loss=loss, + logits=pooled_logits, + past_key_values=transformer_outputs.past_key_values, + hidden_states=transformer_outputs.hidden_states, + attentions=transformer_outputs.attentions, + ) + + +@add_start_docstrings( + """ +The Llama Model transformer with a span classification head on top for extractive question-answering tasks like +SQuAD (a linear layer on top of the hidden-states output to compute `span start logits` and `span end logits`). + """, + LLAMA_START_DOCSTRING, +) +class LlamaForQuestionAnswering(LlamaPreTrainedModel): + base_model_prefix = "transformer" + + # Copied from transformers.models.bloom.modeling_bloom.BloomForQuestionAnswering.__init__ with Bloom->Llama + def __init__(self, config): + super().__init__(config) + self.transformer = LlamaModel(config) + self.qa_outputs = nn.Linear(config.hidden_size, 2) + + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self): + return self.transformer.embed_tokens + + def set_input_embeddings(self, value): + self.transformer.embed_tokens = value + + @add_start_docstrings_to_model_forward(LLAMA_INPUTS_DOCSTRING) + def forward( + self, + input_ids: Optional[torch.LongTensor] = None, + attention_mask: Optional[torch.FloatTensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[Union[Cache, List[torch.FloatTensor]]] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + start_positions: Optional[torch.LongTensor] = None, + end_positions: Optional[torch.LongTensor] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + ) -> Union[Tuple, QuestionAnsweringModelOutput]: + r""" + start_positions (`torch.LongTensor` of shape `(batch_size,)`, *optional*): + Labels for position (index) of the start of the labelled span for computing the token classification loss. + Positions are clamped to the length of the sequence (`sequence_length`). Position outside of the sequence + are not taken into account for computing the loss. + end_positions (`torch.LongTensor` of shape `(batch_size,)`, *optional*): + Labels for position (index) of the end of the labelled span for computing the token classification loss. + Positions are clamped to the length of the sequence (`sequence_length`). Position outside of the sequence + are not taken into account for computing the loss. + """ + return_dict = ( + return_dict if return_dict is not None else self.config.use_return_dict + ) + + outputs = self.transformer( + input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + sequence_output = outputs[0] + + logits = self.qa_outputs(sequence_output) + start_logits, end_logits = logits.split(1, dim=-1) + start_logits = start_logits.squeeze(-1).contiguous() + end_logits = end_logits.squeeze(-1).contiguous() + + total_loss = None + if start_positions is not None and end_positions is not None: + # If we are on multi-GPU, split add a dimension + if len(start_positions.size()) > 1: + start_positions = start_positions.squeeze(-1).to(start_logits.device) + if len(end_positions.size()) > 1: + end_positions = end_positions.squeeze(-1).to(end_logits.device) + # sometimes the start/end positions are outside our model inputs, we ignore these terms + ignored_index = start_logits.size(1) + start_positions = start_positions.clamp(0, ignored_index) + end_positions = end_positions.clamp(0, ignored_index) + + loss_fct = CrossEntropyLoss(ignore_index=ignored_index) + start_loss = loss_fct(start_logits, start_positions) + end_loss = loss_fct(end_logits, end_positions) + total_loss = (start_loss + end_loss) / 2 + + if not return_dict: + output = (start_logits, end_logits) + outputs[2:] + return ((total_loss,) + output) if total_loss is not None else output + + return QuestionAnsweringModelOutput( + loss=total_loss, + start_logits=start_logits, + end_logits=end_logits, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) + + +@add_start_docstrings( + """ + The Llama Model transformer with a token classification head on top (a linear layer on top of the hidden-states + output) e.g. for Named-Entity-Recognition (NER) tasks. + """, + LLAMA_START_DOCSTRING, +) +class LlamaForTokenClassification(LlamaPreTrainedModel): + def __init__(self, config): + super().__init__(config) + self.num_labels = config.num_labels + self.model = LlamaModel(config) + if getattr(config, "classifier_dropout", None) is not None: + classifier_dropout = config.classifier_dropout + elif getattr(config, "hidden_dropout", None) is not None: + classifier_dropout = config.hidden_dropout + else: + classifier_dropout = 0.1 + self.dropout = nn.Dropout(classifier_dropout) + self.score = nn.Linear(config.hidden_size, config.num_labels) + + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self): + return self.model.embed_tokens + + def set_input_embeddings(self, value): + self.model.embed_tokens = value + + @add_start_docstrings_to_model_forward(LLAMA_INPUTS_DOCSTRING) + def forward( + self, + input_ids: Optional[torch.LongTensor] = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[List[torch.FloatTensor]] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + labels: Optional[torch.LongTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + ) -> Union[Tuple, TokenClassifierOutput]: + r""" + labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*): + Labels for computing the sequence classification/regression loss. Indices should be in `[0, ..., + config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If + `config.num_labels > 1` a classification loss is computed (Cross-Entropy). + """ + return_dict = ( + return_dict if return_dict is not None else self.config.use_return_dict + ) + + outputs = self.model( + input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + sequence_output = outputs[0] + sequence_output = self.dropout(sequence_output) + logits = self.score(sequence_output) + + loss = None + if labels is not None: + loss_fct = CrossEntropyLoss() + loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1)) + + if not return_dict: + output = (logits,) + outputs[2:] + return ((loss,) + output) if loss is not None else output + + return TokenClassifierOutput( + loss=loss, + logits=logits, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) diff --git a/archive/kt-sft/ktransformers/models/modeling_mixtral.py b/archive/kt-sft/ktransformers/models/modeling_mixtral.py new file mode 100644 index 0000000..2931069 --- /dev/null +++ b/archive/kt-sft/ktransformers/models/modeling_mixtral.py @@ -0,0 +1,1736 @@ +# coding=utf-8 +''' +Description : +Author : kkk1nak0 +Date : 2024-07-29 02:58:57 +Version : 1.0.0 +LastEditors : kkk1nak0 +LastEditTime : 2024-08-02 06:08:34 +''' + +# Adapted from +# https://github.com/huggingface/transformers/blob/main/src/transformers/models/mixtral/modeling_mixtral.py +# Copyright 2023 Mistral AI and the HuggingFace Inc. team. All rights reserved. +# Copyright (c) 2024 by KVCache.AI, All Rights Reserved. +# +# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX +# and OPT implementations in this library. It has been modified from its +# original forms to accommodate minor architectural differences compared +# to GPT-NeoX and OPT used by the Meta AI team that trained the model. +# +# 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. +"""PyTorch Mixtral model.""" + +import inspect +import math +from typing import List, Optional, Tuple, Union + +import torch +import torch.nn.functional as F +import torch.utils.checkpoint +from torch import nn +from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss + +from transformers.activations import ACT2FN +from transformers.cache_utils import Cache, DynamicCache, StaticCache +from transformers.modeling_attn_mask_utils import ( + AttentionMaskConverter, + _prepare_4d_causal_attention_mask, +) +from transformers.modeling_outputs import ( + MoeCausalLMOutputWithPast, + MoeModelOutputWithPast, + SequenceClassifierOutputWithPast, + TokenClassifierOutput, +) +from transformers.modeling_utils import PreTrainedModel +from transformers.pytorch_utils import is_torch_greater_or_equal_than_1_13 +from transformers.utils import ( + add_start_docstrings, + add_start_docstrings_to_model_forward, + is_flash_attn_2_available, + logging, + replace_return_docstrings, +) +from transformers.utils.import_utils import is_torch_fx_available +from transformers.models.mixtral.configuration_mixtral import MixtralConfig + +from ktransformers.util.grad_wrapper import maybe_no_grad + +if is_flash_attn_2_available(): + from flash_attn import flash_attn_varlen_func, flash_attn_func, flash_attn_with_kvcache + from flash_attn.bert_padding import index_first_axis, pad_input, unpad_input # noqa + + _flash_supports_window_size = "window_size" in list(inspect.signature(flash_attn_func).parameters) + +# This makes `_prepare_4d_causal_attention_mask` a leaf function in the FX graph. +# It means that the function will not be traced through and simply appear as a node in the graph. +if is_torch_fx_available(): + if not is_torch_greater_or_equal_than_1_13: + import torch.fx + + _prepare_4d_causal_attention_mask = torch.fx.wrap(_prepare_4d_causal_attention_mask) + + +logger = logging.get_logger(__name__) + +_CONFIG_FOR_DOC = "MixtralConfig" + + +def load_balancing_loss_func( + gate_logits: torch.Tensor, num_experts: torch.Tensor = None, top_k=2, attention_mask: Optional[torch.Tensor] = None +) -> float: + r""" + Computes auxiliary load balancing loss as in Switch Transformer - implemented in Pytorch. + + See Switch Transformer (https://arxiv.org/abs/2101.03961) for more details. This function implements the loss + function presented in equations (4) - (6) of the paper. It aims at penalizing cases where the routing between + experts is too unbalanced. + + Args: + gate_logits (Union[`torch.Tensor`, Tuple[torch.Tensor]): + Logits from the `gate`, should be a tuple of model.config.num_hidden_layers tensors of + shape [batch_size X sequence_length, num_experts]. + attention_mask (`torch.Tensor`, None): + The attention_mask used in forward function + shape [batch_size X sequence_length] if not None. + num_experts (`int`, *optional*): + Number of experts + + Returns: + The auxiliary loss. + """ + if gate_logits is None or not isinstance(gate_logits, tuple): + return 0 + + if isinstance(gate_logits, tuple): + compute_device = gate_logits[0].device + concatenated_gate_logits = torch.cat([layer_gate.to(compute_device) for layer_gate in gate_logits], dim=0) + + routing_weights = torch.nn.functional.softmax(concatenated_gate_logits, dim=-1) + + _, selected_experts = torch.topk(routing_weights, top_k, dim=-1) + + expert_mask = torch.nn.functional.one_hot(selected_experts, num_experts) + + if attention_mask is None: + # Compute the percentage of tokens routed to each experts + tokens_per_expert = torch.mean(expert_mask.float(), dim=0) + + # Compute the average probability of routing to these experts + router_prob_per_expert = torch.mean(routing_weights, dim=0) + else: + batch_size, sequence_length = attention_mask.shape + num_hidden_layers = concatenated_gate_logits.shape[0] // (batch_size * sequence_length) + + # Compute the mask that masks all padding tokens as 0 with the same shape of expert_mask + expert_attention_mask = ( + attention_mask[None, :, :, None, None] + .expand((num_hidden_layers, batch_size, sequence_length, top_k, num_experts)) + .reshape(-1, top_k, num_experts) + .to(compute_device) + ) + + # Compute the percentage of tokens routed to each experts + tokens_per_expert = torch.sum(expert_mask.float() * expert_attention_mask, dim=0) / torch.sum( + expert_attention_mask, dim=0 + ) + + # Compute the mask that masks all padding tokens as 0 with the same shape of tokens_per_expert + router_per_expert_attention_mask = ( + attention_mask[None, :, :, None] + .expand((num_hidden_layers, batch_size, sequence_length, num_experts)) + .reshape(-1, num_experts) + .to(compute_device) + ) + + # Compute the average probability of routing to these experts + router_prob_per_expert = torch.sum(routing_weights * router_per_expert_attention_mask, dim=0) / torch.sum( + router_per_expert_attention_mask, dim=0 + ) + + overall_loss = torch.sum(tokens_per_expert * router_prob_per_expert.unsqueeze(0)) + return overall_loss * num_experts + + +# Copied from transformers.models.llama.modeling_llama._get_unpad_data +def _get_unpad_data(attention_mask): + seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32) + indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten() + max_seqlen_in_batch = seqlens_in_batch.max().item() + cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0)) + return ( + indices, + cu_seqlens, + max_seqlen_in_batch, + ) + + +# Copied from transformers.models.llama.modeling_llama.LlamaRMSNorm with Llama->Mixtral +class MixtralRMSNorm(nn.Module): + def __init__(self, hidden_size, eps=1e-6): + """ + MixtralRMSNorm is equivalent to T5LayerNorm + """ + super().__init__() + self.weight = nn.Parameter(torch.ones(hidden_size)) + self.variance_epsilon = eps + + def forward(self, hidden_states): + input_dtype = hidden_states.dtype + hidden_states = hidden_states.to(torch.float32) + variance = hidden_states.pow(2).mean(-1, keepdim=True) + hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon) + return self.weight * hidden_states.to(input_dtype) + + def extra_repr(self): + return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}" + + +# copied from transformers.models.mistral.modeling_mistral.MistralRotaryEmbedding with Mistral->Mixtral +# TODO @longjie no longer copied from Mistral after static cache +class MixtralRotaryEmbedding(nn.Module): + def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None): + super().__init__() + + self.dim = dim + self.max_position_embeddings = max_position_embeddings + self.base = base + inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2, dtype=torch.int64).float().to(device) / self.dim)) + self.register_buffer("inv_freq", inv_freq, persistent=False) + + # Build here to make `torch.jit.trace` work. + self.max_seq_len_cached = max_position_embeddings + + @maybe_no_grad() + def forward(self, x, position_ids): + # x: [bs, num_attention_heads, seq_len, head_size] + inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1) + position_ids_expanded = position_ids[:, None, :].float() + + device_type = x.device.type + device_type = device_type if isinstance(device_type, str) and device_type != "mps" else "cpu" + with torch.autocast(device_type=device_type, enabled=False): + freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2) + emb = torch.cat((freqs, freqs), dim=-1) + cos = emb.cos() + sin = emb.sin() + return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype) + + +# Copied from transformers.models.llama.modeling_llama.rotate_half +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) + + +# copied from transformers.models.mistral.modeling_mistral.apply_rotary_pos_emb +# TODO @longjie no longer copied from Mistral after static cache +def apply_rotary_pos_emb(q, k, cos, sin, position_ids, unsqueeze_dim=1): + """Applies Rotary Position Embedding to the query and key tensors. + + Args: + q (`torch.Tensor`): The query tensor. + k (`torch.Tensor`): The key tensor. + cos (`torch.Tensor`): The cosine part of the rotary embedding. + sin (`torch.Tensor`): The sine part of the rotary embedding. + position_ids (`torch.Tensor`): + The position indices of the tokens corresponding to the query and key tensors. For example, this can be + used to pass offsetted position ids when working with a KV-cache. + unsqueeze_dim (`int`, *optional*, defaults to 1): + The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and + sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note + that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and + k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes + cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have + the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2. + Returns: + `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding. + """ + cos = cos.unsqueeze(unsqueeze_dim) + sin = sin.unsqueeze(unsqueeze_dim) + + q_embed = (q * cos) + (rotate_half(q) * sin) + k_embed = (k * cos) + (rotate_half(k) * sin) + return q_embed, k_embed + + +# Copied from transformers.models.llama.modeling_llama.repeat_kv +def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: + """ + This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch, + num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim) + """ + batch, num_key_value_heads, slen, head_dim = hidden_states.shape + if n_rep == 1: + return hidden_states + hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim) + return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim) + + +# copied from transformers.models.mistral.modeling_mistral.MistralAttention with Mistral->Mixtral +# TODO @longjie no longer copied from Mistral after static cache +class MixtralAttention(nn.Module): + """ + Multi-headed attention from 'Attention Is All You Need' paper. Modified to use sliding window attention: Longformer + and "Generating Long Sequences with Sparse Transformers". + """ + + def __init__(self, config: MixtralConfig, layer_idx: Optional[int] = None): + super().__init__() + self.config = config + self.layer_idx = layer_idx + if layer_idx is None: + logger.warning_once( + f"Instantiating {self.__class__.__name__} without passing a `layer_idx` is not recommended and will " + "lead to errors during the forward call if caching is used. Please make sure to provide a `layer_idx` " + "when creating this class." + ) + + self.hidden_size = config.hidden_size + self.num_heads = config.num_attention_heads + self.head_dim = self.hidden_size // self.num_heads + self.num_key_value_heads = config.num_key_value_heads + self.num_key_value_groups = self.num_heads // self.num_key_value_heads + self.max_position_embeddings = config.max_position_embeddings + self.rope_theta = config.rope_theta + self.is_causal = True + self.attention_dropout = config.attention_dropout + + if (self.head_dim * self.num_heads) != self.hidden_size: + raise ValueError( + f"hidden_size must be divisible by num_heads (got `hidden_size`: {self.hidden_size}" + f" and `num_heads`: {self.num_heads})." + ) + self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=False) + self.k_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=False) + self.v_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=False) + self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=False) + + self.rotary_emb = MixtralRotaryEmbedding( + self.head_dim, + max_position_embeddings=self.max_position_embeddings, + base=self.rope_theta, + ) + + def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int): + return tensor.view(bsz, seq_len, self.num_heads, self.head_dim).transpose(1, 2).contiguous() + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_value: Optional[Cache] = None, + output_attentions: bool = False, + use_cache: bool = False, + cache_position: Optional[torch.LongTensor] = None, + ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: + bsz, q_len, _ = hidden_states.size() + + query_states = self.q_proj(hidden_states) + key_states = self.k_proj(hidden_states) + value_states = self.v_proj(hidden_states) + + query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2) + key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) + value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) + + kv_seq_len = key_states.shape[-2] + if past_key_value is not None: + if self.layer_idx is None: + raise ValueError( + f"The cache structure has changed since version v4.36. If you are using {self.__class__.__name__} " + "for auto-regressive decoding with k/v caching, please make sure to initialize the attention class " + "with a layer index." + ) + kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx) + cos, sin = self.rotary_emb(value_states, position_ids) + query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids) + + if past_key_value is not None: + cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position} # Specific to RoPE models + key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs) + + # repeat k/v heads if n_kv_heads < n_heads + key_states = repeat_kv(key_states, self.num_key_value_groups) + value_states = repeat_kv(value_states, self.num_key_value_groups) + + attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim) + + if attn_weights.size() != (bsz, self.num_heads, q_len, kv_seq_len): + raise ValueError( + f"Attention weights should be of size {(bsz, self.num_heads, q_len, kv_seq_len)}, but is" + f" {attn_weights.size()}" + ) + + if attention_mask is not None: # no matter the length, we just slice it + causal_mask = attention_mask[:, :, :, : key_states.shape[-2]] + attn_weights = attn_weights + causal_mask + + # upcast attention to fp32 + attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype) + attn_weights = nn.functional.dropout(attn_weights, p=self.attention_dropout, training=self.training) + attn_output = torch.matmul(attn_weights, value_states) + + if attn_output.size() != (bsz, self.num_heads, q_len, self.head_dim): + raise ValueError( + f"`attn_output` should be of size {(bsz, self.num_heads, q_len, self.head_dim)}, but is" + f" {attn_output.size()}" + ) + + attn_output = attn_output.transpose(1, 2).contiguous() + attn_output = attn_output.reshape(bsz, q_len, self.hidden_size) + + attn_output = self.o_proj(attn_output) + + if not output_attentions: + attn_weights = None + + return attn_output, attn_weights, past_key_value + + +# copied from transformers.models.mistral.modeling_mistral.MistralFlashAttention2 with Mistral->Mixtral +# TODO @longjie no longer copied from Mistral after static cache +class MixtralFlashAttention2(MixtralAttention): + """ + Mixtral flash attention module. This module inherits from `MixtralAttention` as the weights of the module stays + untouched. The only required change would be on the forward pass where it needs to correctly call the public API of + flash attention and deal with padding tokens in case the input contains any of them. + """ + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_value: Optional[Cache] = None, + output_attentions: bool = False, + use_cache: bool = False, + cache_position: Optional[torch.LongTensor] = None, + ): + bsz, q_len, _ = hidden_states.size() + + query_states = self.q_proj(hidden_states) + key_states = self.k_proj(hidden_states) + value_states = self.v_proj(hidden_states) + + query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2) + key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) + value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) + + kv_seq_len = key_states.shape[-2] + if past_key_value is not None: + if self.layer_idx is None: + raise ValueError( + f"The cache structure has changed since version v4.36. If you are using {self.__class__.__name__} " + "for auto-regressive decoding with k/v caching, please make sure to initialize the attention class " + "with a layer index." + ) + kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx) + + cos, sin = self.rotary_emb(value_states, position_ids) + + query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids) + + use_sliding_windows = ( + _flash_supports_window_size + and getattr(self.config, "sliding_window", None) is not None + and kv_seq_len > self.config.sliding_window + and self.config.use_sliding_window + ) + + if not _flash_supports_window_size: + logger.warning_once( + "The current flash attention version does not support sliding window attention, for a more memory efficient implementation" + " make sure to upgrade flash-attn library." + ) + + if past_key_value is not None: + # Activate slicing cache only if the config has a value `sliding_windows` attribute + cache_has_contents = past_key_value.get_seq_length(self.layer_idx) > 0 + if ( + getattr(self.config, "sliding_window", None) is not None + and kv_seq_len > self.config.sliding_window + and cache_has_contents + ): + slicing_tokens = 1 - self.config.sliding_window + + past_key = past_key_value[self.layer_idx][0] + past_value = past_key_value[self.layer_idx][1] + + past_key = past_key[:, :, slicing_tokens:, :].contiguous() + past_value = past_value[:, :, slicing_tokens:, :].contiguous() + + if past_key.shape[-2] != self.config.sliding_window - 1: + raise ValueError( + f"past key must have a shape of (`batch_size, num_heads, self.config.sliding_window-1, head_dim`), got" + f" {past_key.shape}" + ) + + if attention_mask is not None: + attention_mask = attention_mask[:, slicing_tokens:] + attention_mask = torch.cat([attention_mask, torch.ones_like(attention_mask[:, -1:])], dim=-1) + + cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position} # Specific to RoPE models + key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs) + + # we slice the states for static kv cache to be supported in FA2. Not sure it's a must as compile fails + # for bsz == 1, avoid using slice to capture cuda graph + if cache_position is not None and q_len > 1: + key_states = key_states[:, :, : cache_position[-1] + 1, :] + value_states = value_states[:, :, : cache_position[-1] + 1, :] + + # repeat k/v heads if n_kv_heads < n_heads + key_states = repeat_kv(key_states, self.num_key_value_groups) + value_states = repeat_kv(value_states, self.num_key_value_groups) + dropout_rate = 0.0 if not self.training else self.attention_dropout + + # In PEFT, usually we cast the layer norms in float32 for training stability reasons + # therefore the input hidden states gets silently casted in float32. Hence, we need + # cast them back in float16 just to be sure everything works as expected. + input_dtype = query_states.dtype + if input_dtype == torch.float32: + if torch.is_autocast_enabled(): + target_dtype = torch.get_autocast_gpu_dtype() + # Handle the case where the model is quantized + elif hasattr(self.config, "_pre_quantization_dtype"): + target_dtype = self.config._pre_quantization_dtype + else: + target_dtype = self.q_proj.weight.dtype + + logger.warning_once( + f"The input hidden states seems to be silently casted in float32, this might be related to" + f" the fact you have upcasted embedding or layer norm layers in float32. We will cast back the input in" + f" {target_dtype}." + ) + + query_states = query_states.to(target_dtype) + key_states = key_states.to(target_dtype) + value_states = value_states.to(target_dtype) + + # Reashape to the expected shape for Flash Attention + query_states = query_states.transpose(1, 2) + key_states = key_states.transpose(1, 2) + value_states = value_states.transpose(1, 2) + + attn_output = self._flash_attention_forward( + query_states, + key_states, + value_states, + attention_mask, + q_len, + position_ids=position_ids, + dropout=dropout_rate, + sliding_window=getattr(self.config, "sliding_window", None), + is_causal=self.is_causal, + ) + + attn_output = attn_output.reshape(bsz, q_len, self.hidden_size).contiguous() + attn_output = self.o_proj(attn_output) + + if not output_attentions: + attn_weights = None + + return attn_output, attn_weights, past_key_value + + + def _flash_attention_forward( + self, + query_states, + key_states, + value_states, + attention_mask, + q_len, + position_ids, + dropout, + sliding_window, + is_causal, + softmax_scale=None, + ): + """ + Calls the forward method of Flash Attention - if the input hidden states contain at least one padding token + first unpad the input, then computes the attention scores and pad the final attention scores. + + Args: + query_states (`torch.Tensor`): + Input query states to be passed to Flash Attention API + key_states (`torch.Tensor`): + Input key states to be passed to Flash Attention API + value_states (`torch.Tensor`): + Input value states to be passed to Flash Attention API + attention_mask (`torch.Tensor`): + The padding mask - corresponds to a tensor of size `(batch_size, seq_len)` where 0 stands for the + position of padding tokens and 1 for the position of non-padding tokens. + dropout (`float`): + Attention dropout + + """ + + # Decide whether to use SWA or not by layer index. + # if use_sliding_windows and self.layer_idx >= self.config.max_window_layers: + # use_sliding_windows = False + use_sliding_windows = False + + # Contains at least one padding token in the sequence + if attention_mask is not None: + batch_size = query_states.shape[0] + query_states, key_states, value_states, indices_q, cu_seq_lens, max_seq_lens = self._upad_input( + query_states, key_states, value_states, attention_mask, q_len + ) + + cu_seqlens_q, cu_seqlens_k = cu_seq_lens + max_seqlen_in_batch_q, max_seqlen_in_batch_k = max_seq_lens + + if not use_sliding_windows: + attn_output_unpad = flash_attn_varlen_func( + query_states, + key_states, + value_states, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k=cu_seqlens_k, + max_seqlen_q=max_seqlen_in_batch_q, + max_seqlen_k=max_seqlen_in_batch_k, + dropout_p=dropout, + softmax_scale=softmax_scale, + causal=is_causal, + ) + else: + attn_output_unpad = flash_attn_varlen_func( + query_states, + key_states, + value_states, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k=cu_seqlens_k, + max_seqlen_q=max_seqlen_in_batch_q, + max_seqlen_k=max_seqlen_in_batch_k, + dropout_p=dropout, + softmax_scale=softmax_scale, + causal=is_causal, + window_size=(self.config.sliding_window, self.config.sliding_window), + ) + + attn_output = pad_input(attn_output_unpad, indices_q, batch_size, q_len) + else: + if not use_sliding_windows: + if q_len == 1: + position_ids = position_ids.to(dtype=torch.int32).squeeze(1) + attn_output = flash_attn_with_kvcache( + query_states, + key_states, + value_states, + cache_seqlens=position_ids, + softmax_scale=softmax_scale, + causal=is_causal, + ) + else: + attn_output = flash_attn_func( + query_states, + key_states, + value_states, + dropout, + softmax_scale=softmax_scale, + causal=is_causal, + ) + else: + attn_output = flash_attn_func( + query_states, + key_states, + value_states, + dropout, + softmax_scale=softmax_scale, + causal=is_causal, + window_size=(self.config.sliding_window, self.config.sliding_window), + ) + + return attn_output + + # Copied from transformers.models.mistral.modeling_mistral.MistralFlashAttention2._upad_input + def _upad_input(self, query_layer, key_layer, value_layer, attention_mask, query_length): + batch_size, kv_seq_len, num_heads, head_dim = key_layer.shape + + # On the first iteration we need to properly re-create the padding mask + # by slicing it on the proper place + if kv_seq_len != attention_mask.shape[-1]: + attention_mask_num_tokens = attention_mask.shape[-1] + attention_mask = attention_mask[:, attention_mask_num_tokens - kv_seq_len :] + + indices_k, cu_seqlens_k, max_seqlen_in_batch_k = _get_unpad_data(attention_mask) + + key_layer = index_first_axis(key_layer.reshape(batch_size * kv_seq_len, num_heads, head_dim), indices_k) + value_layer = index_first_axis(value_layer.reshape(batch_size * kv_seq_len, num_heads, head_dim), indices_k) + + if query_length == kv_seq_len: + query_layer = index_first_axis( + query_layer.reshape(batch_size * kv_seq_len, num_heads, head_dim), indices_k + ) + cu_seqlens_q = cu_seqlens_k + max_seqlen_in_batch_q = max_seqlen_in_batch_k + indices_q = indices_k + elif query_length == 1: + max_seqlen_in_batch_q = 1 + cu_seqlens_q = torch.arange( + batch_size + 1, dtype=torch.int32, device=query_layer.device + ) # There is a memcpy here, that is very bad. + indices_q = cu_seqlens_q[:-1] + query_layer = query_layer.squeeze(1) + else: + # The -q_len: slice assumes left padding. + attention_mask = attention_mask[:, -query_length:] + query_layer, indices_q, cu_seqlens_q, max_seqlen_in_batch_q = unpad_input(query_layer, attention_mask) + + return ( + query_layer, + key_layer, + value_layer, + indices_q, + (cu_seqlens_q, cu_seqlens_k), + (max_seqlen_in_batch_q, max_seqlen_in_batch_k), + ) + + + +# copied from transformers.models.mistral.modeling_mistral.MistralSdpaAttention with Mistral->Mixtral +# TODO @longjie no longer copied from Mistral after static cache +class MixtralSdpaAttention(MixtralAttention): + """ + Mixtral attention module using torch.nn.functional.scaled_dot_product_attention. This module inherits from + `MixtralAttention` as the weights of the module stays untouched. The only changes are on the forward pass to adapt to + SDPA API. + """ + + # Adapted from MixtralAttention.forward + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_value: Optional[Cache] = None, + output_attentions: bool = False, + use_cache: bool = False, + cache_position: Optional[torch.LongTensor] = None, + ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: + if output_attentions: + # TODO: Improve this warning with e.g. `model.config.attn_implementation = "manual"` once this is implemented. + logger.warning_once( + "MixtralModel is using MixtralSdpaAttention, but `torch.nn.functional.scaled_dot_product_attention` does not support `output_attentions=True`. Falling back to the manual attention implementation, " + 'but specifying the manual implementation will be required from Transformers version v5.0.0 onwards. This warning can be removed using the argument `attn_implementation="eager"` when loading the model.' + ) + return super().forward( + hidden_states=hidden_states, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_value=past_key_value, + output_attentions=output_attentions, + use_cache=use_cache, + ) + + bsz, q_len, _ = hidden_states.size() + + query_states = self.q_proj(hidden_states) + key_states = self.k_proj(hidden_states) + value_states = self.v_proj(hidden_states) + + query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2) + key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) + value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) + + kv_seq_len = key_states.shape[-2] + if past_key_value is not None: + kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx) + cos, sin = self.rotary_emb(value_states, position_ids) + + query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids) + + if past_key_value is not None: + cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position} # Specific to RoPE models + key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs) + + key_states = repeat_kv(key_states, self.num_key_value_groups) + value_states = repeat_kv(value_states, self.num_key_value_groups) + + causal_mask = attention_mask + if attention_mask is not None: # no matter the length, we just slice it + causal_mask = attention_mask[:, :, :, : key_states.shape[-2]] + + # SDPA with memory-efficient backend is currently (torch==2.1.2) bugged with non-contiguous inputs with custom attn_mask, + # Reference: https://github.com/pytorch/pytorch/issues/112577. + if query_states.device.type == "cuda" and attention_mask is not None: + query_states = query_states.contiguous() + key_states = key_states.contiguous() + value_states = value_states.contiguous() + + # We dispatch to SDPA's Flash Attention or Efficient kernels via this `is_causal` if statement instead of an inline conditional assignment + # in SDPA to support both torch.compile's dynamic shapes and full graph options. An inline conditional prevents dynamic shapes from compiling. + # The q_len > 1 is necessary to match with AttentionMaskConverter.to_causal_4d that does not create a causal mask in case q_len == 1. + is_causal = True if causal_mask is None and q_len > 1 else False + + attn_output = torch.nn.functional.scaled_dot_product_attention( + query_states, + key_states, + value_states, + attn_mask=causal_mask, + dropout_p=self.attention_dropout if self.training else 0.0, + is_causal=is_causal, + ) + + attn_output = attn_output.transpose(1, 2).contiguous() + attn_output = attn_output.view(bsz, q_len, self.hidden_size) + + attn_output = self.o_proj(attn_output) + + return attn_output, None, past_key_value + + +MIXTRAL_ATTENTION_CLASSES = { + "eager": MixtralAttention, + "flash_attention_2": MixtralFlashAttention2, + "sdpa": MixtralSdpaAttention, +} + + +class MixtralBlockSparseTop2MLP(nn.Module): + def __init__(self, config: MixtralConfig): + super().__init__() + self.ffn_dim = config.intermediate_size + self.hidden_dim = config.hidden_size + + self.w1 = nn.Linear(self.hidden_dim, self.ffn_dim, bias=False) # gate + self.w2 = nn.Linear(self.ffn_dim, self.hidden_dim, bias=False) # down + self.w3 = nn.Linear(self.hidden_dim, self.ffn_dim, bias=False) # up + + self.act_fn = ACT2FN[config.hidden_act] + + def forward(self, hidden_states): + current_hidden_states = self.act_fn(self.w1(hidden_states)) * self.w3(hidden_states) + current_hidden_states = self.w2(current_hidden_states) + return current_hidden_states + + +class MixtralSparseMoeBlock(nn.Module): + """ + This implementation is + strictly equivalent to standard MoE with full capacity (no + dropped tokens). It's faster since it formulates MoE operations + in terms of block-sparse operations to accomodate imbalanced + assignments of tokens to experts, whereas standard MoE either + (1) drop tokens at the cost of reduced performance or (2) set + capacity factor to number of experts and thus waste computation + and memory on padding. + """ + + def __init__(self, config): + super().__init__() + self.hidden_dim = config.hidden_size + self.ffn_dim = config.intermediate_size + self.num_experts = config.num_local_experts + self.top_k = config.num_experts_per_tok + + # gating + self.gate = nn.Linear(self.hidden_dim, self.num_experts, bias=False) + + self.experts = nn.ModuleList([MixtralBlockSparseTop2MLP(config) for _ in range(self.num_experts)]) + + # Jitter parameters + self.jitter_noise = config.router_jitter_noise + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + """ """ + batch_size, sequence_length, hidden_dim = hidden_states.shape + if self.training and self.jitter_noise > 0: + hidden_states *= torch.empty_like(hidden_states).uniform_(1.0 - self.jitter_noise, 1.0 + self.jitter_noise) + hidden_states = hidden_states.view(-1, hidden_dim) + # router_logits: (batch * sequence_length, n_experts) + router_logits = self.gate(hidden_states) + + routing_weights = F.softmax(router_logits, dim=1, dtype=torch.float) + routing_weights, selected_experts = torch.topk(routing_weights, self.top_k, dim=-1) + routing_weights /= routing_weights.sum(dim=-1, keepdim=True) + # we cast back to the input dtype + routing_weights = routing_weights.to(hidden_states.dtype) + + final_hidden_states = torch.zeros( + (batch_size * sequence_length, hidden_dim), dtype=hidden_states.dtype, device=hidden_states.device + ) + + # One hot encode the selected experts to create an expert mask + # this will be used to easily index which expert is going to be sollicitated + expert_mask = torch.nn.functional.one_hot(selected_experts, num_classes=self.num_experts).permute(2, 1, 0) + + # Loop over all available experts in the model and perform the computation on each expert + for expert_idx in range(self.num_experts): + expert_layer = self.experts[expert_idx] + idx, top_x = torch.where(expert_mask[expert_idx]) + + # Index the correct hidden states and compute the expert hidden state for + # the current expert. We need to make sure to multiply the output hidden + # states by `routing_weights` on the corresponding tokens (top-1 and top-2) + current_state = hidden_states[None, top_x].reshape(-1, hidden_dim) + current_hidden_states = expert_layer(current_state) * routing_weights[top_x, idx, None] + + # However `index_add_` only support torch tensors for indexing so we'll use + # the `top_x` tensor here. + final_hidden_states.index_add_(0, top_x, current_hidden_states.to(hidden_states.dtype)) + final_hidden_states = final_hidden_states.reshape(batch_size, sequence_length, hidden_dim) + return final_hidden_states, router_logits + + +class MixtralDecoderLayer(nn.Module): + def __init__(self, config: MixtralConfig, layer_idx: int): + super().__init__() + self.hidden_size = config.hidden_size + + self.self_attn = MIXTRAL_ATTENTION_CLASSES[config._attn_implementation](config, layer_idx) + + self.block_sparse_moe = MixtralSparseMoeBlock(config) + self.input_layernorm = MixtralRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = MixtralRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_value: Optional[Tuple[torch.Tensor]] = None, + output_attentions: Optional[bool] = False, + output_router_logits: Optional[bool] = False, + use_cache: Optional[bool] = False, + cache_position: Optional[torch.LongTensor] = None, + **kwargs, + ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]: + """ + Args: + hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)` + attention_mask (`torch.FloatTensor`, *optional*): attention mask of size + `(batch, sequence_length)` where padding elements are indicated by 0. + past_key_value (`Tuple(torch.FloatTensor)`, *optional*): cached past key and value projection states + output_attentions (`bool`, *optional*): + Whether or not to return the attentions tensors of all attention layers. See `attentions` under + returned tensors for more detail. + output_router_logits (`bool`, *optional*): + Whether or not to return the logits of all the routers. They are useful for computing the router loss, and + should not be returned during inference. + use_cache (`bool`, *optional*): + If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding + (see `past_key_values`). + cache_position (`torch.LongTensor` of shape `(sequence_length)`, *optional*): + Indices depicting the position of the input sequence tokens in the sequence. + kwargs (`dict`, *optional*): + Arbitrary kwargs to be ignored, used for FSDP and other methods that injects code + into the model + """ + + residual = hidden_states + + hidden_states = self.input_layernorm(hidden_states) + + # Self Attention + hidden_states, self_attn_weights, present_key_value = self.self_attn( + hidden_states=hidden_states, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_value=past_key_value, + output_attentions=output_attentions, + use_cache=use_cache, + cache_position=cache_position, + ) + hidden_states = residual + hidden_states + + # Fully Connected + residual = hidden_states + hidden_states = self.post_attention_layernorm(hidden_states) + hidden_states, router_logits = self.block_sparse_moe(hidden_states) + hidden_states = residual + hidden_states + + outputs = (hidden_states,) + + if output_attentions: + outputs += (self_attn_weights,) + + if use_cache: + outputs += (present_key_value,) + + if output_router_logits: + outputs += (router_logits,) + + return outputs + + +MIXTRAL_START_DOCSTRING = r""" + This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the + library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads + etc.) + + This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass. + Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage + and behavior. + + Parameters: + config ([`MixtralConfig`]): + Model configuration class with all the parameters of the model. Initializing with a config file does not + load the weights associated with the model, only the configuration. Check out the + [`~PreTrainedModel.from_pretrained`] method to load the model weights. +""" + + +@add_start_docstrings( + "The bare Mixtral Model outputting raw hidden-states without any specific head on top.", + MIXTRAL_START_DOCSTRING, +) +# Copied from transformers.models.qwen2.modeling_qwen2.Qwen2PreTrainedModel with Qwen2->Mixtral +class MixtralPreTrainedModel(PreTrainedModel): + config_class = MixtralConfig + base_model_prefix = "model" + supports_gradient_checkpointing = True + _no_split_modules = ["MixtralDecoderLayer"] + _skip_keys_device_placement = "past_key_values" + _supports_flash_attn_2 = True + _supports_sdpa = True + _supports_cache_class = True + + def _init_weights(self, module): + std = self.config.initializer_range + if isinstance(module, nn.Linear): + module.weight.data.normal_(mean=0.0, std=std) + if module.bias is not None: + module.bias.data.zero_() + elif isinstance(module, nn.Embedding): + module.weight.data.normal_(mean=0.0, std=std) + if module.padding_idx is not None: + module.weight.data[module.padding_idx].zero_() + + +MIXTRAL_INPUTS_DOCSTRING = r""" + Args: + input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): + Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide + it. + + Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and + [`PreTrainedTokenizer.__call__`] for details. + + [What are input IDs?](../glossary#input-ids) + attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*): + Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`: + + - 1 for tokens that are **not masked**, + - 0 for tokens that are **masked**. + + [What are attention masks?](../glossary#attention-mask) + + Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and + [`PreTrainedTokenizer.__call__`] for details. + + If `past_key_values` is used, optionally only the last `decoder_input_ids` have to be input (see + `past_key_values`). + + If you want to change padding behavior, you should read [`modeling_opt._prepare_decoder_attention_mask`] + and modify to your needs. See diagram 1 in [the paper](https://arxiv.org/abs/1910.13461) for more + information on the default strategy. + + - 1 indicates the head is **not masked**, + - 0 indicates the head is **masked**. + position_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0, + config.n_positions - 1]`. + + [What are position IDs?](../glossary#position-ids) + past_key_values (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`): + Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of shape + `(batch_size, num_heads, sequence_length, embed_size_per_head)`) and 2 additional tensors of shape + `(batch_size, num_heads, encoder_sequence_length, embed_size_per_head)`. + + Contains pre-computed hidden-states (key and values in the self-attention blocks and in the cross-attention + blocks) that can be used (see `past_key_values` input) to speed up sequential decoding. + + If `past_key_values` are used, the user can optionally input only the last `decoder_input_ids` (those that + don't have their past key value states given to this model) of shape `(batch_size, 1)` instead of all + `decoder_input_ids` of shape `(batch_size, sequence_length)`. + inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*): + Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This + is useful if you want more control over how to convert `input_ids` indices into associated vectors than the + model's internal embedding lookup matrix. + use_cache (`bool`, *optional*): + If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see + `past_key_values`). + output_attentions (`bool`, *optional*): + Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned + tensors for more detail. + output_hidden_states (`bool`, *optional*): + Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for + more detail. + output_router_logits (`bool`, *optional*): + Whether or not to return the logits of all the routers. They are useful for computing the router loss, and + should not be returned during inference. + return_dict (`bool`, *optional*): + Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple. + cache_position (`torch.LongTensor` of shape `(sequence_length)`, *optional*): + Indices depicting the position of the input sequence tokens in the sequence. Contrarily to `position_ids`, + this tensor is not affected by padding. It is used to update the cache in the correct position and to infer + the complete sequence length. +""" + + +@add_start_docstrings( + "The bare Mixtral Model outputting raw hidden-states without any specific head on top.", + MIXTRAL_START_DOCSTRING, +) +# copied from transformers.models.mistral.modeling_mistral.MistralModel with MISTRAL->MIXTRAL,Mistral->Mixtral +# TODO @longjie no longer copied from Mistral after static cache +class MixtralModel(MixtralPreTrainedModel): + """ + Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`MixtralDecoderLayer`] + + Args: + config: MixtralConfig + """ + + def __init__(self, config: MixtralConfig): + super().__init__(config) + self.padding_idx = config.pad_token_id + self.vocab_size = config.vocab_size + + self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx) + self.layers = nn.ModuleList( + [MixtralDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)] + ) + self._attn_implementation = config._attn_implementation + self.norm = MixtralRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + self.gradient_checkpointing = False + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self): + return self.embed_tokens + + def set_input_embeddings(self, value): + self.embed_tokens = value + + # Ignore copy + @add_start_docstrings_to_model_forward(MIXTRAL_INPUTS_DOCSTRING) + def forward( + self, + input_ids: torch.LongTensor = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[List[torch.FloatTensor]] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + output_router_logits: Optional[bool] = None, + return_dict: Optional[bool] = None, + cache_position: Optional[torch.LongTensor] = None, + ) -> Union[Tuple, MoeModelOutputWithPast]: + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_router_logits = ( + output_router_logits if output_router_logits is not None else self.config.output_router_logits + ) + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + use_cache = use_cache if use_cache is not None else self.config.use_cache + + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + if (input_ids is None) ^ (inputs_embeds is not None): + raise ValueError( + "You cannot specify both input_ids and inputs_embeds at the same time, and must specify either one" + ) + + if self.gradient_checkpointing and self.training: + if use_cache: + logger.warning_once( + "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..." + ) + use_cache = False + + use_legacy_cache = False + if use_cache and not isinstance(past_key_values, Cache) and not self.training: + use_legacy_cache = True + past_key_values = DynamicCache.from_legacy_cache(past_key_values) + logger.warning_once( + "We detected that you are passing `past_key_values` as a tuple and this is deprecated and will be removed in v4.43. " + "Please use an appropriate `Cache` class (https://huggingface.co/docs/transformers/v4.41.3/en/internal/generation_utils#transformers.Cache)" + ) + + if inputs_embeds is None: + inputs_embeds = self.embed_tokens(input_ids) + + if cache_position is None: + past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 + cache_position = torch.arange( + past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device + ) + if position_ids is None: + position_ids = cache_position.unsqueeze(0) + + causal_mask = self._update_causal_mask( + attention_mask, inputs_embeds, cache_position, past_key_values, output_attentions + ) + + hidden_states = inputs_embeds + + # decoder layers + all_hidden_states = () if output_hidden_states else None + all_self_attns = () if output_attentions else None + all_router_logits = () if output_router_logits else None + next_decoder_cache = None + + for decoder_layer in self.layers: + if output_hidden_states: + all_hidden_states += (hidden_states,) + + if self.gradient_checkpointing and self.training: + layer_outputs = self._gradient_checkpointing_func( + decoder_layer.__call__, + hidden_states, + causal_mask, + position_ids, + past_key_values, + output_attentions, + output_router_logits, + use_cache, + cache_position, + ) + else: + layer_outputs = decoder_layer( + hidden_states, + attention_mask=causal_mask, + position_ids=position_ids, + past_key_value=past_key_values, + output_attentions=output_attentions, + output_router_logits=output_router_logits, + use_cache=use_cache, + cache_position=cache_position, + ) + + hidden_states = layer_outputs[0] + + if use_cache: + next_decoder_cache = layer_outputs[2 if output_attentions else 1] + + if output_attentions: + all_self_attns += (layer_outputs[1],) + + if output_router_logits: + all_router_logits += (layer_outputs[-1],) + + hidden_states = self.norm(hidden_states) + + # add hidden states from the last decoder layer + if output_hidden_states: + all_hidden_states += (hidden_states,) + + next_cache = None + if use_cache: + next_cache = next_decoder_cache.to_legacy_cache() if use_legacy_cache else next_decoder_cache + + if not return_dict: + return tuple( + v + for v in [hidden_states, next_cache, all_hidden_states, all_self_attns, all_router_logits] + if v is not None + ) + return MoeModelOutputWithPast( + last_hidden_state=hidden_states, + past_key_values=next_cache, + hidden_states=all_hidden_states, + attentions=all_self_attns, + router_logits=all_router_logits, + ) + + # Copied from transformers.models.llama.modeling_llama.LlamaModel._update_causal_mask + def _update_causal_mask( + self, + attention_mask: torch.Tensor, + input_tensor: torch.Tensor, + cache_position: torch.Tensor, + past_key_values: Cache, + output_attentions: bool, + ): + # TODO: As of torch==2.2.0, the `attention_mask` passed to the model in `generate` is 2D and of dynamic length even when the static + # KV cache is used. This is an issue for torch.compile which then recaptures cudagraphs at each decode steps due to the dynamic shapes. + # (`recording cudagraph tree for symint key 13`, etc.), which is VERY slow. A workaround is `@torch.compiler.disable`, but this prevents using + # `fullgraph=True`. See more context in https://github.com/huggingface/transformers/pull/29114 + + if self.config._attn_implementation == "flash_attention_2": + if attention_mask is not None and 0.0 in attention_mask: + return attention_mask + return None + + # For SDPA, when possible, we will rely on its `is_causal` argument instead of its `attn_mask` argument, in + # order to dispatch on Flash Attention 2. This feature is not compatible with static cache, as SDPA will fail + # to infer the attention mask. + past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 + using_static_cache = isinstance(past_key_values, StaticCache) + + # When output attentions is True, sdpa implementation's forward method calls the eager implementation's forward + if self.config._attn_implementation == "sdpa" and not using_static_cache and not output_attentions: + if AttentionMaskConverter._ignore_causal_mask_sdpa( + attention_mask, + inputs_embeds=input_tensor, + past_key_values_length=past_seen_tokens, + is_training=self.training, + ): + return None + + dtype, device = input_tensor.dtype, input_tensor.device + min_dtype = torch.finfo(dtype).min + sequence_length = input_tensor.shape[1] + if using_static_cache: + target_length = past_key_values.get_max_length() + else: + target_length = ( + attention_mask.shape[-1] + if isinstance(attention_mask, torch.Tensor) + else past_seen_tokens + sequence_length + 1 + ) + + if attention_mask is not None and attention_mask.dim() == 4: + # in this case we assume that the mask comes already in inverted form and requires no inversion or slicing + if attention_mask.max() != 0: + raise ValueError("Custom 4D attention mask should be passed in inverted form with max==0`") + causal_mask = attention_mask + else: + causal_mask = torch.full( + (sequence_length, target_length), fill_value=min_dtype, dtype=dtype, device=device + ) + if sequence_length != 1: + causal_mask = torch.triu(causal_mask, diagonal=1) + causal_mask *= torch.arange(target_length, device=device) > cache_position.reshape(-1, 1) + causal_mask = causal_mask[None, None, :, :].expand(input_tensor.shape[0], 1, -1, -1) + if attention_mask is not None: + causal_mask = causal_mask.clone() # copy to contiguous memory for in-place edit + mask_length = attention_mask.shape[-1] + padding_mask = causal_mask[:, :, :, :mask_length] + attention_mask[:, None, None, :] + padding_mask = padding_mask == 0 + causal_mask[:, :, :, :mask_length] = causal_mask[:, :, :, :mask_length].masked_fill( + padding_mask, min_dtype + ) + if ( + self.config._attn_implementation == "sdpa" + and attention_mask is not None + and attention_mask.device.type == "cuda" + and not output_attentions + ): + # Attend to all tokens in fully masked rows in the causal_mask, for example the relevant first rows when + # using left padding. This is required by F.scaled_dot_product_attention memory-efficient attention path. + # Details: https://github.com/pytorch/pytorch/issues/110213 + causal_mask = AttentionMaskConverter._unmask_unattended(causal_mask, min_dtype) + + return causal_mask + + +class MixtralForCausalLM(MixtralPreTrainedModel): + _tied_weights_keys = ["lm_head.weight"] + + def __init__(self, config): + super().__init__(config) + self.model = MixtralModel(config) + self.vocab_size = config.vocab_size + self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) + self.router_aux_loss_coef = config.router_aux_loss_coef + self.num_experts = config.num_local_experts + self.num_experts_per_tok = config.num_experts_per_tok + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self): + return self.model.embed_tokens + + def set_input_embeddings(self, value): + self.model.embed_tokens = value + + def get_output_embeddings(self): + return self.lm_head + + def set_output_embeddings(self, new_embeddings): + self.lm_head = new_embeddings + + def set_decoder(self, decoder): + self.model = decoder + + def get_decoder(self): + return self.model + + @add_start_docstrings_to_model_forward(MIXTRAL_INPUTS_DOCSTRING) + @replace_return_docstrings(output_type=MoeCausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC) + # Ignore copy + def forward( + self, + input_ids: torch.LongTensor = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[List[torch.FloatTensor]] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + labels: Optional[torch.LongTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + output_router_logits: Optional[bool] = None, + return_dict: Optional[bool] = None, + cache_position: Optional[torch.LongTensor] = None, + ) -> Union[Tuple, MoeCausalLMOutputWithPast]: + r""" + Args: + labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Labels for computing the masked language modeling loss. Indices should either be in `[0, ..., + config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored + (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`. + + Returns: + + Example: + + ```python + >>> from transformers import AutoTokenizer, MixtralForCausalLM + + >>> model = MixtralForCausalLM.from_pretrained("mistralai/Mixtral-8x7B-v0.1") + >>> tokenizer = AutoTokenizer.from_pretrained("mistralai/Mixtral-8x7B-v0.1") + + >>> prompt = "Hey, are you conscious? Can you talk to me?" + >>> inputs = tokenizer(prompt, return_tensors="pt") + + >>> # Generate + >>> generate_ids = model.generate(inputs.input_ids, max_length=30) + >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0] + "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you." + ```""" + + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_router_logits = ( + output_router_logits if output_router_logits is not None else self.config.output_router_logits + ) + + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn) + outputs = self.model( + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + output_router_logits=output_router_logits, + return_dict=return_dict, + cache_position=cache_position, + ) + + hidden_states = outputs[0] + logits = self.lm_head(hidden_states) + logits = logits.float() + + loss = None + if labels is not None: + # Shift so that tokens < n predict n + shift_logits = logits[..., :-1, :].contiguous() + shift_labels = labels[..., 1:].contiguous() + # Flatten the tokens + loss_fct = CrossEntropyLoss() + shift_logits = shift_logits.view(-1, self.config.vocab_size) + shift_labels = shift_labels.view(-1) + # Enable model parallelism + shift_labels = shift_labels.to(shift_logits.device) + loss = loss_fct(shift_logits, shift_labels) + + aux_loss = None + if output_router_logits: + aux_loss = load_balancing_loss_func( + outputs.router_logits if return_dict else outputs[-1], + self.num_experts, + self.num_experts_per_tok, + attention_mask, + ) + if labels is not None: + loss += self.router_aux_loss_coef * aux_loss.to(loss.device) # make sure to reside in the same device + + if not return_dict: + output = (logits,) + outputs[1:] + if output_router_logits: + output = (aux_loss,) + output + return (loss,) + output if loss is not None else output + + return MoeCausalLMOutputWithPast( + loss=loss, + aux_loss=aux_loss, + logits=logits, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + router_logits=outputs.router_logits, + ) + + def prepare_inputs_for_generation( + self, + input_ids, + past_key_values=None, + attention_mask=None, + inputs_embeds=None, + cache_position=None, + output_router_logits=False, + position_ids=None, + use_cache=True, + **kwargs, + ): + # If we have cache: let's slice `input_ids` through `cache_position`, to keep only the unprocessed tokens + # Exception 1: when passing input_embeds, input_ids may be missing entries + # Exception 2: some generation methods do special slicing of input_ids, so we don't need to do it here + if past_key_values is not None: + if inputs_embeds is not None: # Exception 1 + input_ids = input_ids[:, -cache_position.shape[0] :] + elif input_ids.shape[1] != cache_position.shape[0]: # Default case (the "else", a no op, is Exception 2) + input_ids = input_ids[:, cache_position] + + if attention_mask is not None and position_ids is None: + # create position_ids on the fly for batch generation + position_ids = attention_mask.long().cumsum(-1) - 1 + position_ids.masked_fill_(attention_mask == 0, 1) + if past_key_values: + position_ids = position_ids[:, -input_ids.shape[1] :] + + # if `inputs_embeds` are passed, we only want to use them in the 1st generation step + if inputs_embeds is not None and cache_position[0] == 0: + model_inputs = {"inputs_embeds": inputs_embeds} + else: + model_inputs = {"input_ids": input_ids.contiguous()} # `contiguous()` needed for compilation use cases + + model_inputs.update( + { + "position_ids": position_ids, + "cache_position": cache_position, + "past_key_values": past_key_values, + "use_cache": use_cache, + "attention_mask": attention_mask, + "output_router_logits": output_router_logits, + } + ) + return model_inputs + + +@add_start_docstrings( + """ + The Mixtral Model transformer with a sequence classification head on top (linear layer). + + [`MixtralForSequenceClassification`] uses the last token in order to do the classification, as other causal models + (e.g. GPT-2) do. + + Since it does classification on the last token, it requires to know the position of the last token. If a + `pad_token_id` is defined in the configuration, it finds the last token that is not a padding token in each row. If + no `pad_token_id` is defined, it simply takes the last value in each row of the batch. Since it cannot guess the + padding tokens when `inputs_embeds` are passed instead of `input_ids`, it does the same (take the last value in + each row of the batch). + """, + MIXTRAL_START_DOCSTRING, +) +# Copied from transformers.models.llama.modeling_llama.LlamaForSequenceClassification with Llama->Mixtral, LLAMA->MIXTRAL +class MixtralForSequenceClassification(MixtralPreTrainedModel): + def __init__(self, config): + super().__init__(config) + self.num_labels = config.num_labels + self.model = MixtralModel(config) + self.score = nn.Linear(config.hidden_size, self.num_labels, bias=False) + + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self): + return self.model.embed_tokens + + def set_input_embeddings(self, value): + self.model.embed_tokens = value + + @add_start_docstrings_to_model_forward(MIXTRAL_INPUTS_DOCSTRING) + def forward( + self, + input_ids: Optional[torch.LongTensor] = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[Union[Cache, List[torch.FloatTensor]]] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + labels: Optional[torch.LongTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + ) -> Union[Tuple, SequenceClassifierOutputWithPast]: + r""" + labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*): + Labels for computing the sequence classification/regression loss. Indices should be in `[0, ..., + config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If + `config.num_labels > 1` a classification loss is computed (Cross-Entropy). + """ + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + transformer_outputs = self.model( + input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + hidden_states = transformer_outputs[0] + logits = self.score(hidden_states) + + if input_ids is not None: + batch_size = input_ids.shape[0] + else: + batch_size = inputs_embeds.shape[0] + + if self.config.pad_token_id is None and batch_size != 1: + raise ValueError("Cannot handle batch sizes > 1 if no padding token is defined.") + if self.config.pad_token_id is None: + sequence_lengths = -1 + else: + if input_ids is not None: + # if no pad token found, use modulo instead of reverse indexing for ONNX compatibility + sequence_lengths = torch.eq(input_ids, self.config.pad_token_id).int().argmax(-1) - 1 + sequence_lengths = sequence_lengths % input_ids.shape[-1] + sequence_lengths = sequence_lengths.to(logits.device) + else: + sequence_lengths = -1 + + pooled_logits = logits[torch.arange(batch_size, device=logits.device), sequence_lengths] + + loss = None + if labels is not None: + labels = labels.to(logits.device) + if self.config.problem_type is None: + if self.num_labels == 1: + self.config.problem_type = "regression" + elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int): + self.config.problem_type = "single_label_classification" + else: + self.config.problem_type = "multi_label_classification" + + if self.config.problem_type == "regression": + loss_fct = MSELoss() + if self.num_labels == 1: + loss = loss_fct(pooled_logits.squeeze(), labels.squeeze()) + else: + loss = loss_fct(pooled_logits, labels) + elif self.config.problem_type == "single_label_classification": + loss_fct = CrossEntropyLoss() + loss = loss_fct(pooled_logits.view(-1, self.num_labels), labels.view(-1)) + elif self.config.problem_type == "multi_label_classification": + loss_fct = BCEWithLogitsLoss() + loss = loss_fct(pooled_logits, labels) + if not return_dict: + output = (pooled_logits,) + transformer_outputs[1:] + return ((loss,) + output) if loss is not None else output + + return SequenceClassifierOutputWithPast( + loss=loss, + logits=pooled_logits, + past_key_values=transformer_outputs.past_key_values, + hidden_states=transformer_outputs.hidden_states, + attentions=transformer_outputs.attentions, + ) + + +@add_start_docstrings( + """ + The Mixtral Model transformer with a token classification head on top (a linear layer on top of the hidden-states + output) e.g. for Named-Entity-Recognition (NER) tasks. + """, + MIXTRAL_START_DOCSTRING, +) +# Copied from transformers.models.llama.modeling_llama.LlamaForTokenClassification with Llama->Mixtral, LLAMA->MIXTRAL +class MixtralForTokenClassification(MixtralPreTrainedModel): + def __init__(self, config): + super().__init__(config) + self.num_labels = config.num_labels + self.model = MixtralModel(config) + if getattr(config, "classifier_dropout", None) is not None: + classifier_dropout = config.classifier_dropout + elif getattr(config, "hidden_dropout", None) is not None: + classifier_dropout = config.hidden_dropout + else: + classifier_dropout = 0.1 + self.dropout = nn.Dropout(classifier_dropout) + self.score = nn.Linear(config.hidden_size, config.num_labels) + + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self): + return self.model.embed_tokens + + def set_input_embeddings(self, value): + self.model.embed_tokens = value + + @add_start_docstrings_to_model_forward(MIXTRAL_INPUTS_DOCSTRING) + def forward( + self, + input_ids: Optional[torch.LongTensor] = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[List[torch.FloatTensor]] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + labels: Optional[torch.LongTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + ) -> Union[Tuple, TokenClassifierOutput]: + r""" + labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*): + Labels for computing the sequence classification/regression loss. Indices should be in `[0, ..., + config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If + `config.num_labels > 1` a classification loss is computed (Cross-Entropy). + """ + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + outputs = self.model( + input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + sequence_output = outputs[0] + sequence_output = self.dropout(sequence_output) + logits = self.score(sequence_output) + + loss = None + if labels is not None: + loss_fct = CrossEntropyLoss() + loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1)) + + if not return_dict: + output = (logits,) + outputs[2:] + return ((loss,) + output) if loss is not None else output + + return TokenClassifierOutput( + loss=loss, + logits=logits, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) \ No newline at end of file diff --git a/archive/kt-sft/ktransformers/models/modeling_qwen2_moe.py b/archive/kt-sft/ktransformers/models/modeling_qwen2_moe.py new file mode 100644 index 0000000..f45577c --- /dev/null +++ b/archive/kt-sft/ktransformers/models/modeling_qwen2_moe.py @@ -0,0 +1,1766 @@ +# coding=utf-8 +''' +Description : +Author : Boxin Zhang +Version : 0.1.0 +''' +# Adapted from +# https://github.com/huggingface/transformers/blob/v4.42.3/src/transformers/models/qwen2_moe/modeling_qwen2_moe.py +# Copyright 2024 The Qwen team, Alibaba Group and the HuggingFace Inc. team. All rights reserved. +# Copyright (c) 2024 by KVCache.AI, All Rights Reserved. +# +# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX +# and OPT implementations in this library. It has been modified from its +# original forms to accommodate minor architectural differences compared +# to GPT-NeoX and OPT used by the Meta AI team that trained the model. +# +# 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. +"""PyTorch Qwen2MoE model.""" + +import inspect +import math +from typing import List, Optional, Tuple, Union + +import torch +import torch.nn.functional as F +import torch.utils.checkpoint +from torch import nn +from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss + +from transformers.activations import ACT2FN +from transformers.cache_utils import Cache, DynamicCache, StaticCache +from transformers.modeling_attn_mask_utils import ( + AttentionMaskConverter, +) +from transformers.modeling_outputs import ( + MoeCausalLMOutputWithPast, + MoeModelOutputWithPast, + SequenceClassifierOutputWithPast, + TokenClassifierOutput, +) +from transformers.modeling_utils import PreTrainedModel +from transformers.utils import ( + add_start_docstrings, + add_start_docstrings_to_model_forward, + is_flash_attn_2_available, + is_flash_attn_greater_or_equal_2_10, + logging, + replace_return_docstrings, +) +from transformers.models.qwen2_moe.configuration_qwen2_moe import Qwen2MoeConfig + +from ktransformers.util.grad_wrapper import maybe_no_grad + +if is_flash_attn_2_available(): + from flash_attn import flash_attn_func, flash_attn_varlen_func, flash_attn_with_kvcache + from flash_attn.bert_padding import index_first_axis, pad_input, unpad_input # noqa + + _flash_supports_window_size = "window_size" in list(inspect.signature(flash_attn_func).parameters) + +logger = logging.get_logger(__name__) + +_CHECKPOINT_FOR_DOC = "Qwen/Qwen1.5-MoE-A2.7B" +_CONFIG_FOR_DOC = "Qwen2MoeConfig" + + +# Copied from transformers.models.mixtral.modeling_mixtral.load_balancing_loss_func +def load_balancing_loss_func( + gate_logits: torch.Tensor, num_experts: torch.Tensor = None, top_k=2, attention_mask: Optional[torch.Tensor] = None +) -> float: + r""" + Computes auxiliary load balancing loss as in Switch Transformer - implemented in Pytorch. + + See Switch Transformer (https://arxiv.org/abs/2101.03961) for more details. This function implements the loss + function presented in equations (4) - (6) of the paper. It aims at penalizing cases where the routing between + experts is too unbalanced. + + Args: + gate_logits (Union[`torch.Tensor`, Tuple[torch.Tensor]): + Logits from the `gate`, should be a tuple of model.config.num_hidden_layers tensors of + shape [batch_size X sequence_length, num_experts]. + attention_mask (`torch.Tensor`, None): + The attention_mask used in forward function + shape [batch_size X sequence_length] if not None. + num_experts (`int`, *optional*): + Number of experts + + Returns: + The auxiliary loss. + """ + if gate_logits is None or not isinstance(gate_logits, tuple): + return 0 + + if isinstance(gate_logits, tuple): + compute_device = gate_logits[0].device + concatenated_gate_logits = torch.cat([layer_gate.to(compute_device) for layer_gate in gate_logits], dim=0) + + routing_weights = torch.nn.functional.softmax(concatenated_gate_logits, dim=-1) + + _, selected_experts = torch.topk(routing_weights, top_k, dim=-1) + + expert_mask = torch.nn.functional.one_hot(selected_experts, num_experts) + + if attention_mask is None: + # Compute the percentage of tokens routed to each experts + tokens_per_expert = torch.mean(expert_mask.float(), dim=0) + + # Compute the average probability of routing to these experts + router_prob_per_expert = torch.mean(routing_weights, dim=0) + else: + batch_size, sequence_length = attention_mask.shape + num_hidden_layers = concatenated_gate_logits.shape[0] // (batch_size * sequence_length) + + # Compute the mask that masks all padding tokens as 0 with the same shape of expert_mask + expert_attention_mask = ( + attention_mask[None, :, :, None, None] + .expand((num_hidden_layers, batch_size, sequence_length, top_k, num_experts)) + .reshape(-1, top_k, num_experts) + .to(compute_device) + ) + + # Compute the percentage of tokens routed to each experts + tokens_per_expert = torch.sum(expert_mask.float() * expert_attention_mask, dim=0) / torch.sum( + expert_attention_mask, dim=0 + ) + + # Compute the mask that masks all padding tokens as 0 with the same shape of tokens_per_expert + router_per_expert_attention_mask = ( + attention_mask[None, :, :, None] + .expand((num_hidden_layers, batch_size, sequence_length, num_experts)) + .reshape(-1, num_experts) + .to(compute_device) + ) + + # Compute the average probability of routing to these experts + router_prob_per_expert = torch.sum(routing_weights * router_per_expert_attention_mask, dim=0) / torch.sum( + router_per_expert_attention_mask, dim=0 + ) + + overall_loss = torch.sum(tokens_per_expert * router_prob_per_expert.unsqueeze(0)) + return overall_loss * num_experts + + +# Copied from transformers.models.llama.modeling_llama._get_unpad_data +def _get_unpad_data(attention_mask): + seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32) + indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten() + max_seqlen_in_batch = seqlens_in_batch.max().item() + cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0)) + return ( + indices, + cu_seqlens, + max_seqlen_in_batch, + ) + + +# Copied from transformers.models.llama.modeling_llama.LlamaRMSNorm with Llama->Qwen2Moe +class Qwen2MoeRMSNorm(nn.Module): + def __init__(self, hidden_size, eps=1e-6): + """ + Qwen2MoeRMSNorm is equivalent to T5LayerNorm + """ + super().__init__() + self.weight = nn.Parameter(torch.ones(hidden_size)) + self.variance_epsilon = eps + + def forward(self, hidden_states): + input_dtype = hidden_states.dtype + hidden_states = hidden_states.to(torch.float32) + variance = hidden_states.pow(2).mean(-1, keepdim=True) + hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon) + return self.weight * hidden_states.to(input_dtype) + +# Copied from transformers.models.mistral.modeling_mistral.MistralRotaryEmbedding with Mistral->Qwen2Moe +class Qwen2MoeRotaryEmbedding(nn.Module): + def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None, scaling_factor=1.0): + super().__init__() + self.scaling_factor = scaling_factor + self.dim = dim + self.max_position_embeddings = max_position_embeddings + self.base = base + inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2, dtype=torch.int64).float().to(device) / self.dim)) + self.register_buffer("inv_freq", inv_freq, persistent=False) + # For BC we register cos and sin cached + self.max_seq_len_cached = max_position_embeddings + + @maybe_no_grad() + def forward(self, x, position_ids): + # x: [bs, num_attention_heads, seq_len, head_size] + inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1) + position_ids_expanded = position_ids[:, None, :].float() + # Force float32 since bfloat16 loses precision on long contexts + # See https://github.com/huggingface/transformers/pull/29285 + device_type = x.device.type + device_type = device_type if isinstance(device_type, str) and device_type != "mps" else "cpu" + with torch.autocast(device_type=device_type, enabled=False): + freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2) + emb = torch.cat((freqs, freqs), dim=-1) + cos = emb.cos() + sin = emb.sin() + return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype) + + +# Copied from transformers.models.llama.modeling_llama.rotate_half +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) + + +# Copied from transformers.models.mistral.modeling_mistral.apply_rotary_pos_emb +def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1): + """Applies Rotary Position Embedding to the query and key tensors. + + Args: + q (`torch.Tensor`): The query tensor. + k (`torch.Tensor`): The key tensor. + cos (`torch.Tensor`): The cosine part of the rotary embedding. + sin (`torch.Tensor`): The sine part of the rotary embedding. + position_ids (`torch.Tensor`): + Deprecated and unused. + unsqueeze_dim (`int`, *optional*, defaults to 1): + The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and + sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note + that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and + k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes + cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have + the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2. + Returns: + `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding. + """ + cos = cos.unsqueeze(unsqueeze_dim) + sin = sin.unsqueeze(unsqueeze_dim) + q_embed = (q * cos) + (rotate_half(q) * sin) + k_embed = (k * cos) + (rotate_half(k) * sin) + return q_embed, k_embed + + +# Modified from transformers.models.mistral.modeling_mistral.MistralMLP with Mistral->Qwen2Moe +class Qwen2MoeMLP(nn.Module): + def __init__(self, config, intermediate_size=None): + super().__init__() + self.config = config + self.hidden_size = config.hidden_size + self.intermediate_size = intermediate_size + self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) + self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) + self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False) + self.act_fn = ACT2FN[config.hidden_act] + + def forward(self, x): + return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x)) + + +# Copied from transformers.models.llama.modeling_llama.repeat_kv +def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: + """ + This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch, + num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim) + """ + batch, num_key_value_heads, slen, head_dim = hidden_states.shape + if n_rep == 1: + return hidden_states + hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim) + return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim) + + +# Copied from transformers.models.qwen2.modeling_qwen2.Qwen2Attention with Qwen2->Qwen2Moe +class Qwen2MoeAttention(nn.Module): + """ + Multi-headed attention from 'Attention Is All You Need' paper. Modified to use sliding window attention: Longformer + and "Generating Long Sequences with Sparse Transformers". + """ + + def __init__(self, config: Qwen2MoeConfig, layer_idx: Optional[int] = None): + super().__init__() + self.config = config + self.layer_idx = layer_idx + if layer_idx is None: + logger.warning_once( + f"Instantiating {self.__class__.__name__} without passing `layer_idx` is not recommended and will " + "to errors during the forward call, if caching is used. Please make sure to provide a `layer_idx` " + "when creating this class." + ) + + self.hidden_size = config.hidden_size + self.num_heads = config.num_attention_heads + self.head_dim = self.hidden_size // self.num_heads + self.num_key_value_heads = config.num_key_value_heads + self.num_key_value_groups = self.num_heads // self.num_key_value_heads + self.max_position_embeddings = config.max_position_embeddings + self.rope_theta = config.rope_theta + self.is_causal = True + self.attention_dropout = config.attention_dropout + + if (self.head_dim * self.num_heads) != self.hidden_size: + raise ValueError( + f"hidden_size must be divisible by num_heads (got `hidden_size`: {self.hidden_size}" + f" and `num_heads`: {self.num_heads})." + ) + self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=True) + self.k_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=True) + self.v_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=True) + self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=False) + + self.rotary_emb = Qwen2MoeRotaryEmbedding( + self.head_dim, + max_position_embeddings=self.max_position_embeddings, + base=self.rope_theta, + ) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_value: Optional[Cache] = None, + output_attentions: bool = False, + use_cache: bool = False, + cache_position: Optional[torch.LongTensor] = None, + ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: + bsz, q_len, _ = hidden_states.size() + + query_states = self.q_proj(hidden_states) + key_states = self.k_proj(hidden_states) + value_states = self.v_proj(hidden_states) + + query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2) + key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) + value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) + + kv_seq_len = key_states.shape[-2] + if past_key_value is not None: + if self.layer_idx is None: + raise ValueError( + f"The cache structure has changed since version v4.36. If you are using {self.__class__.__name__} " + "for auto-regressive decoding with k/v caching, please make sure to initialize the attention class " + "with a layer index." + ) + kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx) + cos, sin = self.rotary_emb(value_states, position_ids) + query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin) + + if past_key_value is not None: + cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position} # Specific to RoPE models + key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs) + + # repeat k/v heads if n_kv_heads < n_heads + key_states = repeat_kv(key_states, self.num_key_value_groups) + value_states = repeat_kv(value_states, self.num_key_value_groups) + + attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim) + + if attn_weights.size() != (bsz, self.num_heads, q_len, kv_seq_len): + raise ValueError( + f"Attention weights should be of size {(bsz, self.num_heads, q_len, kv_seq_len)}, but is" + f" {attn_weights.size()}" + ) + + if attention_mask is not None: # no matter the length, we just slice it + causal_mask = attention_mask[:, :, :, : key_states.shape[-2]] + attn_weights = attn_weights + causal_mask + + # upcast attention to fp32 + attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype) + attn_weights = nn.functional.dropout(attn_weights, p=self.attention_dropout, training=self.training) + attn_output = torch.matmul(attn_weights, value_states) + + if attn_output.size() != (bsz, self.num_heads, q_len, self.head_dim): + raise ValueError( + f"`attn_output` should be of size {(bsz, self.num_heads, q_len, self.head_dim)}, but is" + f" {attn_output.size()}" + ) + + attn_output = attn_output.transpose(1, 2).contiguous() + attn_output = attn_output.reshape(bsz, q_len, self.hidden_size) + + attn_output = self.o_proj(attn_output) + + if not output_attentions: + attn_weights = None + + return attn_output, attn_weights, past_key_value + + +# Copied from transformers.models.qwen2.modeling_qwen2.Qwen2FlashAttention2 with Qwen2->Qwen2Moe +class Qwen2MoeFlashAttention2(Qwen2MoeAttention): + """ + Qwen2Moe flash attention module, following Qwen2Moe attention module. This module inherits from `Qwen2MoeAttention` + as the weights of the module stays untouched. The only required change would be on the forward pass + where it needs to correctly call the public API of flash attention and deal with padding tokens + in case the input contains any of them. Additionally, for sliding window attention, we apply SWA only to the bottom + config.max_window_layers layers. + """ + + # Copied from transformers.models.llama.modeling_llama.LlamaFlashAttention2.__init__ + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + + # TODO: Should be removed once Flash Attention for RoCm is bumped to 2.1. + # flash_attn<2.1 generates top-left aligned causal mask, while what is needed here is bottom-right alignement, that was made default for flash_attn>=2.1. This attribute is used to handle this difference. Reference: https://github.com/Dao-AILab/flash-attention/releases/tag/v2.1.0. + # Beware that with flash_attn<2.1, using q_seqlen != k_seqlen (except for the case q_seqlen == 1) produces a wrong mask (top-left). + self._flash_attn_uses_top_left_mask = not is_flash_attn_greater_or_equal_2_10() + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_value: Optional[Cache] = None, + output_attentions: bool = False, + use_cache: bool = False, + cache_position: Optional[torch.LongTensor] = None, + ): + bsz, q_len, _ = hidden_states.size() + + query_states = self.q_proj(hidden_states) + key_states = self.k_proj(hidden_states) + value_states = self.v_proj(hidden_states) + + query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2) + key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) + value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) + + kv_seq_len = key_states.shape[-2] + if past_key_value is not None: + if self.layer_idx is None: + raise ValueError( + f"The cache structure has changed since version v4.36. If you are using {self.__class__.__name__} " + "for auto-regressive decoding with k/v caching, please make sure to initialize the attention class " + "with a layer index." + ) + kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx) + + cos, sin = self.rotary_emb(value_states, position_ids) + query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin) + + use_sliding_windows = ( + _flash_supports_window_size + and getattr(self.config, "sliding_window", None) is not None + and kv_seq_len > self.config.sliding_window + and self.config.use_sliding_window + ) + + if not _flash_supports_window_size: + logger.warning_once( + "The current flash attention version does not support sliding window attention, for a more memory efficient implementation" + " make sure to upgrade flash-attn library." + ) + + if past_key_value is not None: + # Activate slicing cache only if the config has a value `sliding_windows` attribute + cache_has_contents = past_key_value.get_seq_length(self.layer_idx) > 0 + if ( + getattr(self.config, "sliding_window", None) is not None + and kv_seq_len > self.config.sliding_window + and cache_has_contents + ): + slicing_tokens = 1 - self.config.sliding_window + + past_key = past_key_value[self.layer_idx][0] + past_value = past_key_value[self.layer_idx][1] + + past_key = past_key[:, :, slicing_tokens:, :].contiguous() + past_value = past_value[:, :, slicing_tokens:, :].contiguous() + + if past_key.shape[-2] != self.config.sliding_window - 1: + raise ValueError( + f"past key must have a shape of (`batch_size, num_heads, self.config.sliding_window-1, head_dim`), got" + f" {past_key.shape}" + ) + + if attention_mask is not None: + attention_mask = attention_mask[:, slicing_tokens:] + attention_mask = torch.cat([attention_mask, torch.ones_like(attention_mask[:, -1:])], dim=-1) + + cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position} # Specific to RoPE models + key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs) + # we slice the states for static kv cache to be supported in FA2. Not sure it's a must as compile fails + # for bsz == 1, avoid using slice to capture cuda graph + if cache_position is not None and q_len > 1: + key_states = key_states[:, :, : cache_position[-1] + 1, :] + value_states = value_states[:, :, : cache_position[-1] + 1, :] + + # repeat k/v heads if n_kv_heads < n_heads + key_states = repeat_kv(key_states, self.num_key_value_groups) + value_states = repeat_kv(value_states, self.num_key_value_groups) + dropout_rate = 0.0 if not self.training else self.attention_dropout + + # In PEFT, usually we cast the layer norms in float32 for training stability reasons + # therefore the input hidden states gets silently casted in float32. Hence, we need + # cast them back in float16 just to be sure everything works as expected. + input_dtype = query_states.dtype + if input_dtype == torch.float32: + if torch.is_autocast_enabled(): + target_dtype = torch.get_autocast_gpu_dtype() + # Handle the case where the model is quantized + elif hasattr(self.config, "_pre_quantization_dtype"): + target_dtype = self.config._pre_quantization_dtype + else: + target_dtype = self.q_proj.weight.dtype + + logger.warning_once( + f"The input hidden states seems to be silently casted in float32, this might be related to" + f" the fact you have upcasted embedding or layer norm layers in float32. We will cast back the input in" + f" {target_dtype}." + ) + + query_states = query_states.to(target_dtype) + key_states = key_states.to(target_dtype) + value_states = value_states.to(target_dtype) + + # Reashape to the expected shape for Flash Attention + query_states = query_states.transpose(1, 2) + key_states = key_states.transpose(1, 2) + value_states = value_states.transpose(1, 2) + + attn_output = self._flash_attention_forward( + query_states, + key_states, + value_states, + attention_mask, + q_len, + position_ids=position_ids, + dropout=dropout_rate, + use_sliding_windows=use_sliding_windows, + ) + + attn_output = attn_output.reshape(bsz, q_len, self.hidden_size).contiguous() + attn_output = self.o_proj(attn_output) + + if not output_attentions: + attn_weights = None + + return attn_output, attn_weights, past_key_value + + def _flash_attention_forward( + self, + query_states, + key_states, + value_states, + attention_mask, + query_length, + position_ids, + dropout=0.0, + softmax_scale=None, + use_sliding_windows=False, + ): + """ + Calls the forward method of Flash Attention - if the input hidden states contain at least one padding token + first unpad the input, then computes the attention scores and pad the final attention scores. + + Args: + query_states (`torch.Tensor`): + Input query states to be passed to Flash Attention API + key_states (`torch.Tensor`): + Input key states to be passed to Flash Attention API + value_states (`torch.Tensor`): + Input value states to be passed to Flash Attention API + attention_mask (`torch.Tensor`): + The padding mask - corresponds to a tensor of size `(batch_size, seq_len)` where 0 stands for the + position of padding tokens and 1 for the position of non-padding tokens. + dropout (`float`): + Attention dropout + softmax_scale (`float`, *optional*): + The scaling of QK^T before applying softmax. Default to 1 / sqrt(head_dim) + use_sliding_windows (`bool`, *optional*): + Whether to activate sliding window attention. + """ + if not self._flash_attn_uses_top_left_mask: + causal = self.is_causal + else: + # TODO: Remove the `query_length != 1` check once Flash Attention for RoCm is bumped to 2.1. For details, please see the comment in LlamaFlashAttention2 __init__. + causal = self.is_causal and query_length != 1 + + # Decide whether to use SWA or not by layer index. + if use_sliding_windows and self.layer_idx >= self.config.max_window_layers: + use_sliding_windows = False + + # Contains at least one padding token in the sequence + if attention_mask is not None: + batch_size = query_states.shape[0] + query_states, key_states, value_states, indices_q, cu_seq_lens, max_seq_lens = self._upad_input( + query_states, key_states, value_states, attention_mask, query_length + ) + + cu_seqlens_q, cu_seqlens_k = cu_seq_lens + max_seqlen_in_batch_q, max_seqlen_in_batch_k = max_seq_lens + + if not use_sliding_windows: + attn_output_unpad = flash_attn_varlen_func( + query_states, + key_states, + value_states, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k=cu_seqlens_k, + max_seqlen_q=max_seqlen_in_batch_q, + max_seqlen_k=max_seqlen_in_batch_k, + dropout_p=dropout, + softmax_scale=softmax_scale, + causal=causal, + ) + else: + attn_output_unpad = flash_attn_varlen_func( + query_states, + key_states, + value_states, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k=cu_seqlens_k, + max_seqlen_q=max_seqlen_in_batch_q, + max_seqlen_k=max_seqlen_in_batch_k, + dropout_p=dropout, + softmax_scale=softmax_scale, + causal=causal, + window_size=(self.config.sliding_window, self.config.sliding_window), + ) + + attn_output = pad_input(attn_output_unpad, indices_q, batch_size, query_length) + else: + if not use_sliding_windows: + if query_length == 1: + position_ids = position_ids.to(dtype=torch.int32).squeeze(1) + attn_output = flash_attn_with_kvcache( + query_states, + key_states, + value_states, + cache_seqlens=position_ids, + softmax_scale=softmax_scale, + causal=causal, + ) + else: + attn_output = flash_attn_func( + query_states, + key_states, + value_states, + dropout, + softmax_scale=softmax_scale, + causal=causal, + ) + else: + attn_output = flash_attn_func( + query_states, + key_states, + value_states, + dropout, + softmax_scale=softmax_scale, + causal=causal, + window_size=(self.config.sliding_window, self.config.sliding_window), + ) + + return attn_output + + # Copied from transformers.models.mistral.modeling_mistral.MistralFlashAttention2._upad_input + def _upad_input(self, query_layer, key_layer, value_layer, attention_mask, query_length): + batch_size, kv_seq_len, num_heads, head_dim = key_layer.shape + + # On the first iteration we need to properly re-create the padding mask + # by slicing it on the proper place + if kv_seq_len != attention_mask.shape[-1]: + attention_mask_num_tokens = attention_mask.shape[-1] + attention_mask = attention_mask[:, attention_mask_num_tokens - kv_seq_len :] + + indices_k, cu_seqlens_k, max_seqlen_in_batch_k = _get_unpad_data(attention_mask) + + key_layer = index_first_axis(key_layer.reshape(batch_size * kv_seq_len, num_heads, head_dim), indices_k) + value_layer = index_first_axis(value_layer.reshape(batch_size * kv_seq_len, num_heads, head_dim), indices_k) + + if query_length == kv_seq_len: + query_layer = index_first_axis( + query_layer.reshape(batch_size * kv_seq_len, num_heads, head_dim), indices_k + ) + cu_seqlens_q = cu_seqlens_k + max_seqlen_in_batch_q = max_seqlen_in_batch_k + indices_q = indices_k + elif query_length == 1: + max_seqlen_in_batch_q = 1 + cu_seqlens_q = torch.arange( + batch_size + 1, dtype=torch.int32, device=query_layer.device + ) # There is a memcpy here, that is very bad. + indices_q = cu_seqlens_q[:-1] + query_layer = query_layer.squeeze(1) + else: + # The -q_len: slice assumes left padding. + attention_mask = attention_mask[:, -query_length:] + query_layer, indices_q, cu_seqlens_q, max_seqlen_in_batch_q = unpad_input(query_layer, attention_mask) + + return ( + query_layer, + key_layer, + value_layer, + indices_q, + (cu_seqlens_q, cu_seqlens_k), + (max_seqlen_in_batch_q, max_seqlen_in_batch_k), + ) + + +# Copied from transformers.models.mixtral.modeling_mixtral.MixtralSdpaAttention with Mixtral->Qwen2Moe +class Qwen2MoeSdpaAttention(Qwen2MoeAttention): + """ + Qwen2Moe attention module using torch.nn.functional.scaled_dot_product_attention. This module inherits from + `Qwen2MoeAttention` as the weights of the module stays untouched. The only changes are on the forward pass to adapt to + SDPA API. + """ + + # Adapted from Qwen2MoeAttention.forward + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_value: Optional[Cache] = None, + output_attentions: bool = False, + use_cache: bool = False, + cache_position: Optional[torch.LongTensor] = None, + ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: + if output_attentions: + # TODO: Improve this warning with e.g. `model.config.attn_implementation = "manual"` once this is implemented. + logger.warning_once( + "Qwen2MoeModel is using Qwen2MoeSdpaAttention, but `torch.nn.functional.scaled_dot_product_attention` does not support `output_attentions=True`. Falling back to the manual attention implementation, " + 'but specifying the manual implementation will be required from Transformers version v5.0.0 onwards. This warning can be removed using the argument `attn_implementation="eager"` when loading the model.' + ) + return super().forward( + hidden_states=hidden_states, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_value=past_key_value, + output_attentions=output_attentions, + use_cache=use_cache, + ) + + bsz, q_len, _ = hidden_states.size() + + query_states = self.q_proj(hidden_states) + key_states = self.k_proj(hidden_states) + value_states = self.v_proj(hidden_states) + + query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2) + key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) + value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) + + kv_seq_len = key_states.shape[-2] + if past_key_value is not None: + kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx) + cos, sin = self.rotary_emb(value_states, position_ids) + query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin) + + if past_key_value is not None: + cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position} # Specific to RoPE models + key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs) + + key_states = repeat_kv(key_states, self.num_key_value_groups) + value_states = repeat_kv(value_states, self.num_key_value_groups) + + causal_mask = attention_mask + if attention_mask is not None: # no matter the length, we just slice it + causal_mask = attention_mask[:, :, :, : key_states.shape[-2]] + + # SDPA with memory-efficient backend is currently (torch==2.1.2) bugged with non-contiguous inputs with custom attn_mask, + # Reference: https://github.com/pytorch/pytorch/issues/112577. + if query_states.device.type == "cuda" and attention_mask is not None: + query_states = query_states.contiguous() + key_states = key_states.contiguous() + value_states = value_states.contiguous() + + # We dispatch to SDPA's Flash Attention or Efficient kernels via this `is_causal` if statement instead of an inline conditional assignment + # in SDPA to support both torch.compile's dynamic shapes and full graph options. An inline conditional prevents dynamic shapes from compiling. + # The q_len > 1 is necessary to match with AttentionMaskConverter.to_causal_4d that does not create a causal mask in case q_len == 1. + is_causal = True if causal_mask is None and q_len > 1 else False + + attn_output = torch.nn.functional.scaled_dot_product_attention( + query_states, + key_states, + value_states, + attn_mask=causal_mask, + dropout_p=self.attention_dropout if self.training else 0.0, + is_causal=is_causal, + ) + + attn_output = attn_output.transpose(1, 2).contiguous() + attn_output = attn_output.view(bsz, q_len, self.hidden_size) + + attn_output = self.o_proj(attn_output) + + return attn_output, None, past_key_value + + +QWEN2MOE_ATTENTION_CLASSES = { + "eager": Qwen2MoeAttention, + "flash_attention_2": Qwen2MoeFlashAttention2, + "sdpa": Qwen2MoeSdpaAttention, +} + + +class Qwen2MoeSparseMoeBlock(nn.Module): + def __init__(self, config): + super().__init__() + self.num_experts = config.num_experts + self.top_k = config.num_experts_per_tok + self.norm_topk_prob = config.norm_topk_prob + + # gating + self.gate = nn.Linear(config.hidden_size, config.num_experts, bias=False) + self.experts = nn.ModuleList( + [Qwen2MoeMLP(config, intermediate_size=config.moe_intermediate_size) for _ in range(self.num_experts)] + ) + + self.shared_expert = Qwen2MoeMLP(config, intermediate_size=config.shared_expert_intermediate_size) + self.shared_expert_gate = torch.nn.Linear(config.hidden_size, 1, bias=False) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + """ """ + batch_size, sequence_length, hidden_dim = hidden_states.shape + hidden_states = hidden_states.view(-1, hidden_dim) + # router_logits: (batch * sequence_length, n_experts) + router_logits = self.gate(hidden_states) + + routing_weights = F.softmax(router_logits, dim=1, dtype=torch.float) + routing_weights, selected_experts = torch.topk(routing_weights, self.top_k, dim=-1) + if self.norm_topk_prob: + routing_weights /= routing_weights.sum(dim=-1, keepdim=True) + # we cast back to the input dtype + routing_weights = routing_weights.to(hidden_states.dtype) + + final_hidden_states = torch.zeros( + (batch_size * sequence_length, hidden_dim), dtype=hidden_states.dtype, device=hidden_states.device + ) + + # One hot encode the selected experts to create an expert mask + # this will be used to easily index which expert is going to be sollicitated + expert_mask = torch.nn.functional.one_hot(selected_experts, num_classes=self.num_experts).permute(2, 1, 0) + + # Loop over all available experts in the model and perform the computation on each expert + for expert_idx in range(self.num_experts): + expert_layer = self.experts[expert_idx] + idx, top_x = torch.where(expert_mask[expert_idx]) + + # Index the correct hidden states and compute the expert hidden state for + # the current expert. We need to make sure to multiply the output hidden + # states by `routing_weights` on the corresponding tokens (top-1 and top-2) + current_state = hidden_states[None, top_x].reshape(-1, hidden_dim) + current_hidden_states = expert_layer(current_state) * routing_weights[top_x, idx, None] + + # However `index_add_` only support torch tensors for indexing so we'll use + # the `top_x` tensor here. + final_hidden_states.index_add_(0, top_x, current_hidden_states.to(hidden_states.dtype)) + + shared_expert_output = self.shared_expert(hidden_states) + shared_expert_output = F.sigmoid(self.shared_expert_gate(hidden_states)) * shared_expert_output + + final_hidden_states = final_hidden_states + shared_expert_output + + final_hidden_states = final_hidden_states.reshape(batch_size, sequence_length, hidden_dim) + return final_hidden_states, router_logits + + +class Qwen2MoeDecoderLayer(nn.Module): + def __init__(self, config: Qwen2MoeConfig, layer_idx: int): + super().__init__() + self.hidden_size = config.hidden_size + + self.self_attn = QWEN2MOE_ATTENTION_CLASSES[config._attn_implementation](config, layer_idx) + + if (layer_idx not in config.mlp_only_layers) and ( + config.num_experts > 0 and (layer_idx + 1) % config.decoder_sparse_step == 0 + ): + self.mlp = Qwen2MoeSparseMoeBlock(config) + else: + self.mlp = Qwen2MoeMLP(config, intermediate_size=config.intermediate_size) + + self.input_layernorm = Qwen2MoeRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = Qwen2MoeRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_value: Optional[Tuple[torch.Tensor]] = None, + output_attentions: Optional[bool] = False, + output_router_logits: Optional[bool] = False, + use_cache: Optional[bool] = False, + cache_position: Optional[torch.LongTensor] = None, + **kwargs, + ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]: + """ + Args: + hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)` + attention_mask (`torch.FloatTensor`, *optional*): attention mask of size + `(batch, sequence_length)` where padding elements are indicated by 0. + output_attentions (`bool`, *optional*): + Whether or not to return the attentions tensors of all attention layers. See `attentions` under + returned tensors for more detail. + output_router_logits (`bool`, *optional*): + Whether or not to return the logits of all the routers. They are useful for computing the router loss, + and should not be returned during inference. + use_cache (`bool`, *optional*): + If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding + (see `past_key_values`). + past_key_value (`Tuple(torch.FloatTensor)`, *optional*): cached past key and value projection states + cache_position (`torch.LongTensor` of shape `(sequence_length)`, *optional*): + Indices depicting the position of the input sequence tokens in the sequence. + kwargs (`dict`, *optional*): + Arbitrary kwargs to be ignored, used for FSDP and other methods that injects code + into the model + """ + + residual = hidden_states + + hidden_states = self.input_layernorm(hidden_states) + + # Self Attention + hidden_states, self_attn_weights, present_key_value = self.self_attn( + hidden_states=hidden_states, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_value=past_key_value, + output_attentions=output_attentions, + use_cache=use_cache, + cache_position=cache_position, + ) + hidden_states = residual + hidden_states + + # Fully Connected + residual = hidden_states + hidden_states = self.post_attention_layernorm(hidden_states) + + hidden_states = self.mlp(hidden_states) + if isinstance(hidden_states, tuple): + hidden_states, router_logits = hidden_states + else: + router_logits = None + + hidden_states = residual + hidden_states + + outputs = (hidden_states,) + + if output_attentions: + outputs += (self_attn_weights,) + + if use_cache: + outputs += (present_key_value,) + + if output_router_logits: + outputs += (router_logits,) + + return outputs + + +QWEN2MOE_START_DOCSTRING = r""" + This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the + library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads + etc.) + + This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass. + Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage + and behavior. + + Parameters: + config ([`Qwen2MoeConfig`]): + Model configuration class with all the parameters of the model. Initializing with a config file does not + load the weights associated with the model, only the configuration. Check out the + [`~PreTrainedModel.from_pretrained`] method to load the model weights. +""" + + +@add_start_docstrings( + "The bare Qwen2MoE Model outputting raw hidden-states without any specific head on top.", + QWEN2MOE_START_DOCSTRING, +) +class Qwen2MoePreTrainedModel(PreTrainedModel): + config_class = Qwen2MoeConfig + base_model_prefix = "model" + supports_gradient_checkpointing = True + _no_split_modules = ["Qwen2MoeDecoderLayer"] + _skip_keys_device_placement = "past_key_values" + _supports_flash_attn_2 = True + _supports_sdpa = True + _supports_cache_class = True + _supports_static_cache = True + + def _init_weights(self, module): + std = self.config.initializer_range + if isinstance(module, nn.Linear): + module.weight.data.normal_(mean=0.0, std=std) + if module.bias is not None: + module.bias.data.zero_() + elif isinstance(module, nn.Embedding): + module.weight.data.normal_(mean=0.0, std=std) + if module.padding_idx is not None: + module.weight.data[module.padding_idx].zero_() + + +QWEN2MOE_INPUTS_DOCSTRING = r""" + Args: + input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): + Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide + it. + + Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and + [`PreTrainedTokenizer.__call__`] for details. + + [What are input IDs?](../glossary#input-ids) + attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*): + Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`: + + - 1 for tokens that are **not masked**, + - 0 for tokens that are **masked**. + + [What are attention masks?](../glossary#attention-mask) + + Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and + [`PreTrainedTokenizer.__call__`] for details. + + If `past_key_values` is used, optionally only the last `decoder_input_ids` have to be input (see + `past_key_values`). + + If you want to change padding behavior, you should read [`modeling_opt._prepare_decoder_attention_mask`] + and modify to your needs. See diagram 1 in [the paper](https://arxiv.org/abs/1910.13461) for more + information on the default strategy. + + - 1 indicates the head is **not masked**, + - 0 indicates the head is **masked**. + position_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0, + config.n_positions - 1]`. + + [What are position IDs?](../glossary#position-ids) + past_key_values (`Cache` or `tuple(tuple(torch.FloatTensor))`, *optional*): + Pre-computed hidden-states (key and values in the self-attention blocks and in the cross-attention + blocks) that can be used to speed up sequential decoding. This typically consists in the `past_key_values` + returned by the model at a previous stage of decoding, when `use_cache=True` or `config.use_cache=True`. + + Two formats are allowed: + - a [`~cache_utils.Cache`] instance; + - Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of + shape `(batch_size, num_heads, sequence_length, embed_size_per_head)`). This is also known as the legacy + cache format. + + The model will output the same cache format that is fed as input. If no `past_key_values` are passed, the + legacy cache format will be returned. + + If `past_key_values` are used, the user can optionally input only the last `input_ids` (those that don't + have their past key value states given to this model) of shape `(batch_size, 1)` instead of all `input_ids` + of shape `(batch_size, sequence_length)`. + inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*): + Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This + is useful if you want more control over how to convert `input_ids` indices into associated vectors than the + model's internal embedding lookup matrix. + use_cache (`bool`, *optional*): + If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see + `past_key_values`). + output_attentions (`bool`, *optional*): + Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned + tensors for more detail. + output_hidden_states (`bool`, *optional*): + Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for + more detail. + output_router_logits (`bool`, *optional*): + Whether or not to return the logits of all the routers. They are useful for computing the router loss, and + should not be returned during inference. + return_dict (`bool`, *optional*): + Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple. + cache_position (`torch.LongTensor` of shape `(sequence_length)`, *optional*): + Indices depicting the position of the input sequence tokens in the sequence. Contrarily to `position_ids`, + this tensor is not affected by padding. It is used to update the cache in the correct position and to infer + the complete sequence length. +""" + + +@add_start_docstrings( + "The bare Qwen2MoE Model outputting raw hidden-states without any specific head on top.", + QWEN2MOE_START_DOCSTRING, +) +class Qwen2MoeModel(Qwen2MoePreTrainedModel): + """ + Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`Qwen2MoeDecoderLayer`] + + Args: + config: Qwen2MoeConfig + """ + + def __init__(self, config: Qwen2MoeConfig): + super().__init__(config) + self.padding_idx = config.pad_token_id + self.vocab_size = config.vocab_size + + self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx) + self.layers = nn.ModuleList( + [Qwen2MoeDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)] + ) + self._attn_implementation = config._attn_implementation + self.norm = Qwen2MoeRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + self.gradient_checkpointing = False + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self): + return self.embed_tokens + + def set_input_embeddings(self, value): + self.embed_tokens = value + + @add_start_docstrings_to_model_forward(QWEN2MOE_INPUTS_DOCSTRING) + def forward( + self, + input_ids: torch.LongTensor = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[List[torch.FloatTensor]] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + output_router_logits: Optional[bool] = None, + return_dict: Optional[bool] = None, + cache_position: Optional[torch.LongTensor] = None, + ) -> Union[Tuple, MoeModelOutputWithPast]: + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_router_logits = ( + output_router_logits if output_router_logits is not None else self.config.output_router_logits + ) + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + use_cache = use_cache if use_cache is not None else self.config.use_cache + + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + if (input_ids is None) ^ (inputs_embeds is not None): + raise ValueError( + "You cannot specify both input_ids and inputs_embeds at the same time, and must specify either one" + ) + + if self.gradient_checkpointing and self.training: + if use_cache: + logger.warning_once( + "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..." + ) + use_cache = False + + use_legacy_cache = False + if use_cache and not isinstance(past_key_values, Cache): + use_legacy_cache = True + past_key_values = DynamicCache.from_legacy_cache(past_key_values) + logger.warning_once( + "We detected that you are passing `past_key_values` as a tuple and this is deprecated and will be removed in v4.43. " + "Please use an appropriate `Cache` class (https://huggingface.co/docs/transformers/v4.41.3/en/internal/generation_utils#transformers.Cache)" + ) + + if inputs_embeds is None: + inputs_embeds = self.embed_tokens(input_ids) + + if cache_position is None: + past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 + cache_position = torch.arange( + past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device + ) + if position_ids is None: + position_ids = cache_position.unsqueeze(0) + + causal_mask = self._update_causal_mask( + attention_mask, inputs_embeds, cache_position, past_key_values, output_attentions + ) + + hidden_states = inputs_embeds + + # decoder layers + all_hidden_states = () if output_hidden_states else None + all_self_attns = () if output_attentions else None + all_router_logits = () if output_router_logits else None + next_decoder_cache = None + + for decoder_layer in self.layers: + if output_hidden_states: + all_hidden_states += (hidden_states,) + + if self.gradient_checkpointing and self.training: + layer_outputs = self._gradient_checkpointing_func( + decoder_layer.__call__, + hidden_states, + causal_mask, + position_ids, + past_key_values, + output_attentions, + output_router_logits, + use_cache, + cache_position, + ) + else: + layer_outputs = decoder_layer( + hidden_states, + attention_mask=causal_mask, + position_ids=position_ids, + past_key_value=past_key_values, + output_attentions=output_attentions, + output_router_logits=output_router_logits, + use_cache=use_cache, + cache_position=cache_position, + ) + + hidden_states = layer_outputs[0] + + if use_cache: + next_decoder_cache = layer_outputs[2 if output_attentions else 1] + + if output_attentions: + all_self_attns += (layer_outputs[1],) + + if output_router_logits and layer_outputs[-1] is not None: + all_router_logits += (layer_outputs[-1],) + + hidden_states = self.norm(hidden_states) + + # add hidden states from the last decoder layer + if output_hidden_states: + all_hidden_states += (hidden_states,) + + next_cache = None + if use_cache: + next_cache = next_decoder_cache.to_legacy_cache() if use_legacy_cache else next_decoder_cache + + if not return_dict: + return tuple( + v + for v in [hidden_states, next_cache, all_hidden_states, all_self_attns, all_router_logits] + if v is not None + ) + return MoeModelOutputWithPast( + last_hidden_state=hidden_states, + past_key_values=next_cache, + hidden_states=all_hidden_states, + attentions=all_self_attns, + router_logits=all_router_logits, + ) + + # Copied from transformers.models.llama.modeling_llama.LlamaModel._update_causal_mask + def _update_causal_mask( + self, + attention_mask: torch.Tensor, + input_tensor: torch.Tensor, + cache_position: torch.Tensor, + past_key_values: Cache, + output_attentions: bool, + ): + # TODO: As of torch==2.2.0, the `attention_mask` passed to the model in `generate` is 2D and of dynamic length even when the static + # KV cache is used. This is an issue for torch.compile which then recaptures cudagraphs at each decode steps due to the dynamic shapes. + # (`recording cudagraph tree for symint key 13`, etc.), which is VERY slow. A workaround is `@torch.compiler.disable`, but this prevents using + # `fullgraph=True`. See more context in https://github.com/huggingface/transformers/pull/29114 + + if self.config._attn_implementation == "flash_attention_2": + if attention_mask is not None and 0.0 in attention_mask: + return attention_mask + return None + + # For SDPA, when possible, we will rely on its `is_causal` argument instead of its `attn_mask` argument, in + # order to dispatch on Flash Attention 2. This feature is not compatible with static cache, as SDPA will fail + # to infer the attention mask. + past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 + using_static_cache = isinstance(past_key_values, StaticCache) + + # When output attentions is True, sdpa implementation's forward method calls the eager implementation's forward + if self.config._attn_implementation == "sdpa" and not using_static_cache and not output_attentions: + if AttentionMaskConverter._ignore_causal_mask_sdpa( + attention_mask, + inputs_embeds=input_tensor, + past_key_values_length=past_seen_tokens, + is_training=self.training, + ): + return None + + dtype, device = input_tensor.dtype, input_tensor.device + min_dtype = torch.finfo(dtype).min + sequence_length = input_tensor.shape[1] + if using_static_cache: + target_length = past_key_values.get_max_length() + else: + target_length = ( + attention_mask.shape[-1] + if isinstance(attention_mask, torch.Tensor) + else past_seen_tokens + sequence_length + 1 + ) + + if attention_mask is not None and attention_mask.dim() == 4: + # in this case we assume that the mask comes already in inverted form and requires no inversion or slicing + if attention_mask.max() != 0: + raise ValueError("Custom 4D attention mask should be passed in inverted form with max==0`") + causal_mask = attention_mask + else: + causal_mask = torch.full( + (sequence_length, target_length), fill_value=min_dtype, dtype=dtype, device=device + ) + if sequence_length != 1: + causal_mask = torch.triu(causal_mask, diagonal=1) + causal_mask *= torch.arange(target_length, device=device) > cache_position.reshape(-1, 1) + causal_mask = causal_mask[None, None, :, :].expand(input_tensor.shape[0], 1, -1, -1) + if attention_mask is not None: + causal_mask = causal_mask.clone() # copy to contiguous memory for in-place edit + mask_length = attention_mask.shape[-1] + padding_mask = causal_mask[:, :, :, :mask_length] + attention_mask[:, None, None, :] + padding_mask = padding_mask == 0 + causal_mask[:, :, :, :mask_length] = causal_mask[:, :, :, :mask_length].masked_fill( + padding_mask, min_dtype + ) + if ( + self.config._attn_implementation == "sdpa" + and attention_mask is not None + and attention_mask.device.type == "cuda" + and not output_attentions + ): + # Attend to all tokens in fully masked rows in the causal_mask, for example the relevant first rows when + # using left padding. This is required by F.scaled_dot_product_attention memory-efficient attention path. + # Details: https://github.com/pytorch/pytorch/issues/110213 + causal_mask = AttentionMaskConverter._unmask_unattended(causal_mask, min_dtype) + + return causal_mask + + +class Qwen2MoeForCausalLM(Qwen2MoePreTrainedModel): + _tied_weights_keys = ["lm_head.weight"] + + def __init__(self, config): + super().__init__(config) + self.model = Qwen2MoeModel(config) + self.vocab_size = config.vocab_size + self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) + + self.router_aux_loss_coef = config.router_aux_loss_coef + self.num_experts = config.num_experts + self.num_experts_per_tok = config.num_experts_per_tok + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self): + return self.model.embed_tokens + + def set_input_embeddings(self, value): + self.model.embed_tokens = value + + def get_output_embeddings(self): + return self.lm_head + + def set_output_embeddings(self, new_embeddings): + self.lm_head = new_embeddings + + def set_decoder(self, decoder): + self.model = decoder + + def get_decoder(self): + return self.model + + @add_start_docstrings_to_model_forward(QWEN2MOE_INPUTS_DOCSTRING) + @replace_return_docstrings(output_type=MoeCausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC) + def forward( + self, + input_ids: torch.LongTensor = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[List[torch.FloatTensor]] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + labels: Optional[torch.LongTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + output_router_logits: Optional[bool] = None, + return_dict: Optional[bool] = None, + cache_position: Optional[torch.LongTensor] = None, + ) -> Union[Tuple, MoeCausalLMOutputWithPast]: + r""" + Args: + labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Labels for computing the masked language modeling loss. Indices should either be in `[0, ..., + config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored + (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`. + + Returns: + + Example: + + ```python + >>> from transformers import AutoTokenizer, Qwen2MoeForCausalLM + + >>> model = Qwen2MoeForCausalLM.from_pretrained(PATH_TO_CONVERTED_WEIGHTS) + >>> tokenizer = AutoTokenizer.from_pretrained(PATH_TO_CONVERTED_TOKENIZER) + + >>> prompt = "Hey, are you conscious? Can you talk to me?" + >>> inputs = tokenizer(prompt, return_tensors="pt") + + >>> # Generate + >>> generate_ids = model.generate(inputs.input_ids, max_length=30) + >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0] + "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you." + ```""" + + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_router_logits = ( + output_router_logits if output_router_logits is not None else self.config.output_router_logits + ) + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn) + outputs = self.model( + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + output_router_logits=output_router_logits, + return_dict=return_dict, + cache_position=cache_position, + ) + + hidden_states = outputs[0] + logits = self.lm_head(hidden_states) + logits = logits.float() + + loss = None + if labels is not None: + # Shift so that tokens < n predict n + shift_logits = logits[..., :-1, :].contiguous() + shift_labels = labels[..., 1:].contiguous() + # Flatten the tokens + loss_fct = CrossEntropyLoss() + shift_logits = shift_logits.view(-1, self.config.vocab_size) + shift_labels = shift_labels.view(-1) + # Enable model parallelism + shift_labels = shift_labels.to(shift_logits.device) + loss = loss_fct(shift_logits, shift_labels) + + aux_loss = None + if output_router_logits: + aux_loss = load_balancing_loss_func( + outputs.router_logits if return_dict else outputs[-1], + self.num_experts, + self.num_experts_per_tok, + attention_mask, + ) + if labels is not None: + loss += self.router_aux_loss_coef * aux_loss.to(loss.device) # make sure to reside in the same device + + if not return_dict: + output = (logits,) + outputs[1:] + if output_router_logits: + output = (aux_loss,) + output + return (loss,) + output if loss is not None else output + + return MoeCausalLMOutputWithPast( + loss=loss, + aux_loss=aux_loss, + logits=logits, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + router_logits=outputs.router_logits, + ) + + def prepare_inputs_for_generation( + self, + input_ids, + past_key_values=None, + attention_mask=None, + inputs_embeds=None, + cache_position=None, + use_cache=True, + **kwargs, + ): + past_length = 0 + # Omit tokens covered by past_key_values + if past_key_values is not None: + if isinstance(past_key_values, Cache): + past_length = cache_position[0] if cache_position is not None else past_key_values.get_seq_length() + max_cache_length = ( + torch.tensor(past_key_values.get_max_length(), device=input_ids.device) + if past_key_values.get_max_length() is not None + else None + ) + cache_length = past_length if max_cache_length is None else torch.min(max_cache_length, past_length) + # TODO joao: remove this `else` after `generate` prioritizes `Cache` objects + else: + cache_length = past_length = past_key_values[0][0].shape[2] + max_cache_length = None + + # Keep only the unprocessed tokens: + # 1 - If the length of the attention_mask exceeds the length of input_ids, then we are in a setting where + # some of the inputs are exclusively passed as part of the cache (e.g. when passing input_embeds as + # input) + if attention_mask is not None and attention_mask.shape[1] > input_ids.shape[1]: + input_ids = input_ids[:, -(attention_mask.shape[1] - past_length) :] + # 2 - If the past_length is smaller than input_ids', then input_ids holds all input tokens. We can discard + # input_ids based on the past_length. + elif past_length < input_ids.shape[1]: + input_ids = input_ids[:, past_length:] + # 3 - Otherwise (past_length >= input_ids.shape[1]), let's assume input_ids only has unprocessed tokens. + + # If we are about to go beyond the maximum cache length, we need to crop the input attention mask. + if ( + max_cache_length is not None + and attention_mask is not None + and cache_length + input_ids.shape[1] > max_cache_length + ): + attention_mask = attention_mask[:, -max_cache_length:] + + position_ids = kwargs.get("position_ids", None) + if attention_mask is not None and position_ids is None: + # create position_ids on the fly for batch generation + position_ids = attention_mask.long().cumsum(-1) - 1 + position_ids.masked_fill_(attention_mask == 0, 1) + if past_key_values: + position_ids = position_ids[:, -input_ids.shape[1] :] + + # if `inputs_embeds` are passed, we only want to use them in the 1st generation step + if inputs_embeds is not None and past_length == 0: + model_inputs = {"inputs_embeds": inputs_embeds} + else: + model_inputs = {"input_ids": input_ids} + + input_length = position_ids.shape[-1] if position_ids is not None else input_ids.shape[-1] + if cache_position is None: + cache_position = torch.arange(past_length, past_length + input_length, device=input_ids.device) + elif use_cache: + cache_position = cache_position[-input_length:] + + model_inputs.update( + { + "position_ids": position_ids, + "past_key_values": past_key_values, + "use_cache": use_cache, + "attention_mask": attention_mask, + "cache_position": cache_position, + } + ) + return model_inputs + + @staticmethod + def _reorder_cache(past_key_values, beam_idx): + reordered_past = () + for layer_past in past_key_values: + reordered_past += ( + tuple(past_state.index_select(0, beam_idx.to(past_state.device)) for past_state in layer_past), + ) + return reordered_past + + +@add_start_docstrings( + """ + The Qwen2MoE Model transformer with a sequence classification head on top (linear layer). + + [`Qwen2MoeForSequenceClassification`] uses the last token in order to do the classification, as other causal models + (e.g. GPT-2) do. + + Since it does classification on the last token, it requires to know the position of the last token. If a + `pad_token_id` is defined in the configuration, it finds the last token that is not a padding token in each row. If + no `pad_token_id` is defined, it simply takes the last value in each row of the batch. Since it cannot guess the + padding tokens when `inputs_embeds` are passed instead of `input_ids`, it does the same (take the last value in + each row of the batch). + """, + QWEN2MOE_START_DOCSTRING, +) +# Copied from transformers.models.llama.modeling_llama.LlamaForSequenceClassification with Llama->Qwen2Moe, LLAMA->QWEN2MOE +class Qwen2MoeForSequenceClassification(Qwen2MoePreTrainedModel): + def __init__(self, config): + super().__init__(config) + self.num_labels = config.num_labels + self.model = Qwen2MoeModel(config) + self.score = nn.Linear(config.hidden_size, self.num_labels, bias=False) + + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self): + return self.model.embed_tokens + + def set_input_embeddings(self, value): + self.model.embed_tokens = value + + @add_start_docstrings_to_model_forward(QWEN2MOE_INPUTS_DOCSTRING) + def forward( + self, + input_ids: torch.LongTensor = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[Union[Cache, List[torch.FloatTensor]]] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + labels: Optional[torch.LongTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + ) -> Union[Tuple, SequenceClassifierOutputWithPast]: + r""" + labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*): + Labels for computing the sequence classification/regression loss. Indices should be in `[0, ..., + config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If + `config.num_labels > 1` a classification loss is computed (Cross-Entropy). + """ + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + transformer_outputs = self.model( + input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + hidden_states = transformer_outputs[0] + logits = self.score(hidden_states) + + if input_ids is not None: + batch_size = input_ids.shape[0] + else: + batch_size = inputs_embeds.shape[0] + + if self.config.pad_token_id is None and batch_size != 1: + raise ValueError("Cannot handle batch sizes > 1 if no padding token is defined.") + if self.config.pad_token_id is None: + sequence_lengths = -1 + else: + if input_ids is not None: + # if no pad token found, use modulo instead of reverse indexing for ONNX compatibility + sequence_lengths = torch.eq(input_ids, self.config.pad_token_id).int().argmax(-1) - 1 + sequence_lengths = sequence_lengths % input_ids.shape[-1] + sequence_lengths = sequence_lengths.to(logits.device) + else: + sequence_lengths = -1 + + pooled_logits = logits[torch.arange(batch_size, device=logits.device), sequence_lengths] + + loss = None + if labels is not None: + labels = labels.to(logits.device) + if self.config.problem_type is None: + if self.num_labels == 1: + self.config.problem_type = "regression" + elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int): + self.config.problem_type = "single_label_classification" + else: + self.config.problem_type = "multi_label_classification" + + if self.config.problem_type == "regression": + loss_fct = MSELoss() + if self.num_labels == 1: + loss = loss_fct(pooled_logits.squeeze(), labels.squeeze()) + else: + loss = loss_fct(pooled_logits, labels) + elif self.config.problem_type == "single_label_classification": + loss_fct = CrossEntropyLoss() + loss = loss_fct(pooled_logits.view(-1, self.num_labels), labels.view(-1)) + elif self.config.problem_type == "multi_label_classification": + loss_fct = BCEWithLogitsLoss() + loss = loss_fct(pooled_logits, labels) + if not return_dict: + output = (pooled_logits,) + transformer_outputs[1:] + return ((loss,) + output) if loss is not None else output + + return SequenceClassifierOutputWithPast( + loss=loss, + logits=pooled_logits, + past_key_values=transformer_outputs.past_key_values, + hidden_states=transformer_outputs.hidden_states, + attentions=transformer_outputs.attentions, + ) + + +@add_start_docstrings( + """ + The Qwen2MoE Model transformer with a token classification head on top (a linear layer on top of the hidden-states + output) e.g. for Named-Entity-Recognition (NER) tasks. + """, + QWEN2MOE_START_DOCSTRING, +) +# Copied from transformers.models.llama.modeling_llama.LlamaForTokenClassification with Llama->Qwen2Moe, LLAMA->QWEN2MOE +class Qwen2MoeForTokenClassification(Qwen2MoePreTrainedModel): + def __init__(self, config): + super().__init__(config) + self.num_labels = config.num_labels + self.model = Qwen2MoeModel(config) + if getattr(config, "classifier_dropout", None) is not None: + classifier_dropout = config.classifier_dropout + elif getattr(config, "hidden_dropout", None) is not None: + classifier_dropout = config.hidden_dropout + else: + classifier_dropout = 0.1 + self.dropout = nn.Dropout(classifier_dropout) + self.score = nn.Linear(config.hidden_size, config.num_labels) + + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self): + return self.model.embed_tokens + + def set_input_embeddings(self, value): + self.model.embed_tokens = value + + @add_start_docstrings_to_model_forward(QWEN2MOE_INPUTS_DOCSTRING) + def forward( + self, + input_ids: Optional[torch.LongTensor] = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[List[torch.FloatTensor]] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + labels: Optional[torch.LongTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + ) -> Union[Tuple, TokenClassifierOutput]: + r""" + labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*): + Labels for computing the sequence classification/regression loss. Indices should be in `[0, ..., + config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If + `config.num_labels > 1` a classification loss is computed (Cross-Entropy). + """ + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + outputs = self.model( + input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + sequence_output = outputs[0] + sequence_output = self.dropout(sequence_output) + logits = self.score(sequence_output) + + loss = None + if labels is not None: + loss_fct = CrossEntropyLoss() + loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1)) + + if not return_dict: + output = (logits,) + outputs[2:] + return ((loss,) + output) if loss is not None else output + + return TokenClassifierOutput( + loss=loss, + logits=logits, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) diff --git a/archive/kt-sft/ktransformers/models/modeling_qwen3_moe.py b/archive/kt-sft/ktransformers/models/modeling_qwen3_moe.py new file mode 100644 index 0000000..5881189 --- /dev/null +++ b/archive/kt-sft/ktransformers/models/modeling_qwen3_moe.py @@ -0,0 +1,1473 @@ +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# This file was automatically generated from src/transformers/models/qwen3_moe/modular_qwen3_moe.py. +# Do NOT edit this file manually as any edits will be overwritten by the generation of +# the file from the modular. If any change should be done, please apply the change to the +# modular_qwen3_moe.py file directly. One of our CI enforces this. +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# coding=utf-8 +# Copyright 2025 The Qwen team, Alibaba Group and the HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from typing import Callable, List, Optional, Tuple, Union + +import torch +import torch.nn.functional as F +from torch import nn + +from transformers.activations import ACT2FN +from transformers.cache_utils import Cache, DynamicCache, SlidingWindowCache, StaticCache +from transformers.generation import GenerationMixin +from transformers.modeling_attn_mask_utils import AttentionMaskConverter +# from transformers.modeling_flash_attention_utils import FlashAttentionKwargs +from transformers.modeling_outputs import ( + BaseModelOutputWithPast, + CausalLMOutputWithPast, + MoeCausalLMOutputWithPast, + MoeModelOutputWithPast, + QuestionAnsweringModelOutput, + SequenceClassifierOutputWithPast, + TokenClassifierOutput, +) +from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS +# from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel +from transformers.modeling_utils import PreTrainedModel +# from transformers.processing_utils import Unpack +from transformers.utils import ( + # LossKwargs, + add_code_sample_docstrings, + add_start_docstrings, + add_start_docstrings_to_model_forward, + logging, + replace_return_docstrings, +) +from transformers.utils.deprecation import deprecate_kwarg +from .configuration_qwen3_moe import Qwen3MoeConfig + +from ktransformers.util.grad_wrapper import maybe_no_grad +from ktransformers.models.modeling_qwen2_moe import Qwen2MoeRotaryEmbedding + +logger = logging.get_logger(__name__) + +_CHECKPOINT_FOR_DOC = "Qwen/Qwen3-MoE-15B-A2B" +_CONFIG_FOR_DOC = "Qwen3MoeConfig" + + +def rotate_half(x): + """Rotates half the hidden dims of the input.""" + x1 = x[..., : x.shape[-1] // 2] + x2 = x[..., x.shape[-1] // 2 :] + return torch.cat((-x2, x1), dim=-1) + + +def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1): + """Applies Rotary Position Embedding to the query and key tensors. + + Args: + q (`torch.Tensor`): The query tensor. + k (`torch.Tensor`): The key tensor. + cos (`torch.Tensor`): The cosine part of the rotary embedding. + sin (`torch.Tensor`): The sine part of the rotary embedding. + position_ids (`torch.Tensor`, *optional*): + Deprecated and unused. + unsqueeze_dim (`int`, *optional*, defaults to 1): + The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and + sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note + that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and + k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes + cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have + the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2. + Returns: + `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding. + """ + cos = cos.unsqueeze(unsqueeze_dim) + sin = sin.unsqueeze(unsqueeze_dim) + q_embed = (q * cos) + (rotate_half(q) * sin) + k_embed = (k * cos) + (rotate_half(k) * sin) + return q_embed, k_embed + + +def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: + """ + This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch, + num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim) + """ + batch, num_key_value_heads, slen, head_dim = hidden_states.shape + if n_rep == 1: + return hidden_states + hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim) + return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim) + + +def eager_attention_forward( + module: nn.Module, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attention_mask: Optional[torch.Tensor], + scaling: float, + dropout: float = 0.0, + **kwargs, +): + key_states = repeat_kv(key, module.num_key_value_groups) + value_states = repeat_kv(value, module.num_key_value_groups) + + attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling + if attention_mask is not None: + causal_mask = attention_mask[:, :, :, : key_states.shape[-2]] + attn_weights = attn_weights + causal_mask + + attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype) + attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training) + attn_output = torch.matmul(attn_weights, value_states) + attn_output = attn_output.transpose(1, 2).contiguous() + + return attn_output, attn_weights + + +class Qwen3MoeAttention(nn.Module): + """Multi-headed attention from 'Attention Is All You Need' paper""" + + def __init__(self, config: Qwen3MoeConfig, layer_idx: int): + super().__init__() + self.config = config + self.layer_idx = layer_idx + self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads) + self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads + self.num_heads = config.num_attention_heads + self.num_key_value_heads = config.num_key_value_heads + self.max_position_embeddings = config.max_position_embeddings + self.rope_theta = config.rope_theta + self.scaling = self.head_dim**-0.5 + self.attention_dropout = config.attention_dropout + self.is_causal = True + + self.q_proj = nn.Linear( + config.hidden_size, config.num_attention_heads * self.head_dim, bias=config.attention_bias + ) + self.k_proj = nn.Linear( + config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias + ) + self.v_proj = nn.Linear( + config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias + ) + self.o_proj = nn.Linear( + config.num_attention_heads * self.head_dim, config.hidden_size, bias=config.attention_bias + ) + self.q_norm = Qwen3MoeRMSNorm(self.head_dim, eps=config.rms_norm_eps) # unlike olmo, only on the head dim! + self.k_norm = Qwen3MoeRMSNorm(self.head_dim, eps=config.rms_norm_eps) # thus post q_norm does not need reshape + + self.rotary_emb = Qwen2MoeRotaryEmbedding( + self.head_dim, + max_position_embeddings=self.max_position_embeddings, + base=self.rope_theta, + ) + + self.sliding_window = config.sliding_window + if not ( + self.config.use_sliding_window + and getattr(self.config, "sliding_window", None) is not None + and self.layer_idx >= self.config.max_window_layers + ): + self.sliding_window = None + + def forward( + self, + hidden_states: torch.Tensor, + position_embeddings: Tuple[torch.Tensor, torch.Tensor], + attention_mask: Optional[torch.Tensor], + past_key_value: Optional[Cache] = None, + cache_position: Optional[torch.LongTensor] = None, + # **kwargs: Unpack[FlashAttentionKwargs], + ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: + input_shape = hidden_states.shape[:-1] + hidden_shape = (*input_shape, -1, self.head_dim) + + query_states = self.q_norm(self.q_proj(hidden_states).view(hidden_shape)).transpose(1, 2) + key_states = self.k_norm(self.k_proj(hidden_states).view(hidden_shape)).transpose(1, 2) + value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2) + + cos, sin = position_embeddings + query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin) + + if past_key_value is not None: + # sin and cos are specific to RoPE models; cache_position needed for the static cache + cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position} + key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs) + + attention_interface: Callable = eager_attention_forward + if self.config._attn_implementation != "eager": + if self.config._attn_implementation == "sdpa" and kwargs.get("output_attentions", False): + logger.warning_once( + "`torch.nn.functional.scaled_dot_product_attention` does not support `output_attentions=True`. Falling back to " + 'eager attention. This warning can be removed using the argument `attn_implementation="eager"` when loading the model.' + ) + else: + attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation] + + attn_output, attn_weights = attention_interface( + self, + query_states, + key_states, + value_states, + attention_mask, + dropout=0.0 if not self.training else self.attention_dropout, + scaling=self.scaling, + sliding_window=self.sliding_window, # diff with Llama + # **kwargs, + ) + + attn_output = attn_output.reshape(*input_shape, -1).contiguous() + attn_output = self.o_proj(attn_output) + return attn_output, attn_weights + + +class Qwen3MoeMLP(nn.Module): + def __init__(self, config, intermediate_size=None): + super().__init__() + self.config = config + self.hidden_size = config.hidden_size + self.intermediate_size = intermediate_size if intermediate_size is not None else config.intermediate_size + self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) + self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) + self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False) + self.act_fn = ACT2FN[config.hidden_act] + + def forward(self, x): + down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x)) + return down_proj + + +class Qwen3MoeSparseMoeBlock(nn.Module): + def __init__(self, config): + super().__init__() + self.num_experts = config.num_experts + self.top_k = config.num_experts_per_tok + self.norm_topk_prob = config.norm_topk_prob + + # gating + self.gate = nn.Linear(config.hidden_size, config.num_experts, bias=False) + self.experts = nn.ModuleList( + [Qwen3MoeMLP(config, intermediate_size=config.moe_intermediate_size) for _ in range(self.num_experts)] + ) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + """ """ + batch_size, sequence_length, hidden_dim = hidden_states.shape + hidden_states = hidden_states.view(-1, hidden_dim) + # router_logits: (batch * sequence_length, n_experts) + router_logits = self.gate(hidden_states) + + routing_weights = F.softmax(router_logits, dim=1, dtype=torch.float) + routing_weights, selected_experts = torch.topk(routing_weights, self.top_k, dim=-1) + if self.norm_topk_prob: # only diff with mixtral sparse moe block! + routing_weights /= routing_weights.sum(dim=-1, keepdim=True) + # we cast back to the input dtype + routing_weights = routing_weights.to(hidden_states.dtype) + + final_hidden_states = torch.zeros( + (batch_size * sequence_length, hidden_dim), dtype=hidden_states.dtype, device=hidden_states.device + ) + + # One hot encode the selected experts to create an expert mask + # this will be used to easily index which expert is going to be sollicitated + expert_mask = torch.nn.functional.one_hot(selected_experts, num_classes=self.num_experts).permute(2, 1, 0) + + # Loop over all available experts in the model and perform the computation on each expert + for expert_idx in range(self.num_experts): + expert_layer = self.experts[expert_idx] + idx, top_x = torch.where(expert_mask[expert_idx]) + + # Index the correct hidden states and compute the expert hidden state for + # the current expert. We need to make sure to multiply the output hidden + # states by `routing_weights` on the corresponding tokens (top-1 and top-2) + current_state = hidden_states[None, top_x].reshape(-1, hidden_dim) + current_hidden_states = expert_layer(current_state) * routing_weights[top_x, idx, None] + + # However `index_add_` only support torch tensors for indexing so we'll use + # the `top_x` tensor here. + final_hidden_states.index_add_(0, top_x, current_hidden_states.to(hidden_states.dtype)) + final_hidden_states = final_hidden_states.reshape(batch_size, sequence_length, hidden_dim) + return final_hidden_states, router_logits + + +class Qwen3MoeRMSNorm(nn.Module): + def __init__(self, hidden_size, eps=1e-6): + """ + Qwen3MoeRMSNorm is equivalent to T5LayerNorm + """ + super().__init__() + self.weight = nn.Parameter(torch.ones(hidden_size)) + self.variance_epsilon = eps + self.hidden_size = hidden_size + + def forward(self, hidden_states): + input_dtype = hidden_states.dtype + hidden_states = hidden_states.to(torch.float32) + variance = hidden_states.pow(2).mean(-1, keepdim=True) + hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon) + return self.weight * hidden_states.to(input_dtype) + + def extra_repr(self): + return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}" + + +class Qwen3MoeDecoderLayer(nn.Module): + def __init__(self, config: Qwen3MoeConfig, layer_idx: int): + super().__init__() + self.hidden_size = config.hidden_size + + self.self_attn = Qwen3MoeAttention(config, layer_idx) + self.mlp = Qwen3MoeMLP(config) + + self.self_attn = Qwen3MoeAttention(config, layer_idx) + + if (layer_idx not in config.mlp_only_layers) and ( + config.num_experts > 0 and (layer_idx + 1) % config.decoder_sparse_step == 0 + ): + self.mlp = Qwen3MoeSparseMoeBlock(config) + else: + self.mlp = Qwen3MoeMLP(config, intermediate_size=config.intermediate_size) + + self.input_layernorm = Qwen3MoeRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = Qwen3MoeRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_value: Optional[Tuple[torch.Tensor]] = None, + output_attentions: Optional[bool] = False, + output_router_logits: Optional[bool] = False, + use_cache: Optional[bool] = False, + cache_position: Optional[torch.LongTensor] = None, + position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, # necessary, but kept here for BC + # **kwargs: Unpack[FlashAttentionKwargs], + ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]: + """ + Args: + hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)` + attention_mask (`torch.FloatTensor`, *optional*): attention mask of size + `(batch, sequence_length)` where padding elements are indicated by 0. + output_attentions (`bool`, *optional*): + Whether or not to return the attentions tensors of all attention layers. See `attentions` under + returned tensors for more detail. + output_router_logits (`bool`, *optional*): + Whether or not to return the logits of all the routers. They are useful for computing the router loss, + and should not be returned during inference. + use_cache (`bool`, *optional*): + If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding + (see `past_key_values`). + past_key_value (`Tuple(torch.FloatTensor)`, *optional*): cached past key and value projection states + cache_position (`torch.LongTensor` of shape `(sequence_length)`, *optional*): + Indices depicting the position of the input sequence tokens in the sequence. + position_embeddings (`Tuple[torch.FloatTensor, torch.FloatTensor]`, *optional*): + Tuple containing the cosine and sine positional embeddings of shape `(batch_size, seq_len, head_dim)`, + with `head_dim` being the embedding dimension of each attention head. + kwargs (`dict`, *optional*): + Arbitrary kwargs to be ignored, used for FSDP and other methods that injects code + into the model + """ + residual = hidden_states + + hidden_states = self.input_layernorm(hidden_states) + + # Self Attention + hidden_states, self_attn_weights = self.self_attn( + hidden_states=hidden_states, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_value=past_key_value, + output_attentions=output_attentions, + use_cache=use_cache, + cache_position=cache_position, + position_embeddings=position_embeddings, + ) + hidden_states = residual + hidden_states + + # Fully Connected + residual = hidden_states + hidden_states = self.post_attention_layernorm(hidden_states) + + hidden_states = self.mlp(hidden_states) + if isinstance(hidden_states, tuple): + hidden_states, router_logits = hidden_states + else: + router_logits = None + + hidden_states = residual + hidden_states + + outputs = (hidden_states,) + + if output_attentions: + outputs += (self_attn_weights,) + + if output_router_logits: + outputs += (router_logits,) + + return outputs + + +def _compute_default_rope_parameters( + config: Optional[Qwen3MoeConfig] = None, + device: Optional["torch.device"] = None, + seq_len: Optional[int] = 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 len(rope_kwargs) > 0: + 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 len(rope_kwargs) > 0: + 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 + dim = int(config.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).float().to(device) / dim)) + return inv_freq, attention_factor + +class Qwen3MoeRotaryEmbedding(nn.Module): + def __init__(self, config: Qwen3MoeConfig, device=None): + super().__init__() + # BC: "rope_type" was originally "type" + if hasattr(config, "rope_scaling") and config.rope_scaling is not None: + self.rope_type = config.rope_scaling.get("rope_type", config.rope_scaling.get("type")) + else: + self.rope_type = "default" + self.max_seq_len_cached = config.max_position_embeddings + self.original_max_seq_len = config.max_position_embeddings + + self.config = config + self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type] + + self.scaling_factor = 1.0 + self.dim = config.head_dim + self.max_position_embeddings = config.max_position_embeddings + self.base = config.rope_theta + inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2, dtype=torch.int64).float().to(device) / self.dim)) + + inv_freq, self.attention_scaling = _compute_default_rope_parameters(self.config, device) + self.register_buffer("inv_freq", inv_freq, persistent=False) + self.original_inv_freq = self.inv_freq + + def _dynamic_frequency_update(self, position_ids, device): + """ + dynamic RoPE layers should recompute `inv_freq` in the following situations: + 1 - growing beyond the cached sequence length (allow scaling) + 2 - the current sequence length is in the original scale (avoid losing precision with small sequences) + """ + seq_len = torch.max(position_ids) + 1 + if seq_len > self.max_seq_len_cached: # growth + inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device, seq_len=seq_len) + self.register_buffer("inv_freq", inv_freq, persistent=False) # TODO joao: may break with compilation + self.max_seq_len_cached = seq_len + + if seq_len < self.original_max_seq_len and self.max_seq_len_cached > self.original_max_seq_len: # reset + # This .to() is needed if the model has been moved to a device after being initialized (because + # the buffer is automatically moved, but not the original copy) + self.original_inv_freq = self.original_inv_freq.to(device) + self.register_buffer("inv_freq", self.original_inv_freq, persistent=False) + self.max_seq_len_cached = self.original_max_seq_len + + @maybe_no_grad() + def forward(self, x, position_ids): + if "dynamic" in self.rope_type: + self._dynamic_frequency_update(position_ids, device=x.device) + + # Core RoPE block + inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1) + position_ids_expanded = position_ids[:, None, :].float() + # Force float32 (see https://github.com/huggingface/transformers/pull/29285) + device_type = x.device.type + device_type = device_type if isinstance(device_type, str) and device_type != "mps" else "cpu" + + with torch.autocast(device_type=device_type, enabled=False): + freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2) + emb = torch.cat((freqs, freqs), dim=-1) + cos = emb.cos() + sin = emb.sin() + + # Advanced RoPE types (e.g. yarn) apply a post-processing scaling factor, equivalent to scaling attention + cos = cos * self.attention_scaling + sin = sin * self.attention_scaling + + return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype) + + +QWEN3_MOE_START_DOCSTRING = r""" + This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the + library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads + etc.) + + This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass. + Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage + and behavior. + + Parameters: + config ([`Qwen3MoeConfig`]): + Model configuration class with all the parameters of the model. Initializing with a config file does not + load the weights associated with the model, only the configuration. Check out the + [`~PreTrainedModel.from_pretrained`] method to load the model weights. +""" + + +@add_start_docstrings( + "The bare Qwen3Moe Model outputting raw hidden-states without any specific head on top.", + QWEN3_MOE_START_DOCSTRING, +) +class Qwen3MoePreTrainedModel(PreTrainedModel): + config_class = Qwen3MoeConfig + base_model_prefix = "model" + supports_gradient_checkpointing = True + _no_split_modules = ["Qwen3MoeDecoderLayer"] + _skip_keys_device_placement = ["past_key_values"] + _supports_flash_attn_2 = True + _supports_sdpa = True + _supports_flex_attn = True + _supports_cache_class = True + _supports_quantized_cache = True + _supports_static_cache = False # MoE models don't work with torch.compile (`torch.where(condition)` not supported) + _supports_attention_backend = True + + def _init_weights(self, module): + std = self.config.initializer_range + if isinstance(module, nn.Linear): + module.weight.data.normal_(mean=0.0, std=std) + if module.bias is not None: + module.bias.data.zero_() + elif isinstance(module, nn.Embedding): + module.weight.data.normal_(mean=0.0, std=std) + if module.padding_idx is not None: + module.weight.data[module.padding_idx].zero_() + + +QWEN3_MOE_INPUTS_DOCSTRING = r""" + Args: + input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): + Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide + it. + + Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and + [`PreTrainedTokenizer.__call__`] for details. + + [What are input IDs?](../glossary#input-ids) + attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*): + Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`: + + - 1 for tokens that are **not masked**, + - 0 for tokens that are **masked**. + + [What are attention masks?](../glossary#attention-mask) + + Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and + [`PreTrainedTokenizer.__call__`] for details. + + If `past_key_values` is used, optionally only the last `input_ids` have to be input (see + `past_key_values`). + + If you want to change padding behavior, you should read [`modeling_opt._prepare_decoder_attention_mask`] + and modify to your needs. See diagram 1 in [the paper](https://arxiv.org/abs/1910.13461) for more + information on the default strategy. + + - 1 indicates the head is **not masked**, + - 0 indicates the head is **masked**. + position_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0, + config.n_positions - 1]`. + + [What are position IDs?](../glossary#position-ids) + past_key_values (`Cache`, *optional*): + Pre-computed hidden-states (key and values in the self-attention blocks and in the cross-attention + blocks) that can be used to speed up sequential decoding. This typically consists in the `past_key_values` + returned by the model at a previous stage of decoding, when `use_cache=True` or `config.use_cache=True`. + + It is a [`~cache_utils.Cache`] instance. For more details, see our [kv cache guide](https://huggingface.co/docs/transformers/en/kv_cache). + + If `past_key_values` are used, the user can optionally input only the last `input_ids` (those that don't + have their past key value states given to this model) of shape `(batch_size, 1)` instead of all `input_ids` + of shape `(batch_size, sequence_length)`. + inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*): + Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This + is useful if you want more control over how to convert `input_ids` indices into associated vectors than the + model's internal embedding lookup matrix. + use_cache (`bool`, *optional*): + If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see + `past_key_values`). + output_attentions (`bool`, *optional*): + Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned + tensors for more detail. + output_hidden_states (`bool`, *optional*): + Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for + more detail. + return_dict (`bool`, *optional*): + Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple. + cache_position (`torch.LongTensor` of shape `(sequence_length)`, *optional*): + Indices depicting the position of the input sequence tokens in the sequence. Contrarily to `position_ids`, + this tensor is not affected by padding. It is used to update the cache in the correct position and to infer + the complete sequence length. +""" + + +@add_start_docstrings( + "The bare Qwen3Moe Model outputting raw hidden-states without any specific head on top.", + QWEN3_MOE_START_DOCSTRING, +) +class Qwen3MoeModel(Qwen3MoePreTrainedModel): + """ + Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`Qwen3MoeDecoderLayer`] + + Args: + config: Qwen3MoeConfig + """ + + def __init__(self, config: Qwen3MoeConfig): + super().__init__(config) + self.padding_idx = config.pad_token_id + self.vocab_size = config.vocab_size + + self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx) + self.layers = nn.ModuleList( + [Qwen3MoeDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)] + ) + self.norm = Qwen3MoeRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.rotary_emb = Qwen3MoeRotaryEmbedding(config=config) + self.gradient_checkpointing = False + + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self): + return self.embed_tokens + + def set_input_embeddings(self, value): + self.embed_tokens = value + + @add_start_docstrings_to_model_forward(QWEN3_MOE_INPUTS_DOCSTRING) + def forward( + self, + input_ids: torch.LongTensor = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[List[torch.FloatTensor]] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + output_router_logits: Optional[bool] = None, + return_dict: Optional[bool] = None, + cache_position: Optional[torch.LongTensor] = None, + # **flash_attn_kwargs: Unpack[FlashAttentionKwargs], + ) -> Union[Tuple, BaseModelOutputWithPast]: + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_router_logits = ( + output_router_logits if output_router_logits is not None else self.config.output_router_logits + ) + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + use_cache = use_cache if use_cache is not None else self.config.use_cache + + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + if (input_ids is None) ^ (inputs_embeds is not None): + raise ValueError("You must specify exactly one of input_ids or inputs_embeds") + + if self.gradient_checkpointing and self.training: + if use_cache: + logger.warning_once( + "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..." + ) + use_cache = False + + if use_cache and past_key_values is None: + past_key_values = DynamicCache() + + if inputs_embeds is None: + inputs_embeds = self.embed_tokens(input_ids) + + if cache_position is None: + past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 + cache_position = torch.arange( + past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device + ) + if position_ids is None: + position_ids = cache_position.unsqueeze(0) + + causal_mask = self._update_causal_mask( + attention_mask, inputs_embeds, cache_position, past_key_values, output_attentions + ) + + hidden_states = inputs_embeds + + # create position embeddings to be shared across the decoder layers + position_embeddings = self.rotary_emb(hidden_states, position_ids) + + # decoder layers + all_hidden_states = () if output_hidden_states else None + all_self_attns = () if output_attentions else None + all_router_logits = () if output_router_logits else None + + for decoder_layer in self.layers: + if output_hidden_states: + all_hidden_states += (hidden_states,) + + if self.gradient_checkpointing and self.training: + layer_outputs = self._gradient_checkpointing_func( + decoder_layer.__call__, + hidden_states, + causal_mask, + position_ids, + past_key_values, + output_attentions, + output_router_logits, + use_cache, + cache_position, + position_embeddings, + ) + else: + layer_outputs = decoder_layer( + hidden_states, + attention_mask=causal_mask, + position_ids=position_ids, + past_key_value=past_key_values, + output_attentions=output_attentions, + output_router_logits=output_router_logits, + use_cache=use_cache, + cache_position=cache_position, + position_embeddings=position_embeddings, + # **flash_attn_kwargs, + ) + + hidden_states = layer_outputs[0] + + if output_attentions: + all_self_attns += (layer_outputs[1],) + + if output_router_logits: + all_router_logits += (layer_outputs[-1],) + + hidden_states = self.norm(hidden_states) + + # add hidden states from the last decoder layer + if output_hidden_states: + all_hidden_states += (hidden_states,) + + output = MoeModelOutputWithPast( + last_hidden_state=hidden_states, + past_key_values=past_key_values, + hidden_states=all_hidden_states, + attentions=all_self_attns, + router_logits=all_router_logits, + ) + return output if return_dict else output.to_tuple() + + def _update_causal_mask( + self, + attention_mask: torch.Tensor, + input_tensor: torch.Tensor, + cache_position: torch.Tensor, + past_key_values: Cache, + output_attentions: bool = False, + ): + if self.config._attn_implementation == "flash_attention_2": + if attention_mask is not None and past_key_values is not None: + is_padding_right = attention_mask[:, -1].sum().item() != input_tensor.size()[0] + if is_padding_right: + raise ValueError( + "You are attempting to perform batched generation with padding_side='right'" + " this may lead to unexpected behaviour for Flash Attention version of Qwen3Moe. Make sure to " + " call `tokenizer.padding_side = 'left'` before tokenizing the input. " + ) + if attention_mask is not None and 0.0 in attention_mask: + return attention_mask + return None + + # For SDPA, when possible, we will rely on its `is_causal` argument instead of its `attn_mask` argument, in + # order to dispatch on Flash Attention 2. This feature is not compatible with static cache, as SDPA will fail + # to infer the attention mask. + past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 + using_static_cache = isinstance(past_key_values, StaticCache) + using_sliding_window_cache = isinstance(past_key_values, SlidingWindowCache) + + # When output attentions is True, sdpa implementation's forward method calls the eager implementation's forward + if ( + self.config._attn_implementation == "sdpa" + and not (using_static_cache or using_sliding_window_cache) + and not output_attentions + ): + if AttentionMaskConverter._ignore_causal_mask_sdpa( + attention_mask, + inputs_embeds=input_tensor, + past_key_values_length=past_seen_tokens, + sliding_window=self.config.sliding_window, + is_training=self.training, + ): + return None + + dtype, device = input_tensor.dtype, input_tensor.device + min_dtype = torch.finfo(dtype).min + sequence_length = input_tensor.shape[1] + # SlidingWindowCache or StaticCache + if using_sliding_window_cache or using_static_cache: + target_length = past_key_values.get_max_cache_shape() + # DynamicCache or no cache + else: + target_length = ( + attention_mask.shape[-1] + if isinstance(attention_mask, torch.Tensor) + else past_seen_tokens + sequence_length + 1 + ) + + # In case the provided `attention` mask is 2D, we generate a causal mask here (4D). + causal_mask = self._prepare_4d_causal_attention_mask_with_cache_position( + attention_mask, + sequence_length=sequence_length, + target_length=target_length, + dtype=dtype, + device=device, + cache_position=cache_position, + batch_size=input_tensor.shape[0], + config=self.config, + past_key_values=past_key_values, + ) + + if ( + self.config._attn_implementation == "sdpa" + and attention_mask is not None + and attention_mask.device.type in ["cuda", "xpu"] + and not output_attentions + ): + # Attend to all tokens in fully masked rows in the causal_mask, for example the relevant first rows when + # using left padding. This is required by F.scaled_dot_product_attention memory-efficient attention path. + # Details: https://github.com/pytorch/pytorch/issues/110213 + causal_mask = AttentionMaskConverter._unmask_unattended(causal_mask, min_dtype) + + return causal_mask + + @staticmethod + def _prepare_4d_causal_attention_mask_with_cache_position( + attention_mask: torch.Tensor, + sequence_length: int, + target_length: int, + dtype: torch.dtype, + device: torch.device, + cache_position: torch.Tensor, + batch_size: int, + config: Qwen3MoeConfig, + past_key_values: Cache, + ): + """ + Creates a causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)` from a 2D mask of shape + `(batch_size, key_value_length)`, or if the input `attention_mask` is already 4D, do nothing. + + Args: + attention_mask (`torch.Tensor`): + A 2D attention mask of shape `(batch_size, key_value_length)` or a 4D attention mask of shape `(batch_size, 1, query_length, key_value_length)`. + sequence_length (`int`): + The sequence length being processed. + target_length (`int`): + The target length: when generating with static cache, the mask should be as long as the static cache, to account for the 0 padding, the part of the cache that is not filled yet. + dtype (`torch.dtype`): + The dtype to use for the 4D attention mask. + device (`torch.device`): + The device to place the 4D attention mask on. + cache_position (`torch.Tensor`): + Indices depicting the position of the input sequence tokens in the sequence. + batch_size (`torch.Tensor`): + Batch size. + config (`Qwen3MoeConfig`): + The model's configuration class + past_key_values (`Cache`): + The cache class that is being used currently to generate + """ + if attention_mask is not None and attention_mask.dim() == 4: + # In this case we assume that the mask comes already in inverted form and requires no inversion or slicing. + causal_mask = attention_mask + else: + min_dtype = torch.finfo(dtype).min + causal_mask = torch.full( + (sequence_length, target_length), fill_value=min_dtype, dtype=dtype, device=device + ) + diagonal_attend_mask = torch.arange(target_length, device=device) > cache_position.reshape(-1, 1) + if config.sliding_window is not None: + # if we have sliding window, we should not attend to tokens beyond sliding window length, so we mask them out also + # the check is needed to verify is current checkpoint was trained with sliding window or not + if not isinstance(past_key_values, SlidingWindowCache) or sequence_length > target_length: + sliding_attend_mask = torch.arange(target_length, device=device) <= ( + cache_position.reshape(-1, 1) - config.sliding_window + ) + diagonal_attend_mask.bitwise_or_(sliding_attend_mask) + causal_mask *= diagonal_attend_mask + causal_mask = causal_mask[None, None, :, :].expand(batch_size, 1, -1, -1) + if attention_mask is not None: + causal_mask = causal_mask.clone() # copy to contiguous memory for in-place edit + if attention_mask.shape[-1] > target_length: + attention_mask = attention_mask[:, :target_length] + mask_length = attention_mask.shape[-1] + padding_mask = causal_mask[:, :, :, :mask_length] + attention_mask[:, None, None, :].to( + causal_mask.device + ) + padding_mask = padding_mask == 0 + causal_mask[:, :, :, :mask_length] = causal_mask[:, :, :, :mask_length].masked_fill( + padding_mask, min_dtype + ) + return causal_mask + + +# class KwargsForCausalLM(FlashAttentionKwargs, LossKwargs): ... +class KwargsForCausalLM(): ... + + +def load_balancing_loss_func( + gate_logits: Union[torch.Tensor, Tuple[torch.Tensor], None], + num_experts: Optional[int] = None, + top_k=2, + attention_mask: Optional[torch.Tensor] = None, +) -> Union[torch.Tensor, int]: + r""" + Computes auxiliary load balancing loss as in Switch Transformer - implemented in Pytorch. + + See Switch Transformer (https://arxiv.org/abs/2101.03961) for more details. This function implements the loss + function presented in equations (4) - (6) of the paper. It aims at penalizing cases where the routing between + experts is too unbalanced. + + Args: + gate_logits: + Logits from the `gate`, should be a tuple of model.config.num_hidden_layers tensors of + shape [batch_size X sequence_length, num_experts]. + num_experts: + Number of experts + top_k: + The number of experts to route per-token, can be also interpreted as the `top-k` routing + parameter. + attention_mask (`torch.Tensor`, *optional*): + The attention_mask used in forward function + shape [batch_size X sequence_length] if not None. + + Returns: + The auxiliary loss. + """ + if gate_logits is None or not isinstance(gate_logits, tuple): + return 0 + + if isinstance(gate_logits, tuple): + compute_device = gate_logits[0].device + concatenated_gate_logits = torch.cat([layer_gate.to(compute_device) for layer_gate in gate_logits], dim=0) + + routing_weights = torch.nn.functional.softmax(concatenated_gate_logits, dim=-1) + + _, selected_experts = torch.topk(routing_weights, top_k, dim=-1) + + expert_mask = torch.nn.functional.one_hot(selected_experts, num_experts) + + if attention_mask is None: + # Compute the percentage of tokens routed to each experts + tokens_per_expert = torch.mean(expert_mask.float(), dim=0) + + # Compute the average probability of routing to these experts + router_prob_per_expert = torch.mean(routing_weights, dim=0) + else: + batch_size, sequence_length = attention_mask.shape + num_hidden_layers = concatenated_gate_logits.shape[0] // (batch_size * sequence_length) + + # Compute the mask that masks all padding tokens as 0 with the same shape of expert_mask + expert_attention_mask = ( + attention_mask[None, :, :, None, None] + .expand((num_hidden_layers, batch_size, sequence_length, top_k, num_experts)) + .reshape(-1, top_k, num_experts) + .to(compute_device) + ) + + # Compute the percentage of tokens routed to each experts + tokens_per_expert = torch.sum(expert_mask.float() * expert_attention_mask, dim=0) / torch.sum( + expert_attention_mask, dim=0 + ) + + # Compute the mask that masks all padding tokens as 0 with the same shape of tokens_per_expert + router_per_expert_attention_mask = ( + attention_mask[None, :, :, None] + .expand((num_hidden_layers, batch_size, sequence_length, num_experts)) + .reshape(-1, num_experts) + .to(compute_device) + ) + + # Compute the average probability of routing to these experts + router_prob_per_expert = torch.sum(routing_weights * router_per_expert_attention_mask, dim=0) / torch.sum( + router_per_expert_attention_mask, dim=0 + ) + + overall_loss = torch.sum(tokens_per_expert * router_prob_per_expert.unsqueeze(0)) + return overall_loss * num_experts + + +class Qwen3MoeForCausalLM(Qwen3MoePreTrainedModel, GenerationMixin): + _tied_weights_keys = ["lm_head.weight"] + _tp_plan = {"lm_head": "colwise_rep"} + _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + + def __init__(self, config): + super().__init__(config) + self.model = Qwen3MoeModel(config) + self.vocab_size = config.vocab_size + self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) + self.router_aux_loss_coef = config.router_aux_loss_coef + self.num_experts = config.num_experts + self.num_experts_per_tok = config.num_experts_per_tok + + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self): + return self.model.embed_tokens + + def set_input_embeddings(self, value): + self.model.embed_tokens = value + + def get_output_embeddings(self): + return self.lm_head + + def set_output_embeddings(self, new_embeddings): + self.lm_head = new_embeddings + + def set_decoder(self, decoder): + self.model = decoder + + def get_decoder(self): + return self.model + + @deprecate_kwarg("num_logits_to_keep", version="4.50", new_name="logits_to_keep") + @add_start_docstrings_to_model_forward(QWEN3_MOE_INPUTS_DOCSTRING) + @replace_return_docstrings(output_type=CausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC) + def forward( + self, + input_ids: torch.LongTensor = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[List[torch.FloatTensor]] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + labels: Optional[torch.LongTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + output_router_logits: Optional[bool] = None, + return_dict: Optional[bool] = None, + cache_position: Optional[torch.LongTensor] = None, + logits_to_keep: Union[int, torch.Tensor] = 0, + # **kwargs: Unpack[KwargsForCausalLM], + ) -> Union[Tuple, CausalLMOutputWithPast]: + r""" + labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Labels for computing the masked language modeling loss. Indices should either be in `[0, ..., + config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored + (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`. + + logits_to_keep (`int` or `torch.Tensor`, *optional*): + If an `int`, compute logits for the last `logits_to_keep` tokens. If `0`, calculate logits for all + `input_ids` (special case). Only last token logits are needed for generation, and calculating them only for that + token can save memory, which becomes pretty significant for long sequences or large vocabulary size. + If a `torch.Tensor`, must be 1D corresponding to the indices to keep in the sequence length dimension. + This is useful when using packed tensor format (single dimension for batch and sequence length). + + Returns: + + Example: + + ```python + >>> from transformers import AutoTokenizer, Qwen3MoeForCausalLM + + >>> model = Qwen3MoeForCausalLM.from_pretrained("Qwen/Qwen3-MoE-15B-A2B") + >>> tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-MoE-15B-A2B") + + >>> prompt = "Hey, are you conscious? Can you talk to me?" + >>> inputs = tokenizer(prompt, return_tensors="pt") + + >>> # Generate + >>> generate_ids = model.generate(inputs.input_ids, max_length=30) + >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0] + "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you." + ```""" + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_router_logits = ( + output_router_logits if output_router_logits is not None else self.config.output_router_logits + ) + + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn) + outputs = self.model( + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + output_router_logits=output_router_logits, + return_dict=return_dict, + cache_position=cache_position, + # **kwargs, + ) + + hidden_states = outputs[0] + # Only compute necessary logits, and do not upcast them to float if we are not computing the loss + slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep + logits = self.lm_head(hidden_states[:, slice_indices, :]) + + loss = None + if labels is not None: + loss = self.loss_function(logits, labels, self.vocab_size) + + aux_loss = None + if output_router_logits: + aux_loss = load_balancing_loss_func( + outputs.router_logits if return_dict else outputs[-1], + self.num_experts, + self.num_experts_per_tok, + attention_mask, + ) + if labels is not None: + loss += self.router_aux_loss_coef * aux_loss.to(loss.device) # make sure to reside in the same device + + if not return_dict: + output = (logits,) + outputs[1:] + if output_router_logits: + output = (aux_loss,) + output + return (loss,) + output if loss is not None else output + + return MoeCausalLMOutputWithPast( + loss=loss, + aux_loss=aux_loss, + logits=logits, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + router_logits=outputs.router_logits, + ) + + +@add_start_docstrings( + """ + The Qwen3Moe Model transformer with a sequence classification head on top (linear layer). + + [`Qwen3MoeForSequenceClassification`] uses the last token in order to do the classification, as other causal models + (e.g. GPT-2) do. + + Since it does classification on the last token, it requires to know the position of the last token. If a + `pad_token_id` is defined in the configuration, it finds the last token that is not a padding token in each row. If + no `pad_token_id` is defined, it simply takes the last value in each row of the batch. Since it cannot guess the + padding tokens when `inputs_embeds` are passed instead of `input_ids`, it does the same (take the last value in + each row of the batch). + """, + QWEN3_MOE_START_DOCSTRING, +) +class Qwen3MoeForSequenceClassification(Qwen3MoePreTrainedModel): + def __init__(self, config): + super().__init__(config) + self.num_labels = config.num_labels + self.model = Qwen3MoeModel(config) + self.score = nn.Linear(config.hidden_size, self.num_labels, bias=False) + + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self): + return self.model.embed_tokens + + def set_input_embeddings(self, value): + self.model.embed_tokens = value + + @add_start_docstrings_to_model_forward(QWEN3_MOE_INPUTS_DOCSTRING) + def forward( + self, + input_ids: Optional[torch.LongTensor] = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[Cache] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + labels: Optional[torch.LongTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + ) -> Union[Tuple, SequenceClassifierOutputWithPast]: + r""" + labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*): + Labels for computing the sequence classification/regression loss. Indices should be in `[0, ..., + config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If + `config.num_labels > 1` a classification loss is computed (Cross-Entropy). + """ + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + transformer_outputs = self.model( + input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + hidden_states = transformer_outputs[0] + logits = self.score(hidden_states) + + if input_ids is not None: + batch_size = input_ids.shape[0] + else: + batch_size = inputs_embeds.shape[0] + + if self.config.pad_token_id is None and batch_size != 1: + raise ValueError("Cannot handle batch sizes > 1 if no padding token is defined.") + if self.config.pad_token_id is None: + last_non_pad_token = -1 + elif input_ids is not None: + # To handle both left- and right- padding, we take the rightmost token that is not equal to pad_token_id + non_pad_mask = (input_ids != self.config.pad_token_id).to(logits.device, torch.int32) + token_indices = torch.arange(input_ids.shape[-1], device=logits.device, dtype=torch.int32) + last_non_pad_token = (token_indices * non_pad_mask).argmax(-1) + else: + last_non_pad_token = -1 + logger.warning_once( + f"{self.__class__.__name__} will not detect padding tokens in `inputs_embeds`. Results may be " + "unexpected if using padding tokens in conjunction with `inputs_embeds.`" + ) + + pooled_logits = logits[torch.arange(batch_size, device=logits.device), last_non_pad_token] + + loss = None + if labels is not None: + loss = self.loss_function(logits=logits, labels=labels, pooled_logits=pooled_logits, config=self.config) + + if not return_dict: + output = (pooled_logits,) + transformer_outputs[1:] + return ((loss,) + output) if loss is not None else output + + return SequenceClassifierOutputWithPast( + loss=loss, + logits=pooled_logits, + past_key_values=transformer_outputs.past_key_values, + hidden_states=transformer_outputs.hidden_states, + attentions=transformer_outputs.attentions, + ) + + +@add_start_docstrings( + """ + The Qwen3Moe Model transformer with a token classification head on top (a linear layer on top of the hidden-states + output) e.g. for Named-Entity-Recognition (NER) tasks. + """, + QWEN3_MOE_START_DOCSTRING, +) +class Qwen3MoeForTokenClassification(Qwen3MoePreTrainedModel): + def __init__(self, config): + super().__init__(config) + self.num_labels = config.num_labels + self.model = Qwen3MoeModel(config) + if getattr(config, "classifier_dropout", None) is not None: + classifier_dropout = config.classifier_dropout + elif getattr(config, "hidden_dropout", None) is not None: + classifier_dropout = config.hidden_dropout + else: + classifier_dropout = 0.1 + self.dropout = nn.Dropout(classifier_dropout) + self.score = nn.Linear(config.hidden_size, config.num_labels) + + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self): + return self.model.embed_tokens + + def set_input_embeddings(self, value): + self.model.embed_tokens = value + + @add_start_docstrings_to_model_forward(QWEN3_MOE_INPUTS_DOCSTRING) + @add_code_sample_docstrings( + checkpoint=_CHECKPOINT_FOR_DOC, + output_type=TokenClassifierOutput, + config_class=_CONFIG_FOR_DOC, + ) + def forward( + self, + input_ids: Optional[torch.LongTensor] = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[Cache] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + labels: Optional[torch.LongTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + ) -> Union[Tuple, TokenClassifierOutput]: + r""" + labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*): + Labels for computing the sequence classification/regression loss. Indices should be in `[0, ..., + config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If + `config.num_labels > 1` a classification loss is computed (Cross-Entropy). + """ + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + outputs = self.model( + input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + sequence_output = outputs[0] + sequence_output = self.dropout(sequence_output) + logits = self.score(sequence_output) + + loss = None + if labels is not None: + loss = self.loss_function(logits, labels, self.config) + + if not return_dict: + output = (logits,) + outputs[2:] + return ((loss,) + output) if loss is not None else output + + return TokenClassifierOutput( + loss=loss, + logits=logits, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) + + +@add_start_docstrings( + """ +The Qwen3Moe Model transformer with a span classification head on top for extractive question-answering tasks like +SQuAD (a linear layer on top of the hidden-states output to compute `span start logits` and `span end logits`). + """, + QWEN3_MOE_START_DOCSTRING, +) +class Qwen3MoeForQuestionAnswering(Qwen3MoePreTrainedModel): + base_model_prefix = "transformer" + + def __init__(self, config): + super().__init__(config) + self.transformer = Qwen3MoeModel(config) + self.qa_outputs = nn.Linear(config.hidden_size, 2) + + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self): + return self.transformer.embed_tokens + + def set_input_embeddings(self, value): + self.transformer.embed_tokens = value + + @add_start_docstrings_to_model_forward(QWEN3_MOE_INPUTS_DOCSTRING) + def forward( + self, + input_ids: Optional[torch.LongTensor] = None, + attention_mask: Optional[torch.FloatTensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[Cache] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + start_positions: Optional[torch.LongTensor] = None, + end_positions: Optional[torch.LongTensor] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + **kwargs, + ) -> Union[Tuple, QuestionAnsweringModelOutput]: + r""" + start_positions (`torch.LongTensor` of shape `(batch_size,)`, *optional*): + Labels for position (index) of the start of the labelled span for computing the token classification loss. + Positions are clamped to the length of the sequence (`sequence_length`). Position outside of the sequence + are not taken into account for computing the loss. + end_positions (`torch.LongTensor` of shape `(batch_size,)`, *optional*): + Labels for position (index) of the end of the labelled span for computing the token classification loss. + Positions are clamped to the length of the sequence (`sequence_length`). Position outside of the sequence + are not taken into account for computing the loss. + """ + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + outputs = self.transformer( + input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + sequence_output = outputs[0] + + logits = self.qa_outputs(sequence_output) + start_logits, end_logits = logits.split(1, dim=-1) + start_logits = start_logits.squeeze(-1).contiguous() + end_logits = end_logits.squeeze(-1).contiguous() + + loss = None + if start_positions is not None and end_positions is not None: + loss = self.loss_function(start_logits, end_logits, start_positions, end_positions, **kwargs) + + if not return_dict: + output = (start_logits, end_logits) + outputs[2:] + return ((loss,) + output) if loss is not None else output + + return QuestionAnsweringModelOutput( + loss=loss, + start_logits=start_logits, + end_logits=end_logits, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) + + +__all__ = [ + "Qwen3MoeForCausalLM", + "Qwen3MoeForQuestionAnswering", + "Qwen3MoeModel", + "Qwen3MoePreTrainedModel", + "Qwen3MoeForSequenceClassification", + "Qwen3MoeForTokenClassification", +] \ No newline at end of file diff --git a/archive/kt-sft/ktransformers/moe_test_module.py b/archive/kt-sft/ktransformers/moe_test_module.py new file mode 100644 index 0000000..9423917 --- /dev/null +++ b/archive/kt-sft/ktransformers/moe_test_module.py @@ -0,0 +1,118 @@ +import os +import platform +import sys + +project_dir = os.path.dirname(os.path.dirname(__file__)) +sys.path.insert(0, project_dir) + +from torchviz import make_dot +from torch import nn +import torch +from transformers import ( + AutoTokenizer, + AutoConfig, + AutoModelForCausalLM, + GenerationConfig, + TextStreamer, +) +import unittest +from torch.autograd import gradcheck + +from ktransformers.operators.linear import KLinearTorch, KTransformersLinear +from ktransformers.sft.peft_utils.lora_layer import KTransformersLinearLora +from ktransformers.util.custom_loader import GGUFLoader +from ktransformers.operators.experts import KExpertsTorch +from ktransformers.util.utils import load_weights + +gguf_loader = GGUFLoader(gguf_path="/home/yj/ktransformers/GGUF-DeepSeek-V2-Lite-Chat") +config = AutoConfig.from_pretrained("/home/yj/ktransformers/DeepSeek-V2-Lite-Chat", trust_remote_code=True) +torch.set_default_dtype(config.torch_dtype) + +class TestKExpertsTorch(unittest.TestCase): + def setUp(self): + torch.backends.cudnn.deterministic = True + torch.backends.cudnn.benchmark = False + self.num_experts = 8 + + self.fixed_input = None + self.fixed_expert_ids = None + self.fixed_weights = None + + def _create_fixed_data(self, device, batch_size=2): + """创建固定输入数据""" + if self.fixed_input is None: + with torch.random.fork_rng(): + torch.manual_seed(42) + hidden_size = config.hidden_size + + self.fixed_input = torch.randn(batch_size, hidden_size) + + self.fixed_expert_ids = torch.tensor([[0, 1], [2, 3]], dtype=torch.long) + + self.fixed_weights = torch.tensor([[0.5, 0.5], [0.5, 0.5]], dtype=torch.float32) + + return ( + self.fixed_input.clone().to(device).requires_grad_(True), + self.fixed_expert_ids.clone().to(device), + self.fixed_weights.clone().to(device) + ) + + def _run_single_device_test(self, device, seed=42): + """在指定设备上运行前向反向传播并返回梯度""" + torch.manual_seed(seed) + if device == "cuda": + torch.cuda.manual_seed_all(seed) + + model = KExpertsTorch( + key="blk.1", + gguf_loader=gguf_loader, + config=config, + n_routed_experts=self.num_experts, + device=device + ) + model.load(device=device) + + input_tensor, expert_ids, weights = self._create_fixed_data(device) + + model.to(device) + + with torch.autocast(device_type=device, enabled=False): + output = model(input_tensor, expert_ids, weights) + + loss = output.sum() + loss.backward() + + gradients = { + "input": input_tensor.grad.detach().cpu(), + "loss": loss.detach().cpu(), + "model": [p.grad.detach().cpu() for p in model.parameters() if p.grad is not None] + } + return gradients + + def test_forward_gradient(self): + cpu_gradients = self._run_single_device_test("cpu") + + if torch.cuda.is_available(): + gpu_gradients = self._run_single_device_test("cuda") + + print(f"cpu_gradients:{cpu_gradients}") + print(f"gpu_gradients:{gpu_gradients}") + + input_diff = torch.max(torch.abs(cpu_gradients["input"] - gpu_gradients["input"])) + print(f"input_diff:{input_diff}") + + for i, (cpu_g, gpu_g) in enumerate(zip(cpu_gradients["model"], gpu_gradients["model"])): + param_diff = torch.max(torch.abs(cpu_g - gpu_g)) + print(f"param_diff:{param_diff}") + + for i, (cpu_g, gpu_g) in enumerate(zip(cpu_gradients["model"], gpu_gradients["model"])): + diff = (cpu_g - gpu_g.cpu()).abs().max() + print(f"参数梯度 {i} 最大差异: {diff.item()}") + self.assertTrue(torch.allclose(cpu_g, gpu_g, atol=1e-4, rtol=1e-3), + f"参数梯度 {i} 差异超出阈值，最大差异: {diff.item()}") + + else: + self.skipTest("CUDA不可用，跳过GPU测试") + +if __name__ == '__main__': + unittest.main() \ No newline at end of file diff --git a/archive/kt-sft/ktransformers/moe_test_module_old.py b/archive/kt-sft/ktransformers/moe_test_module_old.py new file mode 100644 index 0000000..13a097d --- /dev/null +++ b/archive/kt-sft/ktransformers/moe_test_module_old.py @@ -0,0 +1,149 @@ +import os +import platform +import sys + +project_dir = os.path.dirname(os.path.dirname(__file__)) +sys.path.insert(0, project_dir) + +from torchviz import make_dot +from torch import nn +import torch +from transformers import ( + AutoTokenizer, + AutoConfig, + AutoModelForCausalLM, + GenerationConfig, + TextStreamer, +) +import unittest +from torch.autograd import gradcheck + +from ktransformers.operators.linear import KLinearTorch, KTransformersLinear +from ktransformers.sft.peft_utils.lora_layer import KTransformersLinearLora +from ktransformers.util.custom_loader import GGUFLoader +from ktransformers.operators.experts import KExpertsTorch +from ktransformers.util.utils import load_weights + +gguf_loader = GGUFLoader(gguf_path="/home/yj/ktransformers/GGUF-DeepSeek-V2-Lite-Chat") +config = AutoConfig.from_pretrained("/home/yj/ktransformers/DeepSeek-V2-Lite-Chat", trust_remote_code=True) +torch.set_default_dtype(config.torch_dtype) + +class TestKExpertsTorch(unittest.TestCase): + def setUp(self): + torch.backends.cudnn.deterministic = True + torch.backends.cudnn.benchmark = False + self.base_device = "cpu" + self.num_experts = 8 + # model = KExpertsTorch( + # key="blk.1", + # gguf_loader=gguf_loader, + # config=config, + # n_routed_experts=self.num_experts, + # device=self.base_device + # ) + # model.load() + + def _run_single_device_test(self, device, seed=42): + """在指定设备上运行前向反向传播并返回梯度""" + torch.manual_seed(seed) + if device == "cuda": + torch.cuda.manual_seed_all(seed) + + model = KExpertsTorch( + key="blk.1", + gguf_loader=gguf_loader, + config=config, + n_routed_experts=self.num_experts, + device=device + ) + model.load(device=device) + + with torch.random.fork_rng(): + torch.manual_seed(seed) + batch_size = 2 + hidden_size = model.config.hidden_size + input_tensor = torch.randn(batch_size, hidden_size, device=device, requires_grad=True) + expert_ids = torch.randint(0, self.num_experts, + (batch_size, model.config.num_experts_per_tok), + device=device) + weights = torch.randn(batch_size, model.config.num_experts_per_tok, device=device) + weights = torch.softmax(weights, dim=-1) + + print(f"input_tensor.device:{input_tensor.device}") + print(f"torch.device(device):{torch.device(device)}") + # assert input_tensor.device == torch.device(device) + for p in model.parameters(): + print(f"p.device:{p.device}") + + for name, param in model.named_parameters(): + print(name, param.size()) + + + model.to(device) + with torch.autocast(device_type=device, enabled=False): + output = model(input_tensor, expert_ids, weights) + + loss = output.sum() + + + # dot = make_dot(output, params=dict(model.named_parameters())) + # dot.render(f"origin_moe_{torch.device(device)}_graph", format="svg") + + loss.backward() + + gradients = { + "input": input_tensor.grad.clone().cpu(), + "loss": loss.clone().cpu(), + "model": [p.grad.clone().cpu() for p in model.parameters() if p.grad is not None] + } + return gradients + + def test_forward_gradient(self): + # for param in model.parameters(): + # self.assertEqual(param.dtype, config.torch_dtype) + + cpu_gradients = self._run_single_device_test("cpu") + print(f"cpu_gradients: {cpu_gradients}") + + self.assertIsNotNone(cpu_gradients["input"]) + self.assertTrue(all(g is not None for g in cpu_gradients["model"])) + + if torch.cuda.is_available(): + gpu_gradients = self._run_single_device_test("cuda") + + print(f"gpu_gradients: {gpu_gradients}") + + + max_diff = (cpu_gradients["input"] - gpu_gradients["input"].cpu()).abs().max() + print(f"Input梯度最大差异: {max_diff.item()}") + + self.assertTrue(torch.allclose(cpu_gradients["input"], gpu_gradients["input"], atol=1e-4, rtol=1e-3), + f"Input梯度差异超出阈值，最大差异: {max_diff.item()}") + + for i, (cpu_g, gpu_g) in enumerate(zip(cpu_gradients["model"], gpu_gradients["model"])): + diff = (cpu_g - gpu_g.cpu()).abs().max() + print(f"参数梯度 {i} 最大差异: {diff.item()}") + self.assertTrue(torch.allclose(cpu_g, gpu_g, atol=1e-4, rtol=1e-3), + f"参数梯度 {i} 差异超出阈值，最大差异: {diff.item()}") + + else: + raise ImportError("NO CUDA FOR TEST!!") + + # def test_detach_effect(self): + # input_tensor = torch.randn(1, model.config.hidden_size, device="cpu", requires_grad=True) + # expert_ids = torch.tensor([[0, 1]], device="cpu") + # weights = torch.tensor([[0.5, 0.5]], device="cpu") + + # output = model(input_tensor, expert_ids, weights) + + # # dot = make_dot(output, params=dict(model.named_parameters())) + # # dot.render("origin_moe_cpu_graph", format="svg") + + # loss = output.sum() + # loss.backward() + + # self.assertIsNotNone(input_tensor.grad) + # self.assertTrue(all(p.grad is not None for p in model.parameters())) + +if __name__ == '__main__': + unittest.main() \ No newline at end of file diff --git a/archive/kt-sft/ktransformers/operators/RoPE.py b/archive/kt-sft/ktransformers/operators/RoPE.py new file mode 100644 index 0000000..0e05563 --- /dev/null +++ b/archive/kt-sft/ktransformers/operators/RoPE.py @@ -0,0 +1,441 @@ +""" +Description : +Author : Boxin Zhang +Version : 0.1.0 +Copyright (c) 2024 by KVCache.AI, All Rights Reserved. +""" + +from torch import nn +from transformers import ROPE_INIT_FUNCTIONS +from ktransformers.models.modeling_llama import ( + LlamaRotaryEmbedding, + LlamaLinearScalingRotaryEmbedding, + LlamaDynamicNTKScalingRotaryEmbedding, +) +from ktransformers.models.modeling_deepseek_v3 import ( + DeepseekV3RotaryEmbedding +) +from ktransformers.models.modeling_deepseek import ( + DeepseekV2YarnRotaryEmbedding, + DeepseekV2RotaryEmbedding, + yarn_get_mscale, + yarn_linear_ramp_mask, + yarn_find_correction_range +) +from ktransformers.operators.base_operator import BaseInjectedModule +from ktransformers.util.custom_loader import GGUFLoader +from ktransformers.util.inference_state import InferenceState +from ktransformers.util.grad_wrapper import maybe_no_grad +from transformers.configuration_utils import PretrainedConfig +import torch + +# Copied from transformers.models.mixtral.modeling_mixtral.MixtralRotaryEmbedding with Mixtral->Qwen2Moe +class RotaryEmbedding(BaseInjectedModule, DeepseekV2RotaryEmbedding): + def __init__( + self, + key: str, + gguf_loader: GGUFLoader, + config: PretrainedConfig, + orig_module: nn.Module, + # device: str = "cuda", + generate_device: str = "cuda", + prefill_device: str = "cuda", + **kwargs, + ): + BaseInjectedModule.__init__( + self, key, gguf_loader, config, orig_module, prefill_device, generate_device, **kwargs + ) + self.orig_module.__init__( + orig_module.dim, orig_module.max_position_embeddings, orig_module.base + ) + self.generate_device = generate_device + self.prefill_device = prefill_device + + def load(self): + self.orig_module.__init__( + self.orig_module.dim, + self.orig_module.max_position_embeddings, + self.orig_module.base, + self.device, + ) + + +class RotaryEmbeddingV3(BaseInjectedModule): + def __init__( + self, + key: str, + gguf_loader: GGUFLoader, + config: PretrainedConfig, + orig_module: nn.Module, + # device: str = "cuda", + generate_device: str = "cuda", + prefill_device: str = "cuda", + **kwargs, + ): + BaseInjectedModule.__init__( + self, key, gguf_loader, config, orig_module, prefill_device, generate_device, **kwargs + ) + self.generate_device = generate_device + self.prefill_device = prefill_device + + @maybe_no_grad() + def forward(self, x, position_ids): + # x: [bs, num_attention_heads, seq_len, head_size] + inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1) + position_ids_expanded = position_ids[:, None, :].float() + # Force float32 since bfloat16 loses precision on long contexts + # See https://github.com/huggingface/transformers/pull/29285 + device_type = x.device.type + device_type = device_type if isinstance(device_type, str) and device_type != "mps" else "cpu" + with torch.autocast(device_type=device_type, enabled=False): + freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2) + emb = torch.cat((freqs, freqs), dim=-1) + cos = emb.cos() + sin = emb.sin() + return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype) + + def load(self): + self._init( + dim=self.config.qk_rope_head_dim, + max_position_embeddings=self.config.max_position_embeddings, + base=self.config.rope_theta, + device=self.device, + ) + def _init(self, dim, max_position_embeddings, base, device, scaling_factor=1.0): + self.scaling_factor = scaling_factor + self.dim = dim + self.max_position_embeddings = max_position_embeddings + self.base = base + self.inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2, dtype=torch.int64).float().to(device) / self.dim)) + # self.register_buffer("inv_freq", inv_freq, persistent=False) + # For BC we register cos and sin cached + self.max_seq_len_cached = max_position_embeddings + +class RotaryEmbeddingV2(BaseInjectedModule, LlamaRotaryEmbedding): + def __init__( + self, + key: str, + gguf_loader: GGUFLoader, + config: PretrainedConfig, + orig_module: nn.Module, + generate_device: str = "cuda", + prefill_device: str = "cuda", + **kwargs, + ): + BaseInjectedModule.__init__( + self, key, gguf_loader, config, orig_module, prefill_device, generate_device, **kwargs + ) + self.orig_module.__init__( + orig_module.dim, + orig_module.max_position_embeddings, + orig_module.base, + None, + orig_module.scaling_factor, + orig_module.rope_type, + orig_module.config, + ) + self.generate_device = generate_device + self.prefill_device = prefill_device + + def load(self): + self.orig_module.__init__( + self.orig_module.dim, + self.orig_module.max_position_embeddings, + self.orig_module.base, + self.device, + self.orig_module.scaling_factor, + self.orig_module.rope_type, + self.orig_module.config, + ) + +class YarnRotaryEmbedding(BaseInjectedModule, DeepseekV2YarnRotaryEmbedding): + def __init__( + self, + key: str, + gguf_loader: GGUFLoader, + config: PretrainedConfig, + orig_module: nn.Module, + # device: str = "cuda", + generate_device: str = "cuda", + prefill_device: str = "cuda", + **kwargs, + ): + BaseInjectedModule.__init__( + self, key, gguf_loader, config, orig_module, prefill_device, generate_device, **kwargs + ) + self.orig_module.__init__( + orig_module.dim, + orig_module.max_position_embeddings, + orig_module.base, + None, # device + orig_module.scaling_factor, + orig_module.original_max_position_embeddings, + orig_module.beta_fast, + orig_module.beta_slow, + orig_module.mscale, + orig_module.mscale_all_dim, + ) + self.generate_device = generate_device + self.prefill_device = prefill_device + + def load(self): + self.orig_module.__init__( + self.orig_module.dim, + self.orig_module.max_position_embeddings, + self.orig_module.base, + self.generate_device, + self.orig_module.scaling_factor, + self.orig_module.original_max_position_embeddings, + self.orig_module.beta_fast, + self.orig_module.beta_slow, + self.orig_module.mscale, + self.orig_module.mscale_all_dim, + ) + +# class DeepSeekV3YarnRotaryEmbedding(BaseInjectedModule, DeepseekV3RotaryEmbedding): +# def __init__( +# self, +# key: str, +# gguf_loader: GGUFLoader, +# config: PretrainedConfig, +# orig_module: nn.Module, +# # device: str = "cuda", +# generate_device: str = "cuda", +# prefill_device: str = "cuda", +# **kwargs, +# ): +# BaseInjectedModule.__init__( +# self, key, gguf_loader, config, orig_module, prefill_device, generate_device, **kwargs +# ) +# self.generate_device = generate_device +# self.prefill_device = prefill_device + +# def load(self): +# # TODO support perlayer prefill +# self.orig_module.__init__( +# self.config, +# device=self.generate_device +# ) +# return + +class YarnRotaryEmbeddingV3(BaseInjectedModule): + def __init__( + self, + key: str, + gguf_loader: GGUFLoader, + config: PretrainedConfig, + orig_module: nn.Module, + # device: str = "cuda", + generate_device: str = "cuda", + prefill_device: str = "cuda", + **kwargs, + ): + BaseInjectedModule.__init__( + self, key, gguf_loader, config, orig_module, prefill_device, generate_device, **kwargs + ) + self.generate_device = generate_device + self.prefill_device = prefill_device + + def load(self): + kwargs = { + key: self.config.rope_scaling[key] + for key in [ + "original_max_position_embeddings", + "beta_fast", + "beta_slow", + "mscale", + "mscale_all_dim", + ] + if key in self.config.rope_scaling + } + self._init( + dim=self.config.qk_rope_head_dim, + max_position_embeddings=self.config.max_position_embeddings, + base=self.config.rope_theta, + device=self.device, + scaling_factor=self.config.rope_scaling["factor"], + **kwargs, + ) + + @maybe_no_grad() + def forward(self, x, position_ids): + # x: [bs, num_attention_heads, seq_len, head_size] + inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1) + position_ids_expanded = position_ids[:, None, :].float() + # Force float32 since bfloat16 loses precision on long contexts + # See https://github.com/huggingface/transformers/pull/29285 + device_type = x.device.type + device_type = device_type if isinstance(device_type, str) and device_type != "mps" else "cpu" + with torch.autocast(device_type=device_type, enabled=False): + freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2) + emb = torch.cat((freqs, freqs), dim=-1) + cos = emb.cos()* self._mscale + sin = emb.sin()* self._mscale + return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype) + + def _init( + self, + dim, + max_position_embeddings=2048, + base=10000, + device=None, + scaling_factor=1.0, + original_max_position_embeddings=4096, + beta_fast=32, + beta_slow=1, + mscale=1, + mscale_all_dim=0, + ): + self.original_max_position_embeddings = original_max_position_embeddings + self.beta_fast = beta_fast + self.beta_slow = beta_slow + self.mscale = mscale + self.mscale_all_dim = mscale_all_dim + self.scaling_factor = scaling_factor + self.dim = dim + self.max_position_embeddings = max_position_embeddings + self.base = base + + freq_extra = 1.0 / ( + self.base + ** (torch.arange(0, dim, 2, dtype=torch.float32, device=device) / dim) + ) + freq_inter = 1.0 / ( + self.scaling_factor + * self.base + ** (torch.arange(0, dim, 2, dtype=torch.float32, device=device) / dim) + ) + + low, high = yarn_find_correction_range( + self.beta_fast, + self.beta_slow, + dim, + self.base, + self.original_max_position_embeddings, + ) + inv_freq_mask = 1.0 - yarn_linear_ramp_mask(low, high, dim // 2).to( + device=device, dtype=torch.float32 + ) + self.inv_freq = freq_inter * (1 - inv_freq_mask) + freq_extra * inv_freq_mask + self._mscale = float( + yarn_get_mscale(self.scaling_factor, self.mscale) + / yarn_get_mscale(self.scaling_factor, self.mscale_all_dim) + ) + # For BC we register cos and sin cached + self.max_seq_len_cached = max_position_embeddings + +class DynamicNTKScalingRotaryEmbedding( + BaseInjectedModule, LlamaDynamicNTKScalingRotaryEmbedding +): + def __init__( + self, + key: str, + gguf_loader: GGUFLoader, + config: PretrainedConfig, + orig_module: nn.Module, + prefill_device: str = "cuda", + generate_device: str = "cuda", + **kwargs, + ): + BaseInjectedModule.__init__( + self, key, gguf_loader, config, orig_module, prefill_device, generate_device, **kwargs + ) + self.orig_module.__init__( + orig_module.dim, + orig_module.max_position_embeddings, + orig_module.base, + None, # device + orig_module.scaling_factor, + orig_module.rope_type, + orig_module.config, + ) + + def load(self): + self.orig_module.__init__( + self.orig_module.dim, + self.orig_module.max_position_embeddings, + self.orig_module.base, + self.orig_module.device, + self.orig_module.scaling_factor, + self.orig_module.rope_type, + self.orig_module.config, + ) + + + +class RotaryEmbeddingV4(BaseInjectedModule): + def __init__( + self, + key: str, + gguf_loader: GGUFLoader, + config: PretrainedConfig, + orig_module: nn.Module, + # device: str = "cuda", + generate_device: str = "cuda", + prefill_device: str = "cuda", + **kwargs, + ): + BaseInjectedModule.__init__( + self, key, gguf_loader, config, orig_module, generate_device, **kwargs + ) + self.generate_device = generate_device + self.prefill_device = prefill_device + + @maybe_no_grad() + def forward(self, x, position_ids): + # x: [bs, num_attention_heads, seq_len, head_size] + inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1) + position_ids_expanded = position_ids[:, None, :].float() + # Force float32 since bfloat16 loses precision on long contexts + # See https://github.com/huggingface/transformers/pull/29285 + device_type = x.device.type + device_type = device_type if isinstance(device_type, str) and device_type != "mps" else "cpu" + with torch.autocast(device_type=device_type, enabled=False): + freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2) + emb = torch.cat((freqs, freqs), dim=-1) + cos = emb.cos() + sin = emb.sin() + return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype) + + def load(self): + self._init( + dim=self.config.qk_rope_head_dim, + max_position_embeddings=self.config.max_position_embeddings, + base=self.config.rope_theta, + device=self.device, + ) + def _init(self, dim, max_position_embeddings, base, device, scaling_factor=1.0): + self.scaling_factor = scaling_factor + self.dim = dim + self.max_position_embeddings = max_position_embeddings + self.base = base + self.inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2, dtype=torch.int64).float().to(device) / self.dim)) + # self.register_buffer("inv_freq", inv_freq, persistent=False) + # For BC we register cos and sin cached + self.max_seq_len_cached = max_position_embeddings + +class KQwen3MoeRotaryEmbedding(BaseInjectedModule, DeepseekV2RotaryEmbedding): + def __init__( + self, + key: str, + gguf_loader: GGUFLoader, + config: PretrainedConfig, + orig_module: nn.Module, + # device: str = "cuda", + generate_device: str = "cuda", + prefill_device: str = "cuda", + **kwargs, + ): + BaseInjectedModule.__init__( + self, key, gguf_loader, config, orig_module, prefill_device, generate_device, **kwargs + ) + self.orig_module.__init__( + config, + ) + self.generate_device = generate_device + self.prefill_device = prefill_device + + def load(self): + self.orig_module.__init__( + self.orig_module.config + ) \ No newline at end of file diff --git a/archive/kt-sft/ktransformers/operators/__init__.py b/archive/kt-sft/ktransformers/operators/__init__.py new file mode 100644 index 0000000..8b13789 --- /dev/null +++ b/archive/kt-sft/ktransformers/operators/__init__.py @@ -0,0 +1 @@ + diff --git a/archive/kt-sft/ktransformers/operators/attention.py b/archive/kt-sft/ktransformers/operators/attention.py new file mode 100644 index 0000000..30e7de9 --- /dev/null +++ b/archive/kt-sft/ktransformers/operators/attention.py @@ -0,0 +1,1083 @@ +''' +Description : +Author : Boxin Zhang +Version : 0.1.0 +Copyright (c) 2024 by KVCache.AI, All Rights Reserved. +''' +import torch +from torch import nn +import warnings +import torch.nn.functional as F +from ktransformers.operators.models import KLlamaModel +from ktransformers.models.configuration_deepseek import DeepseekV2Config +from ktransformers.models.configuration_llama import LlamaConfig +from ktransformers.models.modeling_llama import LlamaRotaryEmbedding +from ktransformers.models.modeling_deepseek import DeepseekV2Attention, apply_rotary_pos_emb +from ktransformers.models.modeling_qwen3_moe import Qwen3MoeAttention, Qwen3MoeRotaryEmbedding +from typing import Optional, Tuple +from ktransformers.operators.base_operator import BaseInjectedModule +from ktransformers.util.custom_loader import GGUFLoader +from ktransformers.util.utils import get_compute_capability +import logging +from transformers.configuration_utils import PretrainedConfig +from transformers.cache_utils import Cache +from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS +from ktransformers.util.vendors import device_manager, get_device, to_device, GPUVendor + +try: + from flash_attn import flash_attn_func +except: + pass +from ktransformers.operators.triton_attention import decode_attention_fwd_grouped +from ktransformers.operators.triton_attention_prefill import context_attention_fwd +import os +from ktransformers.operators.flashinfer_wrapper import flashinfer_enabled +if flashinfer_enabled: + from ktransformers.operators.flashinfer_wrapper import MLAWrapperSingleton + from flashinfer.mla import BatchMLAPagedAttentionWrapper +from ktransformers.models.custom_cache import KDeepSeekV3Cache +logger = logging.getLogger("attention") + +# Copied from transformers.models.llama.modeling_llama.rotate_half +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) + +# V3 MLA is same to V2 +class KDeepseekV2Attention(BaseInjectedModule, DeepseekV2Attention): + """Multi-headed attention from 'Attention Is All You Need' paper""" + attn_mask: Optional[torch.Tensor] = None + + def __init__(self, + key: str, + gguf_loader : GGUFLoader, + config: PretrainedConfig, + orig_module: nn.Module, + prefill_device: str = "cuda", + generate_device: str = "cuda", + chunck_size: int = 1000, + absorb_for_prefill: bool = False, + **kwargs): + BaseInjectedModule.__init__(self, key, gguf_loader, config, orig_module, prefill_device, generate_device, **kwargs) + self.orig_module.__init__(orig_module.config, + orig_module.layer_idx) + self.chunck_size = chunck_size # TODO, generate chunck_size automatically. + self.mla_wrapper = None + self.absorb_for_prefill = absorb_for_prefill + + def get_absorbed(self) -> Tuple[torch.Tensor, torch.Tensor]: + if not (hasattr(self, 'q_absorb') and hasattr(self, 'out_absorb')): + kv_b_proj = self.kv_b_proj.weight.view(self.num_heads, -1, self.kv_lora_rank) + self.q_absorb = kv_b_proj[:, :self.qk_nope_head_dim, :].view(self.num_heads, self.qk_nope_head_dim, self.kv_lora_rank) + self.out_absorb = kv_b_proj[:, self.qk_nope_head_dim:, :].view(self.num_heads, self.v_head_dim, self.kv_lora_rank) + + return self.q_absorb, self.out_absorb + + def forward_chunck( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_value: Optional[Cache] = None, + output_attentions: bool = False, + use_cache: bool = False, + cache_position: Optional[torch.LongTensor] = None, + **kwargs + ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: + bsz, q_len, _ = hidden_states.size() + if self.q_lora_rank is None: + q = self.q_proj(hidden_states) + else: + q = self.q_b_proj(self.q_a_layernorm(self.q_a_proj(hidden_states))) + q = q.view(bsz, q_len, self.num_heads, self.q_head_dim).transpose(1, 2) + q_nope, q_pe = torch.split( + q, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1 + ) + # q_nope [bsz, self.num_heads, q_len, self.qk_nope_head_dim] + # q_pe [bsz, self.num_heads, q_len, self.qk_rope_head_dim] + + compressed_kv = self.kv_a_proj_with_mqa(hidden_states) + compressed_kv, k_pe = torch.split( + compressed_kv, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1 + ) + compressed_kv = self.kv_a_layernorm(compressed_kv) + k_pe = k_pe.view(bsz, q_len, 1, self.qk_rope_head_dim).transpose(1, 2) + + kv_seq_len = k_pe.shape[-2] + if past_key_value is not None: + if self.layer_idx is None: + raise ValueError( + f"The cache structure has changed since transformer version v4.36. If you are using {self.__class__.__name__} " + "for auto-regressive decoding with k/v caching, please make sure to initialize the attention class " + "with a layer index." + ) + kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx) + + cos, sin = self.rotary_emb(q_pe, position_ids) + q_pe, k_pe = apply_rotary_pos_emb(q_pe, k_pe, cos, sin) + + if past_key_value is not None: + cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position} # Specific to RoPE models + + # compressed_kv [bsz, q_len, self.kv_lora_rank] + # k_pe [bsz, 1, q_len, self.qk_rope_head_dim] + k_pe = k_pe.transpose(1,2) + compressed_kv = compressed_kv.unsqueeze(2) + compressed_kv_with_k_pe, _ = past_key_value.update(compressed_kv, k_pe, self.layer_idx, cache_kwargs) + compressed_kv, k_pe = torch.split( + compressed_kv_with_k_pe, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1 + ) + # k_pe [pages, page_size, 1, self.qk_rope_head_dim] + # compressed_kv [pages, page_size, 1, self.kv_lora_rank] + + q_absorb, out_absorb = self.get_absorbed() + + # q_nope [bsz, self.num_heads, q_len, self.qk_nope_head_dim] + # q_pe [bsz, self.num_heads, q_len, self.qk_rope_head_dim] + k_pe = k_pe.view(bsz, 1, -1, self.qk_rope_head_dim)[:,:,:attention_mask.size(-1),:] + compressed_kv = compressed_kv.view(bsz, 1, -1, self.kv_lora_rank)[:,:,:attention_mask.size(-1),:] + # k_pe [bsz, 1, cache_len, self.qk_rope_head_dim] + # compressed_kv [bsz, 1, cache_len,self.kv_lora_rank] + q_nope = torch.matmul(q_nope, q_absorb) + #print(q_pe.shape) + #print(k_pe.shape) + #print(q_nope.shape) + #print(compressed_kv.shape) + + attn_weights = (torch.matmul(q_pe, k_pe.mT) + torch.matmul(q_nope, compressed_kv.mT)) * self.softmax_scale + + #attn_weights [bsz, self.num_heads, q_len, kv_seq_len] + compressed_kv = compressed_kv.squeeze(1) + """ + if attn_weights.size() != (bsz, self.num_heads, q_len, kv_seq_len): + raise ValueError( + f"Attention weights should be of size {(bsz, self.num_heads, q_len, kv_seq_len)}, but is" + f" {attn_weights.size()}" + ) + assert attention_mask is not None + """ + if attention_mask is not None: + """ + if attention_mask.size() != (bsz, 1, q_len, kv_seq_len): + raise ValueError( + f"Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}, but is {attention_mask.size()}" + ) + """ + #causal_mask = attention_mask[:, :, :, : kv_seq_len] + attn_weights = attn_weights + attention_mask + + # upcast attention to fp32 + attn_weights = nn.functional.softmax( + attn_weights, dim=-1, dtype=torch.float32 + ).to(q_pe.dtype) + attn_weights = nn.functional.dropout( + attn_weights, p=self.attention_dropout, training=self.training + ) + + attn_output = torch.einsum('bhql,blc->bhqc', attn_weights, compressed_kv) + + attn_output = torch.matmul(attn_output, out_absorb.mT) + + if attn_output.size() != (bsz, self.num_heads, q_len, self.v_head_dim): + raise ValueError( + f"`attn_output` should be of size {(bsz, self.num_heads, q_len, self.v_head_dim)}, but is" + f" {attn_output.size()}" + ) + + attn_output = attn_output.transpose(1, 2).contiguous() + + attn_output = attn_output.reshape(bsz, q_len, self.num_heads * self.v_head_dim) + + attn_output = self.o_proj(attn_output) + + return attn_output, None, past_key_value + + def forward_linux_triton( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_value: Optional[Cache] = None, + output_attentions: bool = False, + use_cache: bool = False, + cache_position: Optional[torch.LongTensor] = None, + **kwargs, + ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: + + bsz, q_len, _ = hidden_states.size() + + if self.q_lora_rank is None: + q = self.q_proj(hidden_states) + else: + q = self.q_b_proj(self.q_a_layernorm(self.q_a_proj(hidden_states))) + q = q.view(bsz, q_len, self.num_heads, self.q_head_dim) + q_nope, q_pe = torch.split( + q, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1 + ) + + compressed_kv = self.kv_a_proj_with_mqa(hidden_states) + compressed_kv, k_pe = torch.split( + compressed_kv, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1 + ) + compressed_kv = self.kv_a_layernorm(compressed_kv) + k_pe = k_pe.view(bsz, q_len, 1, self.qk_rope_head_dim) + compressed_kv = compressed_kv.view(bsz, q_len, 1, self.kv_lora_rank) + + kv_seq_len = q_len + if past_key_value is not None: + if self.layer_idx is None: + raise ValueError( + f"The cache structure has changed since transformer version v4.36. If you are using {self.__class__.__name__} " + "for auto-regressive decoding with k/v caching, please make sure to initialize the attention class " + "with a layer index." + ) + kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx) + + cos, sin = self.rotary_emb(q_pe, position_ids) + q_pe, k_pe = apply_rotary_pos_emb(q_pe, k_pe, cos, sin, unsqueeze_dim=2) + # q_pe [bsz, q_len, self.num_heads, self.qk_rope_head_dim] k_pe [bsz, q_len, 1, self.qk_rope_head_dim] + + # decode + if q_len == 1: + if past_key_value is not None: + cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position} # Specific to RoPE models + compressed_kv_with_k_pe, page_table = past_key_value.update(compressed_kv, k_pe, self.layer_idx, cache_kwargs) + compressed_kv = compressed_kv_with_k_pe [:, :, :, :self.kv_lora_rank] # for speed + # compressed_kv_with_k_pe [bsz, q_len, 1, self.kv_lora_rank + self.qk_rope_head_dim] + # compressed_kv [bsz, q_len, 1, self.kv_lora_rank] + + # q_nope [bsz, q_len, self.num_heads, self.qk_nope_head_dim] + # q_absorb [self.num_heads, self.qk_nope_head_dim, self.kv_lora_rank] + q_absorb, out_absorb = self.get_absorbed() + q_nope = q_nope.transpose(1, 2) # q_len is 1, no GPU overhead, same below + q_nope = torch.matmul(q_nope, q_absorb) # batched MM + q_nope = q_nope.transpose(1, 2) + #assert q_nope.is_contiguous() + + # q_nope [bsz, q_len, self.num_heads, self.kv_lora_rank] + # q_pe [bsz, q_len, self.num_heads, self.qk_rope_head_dim] + query_states = torch.cat([q_nope, q_pe], dim=-1) + + query_states = query_states.squeeze(1) + attn_output = torch.zeros_like(q_nope) # [bsz, q_len, self.num_heads, self.kv_lora_rank] + + attn_logits = torch.empty( + ( + bsz, + self.num_heads, + 4, #num_kv_splits # follow vLLM, fix it TODO + self.kv_lora_rank + 1, + ), + dtype=torch.float32, + device = attn_output.device + ) + + """ + print("query_states", torch.isnan(query_states).any()) + print("compressed_kv_with_k_pe", torch.isnan(compressed_kv_with_k_pe[:,:,0,:]).any()) + print("compressed_kv", torch.isnan(compressed_kv[:,:,0,:]).any()) + print("position_ids", torch.isnan(position_ids).any()) + """ + + # flash attn doesn't support head_dim bigger than 256 + # use triton attention kernel adapted from vLLM and SGLang for MQA + decode_attention_fwd_grouped(query_states, compressed_kv_with_k_pe, compressed_kv, attn_output, + page_table, + position_ids.squeeze(0).to(torch.int32)+1, attn_logits, + 4, #num_kv_splits # follow vLLM, fix it TODO + self.softmax_scale, + past_key_value.page_size) + + # attn_output [bsz, q_len, self.num_heads, self.kv_lora_rank] + # out_absorb [self.num_heads, self.v_head_dim, self.kv_lora_rank] + attn_output = attn_output.transpose(1, 2) + attn_output = torch.matmul(attn_output, out_absorb.mT) + attn_output = attn_output.transpose(1, 2) + + attn_output = attn_output.reshape(bsz, q_len, self.num_heads * self.v_head_dim) + attn_output = self.o_proj(attn_output) + + #print("attn_output", torch.isnan(attn_output).any()) + return attn_output, None, past_key_value + else: + if past_key_value is not None: + cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position} # Specific to RoPE models + k_pe.squeeze(0) + compressed_kv.squeeze(0) + compressed_kv_with_k_pe, _ = past_key_value.update(compressed_kv, k_pe, self.layer_idx, cache_kwargs) + compressed_kv, k_pe = torch.split( + compressed_kv_with_k_pe, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1 + ) + k_pe = k_pe.view(bsz, -1, self.qk_rope_head_dim) + k_pe = k_pe[:, :kv_seq_len] + compressed_kv = compressed_kv.view(bsz, -1, self.kv_lora_rank) + compressed_kv = compressed_kv[:, :kv_seq_len] + kv = ( + self.kv_b_proj(compressed_kv) + .view(bsz, kv_seq_len, self.num_heads, self.qk_nope_head_dim + self.v_head_dim) + ) + k_nope, value_states = torch.split(kv, [self.qk_nope_head_dim, self.v_head_dim], dim=-1) + query_states = k_pe.new_empty(bsz, q_len, self.num_heads, self.q_head_dim) + query_states[:, :, :, : self.qk_nope_head_dim] = q_nope + query_states[:, :, :, self.qk_nope_head_dim :] = q_pe + + key_states = k_pe.new_empty(bsz, kv_seq_len, self.num_heads, self.q_head_dim) + key_states[:, :, :, :self.qk_nope_head_dim] = k_nope + key_states[:, :, :, self.qk_nope_head_dim:] = k_pe.view(bsz, kv_seq_len, 1, -1) + + value_states = value_states.view(bsz, kv_seq_len, self.num_heads, self.v_head_dim) + value_states_padded = torch.nn.functional.pad(value_states, [0, query_states.shape[-1] - value_states.shape[-1]], value=0) + + attn_output = flash_attn_func( + query_states, + key_states, + value_states_padded, + softmax_scale=self.softmax_scale, + causal=True, + ) + + if self.q_head_dim != self.v_head_dim: + attn_output = attn_output[:, :, :, : self.v_head_dim] + + attn_output = attn_output.reshape( + bsz, q_len, self.num_heads * self.v_head_dim + ).contiguous() + attn_output = self.o_proj(attn_output) + return attn_output, None, past_key_value + + def forward_linux_flashinfer( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.Tensor] = None, + past_key_value: Optional[Cache] = None, + output_attentions: bool = False, + use_cache: bool = False, + cache_position: Optional[torch.Tensor] = None, + **kwargs, + ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: + + bsz, q_len, _ = hidden_states.size() + + if self.q_lora_rank is None: + q = self.q_proj(hidden_states) + else: + q = self.q_b_proj(self.q_a_layernorm(self.q_a_proj(hidden_states))) + q = q.view(bsz, q_len, self.num_heads, self.q_head_dim) + q_nope, q_pe = torch.split( + q, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1 + ) + + compressed_kv = self.kv_a_proj_with_mqa(hidden_states) + compressed_kv, k_pe = torch.split( + compressed_kv, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1 + ) + compressed_kv = self.kv_a_layernorm(compressed_kv) + k_pe = k_pe.view(bsz, q_len, 1, self.qk_rope_head_dim) + compressed_kv = compressed_kv.view(bsz, q_len, 1, self.kv_lora_rank) + + kv_seq_len = q_len + if past_key_value is not None: + if self.layer_idx is None: + raise ValueError( + f"The cache structure has changed since version transformer verision v4.36. If you are using {self.__class__.__name__} " + "for auto-regressive decoding with k/v caching, please make sure to initialize the attention class " + "with a layer index." + ) + kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx) + + cos, sin = self.rotary_emb(q_pe, position_ids) + q_pe, k_pe = apply_rotary_pos_emb(q_pe, k_pe, cos, sin, unsqueeze_dim=2) + # q_pe [bsz, q_len, self.num_heads, self.qk_rope_head_dim] k_pe [bsz, q_len, 1, self.qk_rope_head_dim] + + # decode + if q_len == 1 or self.absorb_for_prefill: + if past_key_value is not None: + cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position} # Specific to RoPE models + compressed_kv_with_k_pe, page_table = past_key_value.update(compressed_kv, k_pe, self.layer_idx, cache_kwargs) + compressed_kv = compressed_kv_with_k_pe [:, :, :, :self.kv_lora_rank].view(-1, past_key_value.page_size, self.kv_lora_rank) + k_pe = compressed_kv_with_k_pe [:, :, :, self.kv_lora_rank:].view(-1, past_key_value.page_size, self.qk_rope_head_dim) + # k_pe [max_pages, page_size, self.qk_rope_head_dim] + # compressed_kv [max_pages, page_size, self.kv_lora_rank] + + # q_nope [bsz, q_len, self.num_heads, self.qk_nope_head_dim] + # q_absorb [self.num_heads, self.qk_nope_head_dim, self.kv_lora_rank] + q_absorb, out_absorb = self.get_absorbed() + q_nope = q_nope.transpose(1, 2) # q_len is 1, no GPU overhead, same below + q_nope = torch.matmul(q_nope, q_absorb) # batched MM + q_nope = q_nope.transpose(1, 2) + q_nope = q_nope.contiguous() + #assert q_nope.is_contiguous() + + # q_nope [bsz, q_len, self.num_heads, self.kv_lora_rank] + # q_pe [bsz, q_len, self.num_heads, self.qk_rope_head_dim] + q_nope.squeeze_(0) + q_pe.squeeze_(0) + + # flash attn doesn't support head_dim bigger than 256, use flashinfer + if self.mla_wrapper is None: + self.mla_wrapper = MLAWrapperSingleton.get_instance(self.device, 1, past_key_value.max_pages, use_cuda_graph = True) + if self.mla_wrapper.need_plan: + self.mla_wrapper.need_plan = False + if q_len == 1: + self.mla_wrapper.plan(None,None,None, + position_ids.squeeze(1)+1, + None, + self.num_heads, + self.kv_lora_rank, + self.qk_rope_head_dim, + past_key_value.page_size, + self.softmax_scale, + q_nope.dtype, + compressed_kv.dtype) + else: + qo_indptr = torch.tensor([0, q_len], dtype=torch.int32, device=self.device) + kv_len_arr = torch.tensor([position_ids[0, -1].item()+1], dtype=torch.int32, device=self.device) + self.mla_wrapper.plan(qo_indptr,None,None, + kv_len_arr, + None, + self.num_heads, + self.kv_lora_rank, + self.qk_rope_head_dim, + past_key_value.page_size, + self.softmax_scale, + q_nope.dtype, + compressed_kv.dtype) + attn_output = self.mla_wrapper.run(q_nope, q_pe, compressed_kv, k_pe).view(bsz, q_len, self.num_heads, self.kv_lora_rank) + """ + k = ( + torch.cat([compressed_kv, k_pe], dim=-1) + .view(-1, 1, 512 + 64) + .repeat_interleave(self.num_heads, dim=1) + ) + v = compressed_kv.view(-1, 1, 512).repeat_interleave(self.num_heads, dim=1) + lens = position_ids.item() + 1 + #print("lens", lens) + attn_ref, lse_ref = attention_ref( + 1, + torch.cat([q_nope, q_pe], dim=-1), + k[:lens], + v[:lens], + False, + self.softmax_scale + ) + attn_output = attn_ref.view(bsz, q_len, self.num_heads, self.kv_lora_rank) + """ + + # mla_wrapper run output: [tokens, self.num_heads, self.kv_lora_rank] + # attn_output [bsz, q_len, self.num_heads, self.kv_lora_rank] + # out_absorb [self.num_heads, self.v_head_dim, self.kv_lora_rank] + attn_output = attn_output.transpose(1, 2) # [bsz, self.num_heads, q_len, self.kv_lora_rank] + attn_output = torch.matmul(attn_output, out_absorb.mT) # [bsz, self.num_heads, q_len, self.v_head_dim] + attn_output = attn_output.transpose(1, 2).contiguous() # [bsz, q_len, self.num_heads, self.kv_lora_rank] + + attn_output = attn_output.reshape(bsz, q_len, self.num_heads * self.v_head_dim) # [bsz, q_len, self.num_heads * self.v_head_dim] + attn_output = self.o_proj(attn_output) + + return attn_output, None, past_key_value + else: + if past_key_value is not None: + cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position} # Specific to RoPE models + k_pe.squeeze(0) + compressed_kv.squeeze(0) + compressed_kv_with_k_pe, _ = past_key_value.update(compressed_kv, k_pe, self.layer_idx, cache_kwargs) + compressed_kv, k_pe = torch.split( + compressed_kv_with_k_pe, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1 + ) + k_pe = k_pe.view(bsz, -1, self.qk_rope_head_dim) + k_pe = k_pe[:, :kv_seq_len] + compressed_kv = compressed_kv.view(bsz, -1, self.kv_lora_rank) + compressed_kv = compressed_kv[:, :kv_seq_len] + kv = ( + self.kv_b_proj(compressed_kv) + .view(bsz, kv_seq_len, self.num_heads, self.qk_nope_head_dim + self.v_head_dim) + ) + k_nope, value_states = torch.split(kv, [self.qk_nope_head_dim, self.v_head_dim], dim=-1) + query_states = k_pe.new_empty(bsz, q_len, self.num_heads, self.q_head_dim) + query_states[:, :, :, : self.qk_nope_head_dim] = q_nope + query_states[:, :, :, self.qk_nope_head_dim :] = q_pe + + key_states = k_pe.new_empty(bsz, kv_seq_len, self.num_heads, self.q_head_dim) + key_states[:, :, :, :self.qk_nope_head_dim] = k_nope + key_states[:, :, :, self.qk_nope_head_dim:] = k_pe.view(bsz, kv_seq_len, 1, -1) + + value_states = value_states.view(bsz, kv_seq_len, self.num_heads, self.v_head_dim) + value_states_padded = torch.nn.functional.pad(value_states, [0, query_states.shape[-1] - value_states.shape[-1]], value=0) + + attn_output = flash_attn_func( + query_states, + key_states, + value_states_padded, + softmax_scale=self.softmax_scale, + causal=True, + ) + + if self.q_head_dim != self.v_head_dim: + attn_output = attn_output[:, :, :, : self.v_head_dim] + + attn_output = attn_output.reshape( + bsz, q_len, self.num_heads * self.v_head_dim + ).contiguous() + attn_output = self.o_proj(attn_output) + return attn_output, None, past_key_value + + def forward_windows( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_value: Optional[Cache] = None, + output_attentions: bool = False, + use_cache: bool = False, + cache_position: Optional[torch.LongTensor] = None, + **kwargs, + ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: + if "padding_mask" in kwargs: + warnings.warn( + "Passing `padding_mask` is deprecated and will be removed in v4.37. Please make sure use `attention_mask` instead.`" + ) + bsz, q_len, _ = hidden_states.size() + + if q_len <= self.chunck_size: + return self.forward_chunck( + hidden_states, + attention_mask, + position_ids, + past_key_value, + output_attentions, + use_cache, + cache_position, + **kwargs + ) + + assert output_attentions == False, "output_attentions is not supported when using chunked attention" + attn_output = None + cur_idx = 0 + while cur_idx < q_len: + if attention_mask is not None: + chunk_mask = attention_mask[:, :, cur_idx:min(cur_idx + self.chunck_size, q_len), ...] + else: + # generate chunk_mask automatically. + self.attn_mask = \ + torch.zeros(1, 1, self.chunck_size, past_key_value.max_cache_len, device=hidden_states.device) \ + if self.attn_mask is None \ + else self.attn_mask + self.attn_mask[:, :, :, cur_idx:min(cur_idx+self.chunck_size, past_key_value.max_cache_len)] = \ + -1e+38 * torch.triu(torch.ones(self.chunck_size, self.chunck_size, device=hidden_states.device), diagonal=1)\ + [:,:min(self.chunck_size, min(past_key_value.max_cache_len-cur_idx, self.chunck_size))] + self.attn_mask[:, :, :, cur_idx+self.chunck_size:] = -1e+38 + self.attn_mask[:, :, :, :cur_idx] = 0 + chunk_mask = torch.narrow(self.attn_mask, 2, 0, min(self.chunck_size, q_len-cur_idx)) + + cur_output, _, _ = self.forward_chunck( + hidden_states[:, cur_idx:min(cur_idx + self.chunck_size, q_len), ...], + chunk_mask, + position_ids[:, cur_idx:min(cur_idx + self.chunck_size, q_len)], + past_key_value, + output_attentions, + use_cache, + cache_position[cur_idx:min(cur_idx + self.chunck_size, q_len)], + **kwargs + ) + cur_idx += self.chunck_size + if attn_output is None: + attn_output = cur_output + else: + attn_output = torch.cat((attn_output, cur_output), dim=-2) + + return attn_output, None, past_key_value + + def forward_xpu( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_value: Optional[Cache] = None, + output_attentions: bool = False, + use_cache: bool = False, + cache_position: Optional[torch.LongTensor] = None, + **kwargs, + ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: + if "padding_mask" in kwargs: + warnings.warn( + "Passing `padding_mask` is deprecated and will be removed in v4.37. Please make sure use `attention_mask` instead.`" + ) + bsz, q_len, _ = hidden_states.size() + + if self.q_lora_rank is None: + q = self.q_proj(hidden_states) + else: + q = self.q_b_proj(self.q_a_layernorm(self.q_a_proj(hidden_states))) + query_states = q.view(bsz, q_len, self.num_heads, self.q_head_dim).transpose(1, 2) + + compressed_kv = self.kv_a_proj_with_mqa(hidden_states) + compressed_kv, k_pe = torch.split( + compressed_kv, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1 + ) + k_pe = k_pe.view(bsz, q_len, 1, self.qk_rope_head_dim).transpose(1, 2) + kv = ( + self.kv_b_proj(self.kv_a_layernorm(compressed_kv)) + .view(bsz, q_len, self.num_heads, self.qk_nope_head_dim + self.v_head_dim) + .transpose(1, 2) + ) + + k_nope, value_states = torch.split( + kv, [self.qk_nope_head_dim, self.v_head_dim], dim=-1 + ) + kv_seq_len = value_states.shape[-2] + if past_key_value is not None: + if self.layer_idx is None: + raise ValueError( + f"The cache structure has changed since version v4.36. If you are using {self.__class__.__name__} " + "for auto-regressive decoding with k/v caching, please make sure to initialize the attention class " + "with a layer index." + ) + kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx) + + position_embeddings = kwargs.get("position_embeddings", None) + if position_embeddings is not None: + cos, sin = position_embeddings + key_states = torch.cat( + [k_nope, k_pe.expand([-1, self.num_heads, -1, -1])], + dim=-1 + ) + from ipex_llm.transformers.models.common import rotary_two_with_cache_inplaced + rotary_two_with_cache_inplaced(query_states[:, :, :, self.qk_nope_head_dim :], + key_states[:, :, :, self.qk_nope_head_dim:], + cos, sin, True) + else: + q_nope, q_pe = torch.split( + query_states, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1 + ) + cos, sin = self.rotary_emb(q_pe, position_ids) + q_pe, k_pe = apply_rotary_pos_emb(q_pe, k_pe, cos, sin) + query_states = k_pe.new_empty(bsz, self.num_heads, q_len, self.q_head_dim) + query_states[:, :, :, : self.qk_nope_head_dim] = q_nope + query_states[:, :, :, self.qk_nope_head_dim :] = q_pe + + key_states = k_pe.new_empty(bsz, self.num_heads, q_len, self.q_head_dim) + key_states[:, :, :, : self.qk_nope_head_dim] = k_nope + key_states[:, :, :, self.qk_nope_head_dim :] = k_pe + + if past_key_value is not None: + cache_kwargs = {"sin": sin, "cos": cos} # Specific to RoPE models + key_states, value_states = past_key_value.update( + key_states.half(), value_states.half(), self.layer_idx, cache_kwargs + ) + + attn_weights = None + from ipex_llm.transformers.models.common import scaled_dot_product_attention + attn_output = scaled_dot_product_attention( + query_states.half(), key_states, value_states, + attention_mask.half(), q_len == kv_seq_len, self.softmax_scale + ) + + attn_output = attn_output.transpose(1, 2).contiguous() + attn_output = attn_output.reshape(bsz, q_len, self.num_heads * self.v_head_dim) + attn_output = self.o_proj(attn_output).to(hidden_states.dtype) + + if not output_attentions: + attn_weights = None + + return attn_output, attn_weights, past_key_value + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_value: Optional[Cache] = None, + output_attentions: bool = False, + use_cache: bool = False, + cache_position: Optional[torch.LongTensor] = None, + **kwargs, + ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: + if torch.xpu.is_available(): + return self.forward_xpu( + hidden_states, + attention_mask, + position_ids, + past_key_value, + output_attentions, + use_cache, + cache_position, + **kwargs, + ) + elif (os.name == 'nt' + or get_compute_capability() < 8 + or hidden_states.device.type == 'cpu' + or device_manager.gpu_vendor != GPUVendor.NVIDIA): + return self.forward_windows( + hidden_states, + attention_mask, + position_ids, + past_key_value, + output_attentions, + use_cache, + cache_position, + **kwargs, + ) + else: + if flashinfer_enabled: + return self.forward_linux_flashinfer( + hidden_states, + attention_mask, + position_ids, + past_key_value, + output_attentions, + use_cache, + cache_position, + **kwargs, + ) + else: + return self.forward_linux_triton( + hidden_states, + attention_mask, + position_ids, + past_key_value, + output_attentions, + use_cache, + cache_position, + **kwargs, + ) + + +class KLlamaAttention(BaseInjectedModule): + """Multi-headed attention from 'Attention Is All You Need' paper""" + + def __init__(self, + key: str, + gguf_loader : GGUFLoader, + config: PretrainedConfig, + orig_module: nn.Module, + prefill_device: str = "cuda", + generate_device: str = "cuda", + **kwargs): + BaseInjectedModule.__init__(self, key, gguf_loader, config, orig_module, prefill_device, generate_device, **kwargs) + self.orig_module.__init__(orig_module.config, + orig_module.layer_idx) + def apply_rotary_pos_emb(self, q, k, cos, sin, position_ids=None, unsqueeze_dim=1): + """Applies Rotary Position Embedding to the query and key tensors. + + Args: + q (`torch.Tensor`): The query tensor. + k (`torch.Tensor`): The key tensor. + cos (`torch.Tensor`): The cosine part of the rotary embedding. + sin (`torch.Tensor`): The sine part of the rotary embedding. + position_ids (`torch.Tensor`, *optional*): + Deprecated and unused. + unsqueeze_dim (`int`, *optional*, defaults to 1): + The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and + sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note + that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and + k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes + cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have + the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2. + Returns: + `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding. + """ + cos = cos.unsqueeze(unsqueeze_dim) + sin = sin.unsqueeze(unsqueeze_dim) + q_embed = (q * cos) + (rotate_half(q) * sin) + k_embed = (k * cos) + (rotate_half(k) * sin) + return q_embed, k_embed + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_value: Optional[Cache] = None, + output_attentions: bool = False, + use_cache: bool = False, + cache_position: Optional[torch.LongTensor] = None, + position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, # will become mandatory in v4.45 + **kwargs, + ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: + bsz, q_len, _ = hidden_states.size() + + if self.config.pretraining_tp > 1: + key_value_slicing = (self.num_key_value_heads * self.head_dim) // self.config.pretraining_tp + query_slices = self.q_proj.weight.split( + (self.num_heads * self.head_dim) // self.config.pretraining_tp, dim=0 + ) + key_slices = self.k_proj.weight.split(key_value_slicing, dim=0) + value_slices = self.v_proj.weight.split(key_value_slicing, dim=0) + + query_states = [F.linear(hidden_states, query_slices[i]) for i in range(self.config.pretraining_tp)] + query_states = torch.cat(query_states, dim=-1) + + key_states = [F.linear(hidden_states, key_slices[i]) for i in range(self.config.pretraining_tp)] + key_states = torch.cat(key_states, dim=-1) + + value_states = [F.linear(hidden_states, value_slices[i]) for i in range(self.config.pretraining_tp)] + value_states = torch.cat(value_states, dim=-1) + + else: + query_states = self.q_proj(hidden_states) + key_states = self.k_proj(hidden_states) + value_states = self.v_proj(hidden_states) + + query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2) + key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) + value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) + + if position_embeddings is None: + + logger.warning( + "The attention layers in this model are transitioning from computing the RoPE embeddings internally " + "through `position_ids` (2D tensor with the indexes of the tokens), to using externally computed " + "`position_embeddings` (Tuple of tensors, containing cos and sin). In v4.45 `position_ids` will be " + "removed and `position_embeddings` will be mandatory." + ) + cos, sin = self.rotary_emb(value_states, position_ids) + else: + cos, sin = position_embeddings + query_states, key_states = self.apply_rotary_pos_emb(query_states, key_states, cos, sin) + if q_len == 1: + position_ids = position_ids[0][-1].unsqueeze(0).unsqueeze(0) + query_states = query_states[:, :, -1:] + key_states = key_states[:, :, -1:] + + attn_output = KLlamaModel.dynamic_sdpa.apply( + self.layer_idx, + bsz, + position_ids[0][0], + query_states.transpose(1, 2).to(torch.float16), + key_states.transpose(1, 2).to(torch.float16), + value_states.transpose(1, 2).to(torch.float16), + mode="prefill" if q_len > 1 else "generate", + ) + + + if attn_output.size() != (bsz, self.num_heads, q_len, self.head_dim): + raise ValueError( + f"`attn_output` should be of size {(bsz, self.num_heads, q_len, self.head_dim)}, but is" + f" {attn_output.size()}" + ) + + attn_output = attn_output.transpose(1, 2).contiguous() + + attn_output = attn_output.reshape(bsz, q_len, -1) + + if self.config.pretraining_tp > 1: + attn_output = attn_output.split(self.hidden_size // self.config.pretraining_tp, dim=2) + o_proj_slices = self.o_proj.weight.split(self.hidden_size // self.config.pretraining_tp, dim=1) + attn_output = sum([F.linear(attn_output[i], o_proj_slices[i]) for i in range(self.config.pretraining_tp)]) + else: + attn_output = self.o_proj(attn_output) + + if not output_attentions: + attn_weights = None + + return attn_output, attn_weights, past_key_value + + +class KQwen3MoeAttentionIPEXLLM(BaseInjectedModule, Qwen3MoeAttention): + def __init__(self, + key: str, + gguf_loader : GGUFLoader, + config: PretrainedConfig, + orig_module: nn.Module, + prefill_device: str = "xpu", + generate_device: str = "xpu", + chunck_size: int = 1000, + **kwargs): + BaseInjectedModule.__init__(self, key, gguf_loader, config, orig_module, prefill_device, **kwargs) + self.orig_module.__init__(orig_module.config, + orig_module.layer_idx) + self.chunck_size = chunck_size # TODO, generate chunck_size automatically. + assert prefill_device.lower()[:3] == "xpu", "KQwen3MoeAttentionIPEXLLM only supports XPU device" + assert generate_device.lower()[:3] == "xpu", "KQwen3MoeAttentionIPEXLLM only supports XPU device" + + def forward( + self, + hidden_states: torch.Tensor, + position_ids: Optional[torch.Tensor], + position_embeddings: Tuple[torch.Tensor, torch.Tensor], + attention_mask: Optional[torch.Tensor], + past_key_value: Optional[Cache] = None, + cache_position: Optional[torch.LongTensor] = None, + **kwargs + ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: + input_shape = hidden_states.shape[:-1] + bsz, q_len, _ = hidden_states.size() + input_dtype = hidden_states.dtype + hidden_shape = (*input_shape, -1, self.head_dim) + + if not hasattr(self, 'qkv_proj'): + from ipex_llm.transformers.models.common import merge_quantized_qkv + merge_quantized_qkv(self.q_proj.generate_linear, self.k_proj.generate_linear, self.v_proj.generate_linear, self.orig_module) + + qkv = self.qkv_proj(hidden_states) + qkv = qkv.view(bsz, q_len, -1, self.head_dim) + qkv = qkv.transpose(1, 2) + query_states, key_states, value_states = qkv.split([self.config.num_attention_heads, + self.config.num_key_value_heads, + self.config.num_key_value_heads], dim=1) + query_states = self.q_norm(query_states) + key_states = self.k_norm(key_states) + + if position_embeddings is None: + position_embeddings = self.rotary_emb(hidden_states, position_ids) + + cos, sin = position_embeddings + + from ipex_llm.transformers.models.common import rotary_half_with_cache_inplaced + rotary_half_with_cache_inplaced(query_states, key_states, cos, sin) + + if past_key_value is not None: + # sin and cos are specific to RoPE models; cache_position needed for the static cache + cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position} + key_states, value_states = past_key_value.update(key_states.half(), value_states.half(), + self.layer_idx, cache_kwargs) + + attn_weights = None + from ipex_llm.transformers.models.common import scaled_dot_product_attention + attn_output = scaled_dot_product_attention( + query_states.half(), key_states, value_states, + attention_mask.half(), q_len == key_states.size(2), self.scaling + ) + attn_output = attn_output.transpose(1, 2).contiguous() + + attn_output = attn_output.reshape(*input_shape, -1).contiguous() + attn_output = self.o_proj(attn_output).to(input_dtype) + return attn_output, attn_weights + + +def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: + """ + This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch, + num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim) + """ + batch, num_key_value_heads, slen, head_dim = hidden_states.shape + if n_rep == 1: + return hidden_states + hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim) + return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim) + + +def eager_attention_forward( + module: nn.Module, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attention_mask: Optional[torch.Tensor], + scaling: float, + dropout: float = 0.0, + **kwargs, +): + key_states = repeat_kv(key, module.num_key_value_groups) + value_states = repeat_kv(value, module.num_key_value_groups) + + attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling + if attention_mask is not None: + causal_mask = attention_mask[:, :, :, : key_states.shape[-2]] + attn_weights = attn_weights + causal_mask + + attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype) + attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training) + attn_output = torch.matmul(attn_weights, value_states) + attn_output = attn_output.transpose(1, 2).contiguous() + + return attn_output, attn_weights + + +class KQwen3MoeAttention(BaseInjectedModule, Qwen3MoeAttention ): + def __init__(self, + key: str, + gguf_loader: GGUFLoader, + config: PretrainedConfig, + orig_module: nn.Module, + prefill_device: str = "cuda", + generate_device: str = "cuda", + chunck_size: int = 1000, + **kwargs): + BaseInjectedModule.__init__(self, key, gguf_loader, config, orig_module, prefill_device, generate_device, + **kwargs) + self.orig_module.__init__(self.orig_module.config, + orig_module.layer_idx) + self.chunck_size = chunck_size # TODO, generate chunck_size automatically. + + # Copied from transformers.models.mistral.modeling_mistral.apply_rotary_pos_emb + def apply_rotary_pos_emb(self, q, k, cos, sin, position_ids=None, unsqueeze_dim=1): + """Applies Rotary Position Embedding to the query and key tensors. + + Args: + q (`torch.Tensor`): The query tensor. + k (`torch.Tensor`): The key tensor. + cos (`torch.Tensor`): The cosine part of the rotary embedding. + sin (`torch.Tensor`): The sine part of the rotary embedding. + position_ids (`torch.Tensor`): + Deprecated and unused. + unsqueeze_dim (`int`, *optional*, defaults to 1): + The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and + sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note + that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and + k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes + cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have + the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2. + Returns: + `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding. + """ + cos = cos.unsqueeze(unsqueeze_dim) + sin = sin.unsqueeze(unsqueeze_dim) + q_embed = (q * cos) + (rotate_half(q) * sin) + k_embed = (k * cos) + (rotate_half(k) * sin) + return q_embed, k_embed + + def forward(self, + hidden_states: torch.Tensor, + position_ids: Optional[torch.Tensor], + position_embeddings: Tuple[torch.Tensor, torch.Tensor], + attention_mask: Optional[torch.Tensor], + past_key_value: Optional[Cache] = None, + cache_position: Optional[torch.LongTensor] = None, + **kwargs + ): + input_shape = hidden_states.shape[:-1] + hidden_shape = (*input_shape, -1, self.head_dim) + + query_states = self.q_norm(self.q_proj(hidden_states).view(hidden_shape)).transpose(1, 2) + key_states = self.k_norm(self.k_proj(hidden_states).view(hidden_shape)).transpose(1, 2) + value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2) + + if position_embeddings is None: + position_embeddings = self.rotary_emb(hidden_states, position_ids) + + cos, sin = position_embeddings + + query_states, key_states = self.apply_rotary_pos_emb(query_states, key_states, cos, sin) + + + if past_key_value is not None: + # sin and cos are specific to RoPE models; cache_position needed for the static cache + cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position} + key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs) + + attention_interface: Callable = eager_attention_forward + if self.config._attn_implementation != "eager": + if self.config._attn_implementation == "sdpa" and kwargs.get("output_attentions", False): + logger.warning_once( + "`torch.nn.functional.scaled_dot_product_attention` does not support `output_attentions=True`. Falling back to " + 'eager attention. This warning can be removed using the argument `attn_implementation="eager"` when loading the model.' + ) + else: + attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation] + + attn_output, attn_weights = attention_interface( + self, + query_states, + key_states, + value_states, + attention_mask, + dropout=0.0 if not self.training else self.attention_dropout, + scaling=self.scaling, + sliding_window=self.sliding_window, # diff with Llama + **kwargs, + ) + + attn_output = attn_output.reshape(*input_shape, -1).contiguous() + attn_output = self.o_proj(attn_output) + return attn_output, attn_weights diff --git a/archive/kt-sft/ktransformers/operators/balance_serve_attention.py b/archive/kt-sft/ktransformers/operators/balance_serve_attention.py new file mode 100644 index 0000000..51695f3 --- /dev/null +++ b/archive/kt-sft/ktransformers/operators/balance_serve_attention.py @@ -0,0 +1,457 @@ +''' +Description : +Author : Boxin Zhang +Version : 0.2.5 +Copyright (c) 2024 by KVCache.AI, All Rights Reserved. +''' +import torch +from torch import nn +from ktransformers.models.modeling_deepseek import DeepseekV2Attention, apply_rotary_pos_emb +from ktransformers.models.modeling_qwen2_moe import Qwen2MoeAttention +from ktransformers.models.modeling_qwen3_moe import Qwen3MoeAttention +from typing import Optional, Tuple +from ktransformers.operators.base_operator import BaseInjectedModule +from ktransformers.util.custom_loader import GGUFLoader +import logging +from transformers.configuration_utils import PretrainedConfig +from flashinfer import BatchMLAPagedAttentionWrapper +from ktransformers.operators.flashinfer_batch_prefill_wrapper import flashInferAttn +from ktransformers.models.custom_cache import KDeepSeekV3Cache, KGQACache +logger = logging.getLogger("attention") + +# Copied from transformers.models.llama.modeling_llama.rotate_half +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) + +class flashinfer_attn(BaseInjectedModule, DeepseekV2Attention): + def __init__(self, + key: str, + gguf_loader : GGUFLoader, + config: PretrainedConfig, + orig_module: nn.Module, + prefill_device: str = "cuda", + generate_device: str = "cuda", + chunck_size: int = 1000, + **kwargs): + BaseInjectedModule.__init__(self, key, gguf_loader, config, orig_module, prefill_device, **kwargs) + self.orig_module.__init__(orig_module.config, + orig_module.layer_idx) + self.chunck_size = chunck_size # TODO, generate chunck_size automatically. + + + def get_absorbed(self) -> Tuple[torch.Tensor, torch.Tensor]: + if not (hasattr(self, 'q_absorb') and hasattr(self, 'out_absorb')): + kv_b_proj = self.kv_b_proj.weight.view(self.num_heads, -1, self.kv_lora_rank) + q_absorb = kv_b_proj[:, :self.qk_nope_head_dim, :].reshape(-1, self.kv_lora_rank) + out_absorb = kv_b_proj[:, self.qk_nope_head_dim:, :].reshape(-1, self.kv_lora_rank) + self.q_absorb = nn.Linear(self.kv_lora_rank, self.num_heads * self.qk_nope_head_dim, + bias=False, dtype=q_absorb.dtype, device=q_absorb.device) + self.q_absorb.weight.data = q_absorb + self.out_absorb = nn.Linear(self.kv_lora_rank, self.num_heads * self.v_head_dim, + bias=False, dtype=out_absorb.dtype, device=out_absorb.device) + self.out_absorb.weight.data = out_absorb + #del self.orig_module.kv_b_proj + q_absorb = self.q_absorb.weight.view(self.num_heads, self.qk_nope_head_dim, self.kv_lora_rank) + out_absorb = self.out_absorb.weight.view(self.num_heads, self.v_head_dim, self.kv_lora_rank) + return q_absorb, out_absorb + + + def forward(self, + hidden_states: torch.Tensor, + kv_cache: KDeepSeekV3Cache, + position_ids: torch.Tensor, + wrapper: BatchMLAPagedAttentionWrapper, + num_tokens_tensors: torch.Tensor, + page_idx: torch.Tensor, + page_offset: torch.Tensor, + ): + q_len, _ = hidden_states.size() + + if self.q_lora_rank is None: + q = self.q_proj(hidden_states, num_tokens_tensors) + else: + q = self.q_b_proj(self.q_a_layernorm(self.q_a_proj(hidden_states, num_tokens_tensors), num_tokens_tensors), num_tokens_tensors) + q = q.view(q_len, self.num_heads, self.q_head_dim) + q_nope, q_pe = torch.split( + q, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1 + ) + + compressed_kv = self.kv_a_proj_with_mqa(hidden_states, num_tokens_tensors) + compressed_kv, k_pe = torch.split( + compressed_kv, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1 + ) + compressed_kv = compressed_kv.contiguous() + compressed_kv = self.kv_a_layernorm(compressed_kv, num_tokens_tensors) + k_pe = k_pe.view(q_len, 1, self.qk_rope_head_dim) + compressed_kv = compressed_kv.view(q_len, 1, self.kv_lora_rank) + + cos, sin = self.rotary_emb(q_pe, position_ids.unsqueeze(0)) + q_pe, k_pe = apply_rotary_pos_emb(q_pe.unsqueeze(0), k_pe.unsqueeze(0), cos, sin, unsqueeze_dim=2) + q_pe = q_pe.squeeze(0) + if kv_cache is not None: + + # page_idx, page_offset = kv_cache.get_page_table(position_ids, q_indptr, kv_indptr, kv_indices) + cache_kwargs = {"sin": sin, "cos": cos, "page_idx": page_idx, "page_offset": page_offset} # Specific to RoPE models + compressed_kv_with_k_pe = kv_cache.update(compressed_kv.unsqueeze(0), k_pe, self.layer_idx, page_idx, page_offset, cache_kwargs) + compressed_kv = compressed_kv_with_k_pe [:, :, :, :self.kv_lora_rank].view(-1, kv_cache.page_size, self.kv_lora_rank) + k_pe = compressed_kv_with_k_pe [:, :, :, self.kv_lora_rank:].view(-1, kv_cache.page_size, self.qk_rope_head_dim) + + q_absorb, out_absorb = self.get_absorbed() + q_nope = q_nope.transpose(0, 1) # q_len is 1, no GPU overhead, same below + q_nope = torch.matmul(q_nope, q_absorb) # batched MM + q_nope = q_nope.transpose(0, 1) + # q_nope.squeeze_(1) + # q_pe.squeeze_(1) + + attn_output = wrapper.run(q_nope, q_pe, compressed_kv, k_pe).view(q_len, self.num_heads, self.kv_lora_rank) + attn_output = attn_output.transpose(0, 1) + attn_output = torch.matmul(attn_output, out_absorb.mT) # [self.num_heads, q_len, self.v_head_dim] + attn_output = attn_output.transpose(0, 1) + attn_output = attn_output.reshape(q_len, self.num_heads * self.v_head_dim) + attn_output = self.o_proj(attn_output, num_tokens_tensors) + return attn_output + +class KQwen2MoeAttention(BaseInjectedModule, Qwen2MoeAttention): + def __init__(self, + key: str, + gguf_loader : GGUFLoader, + config: PretrainedConfig, + orig_module: nn.Module, + prefill_device: str = "cuda", + generate_device: str = "cuda", + chunck_size: int = 1000, + **kwargs): + BaseInjectedModule.__init__(self, key, gguf_loader, config, orig_module, prefill_device, **kwargs) + self.orig_module.__init__(orig_module.config, + orig_module.layer_idx) + self.chunck_size = chunck_size # TODO, generate chunck_size automatically. + + + # Copied from transformers.models.mistral.modeling_mistral.apply_rotary_pos_emb + def apply_rotary_pos_emb(self, q, k, cos, sin, position_ids=None, unsqueeze_dim=1): + """Applies Rotary Position Embedding to the query and key tensors. + + Args: + q (`torch.Tensor`): The query tensor. + k (`torch.Tensor`): The key tensor. + cos (`torch.Tensor`): The cosine part of the rotary embedding. + sin (`torch.Tensor`): The sine part of the rotary embedding. + position_ids (`torch.Tensor`): + Deprecated and unused. + unsqueeze_dim (`int`, *optional*, defaults to 1): + The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and + sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note + that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and + k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes + cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have + the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2. + Returns: + `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding. + """ + cos = cos.unsqueeze(unsqueeze_dim) + sin = sin.unsqueeze(unsqueeze_dim) + q_embed = (q * cos) + (rotate_half(q) * sin) + k_embed = (k * cos) + (rotate_half(k) * sin) + return q_embed, k_embed + + + def forward(self, + hidden_states: torch.Tensor, + kv_cache: KGQACache, + position_ids: torch.Tensor, + wrapper: flashInferAttn, + bsz_tensors: torch.Tensor, + page_idx: torch.Tensor, + page_offset: torch.Tensor, + ): + q_len, _ = hidden_states.size() + + query_states = self.q_proj(hidden_states, bsz_tensors) + key_states = self.k_proj(hidden_states, bsz_tensors) + value_states = self.v_proj(hidden_states, bsz_tensors) + + + query_states = query_states.view(q_len, self.num_heads, self.head_dim) + key_states = key_states.view(q_len, self.num_key_value_heads, self.head_dim) + value_states = value_states.view(q_len, self.num_key_value_heads, self.head_dim) + + cos, sin = self.rotary_emb(value_states.unsqueeze(0), position_ids.unsqueeze(0)) + query_states, key_states = self.apply_rotary_pos_emb(query_states.unsqueeze(0), key_states.unsqueeze(0), cos, sin, unsqueeze_dim=2) + + query_states = query_states.view(q_len, self.num_heads, self.head_dim) + key_states = key_states.view( + q_len, self.num_key_value_heads, self.head_dim + ) + value_states = value_states.view( + q_len, self.num_key_value_heads, self.head_dim + ) + + k_cache = kv_cache.get_k_cache(self.layer_idx) + v_cache = kv_cache.get_v_cache(self.layer_idx) + + + attn_output = wrapper.forward(query_states, k_cache, v_cache, key_states, value_states) + + + attn_output = self.o_proj(attn_output.view(q_len, self.num_heads * self.head_dim), bsz_tensors) + + return attn_output + +class KQwen3MoeAttention(BaseInjectedModule, Qwen3MoeAttention): + def __init__(self, + key: str, + gguf_loader : GGUFLoader, + config: PretrainedConfig, + orig_module: nn.Module, + prefill_device: str = "cuda", + generate_device: str = "cuda", + chunck_size: int = 1000, + **kwargs): + BaseInjectedModule.__init__(self, key, gguf_loader, config, orig_module, prefill_device, **kwargs) + self.orig_module.__init__(orig_module.config, + orig_module.layer_idx) + self.chunck_size = chunck_size # TODO, generate chunck_size automatically. + + + # Copied from transformers.models.mistral.modeling_mistral.apply_rotary_pos_emb + def apply_rotary_pos_emb(self, q, k, cos, sin, position_ids=None, unsqueeze_dim=1): + """Applies Rotary Position Embedding to the query and key tensors. + + Args: + q (`torch.Tensor`): The query tensor. + k (`torch.Tensor`): The key tensor. + cos (`torch.Tensor`): The cosine part of the rotary embedding. + sin (`torch.Tensor`): The sine part of the rotary embedding. + position_ids (`torch.Tensor`): + Deprecated and unused. + unsqueeze_dim (`int`, *optional*, defaults to 1): + The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and + sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note + that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and + k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes + cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have + the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2. + Returns: + `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding. + """ + cos = cos.unsqueeze(unsqueeze_dim) + sin = sin.unsqueeze(unsqueeze_dim) + q_embed = (q * cos) + (rotate_half(q) * sin) + k_embed = (k * cos) + (rotate_half(k) * sin) + return q_embed, k_embed + + + def forward(self, + hidden_states: torch.Tensor, + kv_cache: KGQACache, + position_ids: torch.Tensor, + wrapper: flashInferAttn, + bsz_tensors: torch.Tensor, + page_idx: torch.Tensor, + page_offset: torch.Tensor, + ): + q_len, _ = hidden_states.size() + + bsz_tensors_q = bsz_tensors * self.num_heads + bsz_tensors_kv = bsz_tensors * self.num_key_value_heads + + query_states = self.q_norm(self.q_proj(hidden_states, bsz_tensors), bsz_tensors_q) + key_states = self.k_norm(self.k_proj(hidden_states, bsz_tensors), bsz_tensors_kv) + value_states = self.v_proj(hidden_states, bsz_tensors) + + + query_states = query_states.view(q_len, self.num_heads, self.head_dim) + key_states = key_states.view(q_len, self.num_key_value_heads, self.head_dim) + value_states = value_states.view(q_len, self.num_key_value_heads, self.head_dim) + + cos, sin = self.rotary_emb(value_states.unsqueeze(0), position_ids.unsqueeze(0)) + query_states, key_states = self.apply_rotary_pos_emb(query_states.unsqueeze(0), key_states.unsqueeze(0), cos, sin, unsqueeze_dim=2) + + query_states = query_states.view(q_len, self.num_heads, self.head_dim) + key_states = key_states.view( + q_len, self.num_key_value_heads, self.head_dim + ) + value_states = value_states.view( + q_len, self.num_key_value_heads, self.head_dim + ) + + k_cache = kv_cache.get_k_cache(self.layer_idx) + v_cache = kv_cache.get_v_cache(self.layer_idx) + + + attn_output = wrapper.forward(query_states, k_cache, v_cache, key_states, value_states) + + + attn_output = self.o_proj(attn_output.view(q_len, self.num_heads * self.head_dim), bsz_tensors) + + return attn_output + + +class deepseek_torch_attn(BaseInjectedModule, DeepseekV2Attention): + def __init__(self, + key: str, + gguf_loader : GGUFLoader, + config: PretrainedConfig, + orig_module: nn.Module, + prefill_device: str = "cuda", + generate_device: str = "cuda", + chunck_size: int = 1000, + **kwargs): + BaseInjectedModule.__init__(self, key, gguf_loader, config, orig_module, prefill_device, **kwargs) + self.orig_module.__init__(orig_module.config, + orig_module.layer_idx) + self.chunck_size = chunck_size # TODO, generate chunck_size automatically. + + + def get_absorbed(self) -> Tuple[torch.Tensor, torch.Tensor]: + if not (hasattr(self, 'q_absorb') and hasattr(self, 'out_absorb')): + kv_b_proj = self.kv_b_proj.weight.view(self.num_heads, -1, self.kv_lora_rank) + q_absorb = kv_b_proj[:, :self.qk_nope_head_dim, :].reshape(-1, self.kv_lora_rank) + out_absorb = kv_b_proj[:, self.qk_nope_head_dim:, :].reshape(-1, self.kv_lora_rank) + self.q_absorb = nn.Linear(self.kv_lora_rank, self.num_heads * self.qk_nope_head_dim, + bias=False, dtype=q_absorb.dtype, device=q_absorb.device) + self.q_absorb.weight.data = q_absorb + self.out_absorb = nn.Linear(self.kv_lora_rank, self.num_heads * self.v_head_dim, + bias=False, dtype=out_absorb.dtype, device=out_absorb.device) + self.out_absorb.weight.data = out_absorb + #del self.orig_module.kv_b_proj + q_absorb = self.q_absorb.weight.view(self.num_heads, self.qk_nope_head_dim, self.kv_lora_rank) + out_absorb = self.out_absorb.weight.view(self.num_heads, self.v_head_dim, self.kv_lora_rank) + return q_absorb, out_absorb + + + + def forward(self, + hidden_states: torch.Tensor, + kv_cache: KDeepSeekV3Cache, + position_ids: torch.Tensor, + wrapper: None, + num_tokens_tensors: torch.Tensor, + page_idx: torch.Tensor, + page_offset: torch.Tensor, + attention_masks: Optional[list[torch.Tensor]] = None, + q_indptr: Optional[torch.Tensor] = None, + kv_indices: Optional[torch.Tensor] = None, + kv_indptr: Optional[torch.Tensor] = None, + bsz_tensors: Optional[torch.Tensor] = None, + last_page_len: Optional[torch.Tensor] = None, + ): + # range bsz_tensors + final_attention_output = torch.tensor([], device=hidden_states.device) + for i in range(bsz_tensors[0]): + batch_num_tokens_tensors = q_indptr[i+1] - q_indptr[i] + batch_last_page_len = last_page_len[i] + # kv_total_len is kv_len, batch_compressed_kv is compressed_kv, batch_k_pe is k_pe + batch_page_idx = page_idx[q_indptr[i]:q_indptr[i+1]] + batch_page_offset = page_offset[q_indptr[i]:q_indptr[i+1]] + # kv_page_nums is the number of pages for the current batch + kv_page_nums = kv_indptr[i+1] - kv_indptr[i] + # kv_total_len is the total length of the kv cache for the current batch (kv_len for algorithm) + kv_total_len = kv_page_nums * kv_cache.page_size + if batch_last_page_len is not None: + kv_total_len = kv_total_len - (kv_cache.page_size - batch_last_page_len) + # print(f"kv_total_len's shape {kv_total_len.shape}") + # kv_index is the index of the kv cache pages for the current batch + kv_index = kv_indices[kv_indptr[i]:kv_indptr[i+1]] + # we can index [kv_index, page_offset_indices] to get the kv cache for the current batch + # from q_indptr[i] to q_indptr[i+1] is the range of the current batch + batch_hidden_states = hidden_states[q_indptr[i]:q_indptr[i+1]] + batch_position_ids = position_ids[q_indptr[i]:q_indptr[i+1]] + q_len, _ = batch_hidden_states.size() + # print("q_len -> ", q_len) + + if self.q_lora_rank is None: + q = self.q_proj(batch_hidden_states, batch_num_tokens_tensors) + else: + q = self.q_b_proj(self.q_a_layernorm(self.q_a_proj(batch_hidden_states, batch_num_tokens_tensors), batch_num_tokens_tensors), batch_num_tokens_tensors) + # for v3, bsz, q_len, num_heads(128), qk_head_dim(192=128(nope)+64(rope)) + q = q.view(q_len, self.num_heads, self.q_head_dim) + # q_nope is [q_len, num_heads(128), qk_nope_head_dim(128)] + # q_pe is [q_len, num_heads(128), qk_rope_head_dim(64)] + q_nope, q_pe = torch.split( + q, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1 + ) + # compressed_kv is [q_len, kv_lora_rank(512) + rope(64)] + compressed_kv = self.kv_a_proj_with_mqa(batch_hidden_states, batch_num_tokens_tensors) + # compressed_kv is [q_len, kv_lora_rank(512)], k_pe is [q_len, rope(64)] + compressed_kv, k_pe = torch.split( + compressed_kv, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1 + ) + compressed_kv = compressed_kv.contiguous() + compressed_kv = self.kv_a_layernorm(compressed_kv, batch_num_tokens_tensors) + # k_pe is [q_len, 1, qk_rope_head_dim(64)] + k_pe = k_pe.view(q_len, 1, self.qk_rope_head_dim) + # compressed_kv is [q_len, 1, kv_lora_rank(512)] + compressed_kv = compressed_kv.view(q_len, 1, self.kv_lora_rank) + + cos, sin = self.rotary_emb(q_pe, batch_position_ids.unsqueeze(0)) + # print(f"q_pe shape{q_pe.shape}, k_pe shape {k_pe.shape}") + q_pe, k_pe = apply_rotary_pos_emb(q_pe.unsqueeze(0), k_pe.unsqueeze(0), cos, sin, unsqueeze_dim=2) + q_pe = q_pe.squeeze(0) + # q_pe is [num_heads(128), q_len, qk_rope_head_dim(64)] + q_pe.transpose_(0, 1) + if kv_cache is not None: + cache_kwargs = {"sin": sin, "cos": cos, "page_idx": batch_page_idx, "page_offset": batch_page_offset} # Specific to RoPE models + compressed_kv_with_k_pe = kv_cache.update(compressed_kv.unsqueeze(0), k_pe, self.layer_idx, batch_page_idx, batch_page_offset, cache_kwargs) + compressed_kv = compressed_kv_with_k_pe [:, :, :, :self.kv_lora_rank].view(-1, kv_cache.page_size, self.kv_lora_rank) + k_pe = compressed_kv_with_k_pe [:, :, :, self.kv_lora_rank:].view(-1, kv_cache.page_size, self.qk_rope_head_dim) + # q_absorb is [num_heads(128), qk_nope_head_dim(128), kv_lora_rank(512)] + # out_absorb is [num_heads(128), kv_lora_rank(512), v_head_dim(128)] v_head_dim is also the nope dim + q_absorb, out_absorb = self.get_absorbed() + # q_nope is [num_heads(128), q_len, qk_nope_head_dim(128)] + q_nope = q_nope.transpose(0, 1) # q_len is 1, no GPU overhead, same below + # q_nope is [num_heads(128), q_len, kv_lora_rank(512)] + q_nope = torch.matmul(q_nope, q_absorb) # batched MM + + # # q_nope is [q_len, num_heads(128), kv_lora_rank(512)] + # q_nope = q_nope.transpose(0, 1) + + # we need to index out the compressed_kv and k_pe for the current batch + batch_compressed_kv = None + batch_k_pe = None + for page_index in kv_index: + if kv_total_len > kv_cache.page_size: + tmp_compressed_kv = compressed_kv[page_index, 0:kv_cache.page_size, :] + tmp_k_pe = k_pe[page_index, 0:kv_cache.page_size, :] + if batch_compressed_kv is None or batch_k_pe is None: + batch_compressed_kv = tmp_compressed_kv + batch_k_pe = tmp_k_pe + else: + batch_compressed_kv = torch.cat((batch_compressed_kv, tmp_compressed_kv), dim=0) + batch_k_pe = torch.cat((batch_k_pe, tmp_k_pe), dim=0) + kv_total_len -= kv_cache.page_size + else: + tmp_compressed_kv = compressed_kv[page_index, 0:kv_total_len, :] + tmp_k_pe = k_pe[page_index, 0:kv_total_len, :] + if batch_compressed_kv is None or batch_k_pe is None: + batch_compressed_kv = tmp_compressed_kv + batch_k_pe = tmp_k_pe + else: + batch_compressed_kv = torch.cat((batch_compressed_kv, tmp_compressed_kv), dim=0) + batch_k_pe = torch.cat((batch_k_pe, tmp_k_pe), dim=0) + break + # batch_compressed_kv is [kv_total_len(k_len), kv_lora_rank(512)] + # batch_k_pe is [kv_total_len(k_len), qk_rope_head_dim(64)] + attention_weights = (torch.matmul(q_pe,batch_k_pe.mT) + torch.matmul(q_nope, batch_compressed_kv.mT)) * self.softmax_scale + # attention_weights is [num_heads(128), q_len, k_len] + + # attention_weights = attention_weights.transpose(0,1).unsqueeze(0).squeeze(-1).expand(q_len,-1,-1).transpose(0,1) + + # attention_masks[i] is [q_len, k_len] + + attention_weights = (attention_weights + attention_masks[i]) + # attention_weights shape is [num_heads(128), q_len, k_len] + attention_weights = nn.functional.softmax(attention_weights,dim=-1,dtype=torch.float32).to(q_pe.dtype) + attn_output = torch.matmul(attention_weights, batch_compressed_kv) # [num_heads(128),q_len, lora_rank(512)] + # out_absorb shape is [num_heads(128), kv_lora_rank(512), v_head_dim(128)] + out_absorb = out_absorb.transpose(1,2) + # q for q_len, n for num_heads, h for v_head_dim, v for kv_lora_rank + attn_output = torch.matmul(attn_output, out_absorb) # [num_heads(128), q_len, v_head_dim(128)] + attn_output = attn_output.transpose(0, 1) # [q_len, num_heads(128), v_head_dim(128)] + attn_output = attn_output.reshape(q_len, self.num_heads * self.v_head_dim) + attn_output = self.o_proj(attn_output, batch_num_tokens_tensors) + final_attention_output = torch.cat((final_attention_output, attn_output), dim=0) + return final_attention_output \ No newline at end of file diff --git a/archive/kt-sft/ktransformers/operators/base_operator.py b/archive/kt-sft/ktransformers/operators/base_operator.py new file mode 100644 index 0000000..43f9510 --- /dev/null +++ b/archive/kt-sft/ktransformers/operators/base_operator.py @@ -0,0 +1,68 @@ +''' +Description : +Author : Boxin Zhang +Version : 0.1.0 +Copyright (c) 2024 by KVCache.AI, All Rights Reserved. +''' +from typing import Any +from torch import nn, Tensor +from ktransformers.util.custom_loader import GGUFLoader +from transformers.configuration_utils import PretrainedConfig +import ktransformers.util.utils as utils +class BaseInjectedModule(nn.Module): + + def __init__(self, + key: str, + gguf_loader : GGUFLoader, + config: PretrainedConfig, + orig_module: nn.Module, + prefill_device: str = "cuda", + generate_device: str = "cuda", + **kwargs): + nn.Module.__init__(self) + nn.Module.__setattr__(self, "orig_module", orig_module) + object.__setattr__(self, "key", key) + object.__setattr__(self, "gguf_loader", gguf_loader) + object.__setattr__(self, "config", config) + object.__setattr__(self, "prefill_device", prefill_device) + object.__setattr__(self, "generate_device", generate_device) + object.__setattr__(self, "device", generate_device) + + def __getattr__(self, name: str) -> Any: + # __getattr__ in nn.Module doesn't call super().__getattribute__ when name is not in nn.Module.__dict__, + # but __setattr__ in nn.Module call super().__setattr__ in that case, there may be some attribute set + # but can't get using __getattr__, typically these attr is build in attr of the class, so class.attr does not + # call __getattr__. + # Example: + # ...import torch + # ...l=torch.nn.Linear(100,200) + # ...l.out_features # 200 + # ...l.__getattr__("out_features") # AttributeError: 'Linear' object has no attribute 'out_features' + try: + return object.__getattribute__(self, name) # if this attr belongs to BaseInjectedModule + except: + if name == "orig_module": + return nn.Module.__getattr__(self, "orig_module") + try: + return nn.Module.__getattr__(self, "orig_module").__getattr__(name) # if this attr belongs to orig_module + except: + return super(nn.Module, nn.Module.__getattr__(self, "orig_module")).__getattribute__(name) # if this attr belongs to orig_module but not in nn.Module.__dict__ + + def __setattr__(self, name: str, value: Tensor | nn.Module) -> None: + if name == "orig_module": + return nn.Module.__setattr__(self, "orig_module", value) + # elif name == "base_layer": + # return nn.Module.__setattr__(self, "base_layer", value) + elif hasattr(self, name): + return object.__setattr__(self, name, value) + return nn.Module.__getattr__(self, "orig_module").__setattr__(name, value) + + def forward(self, *args, **kwargs): + return self.orig_module.forward(*args, **kwargs) + + def load(self, gguf_loader=None, adapter_gguf : bool = False): + for name, child in self._modules.items(): + if gguf_loader==None: + utils.load_weights(child, self.gguf_loader, self.key+".", adapter_gguf=adapter_gguf) + else: + utils.load_weights(child, gguf_loader, self.key+".", adapter_gguf=adapter_gguf) diff --git a/archive/kt-sft/ktransformers/operators/cpuinfer.py b/archive/kt-sft/ktransformers/operators/cpuinfer.py new file mode 100644 index 0000000..bf66efb --- /dev/null +++ b/archive/kt-sft/ktransformers/operators/cpuinfer.py @@ -0,0 +1,751 @@ +#!/usr/bin/env python +# coding=utf-8 +""" +Description : This script defines the `CPUInferKVCache` and `CPUInfer` classes for performing inference + with a Key-Value Cache on the CPU. The `CPUInferKVCache` class is responsible for configuring + and managing key-value caches, updating and retrieving cache data, and handling attention + operations. It supports different cache types (e.g., Q4_0, FP16) and retrieval strategies + (e.g., shared, separate). The `CPUInfer` class handles task submission and synchronization + on the CPU, with optional CUDA stream integration for tasks involving GPU acceleration. + These classes facilitate efficient caching and memory management for deep learning models + that leverage key-value attention mechanisms, particularly on CPU-based systems. +Author : djw +Date : 2024-08-26 23:25:24 +Version : 1.0.0 +LastEditors : djw +LastEditTime : 2024-08-26 23:25:24 +Copyright (c) 2024 by KVCache.AI, All Rights Reserved. +""" +import sys, os +from typing import Any +import torch +sys.path.append(os.path.join(os.path.dirname(__file__), "..", "ktransformers_ext", "build")) +sys.path.append(os.path.join(os.path.dirname(__file__), "..", "ktransformers_ext", "build", "Release")) +sys.path.append(os.path.join(os.path.dirname(__file__), "..", "ktransformers_ext", "build", "Debug")) +import cpuinfer_ext +from ktransformers.server.config.config import Config + + +class CPUInferKVCache: + def __init__( + self, + layer_num: int = 32, + kv_head_num: int = 8, + q_head_num: int = 32, + head_dim: int = 128, + block_len: int = 256, + anchor_num: int = 4, + anchor_type: str = "FIXED", + kv_type: str = "Q4_0", + retrieval_type: str = "SHARED", + layer_step: int = 1, + token_step: int = 1, + layer_offset: int = 0, + max_thread_num: int = 32, + max_batch_size: int = 4, + max_block_num: int = 512, + ): + + if anchor_type == "FIXED": + anchor_type = cpuinfer_ext.kvcache.AnchorType.FIXED + elif anchor_type == "QUEST": + anchor_type = cpuinfer_ext.kvcache.AnchorType.QUEST + elif anchor_type == "DYNAMIC": + anchor_type = cpuinfer_ext.kvcache.AnchorType.DYNAMIC + elif anchor_type == "BLOCK_MEAN": + anchor_type = cpuinfer_ext.kvcache.AnchorType.BLOCK_MEAN + elif anchor_type == "BLOCK_MAX": + anchor_type = cpuinfer_ext.kvcache.AnchorType.BLOCK_MAX + else: + raise ValueError(f"Unknown anchor type: {anchor_type}") + + if kv_type == "FP16": + kv_type = cpuinfer_ext.kvcache.ggml_type.FP16 + elif kv_type == "FP32": + assert False, "FP32 is not supported yet." + kv_type = cpuinfer_ext.kvcache.ggml_type.FP32 + elif kv_type == "Q4_0": + kv_type = cpuinfer_ext.kvcache.ggml_type.Q4_0 + elif kv_type == "Q8_0": + kv_type = cpuinfer_ext.kvcache.ggml_type.Q8_0 + else: + raise ValueError(f"Unknown kv type: {kv_type}") + + if retrieval_type == "SHARED": + retrieval_type = cpuinfer_ext.kvcache.RetrievalType.LAYER + elif retrieval_type == "INDIVIDUAL": + retrieval_type = cpuinfer_ext.kvcache.RetrievalType.QHEAD + elif retrieval_type == "SEPARATE": + retrieval_type = cpuinfer_ext.kvcache.RetrievalType.KVHEAD + + self.config = cpuinfer_ext.kvcache.KVCacheConfig( + layer_num, + kv_head_num, + q_head_num, + head_dim, + block_len, + anchor_num, + anchor_type, + kv_type, + retrieval_type, + layer_step, + token_step, + layer_offset, + max_block_num, + max_batch_size, + max_thread_num, + ) + self.kvcache = cpuinfer_ext.kvcache.KVCache(self.config) + + def load_kvcache(self, tensor_file_path: str): + if not os.path.exists(tensor_file_path): + raise FileNotFoundError(f"The file {tensor_file_path} does not exist.") + return self.kvcache.load_kvcache(tensor_file_path,) + + def dump_kvcache( + self, block_table: torch.Tensor, cache_total_len: int, tensor_file_path: str + ): + assert ( + block_table.dim() == 1 + and block_table.dtype == torch.int + and block_table.is_contiguous() + and block_table.device == torch.device("cpu") + ), "block_table dim: {}, size: {}, dtype: {}, contiguous: {}, device: {}".format( + block_table.dim(), + block_table.size(), + block_table.dtype, + block_table.is_contiguous(), + block_table.device, + ) + + assert ( + cache_total_len > 0 + and cache_total_len <= self.config.block_len * block_table.size(0) + ), "cache_total_len: {}".format(cache_total_len) + + if not os.path.exists(os.path.dirname(tensor_file_path)): + os.makedirs(os.path.dirname(tensor_file_path)) + + return self.kvcache.dump_kvcache( + block_table.data_ptr(), + cache_total_len, + tensor_file_path, + ) + + def update_cache_total_len(self, cache_total_len: int): + assert cache_total_len > 0, "cache_total_len: {}".format(cache_total_len) + self.kvcache.update_cache_total_len(cache_total_len) + + # q_in: (bsz, q_len, q_head_num, head_dim) + # output: (bsz, q_len, q_head_num, head_dim) + # attn_lse: (bsz, q_len, q_head_num) + # block_table: (bsz, max_block_num) + def attn( + self, + q_in: torch.Tensor, + output: torch.Tensor, + attn_lse: torch.Tensor, + layer_idx: int, + generate_token_idx: int, + block_table: torch.Tensor | None = None, + cache_seqlens: torch.Tensor | None = None, + pick_block_num: int | None = None, + init_block_num: int | None = None, + local_block_num: int | None = None, + ): + + assert ( + q_in.dim() == 4 + and q_in.size(2) == self.config.q_head_num + and q_in.size(3) == self.config.head_dim + and q_in.dtype == torch.float16 + and q_in.is_contiguous() + and q_in.device == torch.device("cpu") + ), "q_in dim: {}, size: {}, dtype: {}, contiguous: {}, device: {}".format( + q_in.dim(), q_in.size(), q_in.dtype, q_in.is_contiguous(), q_in.device + ) + + batch_size = q_in.size(0) + q_len = q_in.size(1) + + assert (block_table is None) or ( + block_table.dim() == 2 + and block_table.size(0) == batch_size + and block_table.dtype == torch.int + and block_table.is_contiguous() + and block_table.device == torch.device("cpu") + ), "block_table dim: {}, size: {}, dtype: {}, contiguous: {}, device: {}".format( + block_table.dim(), + block_table.size(), + block_table.dtype, + block_table.is_contiguous(), + block_table.device, + ) + + max_block_num = block_table.size(1) if block_table is not None else 0 + + assert ( + output.dim() == 4 + and output.size(0) == batch_size + and output.size(2) == self.config.q_head_num + and output.size(1) == q_len + and output.size(3) == self.config.head_dim + and output.dtype == torch.float16 + and output.is_contiguous() + and output.device == torch.device("cpu") + ), "output dim: {}, size: {}, dtype: {}, contiguous: {}, device: {}".format( + output.dim(), + output.size(), + output.dtype, + output.is_contiguous(), + output.device, + ) + + assert ( + attn_lse.dim() == 3 + and attn_lse.size(0) == batch_size + and attn_lse.size(1) == q_len + and attn_lse.size(2) == self.config.q_head_num + and attn_lse.dtype == torch.float32 + and attn_lse.is_contiguous() + and attn_lse.device == torch.device("cpu") + ), "attn_lse dim: {}, size: {}, dtype: {}, contiguous: {}, device: {}".format( + attn_lse.dim(), + attn_lse.size(), + attn_lse.dtype, + attn_lse.is_contiguous(), + attn_lse.device, + ) + + assert ( + layer_idx >= 0 and layer_idx < self.config.layer_num + ), "layer_idx: {}".format(layer_idx) + + assert (cache_seqlens is None) or ( + cache_seqlens.dim() == 1 + and cache_seqlens.size(0) == batch_size + and cache_seqlens.dtype == torch.int + and cache_seqlens.is_contiguous() + and cache_seqlens.device == torch.device("cpu") + ), "cache_seqlens dim: {}, size: {}, dtype: {}, contiguous: {}, device: {}".format( + cache_seqlens.dim(), + cache_seqlens.size(), + cache_seqlens.dtype, + cache_seqlens.is_contiguous(), + cache_seqlens.device, + ) + + return self.kvcache.attn( + q_in.data_ptr(), + output.data_ptr(), + attn_lse.data_ptr(), + layer_idx, + generate_token_idx, + q_len, + batch_size, + max_block_num, + block_table.data_ptr() if block_table is not None else 0, + cache_seqlens.data_ptr() if cache_seqlens is not None else 0, + pick_block_num, + init_block_num, + local_block_num, + ) + + # k_in: (block_len, kv_head_num, head_dim) + # v_in: (block_len, kv_head_num, head_dim) + def update_kvcache_one_block_fp16( + self, k_in: torch.Tensor, v_in: torch.Tensor, layer_id: int, block_idx: int + ): + assert ( + k_in.dim() == 3 + and k_in.size(1) == self.config.block_len + and k_in.size(0) == self.config.kv_head_num + and k_in.size(2) == self.config.head_dim + and k_in.dtype == torch.float16 + and k_in.is_contiguous() + and k_in.device == torch.device("cpu") + ), "k_in dim: {}, size: {}, dtype: {}, contiguous: {}, device: {}".format( + k_in.dim(), k_in.size(), k_in.dtype, k_in.is_contiguous(), k_in.device + ) + assert ( + v_in.dim() == 3 + and v_in.size(1) == self.config.block_len + and v_in.size(0) == self.config.kv_head_num + and v_in.size(2) == self.config.head_dim + and v_in.dtype == torch.float16 + and v_in.is_contiguous() + and v_in.device == torch.device("cpu") + ), "v_in dim: {}, size: {}, dtype: {}, contiguous: {}, device: {}".format( + v_in.dim(), v_in.size(), v_in.dtype, v_in.is_contiguous(), v_in.device + ) + assert ( + layer_id >= 0 and layer_id < self.config.layer_num + ), "layer_id: {}".format(layer_id) + assert block_idx >= 0, "block_idx: {}".format(block_idx) + return self.kvcache.update_one_block_fp16( + k_in.data_ptr(), + v_in.data_ptr(), + layer_id, + block_idx, + ) + + def get_kvcache_one_block_fp16( + self, k_in: torch.Tensor, v_in: torch.Tensor, layer_id: int, block_idx: int + ): + assert ( + k_in.dim() == 3 + and k_in.size(1) == self.config.block_len + and k_in.size(0) == self.config.kv_head_num + and k_in.size(2) == self.config.head_dim + and k_in.dtype == torch.float16 + and k_in.is_contiguous() + and k_in.device == torch.device("cpu") + ), "k_in dim: {}, size: {}, dtype: {}, contiguous: {}, device: {}".format( + k_in.dim(), k_in.size(), k_in.dtype, k_in.is_contiguous(), k_in.device + ) + assert ( + v_in.dim() == 3 + and v_in.size(1) == self.config.block_len + and v_in.size(0) == self.config.kv_head_num + and v_in.size(2) == self.config.head_dim + and v_in.dtype == torch.float16 + and v_in.is_contiguous() + and v_in.device == torch.device("cpu") + ), "v_in dim: {}, size: {}, dtype: {}, contiguous: {}, device: {}".format( + v_in.dim(), v_in.size(), v_in.dtype, v_in.is_contiguous(), v_in.device + ) + assert ( + layer_id >= 0 and layer_id < self.config.layer_num + ), "layer_id: {}".format(layer_id) + assert block_idx >= 0, "block_idx: {}".format(block_idx) + return self.kvcache.get_one_block_fp16( + k_in.data_ptr(), + v_in.data_ptr(), + layer_id, + block_idx, + ) + + def update_importance_one_block( + self, importance: torch.Tensor, layer_id: int, block_idx: int + ): + assert ( + importance.dim() == 1 + and importance.size(0) == self.config.block_len + and importance.dtype == torch.float16 + and importance.is_contiguous() + and importance.device == torch.device("cpu") + ), "importance dim: {}, size: {}, dtype: {}, contiguous: {}, device: {}".format( + importance.dim(), + importance.size(), + importance.dtype, + importance.is_contiguous(), + importance.device, + ) + assert ( + layer_id >= 0 and layer_id < self.config.layer_num + ), "layer_id: {}".format(layer_id) + assert block_idx >= 0, "block_idx: {}".format(block_idx) + return self.kvcache.update_importance_one_block( + importance.data_ptr(), + layer_id, + block_idx, + ) + + def get_importance_one_block( + self, importance: torch.Tensor, layer_id: int, block_idx: int + ): + assert ( + importance.dim() == 1 + and importance.size(0) == self.config.block_len + and importance.dtype == torch.float16 + and importance.is_contiguous() + and importance.device == torch.device("cpu") + ), "importance dim: {}, size: {}, dtype: {}, contiguous: {}, device: {}".format( + importance.dim(), + importance.size(), + importance.dtype, + importance.is_contiguous(), + importance.device, + ) + assert ( + layer_id >= 0 and layer_id < self.config.layer_num + ), "layer_id: {}".format(layer_id) + assert block_idx >= 0, "block_idx: {}".format(block_idx) + return self.kvcache.get_importance_one_block( + importance.data_ptr(), + layer_id, + block_idx, + ) + + def get_anchor_one_block(self, anchor: torch.Tensor, layer_id: int, block_idx: int): + assert ( + anchor.dim() == 3 + and anchor.size(0) == self.config.kv_head_num + and anchor.size(1) == self.config.anchor_num + and anchor.size(2) == self.config.head_dim + and anchor.dtype == torch.float16 + and anchor.is_contiguous() + and anchor.device == torch.device("cpu") + ), "anchor dim: {}, size: {}, dtype: {}, contiguous: {}, device: {}".format( + anchor.dim(), + anchor.size(), + anchor.dtype, + anchor.is_contiguous(), + anchor.device, + ) + assert ( + layer_id >= 0 and layer_id < self.config.layer_num + ), "layer_id: {}".format(layer_id) + assert block_idx >= 0, "block_idx: {}".format(block_idx) + return self.kvcache.get_anchor_one_block( + anchor.data_ptr(), + layer_id, + block_idx, + ) + + def update_anchor_one_block( + self, anchor: torch.Tensor, layer_id: int, block_idx: int + ): + assert ( + anchor.dim() == 3 + and anchor.size(0) == self.config.kv_head_num + and anchor.size(1) == self.config.anchor_num + and anchor.size(2) == self.config.head_dim + and anchor.dtype == torch.float16 + and anchor.is_contiguous() + and anchor.device == torch.device("cpu") + ), "anchor dim: {}, size: {}, dtype: {}, contiguous: {}, device: {}".format( + anchor.dim(), + anchor.size(), + anchor.dtype, + anchor.is_contiguous(), + anchor.device, + ) + assert ( + layer_id >= 0 and layer_id < self.config.layer_num + ), "layer_id: {}".format(layer_id) + assert block_idx >= 0, "block_idx: {}".format(block_idx) + return self.kvcache.update_anchor_one_block( + anchor.data_ptr(), + layer_id, + block_idx, + ) + + def calc_anchor_all_layers( + self, + block_table: torch.Tensor, + cache_seqlens: torch.Tensor, + ): + assert ( + block_table.dim() == 2 + and block_table.size(0) == cache_seqlens.size(0) + and block_table.dtype == torch.int + and block_table.is_contiguous() + and block_table.device == torch.device("cpu") + ), "block_table dim: {}, size: {}, dtype: {}, contiguous: {}, device: {}".format( + block_table.dim(), + block_table.size(), + block_table.dtype, + block_table.is_contiguous(), + block_table.device, + ) + assert ( + cache_seqlens.dim() == 1 + and cache_seqlens.dtype == torch.int + and cache_seqlens.is_contiguous() + and cache_seqlens.device == torch.device("cpu") + ), "cache_seqlens dim: {}, size: {}, dtype: {}, contiguous: {}, device: {}".format( + cache_seqlens.dim(), + cache_seqlens.size(), + cache_seqlens.dtype, + cache_seqlens.is_contiguous(), + cache_seqlens.device, + ) + batch_size = block_table.size(0) + max_block_num = block_table.size(1) + return self.kvcache.calc_anchor_all_layers( + block_table.data_ptr(), + cache_seqlens.data_ptr(), + batch_size, + max_block_num, + ) + + def clear_importance_all_layers( + self, + block_table: torch.Tensor, + cache_seqlens: torch.Tensor, + ): + assert ( + block_table.dim() == 2 + and block_table.size(0) == cache_seqlens.size(0) + and block_table.dtype == torch.int + and block_table.is_contiguous() + and block_table.device == torch.device("cpu") + ), "block_table dim: {}, size: {}, dtype: {}, contiguous: {}, device: {}".format( + block_table.dim(), + block_table.size(), + block_table.dtype, + block_table.is_contiguous(), + block_table.device, + ) + assert ( + cache_seqlens.dim() == 1 + and cache_seqlens.dtype == torch.int + and cache_seqlens.is_contiguous() + and cache_seqlens.device == torch.device("cpu") + ), "cache_seqlens dim: {}, size: {}, dtype: {}, contiguous: {}, device: {}".format( + cache_seqlens.dim(), + cache_seqlens.size(), + cache_seqlens.dtype, + cache_seqlens.is_contiguous(), + cache_seqlens.device, + ) + batch_size = block_table.size(0) + max_block_num = block_table.size(1) + return self.kvcache.clear_importance_all_layers( + block_table.data_ptr(), + cache_seqlens.data_ptr(), + batch_size, + max_block_num, + ) + + def get_cache_total_len(self): + return self.kvcache.get_cache_total_len() + + def update_kvcache_q4( + self, + k_in: torch.Tensor, + k_scales: torch.Tensor, + v_in: torch.Tensor, + v_scales: torch.Tensor, + layer_id: int, + seq_offset: int | None = None, + seq_len: int | None = None, + block_table: torch.Tensor | None = None, + ): + raise NotImplementedError + + def update_kvcache_fp16( + self, + k_in: torch.Tensor, + v_in: torch.Tensor, + layer_idx, + block_table: torch.Tensor, + max_block_num, + past_len: torch.Tensor, + q_len, + ): + batch_size = block_table.size(0) + return self.kvcache.get_kvcache_fp16( + k_in.data_ptr(), + v_in.data_ptr(), + layer_idx, + block_table.data_ptr(), + batch_size, + max_block_num, + past_len.data_ptr(), + q_len + ) + + def get_kvcache_q4( + self, + k_in: torch.Tensor, + k_scales: torch.Tensor, + v_in: torch.Tensor, + v_scales: torch.Tensor, + layer_id: int, + seq_offset: int | None = None, + seq_len: int | None = None, + block_table: torch.Tensor | None = None, + ): + raise NotImplementedError + + def get_kvcache_fp16( + self, + k_in: torch.Tensor, + v_in: torch.Tensor, + layer_id: int, + layer_idx, + block_table: torch.Tensor, + max_block_num, + past_len: torch.Tensor, + ): + batch_size = block_table.size(0) + return self.kvcache.get_kvcache_fp16( + k_in.data_ptr(), + v_in.data_ptr(), + layer_idx, + block_table.data_ptr(), + batch_size, + max_block_num, + past_len.data_ptr(), + ) + + def get_and_update_kvcache_fp16( + self, + k_cache_cpu: torch.Tensor, + v_cache_cpu: torch.Tensor, + layer_idx, + block_table: torch.Tensor, + max_block_num, + past_len: torch.Tensor, + q_len, + ): + batch_size = block_table.size(0) + return self.kvcache.get_and_update_kvcache_fp16( + k_cache_cpu.data_ptr(), + v_cache_cpu.data_ptr(), + layer_idx, + block_table.data_ptr(), + batch_size, + max_block_num, + past_len.data_ptr(), + q_len, + ) + + def update_importance( + self, + importance_cache: torch.Tensor, + layer_idx, + block_table: torch.Tensor, + max_block_num, + offset: torch.Tensor, + width, + ): + batch_size = block_table.size(0) + return self.kvcache.update_importance( + importance_cache.data_ptr(), + layer_idx, + block_table.data_ptr(), + batch_size, + max_block_num, + offset.data_ptr(), + width, + ) + + # attn_sparsity: ((bsz, q_len, q_head_num), dtype = torch.float32) + def get_attn_sparsity( + self, + q_in: torch.Tensor, + attn_sparsity: torch.Tensor, + layer_idx: int, + block_table: torch.Tensor, + cache_seqlens: torch.Tensor, + block_table_origin: torch.Tensor, + cache_seqlens_origin: torch.Tensor, + generate_token_idx: int = 0, + topk: int | None = None, + local: int | None = None, + ): + batch_size = block_table.size(0) + max_block_num = block_table.size(1) + max_block_num_origin = block_table_origin.size(1) + q_len = q_in.size(1) + + if topk is None or local is None or topk + local >= max_block_num: + topk = -1 + local = -1 + return self.kvcache.get_attn_sparsity( + q_in.data_ptr(), + attn_sparsity.data_ptr(), + layer_idx, + generate_token_idx, + q_len, + batch_size, + max_block_num, + block_table.data_ptr(), + cache_seqlens.data_ptr(), + block_table_origin.data_ptr(), + cache_seqlens_origin.data_ptr(), + max_block_num_origin, + topk, + local, + ) + + def attn_with_kvcache( + self, + q_in: torch.Tensor, + k_in: torch.Tensor, + v_in: torch.Tensor, + output: torch.Tensor, + attn_lse: torch.Tensor, + layer_idx: int, + block_table: torch.Tensor, + cache_seqlens: torch.Tensor, + generate_token_idx: int = 0, + topk: int | None = None, + local: int | None = None, + ): + + batch_size = block_table.size(0) + max_block_num = block_table.size(1) + q_len = q_in.size(1) + + if topk is None or local is None or topk + local >= max_block_num: + topk = -1 + local = -1 + return self.kvcache.attn_with_kvcache( + q_in.data_ptr(), + k_in.data_ptr(), + v_in.data_ptr(), + output.data_ptr(), + attn_lse.data_ptr(), + layer_idx, + generate_token_idx, + q_len, + batch_size, + max_block_num, + block_table.data_ptr(), + cache_seqlens.data_ptr(), + topk, + local, + ) + + def get_all_kvcache_one_layer( + self, k_in: torch.Tensor, v_in: torch.Tensor, layer_id: int + ): + return self.kvcache.get_all_kvcache_one_layer( + k_in.data_ptr(), + v_in.data_ptr(), + layer_id, + ) + + def get_importance( + self, + importance: torch.Tensor, + block_table: torch.Tensor, + ): + raise NotImplementedError + + def get_anchor( + self, + anchor: torch.Tensor, + block_table: torch.Tensor, + ): + raise NotImplementedError + + +class CPUInfer: + cpuinfer = None + cur_backend_thread_num = 0 + + def __init__(self, thread_num): + if thread_num > CPUInfer.cur_backend_thread_num: + CPUInfer.cur_backend_thread_num = thread_num + del CPUInfer.cpuinfer + CPUInfer.cpuinfer = cpuinfer_ext.CPUInfer(thread_num) + + def submit(self, task): + CPUInfer.cpuinfer.submit(task) + + def submit_with_cuda_stream(self, current_cuda_stream, task): + CPUInfer.cpuinfer.submit_with_cuda_stream(current_cuda_stream, task) + + def sync(self): + CPUInfer.cpuinfer.sync() + + def sync_with_cuda_stream(self, current_cuda_stream): + CPUInfer.cpuinfer.sync_with_cuda_stream(current_cuda_stream) + + + diff --git a/archive/kt-sft/ktransformers/operators/dynamic_attention.py b/archive/kt-sft/ktransformers/operators/dynamic_attention.py new file mode 100644 index 0000000..f64e374 --- /dev/null +++ b/archive/kt-sft/ktransformers/operators/dynamic_attention.py @@ -0,0 +1,781 @@ +#!/usr/bin/env python +# coding=utf-8 +""" +Description : +Author : Jianwei Dong +Date : 2024-08-26 23:25:24 +Version : 1.0.0 +LastEditors : Jianwei Dong +LastEditTime : 2024-08-26 23:25:24 +Copyright (c) 2024 by KVCache.AI, All Rights Reserved. +""" + +import torch +from transformers import AutoConfig +import sys, os +import logging +logger = logging.getLogger("dynamic_attention") +sys.path.append(os.path.dirname(__file__) + "/../ktransformers_ext/cpu_backend") +from ktransformers.operators.cpuinfer import CPUInfer, CPUInferKVCache +try: + from flash_attn import flash_attn_func, flash_attn_with_kvcache +except: + print("falsh attn not found") + + +import math +import json + + +class DynamicScaledDotProductAttention: + remaining_length: int + cpu_infer = None + + def __init__( + self, + max_seq_len: int, + block_size: int, + config: AutoConfig, + device: torch.device, + local_windows_len: int, + topk: int, + threads_num: int, + anchor_type: str = "DYNAMIC", + kv_type: str = "FP16", + dense_layer_num: int = 0, + anchor_num: int = 1, + block_selection_mode: str = "SHARED", + layer_step: int = 1, + token_step: int = 1, + preselect_block: bool = False, + preselect_block_count: int = 96, + prefill_chunk_size: int = 20480, + use_attn_sparsity: bool = False, + ): + # assert anchor_num == 1 + # assert anchor_type == "DYNAMIC" + self.remaining_length = 0 + valid_anchor_types = ["DYNAMIC", "FIXED", "BLOCK_MEAN", "BLOCK_MAX", "QUEST"] + assert anchor_type in valid_anchor_types + if anchor_type == "QUEST": + assert anchor_num == 2 + elif anchor_type != "FIXED" and anchor_type != "DYNAMIC": + assert anchor_num == 1 + + valid_kv_types = ["FP16", "FP32", "Q4_0", "Q8_0"] + assert kv_type in valid_kv_types + if kv_type != "FP16" and kv_type != "FP32": + assert block_size % 32 == 0 + + valid_block_selection_modes = ["SHARED", "SEPARATE"] # individual + assert block_selection_mode in valid_block_selection_modes + + self.max_seq_len = max_seq_len + self.block_num = max_seq_len // block_size + self.block_size = block_size + self.anchor_type = anchor_type + self.kv_type = kv_type + self.anchor_num = anchor_num + self.threads_num = threads_num + self.layer_step = layer_step + self.token_step = token_step + self.preselect_block = preselect_block + self.preselect_block_count = preselect_block_count + self.block_selection_mode = block_selection_mode + self.use_attn_sparsity = use_attn_sparsity + + # model config + self.kv_head_num = config.num_key_value_heads + self.q_head_num = config.num_attention_heads + self.head_dim = config.hidden_size // config.num_attention_heads + self.layer_num = config.num_hidden_layers + + self.device = device + self.local_windows_len = local_windows_len + self.local_block_num = self.local_windows_len // self.block_size + 1 + self.prefill_chunk_size = prefill_chunk_size + + self.topk = topk + self.dense_layer_num = dense_layer_num + # self.dense_layer_num = 32 + self.cache_key_states = torch.zeros( + (self.block_num, block_size, self.kv_head_num, self.head_dim), + device=device, + dtype=torch.float16, + ) + self.cache_value_states = torch.zeros( + (self.block_num, block_size, self.kv_head_num, self.head_dim), + device=device, + dtype=torch.float16, + ) + # [max_num_block, block_size, head_num] + self.cache_importance = torch.zeros( + (self.block_num, block_size, self.q_head_num), + device=device, + dtype=torch.float16, + ) + + # key_states: [bsz, q_len, kv_head_num, head_dim] + # value_states: [bsz, q_len, kv_head_num, head_dim] + # query_states: [bsz, q_len, q_head_num, head_dim] + self.q_in_cpu = torch.zeros( + (1, 1, self.q_head_num, self.head_dim), + device="cpu", + dtype=torch.float16, + pin_memory=True, + ) + self.k_in_cpu = torch.zeros( + (1, 1, self.kv_head_num, self.head_dim), + device="cpu", + dtype=torch.float16, + pin_memory=True, + ) + self.v_in_cpu = torch.zeros( + (1, 1, self.kv_head_num, self.head_dim), + device="cpu", + dtype=torch.float16, + pin_memory=True, + ) + + self.cache_seqlens_cpu = torch.empty( + (1,), device="cpu", dtype=torch.int32, pin_memory=True + ) + + self.cache_seqlens_cuda = torch.empty((1,), device=device, dtype=torch.int32) + + self.prefix_block_table = torch.arange( + self.block_num, device="cpu", dtype=torch.int32, pin_memory=True + ).view(1, -1) + + self.block_table_cpu = torch.arange( + self.block_num, device="cpu", dtype=torch.int32, pin_memory=True + ).view(1, -1) + + # assert ( + # self.local_windows_len // self.block_size + 1 + self.preselect_block_count + # <= self.block_num + # ) + + self.output_cpu = torch.empty( + (1, 1, self.q_head_num, self.head_dim), + device="cpu", + dtype=torch.float16, + pin_memory=True, + ) + self.lse_cpu = torch.empty( + (1, 1, self.q_head_num), device="cpu", dtype=torch.float32, pin_memory=True + ) + + self.output_cuda = torch.empty( + (1, 1, self.q_head_num, self.head_dim), device=device, dtype=torch.float16 + ) + + self.attn_sparsity = torch.zeros( + (1, 1, self.q_head_num), device="cpu", dtype=torch.float32, pin_memory=True + ) + + if preselect_block == True: + self.preselect_block_table = torch.zeros( + self.layer_num, + self.preselect_block_count, + device=device, + dtype=torch.int32, + ) + self.preselect_block_num = 0 # block_num before preselect + self.evict_tokens = 0 + + if DynamicScaledDotProductAttention.cpu_infer is None: + DynamicScaledDotProductAttention.cpu_infer = CPUInfer(threads_num) + self.cpu_infer = DynamicScaledDotProductAttention.cpu_infer + self.local_thread = CPUInferKVCache( + self.layer_num, + self.kv_head_num, + self.q_head_num, + self.head_dim, + self.block_size, + anchor_num=self.anchor_num, + anchor_type=anchor_type, + kv_type=self.kv_type, + retrieval_type=self.block_selection_mode, + layer_step=self.layer_step, + token_step=self.token_step, + layer_offset=self.dense_layer_num % self.layer_step, + max_batch_size=1, + max_block_num=self.block_num, + max_thread_num=self.threads_num, + ) + + print( + f"local_windows_len: {local_windows_len}, topk: {topk}, dense_layer_num: {dense_layer_num}, kv_type: {self.kv_type}, anchor_type: {self.anchor_type}, preselect_block: {self.preselect_block}, preselect_block_count: {self.preselect_block_count}, token_step: {self.token_step}, layer_step: {self.layer_step}" + ) + + self.shape_mask = ( + self.q_head_num, + self.block_size, + self.block_size, + ) + + mask = torch.zeros( + self.shape_mask, dtype=torch.uint8, device=device + ).contiguous() + elm_idx = torch.arange(self.block_size, device=device) + + for i in range(mask.size(-2)): + idx = i + mask.size(-1) - mask.size(-2) - elm_idx + idx = idx[idx >= 0] + mask[..., i, idx] = 1 + + self.tril_mask = mask + self.triu_mask = mask ^ 1 + + self.generate_token_idx = 0 + + def get_attn_score_one_block( + self, + batch_idx: int, + max_block_num: int, + query: torch.Tensor, + key: torch.Tensor, + offset: int, + width: int, + mask_mode: str | None = None, + use_softmax: bool = True, + ): + n_rep = self.q_head_num // self.kv_head_num + importance = self.cache_importance.view(-1, self.q_head_num) + importance = importance.narrow(0, batch_idx * max_block_num + offset, width) + n_gqa_ = self.q_head_num // self.kv_head_num + for head_idx in range(self.q_head_num): + key_item = key[..., head_idx // n_gqa_, :].view(key.size(0), -1) + qk = torch.einsum( + "qd,kd->qk", query[:,head_idx,:], key_item + ) # (num_attention_heads, len_q, len_k) + + if mask_mode == "tril": + mask = self.tril_mask + mask = mask[0, -qk.size(-2) :, -qk.size(-1) :] + qk = qk * mask + elif mask_mode == "triu": + mask = self.triu_mask + mask = mask[0, -qk.size(-2) :, -qk.size(-1) :] + qk = qk * mask + + if use_softmax: + qk = torch.nn.functional.softmax( + qk / math.sqrt(self.head_dim), dim=-1, dtype=torch.float32 + ).to(torch.float16) + + qk = torch.sum(qk, dim=-2) + importance[...,head_idx] += qk + + def get_preselect_block_table_and_attn_score( + self, + layer_idx: int, + batch_size: int, + offset: torch.Tensor, + width: int, + query: torch.Tensor, + key: torch.Tensor, + union_with_last_layer: bool = True, + ): + max_seqs_len = offset.max().item() + width + max_block_num = (max_seqs_len + self.block_size - 1) // self.block_size + + for batch_idx in range(batch_size): + query_cur = query[batch_idx][-128:] + self.get_attn_score_one_block( + batch_idx, + max_block_num, + query_cur, + key[batch_idx][: offset[batch_idx].item() + width], + 0, + offset[batch_idx].item() + width, + mask_mode=None, + ) + + if self.preselect_block: + self.prefill_block_num = max( + 0, max_block_num - self.local_windows_len // self.block_size + ) + self.evict_tokens = ( + max(self.prefill_block_num - self.preselect_block_count, 0) + * self.block_size + ) + + if self.prefill_block_num != 0: + importance_cache = self.cache_importance.narrow( + 0, 0, self.prefill_block_num * batch_size + ).view( + batch_size, self.prefill_block_num, self.block_size, self.q_head_num + ) + + importance_r = importance_cache[:, 1:, : self.block_size // 4] + pad_r = torch.zeros_like(importance_r[:, :1]) + importance_r = torch.cat((importance_r, pad_r), dim=1) + importance_l = importance_cache[:, :-1, -self.block_size // 4 :] + pad_l = torch.zeros_like(importance_l[:, :1]) + importance_l = torch.cat((pad_l, importance_l), dim=1) + importance = torch.cat( + (importance_l, importance_cache, importance_r), dim=2 + ) + importance = importance.mean(dim=-1) + importance = importance.mean(dim=-1) + # importance: (batch_size, max_block_num) + topk = min(self.preselect_block_count, self.prefill_block_num) + values, indices = torch.topk( + importance, + k=topk, + dim=1, + ) + + self.preselect_block_table[ + layer_idx : layer_idx + 1, + :topk, + ].copy_(indices) + + if union_with_last_layer and layer_idx == 31: + for tmp_layer_idx in range(self.layer_num - 1): + for i in range(1, min(topk, 6)): + x = self.preselect_block_table[-1, i] + if x not in self.preselect_block_table[tmp_layer_idx]: + self.preselect_block_table[tmp_layer_idx, topk - i] = x + if self.anchor_type == "DYNAMIC": + importance_cache = self.cache_importance.narrow( + 0, 0, max_block_num * batch_size + ).view(batch_size, max_block_num * self.block_size, self.q_head_num) + importance_cache_cpu = torch.empty_like( + importance_cache, device="cpu", pin_memory=True + ) + + importance_cache_cpu.copy_(importance_cache) + + block_table_cpu = self.prefix_block_table[:, :max_block_num].to("cpu") + offset_cpu = offset.contiguous().to("cpu") + + self.cpu_infer.submit( + self.local_thread.update_importance( + importance_cache_cpu, + layer_idx, + block_table_cpu, + max_block_num, + offset_cpu, + width, + ) + ) + self.cpu_infer.sync() + + importance_cache = self.cache_importance.narrow( + 0, 0, max_block_num * batch_size + ).view(batch_size, max_block_num * self.block_size, self.q_head_num) + importance_cache.zero_() + + # key: [bsz, past_len, head_num, head_dim] float16 + # query: [bsz, q_len, q_head_num, head_dim] float16 + def get_attn_score( + self, + layer_idx: int, + batch_size: int, + offset: torch.Tensor, + width: int, + query: torch.Tensor, + key: torch.Tensor, + ): + max_seqs_len = offset.max().item() + width + max_block_num = (max_seqs_len + self.block_size - 1) // self.block_size + + for batch_idx in range(batch_size): + for idx in range(width // self.block_size): + offset_cur = idx * self.block_size + query_cur = query[batch_idx, offset_cur : offset_cur + self.block_size] + self.get_attn_score_one_block( + batch_idx, + max_block_num, + query_cur, + key[ + batch_idx, + offset[batch_idx] + + offset_cur : offset[batch_idx] + + offset_cur + + self.block_size, + ], + offset[batch_idx].item() + offset_cur, + self.block_size, + mask_mode="tril", + use_softmax=False, + ) + + offset_key = ( + offset[batch_idx].item() + + idx * self.block_size + - self.local_windows_len + ) + if offset_key >= 0: + self.get_attn_score_one_block( + batch_idx, + max_block_num, + query_cur, + key[batch_idx, offset_key : offset_key + self.block_size], + offset_key, + self.block_size, + mask_mode="triu", + use_softmax=False, + ) + + offset_key = max(0, offset_key + self.block_size) + width_key = ( + offset[batch_idx].item() + idx * self.block_size - offset_key + ) + if width_key > 0: + self.get_attn_score_one_block( + batch_idx, + max_block_num, + query_cur, + key[batch_idx, offset_key : offset_key + width_key], + offset_key, + width_key, + mask_mode=None, + use_softmax=False, + ) + + importance_cache = self.cache_importance.narrow( + 0, 0, max_block_num * batch_size + ).view(batch_size, max_block_num * self.block_size, self.q_head_num) + importance_cache_cpu = torch.empty_like( + importance_cache, device="cpu", pin_memory=True + ) + + importance_cache_cpu.copy_(importance_cache) + + block_table_cpu = self.prefix_block_table[:, :max_block_num].to("cpu") + offset_cpu = offset.contiguous().to("cpu") + + self.cpu_infer.submit( + self.local_thread.update_importance( + importance_cache_cpu, + layer_idx, + block_table_cpu, + max_block_num, + offset_cpu, + width, + ) + ) + self.cpu_infer.sync() + importance_cache.zero_() + + # key: [bsz, q_len, head_num, head_dim] float16 + # value: [bsz, q_len, head_num, head_dim] float16 + def swap_in_and_swap_out(self, layer_idx, past_len, q_len, key, value): + batch_size = 1 + max_seqs_len = past_len.max().item() + q_len + max_block_num = (max_seqs_len + self.block_size - 1) // self.block_size + k_cache = self.cache_key_states.narrow(0, 0, max_block_num * batch_size).view( + batch_size, max_block_num * self.block_size, self.kv_head_num, self.head_dim + ) + v_cache = self.cache_value_states.narrow(0, 0, max_block_num * batch_size).view( + batch_size, max_block_num * self.block_size, self.kv_head_num, self.head_dim + ) + + for batch_idx in range(batch_size): + offset = past_len[batch_idx] + width = q_len + k_cache[batch_idx][offset : offset + width].copy_( + key[batch_idx].view(-1, self.kv_head_num, self.head_dim) + ) + v_cache[batch_idx][offset : offset + width].copy_( + value[batch_idx].view(-1, self.kv_head_num, self.head_dim) + ) + + k_cache_cpu = torch.empty_like(k_cache, device="cpu", pin_memory=True) + v_cache_cpu = torch.empty_like(v_cache, device="cpu", pin_memory=True) + + k_cache_cpu.copy_(k_cache) + v_cache_cpu.copy_(v_cache) + + cur_block_num = ( + q_len + past_len[0].item() + self.block_size - 1 + ) // self.block_size + block_table_cpu = self.prefix_block_table[:, :cur_block_num].to("cpu") + past_len_cpu = past_len.contiguous().to("cpu") + + self.cpu_infer.submit( + self.local_thread.get_and_update_kvcache_fp16( + k_cache_cpu, + v_cache_cpu, + layer_idx, + block_table_cpu, + max_block_num, + past_len_cpu, + q_len, + ) + ) + + self.cpu_infer.sync() + k_cache.copy_(k_cache_cpu) + v_cache.copy_(v_cache_cpu) + + return k_cache, v_cache + + def calc_anchor(self, cache_seqlens: int): + cur_block_num = (cache_seqlens + self.block_size - 1) // self.block_size + block_table_cpu = self.prefix_block_table[:, :cur_block_num].to("cpu") + cache_seqlens_cpu = torch.tensor( + [cache_seqlens], device="cpu", dtype=torch.int32 + ) + + self.cpu_infer.submit( + self.local_thread.calc_anchor_all_layers( + block_table_cpu, + cache_seqlens_cpu, + ) + ) + self.cpu_infer.sync() + + def clear_importance(self, cache_seqlens: int): + print(f"clear importance: {cache_seqlens}") + cur_block_num = (cache_seqlens + self.block_size - 1) // self.block_size + block_table_cpu = self.prefix_block_table[:, :cur_block_num].to("cpu") + cache_seqlens_cpu = torch.tensor( + [cache_seqlens], device="cpu", dtype=torch.int32 + ) + + self.cpu_infer.submit( + self.local_thread.clear_importance_all_layers( + block_table_cpu, + cache_seqlens_cpu, + ) + ) + self.cpu_infer.sync() + + def clear_kvcache(self, cache_seqlens: int): + cur_block_num = (cache_seqlens + self.block_size - 1) // self.block_size + block_table_cpu = self.prefix_block_table[:, :cur_block_num].to("cpu") + cache_seqlens_cpu = torch.tensor( + [cache_seqlens], device="cpu", dtype=torch.int32 + ) + + self.cpu_infer.submit( + self.local_thread.clear_kvcache_all_layers( + block_table_cpu, + cache_seqlens_cpu, + ) + ) + self.cpu_infer.sync() + + def get_attn_sparsity( + self, + q_in: torch.Tensor, + layer_idx: int, + block_table: torch.Tensor, + cache_seqlens: torch.Tensor, + block_table_origin: torch.Tensor, + cache_seqlens_origin: torch.Tensor, + generate_token_idx: int = 0, + topk: int | None = None, + local: int | None = None, + output_path: str = "./attn_sparsity.json", + ): + self.attn_sparsity.zero_() + self.pcinfer.submit( + self.local_thread.get_attn_sparsity( + q_in, + self.attn_sparsity, + layer_idx, + block_table, + cache_seqlens, + block_table_origin, + cache_seqlens_origin, + generate_token_idx, + topk, + local, + ) + ) + self.cpu_infer.sync() + with open(output_path, "a") as file: + for head_idx in range(self.q_head_num): + sparsity = self.attn_sparsity[0][0][head_idx].item() + json_obj = { + "token_idx": generate_token_idx, + "layer_idx": layer_idx, + "head_idx": head_idx, + "sparsity": sparsity, + } + json.dump(json_obj, file) + file.write("\n") + + def apply( + self, + layer_idx: int, + bsz: int, + past_len: int, + query_states: torch.Tensor, + key_states: torch.Tensor, + value_states: torch.Tensor, + mode: str = "prefill", + generate_token_idx: int = -1, + ): + + # key_states: [bsz, q_len, kv_head_num, head_dim] + # value_states: [bsz, q_len, kv_head_num, head_dim] + # query_states: [bsz, q_len, q_head_num, head_dim] + assert query_states.dtype == torch.float16 + assert key_states.dtype == torch.float16 + assert value_states.dtype == torch.float16 + + assert key_states.size(2) == self.kv_head_num + assert value_states.size(2) == self.kv_head_num + assert query_states.size(2) == self.q_head_num + + q_len = query_states.size(1) + batch_size = query_states.size(0) + self.cache_seqlens_cuda.fill_(past_len) + last_chunk = False + if self.remaining_length <= self.prefill_chunk_size and q_len != 1: + last_chunk = True + device = query_states.device + if layer_idx == 0: + if q_len == 1: + self.generate_token_idx += 1 + elif last_chunk: + self.generate_token_idx = -1 + + if mode == "prefill": + key, value = self.swap_in_and_swap_out( + layer_idx, + self.cache_seqlens_cuda, + q_len, + key_states, + value_states, + ) + + if last_chunk and (self.anchor_type == "DYNAMIC" or self.preselect_block): + self.get_preselect_block_table_and_attn_score( + layer_idx, + bsz, + self.cache_seqlens_cuda, + q_len, + query_states, + key, + ) + output = flash_attn_with_kvcache( + q=query_states, + k_cache=key, + v_cache=value, + cache_seqlens=self.cache_seqlens_cuda + q_len, + causal=True, + ) + return output.transpose(1, 2) + + elif mode == "generate": + assert self.generate_token_idx >= 0 + self.q_in_cpu.copy_(query_states, non_blocking=True) + self.k_in_cpu.copy_(key_states, non_blocking=True) + self.v_in_cpu.copy_(value_states, non_blocking=True) + self.cache_seqlens_cpu.copy_(self.cache_seqlens_cuda, non_blocking=True) + # print(layer_idx) + if layer_idx < self.dense_layer_num: + self.block_table_cpu.copy_(self.prefix_block_table, non_blocking=True) + self.cpu_infer.submit_with_cuda_stream( + torch.cuda.current_stream("cuda").cuda_stream, + self.local_thread.attn_with_kvcache( + q_in=self.q_in_cpu, + k_in=self.k_in_cpu, + v_in=self.v_in_cpu, + output=self.output_cpu, + attn_lse=self.lse_cpu, + layer_idx=layer_idx, + block_table=self.block_table_cpu, + cache_seqlens=self.cache_seqlens_cpu, + ), + ) + else: + if self.preselect_block: + self.cache_seqlens_cpu.copy_( + self.cache_seqlens_cuda - self.evict_tokens, non_blocking=True + ) + if self.preselect_block_count < self.prefill_block_num: + self.block_table_cpu[:, : self.preselect_block_count].copy_( + self.preselect_block_table[layer_idx : layer_idx + 1], + non_blocking=True, + ) + + self.block_table_cpu[ + :, + self.preselect_block_count : self.preselect_block_count + + self.local_block_num, + ].copy_( + self.prefix_block_table[ + :, + self.prefill_block_num : self.prefill_block_num + + self.local_block_num, + ], + non_blocking=True, + ) + # print("submit_with_cuda_stream") + self.cpu_infer.submit_with_cuda_stream( + torch.cuda.current_stream("cuda").cuda_stream, + self.local_thread.attn_with_kvcache( + q_in=self.q_in_cpu, + k_in=self.k_in_cpu, + v_in=self.v_in_cpu, + output=self.output_cpu, + attn_lse=self.lse_cpu, + layer_idx=layer_idx, + generate_token_idx=self.generate_token_idx, + block_table=self.block_table_cpu, + cache_seqlens=self.cache_seqlens_cpu, + topk=( + self.topk + if self.topk <= self.preselect_block_count + else None + ), + local=self.local_windows_len // self.block_size, + ), + ) + # print("submit_with_cuda_stream enqueue\n") + else: + self.block_table_cpu.copy_( + self.prefix_block_table, non_blocking=True + ) + self.cpu_infer.submit_with_cuda_stream( + torch.cuda.current_stream("cuda").cuda_stream, + self.local_thread.attn_with_kvcache( + q_in=self.q_in_cpu, + k_in=self.k_in_cpu, + v_in=self.v_in_cpu, + output=self.output_cpu, + attn_lse=self.lse_cpu, + layer_idx=layer_idx, + generate_token_idx=self.generate_token_idx, + block_table=self.block_table_cpu, + cache_seqlens=self.cache_seqlens_cpu, + topk=self.topk, + local=self.local_windows_len // self.block_size, + ), + ) + self.cpu_infer.sync_with_cuda_stream( + torch.cuda.current_stream("cuda").cuda_stream + ) + # print("submit_with_cuda_stream finished\n") + self.output_cuda.copy_(self.output_cpu, non_blocking=True) + return self.output_cuda.transpose(1, 2) + + def save(self, path: str, length: int): + cur_block_num = (length + self.block_size - 1) // self.block_size + block_table_cpu = self.prefix_block_table[0, :cur_block_num].to("cpu") + cache_seqlens_cpu = torch.tensor([length], device="cpu", dtype=torch.int32) + self.cpu_infer.submit( + self.local_thread.dump_kvcache( + block_table_cpu, + cache_seqlens_cpu, + path, + ) + ) + self.cpu_infer.sync() + + def load(self, path: str, length: int): + self.cpu_infer.submit( + self.local_thread.load_kvcache( + path, + ) + ) + self.cpu_infer.sync() diff --git a/archive/kt-sft/ktransformers/operators/experts.py b/archive/kt-sft/ktransformers/operators/experts.py new file mode 100644 index 0000000..0e80bf1 --- /dev/null +++ b/archive/kt-sft/ktransformers/operators/experts.py @@ -0,0 +1,2310 @@ +#!/usr/bin/env python +# coding=utf-8 +''' +Description : +Author : Azure-Tang, Boxin Zhang, chenht2022 +Date : 2024-07-25 11:25:24 +Version : 0.1.0 +LastEditors : Azure +LastEditTime : 2024-08-29 09:41:10 +Copyright (c) 2024 by KVCache.AI, All Rights Reserved. +''' + +from typing import Any, Union +import numpy as np +import numpy.typing as npt +from torch import Tensor, nn +import torch.nn.functional as F +import torch +import sys, os +from ktransformers.operators.base_operator import BaseInjectedModule +from tqdm import tqdm +import time +import logging +from tqdm.auto import tqdm +import re + +sys.path.append(os.path.join(os.path.dirname(__file__), "..", "ktransformers_ext", "build")) +sys.path.append(os.path.join(os.path.dirname(__file__), "..", "ktransformers_ext", "build", "Release")) +sys.path.append(os.path.join(os.path.dirname(__file__), "..", "ktransformers_ext", "build", "Debug")) +import cpuinfer_ext +from cpuinfer_ext.moe import MOEConfig, MOE +from cpuinfer_ext.sft_moe import SFT_MOEConfig, SFT_MOE +import ctypes +from ktransformers.util.custom_loader import GGUFLoader +from ktransformers.util.inference_state import InferenceState +from ktransformers.util.custom_gguf import GGMLQuantizationType +from ktransformers.util.custom_loader import GGUFLoader, SafeTensorLoader, ModelLoader +from ktransformers.server.config.config import Config +from transformers.activations import ACT2FN +from transformers.configuration_utils import PretrainedConfig +from abc import ABC, abstractmethod +from ktransformers.operators.linear import KLinearMarlin, KLinearTorch, KTransformersLinear +import time +from ktransformers.operators.cpuinfer import CPUInfer +from ktransformers.util.grad_wrapper import maybe_no_grad + +H_FIXED = 7168 +M_FIXED = 2048 + +def deduplicate_and_sort(lst): + return sorted(set(lst)) +def generate_cuda_graphs(chunk_size: int) -> list: + assert chunk_size <= 1024 or chunk_size % 1024 == 0, "chunk_size must <= 1024 or a multiple of 1024" + base_list = [1, 2, 3, Config().max_batch_size, 64, 256, 512, chunk_size] + + if chunk_size <= 1024: + return deduplicate_and_sort(base_list) + + multiples = [i for i in range(1024, chunk_size + 1, 1024)] + + return deduplicate_and_sort(base_list + multiples) +#cuda_graphs = [Config().chunk_size] +if torch.cuda.is_available(): + cuda_graphs = generate_cuda_graphs(Config().chunk_size) +else: + cuda_graphs = 1 +# class Base(BaseInjectedModule, ABC): +class KExpertsBase(ABC): + def __init__(self, key: str, gguf_loader: GGUFLoader, config: PretrainedConfig, orig_module: nn.Module, device: str = "cuda", **kwargs): + # super().__init__(key, gguf_loader, config, orig_module, device, **kwargs) + self.key = key + self.gguf_loader = gguf_loader + self.config = config + self.device = device + + @abstractmethod + def forward(self, input_tensor, expert_ids, weights): + pass + + @abstractmethod + def load(self, w: dict | nn.Parameter | tuple | None = None, device: str = "cpu", warmup: bool = False): + pass + + @abstractmethod + def unload(): + pass + + def load_weights(self, override_key: str | None = None, device: str = "cpu"): + res = {} + if override_key is not None: + keys = override_key + else: + keys = [self.key] + + gate = None + up = None + down = None + gate_type = None + up_type = None + down_type = None + + for key in keys: + # if key + ".ffn_gate_exps.weight" in self.gguf_loader.tensor_info: # TODO: maybe problem in merge (this is origin one) + if self.gguf_loader.has_tensor(key + ".ffn_gate_exps.weight"): + targets = [".ffn_gate_exps.weight", ".ffn_up_exps.weight", ".ffn_down_exps.weight" ] + tensors = self.load_multi(key, targets, device=device) + gate = tensors[".ffn_gate_exps.weight"] + up = tensors[".ffn_up_exps.weight"] + down = tensors[".ffn_down_exps.weight"] + gate_type = self.gguf_loader.tensor_info[key + ".ffn_gate_exps.weight"]["ggml_type"] + up_type = self.gguf_loader.tensor_info[key + ".ffn_up_exps.weight"]["ggml_type"] + down_type = self.gguf_loader.tensor_info[key + ".ffn_down_exps.weight"]["ggml_type"] + # elif key + ".ffn_down.0.weight" in self.gguf_loader.tensor_info: # TODO: maybe problem in merge (this is origin one) + elif self.gguf_loader.has_tensor(key + ".ffn_down.0.weight"): + # for supporting Mixtral-8x7B-Instuct + gate = [] + up = [] + down = [] + for i in range(8): + gatei, upi, downi = f".ffn_gate.{i}.weight", f".ffn_up.{i}.weight", f".ffn_down.{i}.weight" + targets = [gatei, upi, downi] + tensors = self.load_multi(key, targets, device=device) + gate_it, up_it, down_it = tensors[gatei], tensors[upi], tensors[downi] + gate.append(gate_it) + up.append(up_it) + down.append(down_it) + gate = torch.stack(gate) + up = torch.stack(up) + down = torch.stack(down) + gate_type = self.gguf_loader.tensor_info[key + ".ffn_gate.0.weight"]["ggml_type"] + up_type = self.gguf_loader.tensor_info[key + ".ffn_up.0.weight"]["ggml_type"] + down_type = self.gguf_loader.tensor_info[key + ".ffn_down.0.weight"]["ggml_type"] + else: + raise ValueError(f"Experts {key} not found in gguf_loader") + res = {key:{"gate": gate, "up": up, "down": down, "gate_type": gate_type, "up_type": up_type, "down_type": down_type}} + return res + + def load_multi(self, key: str, keys: list[str], device: str = "cpu"): + tensors = {} + for k in keys: + tensors[k] = self.gguf_loader.load_gguf_tensor(key + k, device=device) + return tensors +class KExpertsCPU(KExpertsBase): + input_tensor_cpu:Tensor = None + expert_ids_cpu:Tensor = None + weights_cpu:Tensor = None + output_cpu:Tensor = None + output_gpu_map:dict = {} # Manage output tensor buffer on different gpu + #stream_map:dict = {} # Manage cuda stream on different gpu + # @TODO add yaml + CPU_INFER = CPUInfer(Config().cpu_infer) + def __init__( + self, + key: str, + gguf_loader: GGUFLoader, + config: PretrainedConfig, + n_routed_experts: int, + orig_module: nn.Module = None, + device: str = "cpu", + out_device: str = "cuda", # this device mean which device the output should on. TODO: support cpu. + **kwargs + ): + super().__init__(key, gguf_loader, config, orig_module, device, **kwargs) + assert device.lower() == "cpu", "KExpertsCPU can only be loaded on CPU" + self.n_routed_experts = n_routed_experts + self.out_device = out_device + self.backend = kwargs.get("backend", "llamafile") + + def load(self, w: dict | nn.Parameter | tuple | None = None, device:str|None = None, warmup:bool = False): + if device: + assert device.lower() == "cpu", "KExpertsCPU can only be loaded on CPU, Parameter \"device\" can be cpu or None." + if w is None: w = self.load_weights()[self.key] + self.gate = w["gate"] + self.up = w["up"] + self.down = w["down"] + self.gate_type = w["gate_type"] + self.up_type = w["up_type"] + self.down_type = w["down_type"] + gate_ptr = ctypes.addressof( + ctypes.cast(self.gate.ctypes.data, ctypes.POINTER(ctypes.c_uint64)).contents + ) + up_ptr = ctypes.addressof( + ctypes.cast(self.up.ctypes.data, ctypes.POINTER(ctypes.c_uint64)).contents + ) + down_ptr = ctypes.addressof( + ctypes.cast(self.down.ctypes.data, ctypes.POINTER(ctypes.c_uint64)).contents + ) + # print(self.gate_qtype, self.up_qtype, self.down_qtype) + n_routed_experts = self.n_routed_experts + self.cpu_infer = KExpertsCPU.CPU_INFER + # n_routed_experts = len(self.orig_module) + model_dtype = torch.get_default_dtype() + if torch.xpu.is_available() and model_dtype == torch.float16: + hidden_type = 1 # fp16 + else: + hidden_type = 30 # bf16 + if self.backend == "llamafile": + moe_config = MOEConfig( + n_routed_experts, + self.config.num_experts_per_tok, + self.config.hidden_size, + self.config.moe_intermediate_size, + 64, + 10, + 1024, + gate_ptr, + up_ptr, + down_ptr, + self.gate_type, + self.up_type, + self.down_type, + hidden_type, # TODO: get from model.dtype + ) + self.moe = MOE(moe_config) + elif self.backend == "AMXBF16": + from cpuinfer_ext.moe import AMX_MOEConfig, AMXBF16_MOE + assert self.gate_type == GGMLQuantizationType.BF16 + assert self.up_type == GGMLQuantizationType.BF16 + assert self.down_type == GGMLQuantizationType.BF16 + moe_config = AMX_MOEConfig( + n_routed_experts, + self.config.num_experts_per_tok, + self.config.hidden_size, + self.config.moe_intermediate_size, + max(cuda_graphs) if isinstance(cuda_graphs, list) else Config().chunk_size, + gate_ptr, + up_ptr, + down_ptr, + ) + self.moe = AMXBF16_MOE(moe_config) + self.cpu_infer.submit(self.moe.load_weights()) + self.cpu_infer.sync() + elif self.backend == "AMXInt8": + from cpuinfer_ext.moe import AMX_MOEConfig, AMXInt8_MOE + assert self.gate_type == GGMLQuantizationType.BF16 + assert self.up_type == GGMLQuantizationType.BF16 + assert self.down_type == GGMLQuantizationType.BF16 + moe_config = AMX_MOEConfig( + n_routed_experts, + self.config.num_experts_per_tok, + self.config.hidden_size, + self.config.moe_intermediate_size, + max(cuda_graphs) if isinstance(cuda_graphs, list) else Config().chunk_size, + gate_ptr, + up_ptr, + down_ptr, + ) + self.moe = AMXInt8_MOE(moe_config) + self.cpu_infer.submit(self.moe.load_weights()) + self.cpu_infer.sync() + # print(n_routed_experts, hidden_size, moe_intermediate_size) + num_experts_per_tok = self.config.num_experts_per_tok + if warmup: + self.cpu_infer.submit(self.moe.warm_up()) + self.cpu_infer.sync() + if self.out_device not in KExpertsCPU.output_gpu_map: + if isinstance(cuda_graphs, list): + KExpertsCPU.output_gpu_map[self.out_device] = [torch.zeros((cuda_graphs[i], self.config.hidden_size), device=self.out_device) for i in range(len(cuda_graphs))] + else: + KExpertsCPU.output_gpu_map[self.out_device] = torch.zeros((cuda_graphs, self.config.hidden_size), device=self.out_device) + if KExpertsCPU.input_tensor_cpu == None: + if isinstance(cuda_graphs, list): + KExpertsCPU.input_tensor_cpu = [torch.zeros((cuda_graphs[i], self.config.hidden_size), device="cpu", pin_memory=True) for i in range(len(cuda_graphs))] + KExpertsCPU.expert_ids_cpu = [torch.zeros((cuda_graphs[i], num_experts_per_tok), device="cpu", dtype=torch.long, pin_memory=True) for i in range(len(cuda_graphs))] + KExpertsCPU.weights_cpu = [torch.zeros((cuda_graphs[i], num_experts_per_tok), device="cpu", dtype=torch.float32, pin_memory=True) for i in range(len(cuda_graphs))] + KExpertsCPU.output_cpu = [torch.zeros((cuda_graphs[i], self.config.hidden_size), device="cpu", pin_memory=True, dtype=torch.bfloat16) for i in range(len(cuda_graphs))] + KExpertsCPU.bsz_tensor_cpu = [torch.zeros((1), device="cpu", dtype=torch.int32, pin_memory=True) for i in range(len(cuda_graphs))] + else: + KExpertsCPU.input_tensor_cpu = torch.zeros((cuda_graphs, self.config.hidden_size), device="cpu", pin_memory=True) + KExpertsCPU.expert_ids_cpu = torch.zeros((cuda_graphs, num_experts_per_tok), device="cpu", dtype=torch.long, pin_memory=True) + KExpertsCPU.weights_cpu = torch.zeros((cuda_graphs, num_experts_per_tok), device="cpu", dtype=torch.float32, pin_memory=True) + if torch.xpu.is_available(): + KExpertsCPU.output_cpu = torch.zeros((cuda_graphs, self.config.hidden_size), device="cpu", pin_memory=True, dtype=model_dtype) + KExpertsCPU.bsz_tensor_cpu = torch.ones((1), device="cpu", dtype=torch.int32, pin_memory=True) + else: + KExpertsCPU.output_cpu = torch.zeros((cuda_graphs, self.config.hidden_size), device="cpu", pin_memory=True, dtype=torch.bfloat16) + KExpertsCPU.bsz_tensor_cpu = torch.zeros((1), device="cpu", dtype=torch.int32, pin_memory=True) + + def submit_for_one_decode(self, input_tensor, expert_ids, weights, bsz_tensor=None, cuda_graph_idx=0): + if bsz_tensor is None: + bsz_tensor = torch.ones(1, device=input_tensor.device, dtype=torch.int32) + if cuda_graph_idx != -1: + KExpertsCPU.input_tensor_cpu[cuda_graph_idx].copy_(input_tensor, non_blocking=True) + KExpertsCPU.expert_ids_cpu[cuda_graph_idx].copy_(expert_ids, non_blocking=True) + KExpertsCPU.weights_cpu[cuda_graph_idx].copy_(weights, non_blocking=True) + KExpertsCPU.bsz_tensor_cpu[cuda_graph_idx].copy_(bsz_tensor, non_blocking=True) + self.cpu_infer.submit_with_cuda_stream(torch.cuda.current_stream(self.out_device).cuda_stream, self.moe.forward(1, expert_ids.size(-1), KExpertsCPU.expert_ids_cpu[cuda_graph_idx].data_ptr(), KExpertsCPU.weights_cpu[cuda_graph_idx].data_ptr(), KExpertsCPU.input_tensor_cpu[cuda_graph_idx].data_ptr(), KExpertsCPU.output_cpu[cuda_graph_idx].data_ptr(), KExpertsCPU.bsz_tensor_cpu[cuda_graph_idx].data_ptr())) + else: + KExpertsCPU.input_tensor_cpu.copy_(input_tensor, non_blocking=True) + KExpertsCPU.expert_ids_cpu.copy_(expert_ids, non_blocking=True) + KExpertsCPU.weights_cpu.copy_(weights, non_blocking=True) + KExpertsCPU.bsz_tensor_cpu.copy_(bsz_tensor, non_blocking=True) + self.cpu_infer.submit_with_cuda_stream(torch.cuda.current_stream(self.out_device).cuda_stream, self.moe.forward(1, expert_ids.size(-1), KExpertsCPU.expert_ids_cpu.data_ptr(), KExpertsCPU.weights_cpu.data_ptr(), KExpertsCPU.input_tensor_cpu.data_ptr(), KExpertsCPU.output_cpu.data_ptr(), KExpertsCPU.bsz_tensor_cpu.data_ptr())) + + + def sync_for_one_decode(self, cuda_graph_idx=0): + if cuda_graph_idx != -1: + self.cpu_infer.sync_with_cuda_stream(torch.cuda.current_stream(self.out_device).cuda_stream) + KExpertsCPU.output_gpu_map[self.out_device][cuda_graph_idx].copy_(KExpertsCPU.output_cpu[cuda_graph_idx], non_blocking=True) + return KExpertsCPU.output_gpu_map[self.out_device][cuda_graph_idx] + else: + self.cpu_infer.sync_with_cuda_stream(torch.cuda.current_stream(self.out_device).cuda_stream) + KExpertsCPU.output_gpu_map[self.out_device].copy_(KExpertsCPU.output_cpu, non_blocking=True) + return KExpertsCPU.output_gpu_map[self.out_device] + + def forward(self, input_tensor, expert_ids, weights, bsz_tensor=None, cuda_graph_idx=0): + # generate, capture and run cuda graph + # print(expert_ids) + if bsz_tensor is None and (not torch.xpu.is_available() or input_tensor.size(0) > 1): + bsz_tensor = torch.tensor([input_tensor.size(0)], device=input_tensor.device, dtype=torch.int32) + if torch.cuda.is_available() and torch.cuda.is_current_stream_capturing(): + if cuda_graph_idx != -1: + KExpertsCPU.input_tensor_cpu[cuda_graph_idx].copy_(input_tensor, non_blocking=True) + KExpertsCPU.expert_ids_cpu[cuda_graph_idx].copy_(expert_ids, non_blocking=True) + KExpertsCPU.weights_cpu[cuda_graph_idx].copy_(weights, non_blocking=True) + KExpertsCPU.bsz_tensor_cpu[cuda_graph_idx].copy_(bsz_tensor, non_blocking=True) + self.cpu_infer.submit_with_cuda_stream(torch.cuda.current_stream().cuda_stream, self.moe.forward(expert_ids.size(0), expert_ids.size(-1), KExpertsCPU.expert_ids_cpu[cuda_graph_idx].data_ptr(), KExpertsCPU.weights_cpu[cuda_graph_idx].data_ptr(), KExpertsCPU.input_tensor_cpu[cuda_graph_idx].data_ptr(), KExpertsCPU.output_cpu[cuda_graph_idx].data_ptr(), KExpertsCPU.bsz_tensor_cpu[cuda_graph_idx].data_ptr())) + self.cpu_infer.sync_with_cuda_stream(torch.cuda.current_stream().cuda_stream) + KExpertsCPU.output_gpu_map[self.out_device][cuda_graph_idx].copy_(KExpertsCPU.output_cpu[cuda_graph_idx], non_blocking=True) + return KExpertsCPU.output_gpu_map[self.out_device][cuda_graph_idx] + + else: + KExpertsCPU.input_tensor_cpu.copy_(input_tensor, non_blocking=True) + KExpertsCPU.expert_ids_cpu.copy_(expert_ids, non_blocking=True) + KExpertsCPU.weights_cpu.copy_(weights, non_blocking=True) + KExpertsCPU.bsz_tensor_cpu.copy_(bsz_tensor, non_blocking=True) + self.cpu_infer.submit_with_cuda_stream(torch.cuda.current_stream().cuda_stream, self.moe.forward(expert_ids.size(0), expert_ids.size(-1), KExpertsCPU.expert_ids_cpu.data_ptr(), KExpertsCPU.weights_cpu.data_ptr(), KExpertsCPU.input_tensor_cpu.data_ptr(), KExpertsCPU.output_cpu.data_ptr(), KExpertsCPU.bsz_tensor_cpu.data_ptr())) + self.cpu_infer.sync_with_cuda_stream(torch.cuda.current_stream().cuda_stream) + KExpertsCPU.output_gpu_map[self.out_device].copy_(KExpertsCPU.output_cpu, non_blocking=True) + return KExpertsCPU.output_gpu_map[self.out_device] + elif input_tensor.size(0)==1 and torch.xpu.is_available(): + KExpertsCPU.input_tensor_cpu.copy_(input_tensor.view(-1), non_blocking=True) + KExpertsCPU.expert_ids_cpu.copy_(expert_ids.view(-1), non_blocking=True) + KExpertsCPU.weights_cpu.copy_(weights.view(-1), non_blocking=True) + # KExpertsCPU.bsz_tensor_cpu.copy_(bsz_tensor.view(-1), non_blocking=True) + self.cpu_infer.submit(self.moe.forward(expert_ids.size(0), expert_ids.size(1), KExpertsCPU.expert_ids_cpu.data_ptr(), KExpertsCPU.weights_cpu.data_ptr(), KExpertsCPU.input_tensor_cpu.data_ptr(), KExpertsCPU.output_cpu.data_ptr(), KExpertsCPU.bsz_tensor_cpu.data_ptr())) + self.cpu_infer.sync() + KExpertsCPU.output_gpu_map[self.out_device].copy_(KExpertsCPU.output_cpu, non_blocking=True) + return KExpertsCPU.output_gpu_map[self.out_device].view(1, -1) + else: + input_tensor = input_tensor.contiguous().cpu() + expert_ids = expert_ids.contiguous().cpu() + weights = weights.contiguous().to(torch.float32).cpu() + bsz_tensor = bsz_tensor.contiguous().cpu() + output = torch.empty_like(input_tensor).contiguous() + self.cpu_infer.submit(self.moe.forward(expert_ids.size(0), expert_ids.size(1), expert_ids.data_ptr(), weights.data_ptr(), input_tensor.data_ptr(), output.data_ptr(), bsz_tensor.data_ptr())) + self.cpu_infer.sync() + return output.to(device=object.__getattribute__(self, "out_device")) + + def unload(self): + return + + def load_weights(self, override_key: str | None = None, device: str = "cpu"): + # TODO: support Bias + res = {} + if override_key is not None: + keys = override_key + else: + keys = [self.key] + + gate = None + up = None + down = None + gate_type = None + up_type = None + down_type = None + + for key in keys: + if isinstance(self.gguf_loader, SafeTensorLoader): + res = self.gguf_loader.load_experts(key) + return {key: res} + elif self.gguf_loader.has_tensor(key + ".ffn_gate_exps.weight"): + gate = self.gguf_loader.get_mmap_tensor(key + ".ffn_gate_exps.weight") + up = self.gguf_loader.get_mmap_tensor(key + ".ffn_up_exps.weight") + down = self.gguf_loader.get_mmap_tensor(key + ".ffn_down_exps.weight") + # gate_type = self.gguf_loader.tensor_info[key + ".ffn_gate_exps.weight"]["ggml_type"] + # up_type = self.gguf_loader.tensor_info[key + ".ffn_up_exps.weight"]["ggml_type"] + # down_type = self.gguf_loader.tensor_info[key + ".ffn_down_exps.weight"]["ggml_type"] + gate_type = self.gguf_loader.get_ggml_type(key + ".ffn_gate_exps.weight") + up_type = self.gguf_loader.get_ggml_type(key + ".ffn_up_exps.weight") + down_type = self.gguf_loader.get_ggml_type(key + ".ffn_down_exps.weight") + + elif key + ".ffn_gate_exps.weight" in self.gguf_loader.tensor_info: + gate = self.gguf_loader.get_mmap_tensor(key + ".ffn_gate_exps.weight") + up = self.gguf_loader.get_mmap_tensor(key + ".ffn_up_exps.weight") + down = self.gguf_loader.get_mmap_tensor(key + ".ffn_down_exps.weight") + gate_type = self.gguf_loader.tensor_info[key + ".ffn_gate_exps.weight"]["ggml_type"] + up_type = self.gguf_loader.tensor_info[key + ".ffn_up_exps.weight"]["ggml_type"] + down_type = self.gguf_loader.tensor_info[key + ".ffn_down_exps.weight"]["ggml_type"] + elif key + ".ffn_down.0.weight" in self.gguf_loader.tensor_info: + # for supporting Mixtral-8x7B-Instuct + gate = [] + up = [] + down = [] + for i in range(8): + gate_it = self.gguf_loader.get_mmap_tensor(f"{key}.ffn_gate.{i}.weight") + up_it = self.gguf_loader.get_mmap_tensor(f"{key}.ffn_up.{i}.weight") + down_it = self.gguf_loader.get_mmap_tensor(f"{key}.ffn_down.{i}.weight") + gate.append(gate_it) + up.append(up_it) + down.append(down_it) + gate = np.stack(gate) + up = np.stack(up) + down = np.stack(down) + gate_type = self.gguf_loader.get_ggml_type(key + ".ffn_gate.0.weight") + up_type = self.gguf_loader.get_ggml_type(key + ".ffn_up.0.weight") + down_type = self.gguf_loader.get_ggml_type(key + ".ffn_down.0.weight") + else: + raise ValueError(f"Experts {key} not found in gguf_loader") + res = {key:{"gate": gate, "up": up, "down": down, "gate_type": gate_type, "up_type": up_type, "down_type": down_type}} + return res +class KSFTExpertsCPU(torch.autograd.Function): + input_tensor_cpu:Tensor = None + expert_ids_cpu:Tensor = None + weights_cpu:Tensor = None + output_cpu:Tensor = None + output_gpu_map:dict = {} # Manage output tensor buffer on different gpu + #stream_map:dict = {} # Manage cuda stream on different gpu + #gguf_loader:GGUFLoader = None + CPU_INFER = CPUInfer(Config().cpu_infer) + + # Pinned memory buffers for training (batch mode) + # These are used for efficient CPU-GPU data transfer + _pinned_input_buf: Tensor = None # [max_tokens, hidden_size] + _pinned_output_buf: Tensor = None # [max_tokens, hidden_size] + _pinned_expert_ids_buf: Tensor = None # [max_tokens, num_experts_per_tok] + _pinned_weights_buf: Tensor = None # [max_tokens, num_experts_per_tok] + _pinned_grad_out_buf: Tensor = None # [max_tokens, hidden_size] for backward + _pinned_grad_in_buf: Tensor = None # [max_tokens, hidden_size] for backward + _pinned_buf_size: int = 0 # current buffer capacity (max_tokens) + _hidden_size: int = 0 + _num_experts_per_tok: int = 0 + def __init__( + self, + key: str, + gguf_loader: GGUFLoader, + config: PretrainedConfig, + n_routed_experts: int, + orig_module: nn.Module = None, + device: str = "cpu", + out_device: str = "cuda", # this device mean which device the output should on. TODO: support cpu. + **kwargs + ): + super().__init__(key, gguf_loader, config, orig_module, device, **kwargs) + #if KExpertsCPU.gguf_loader is None: + # KExpertsCPU.gguf_loader = GGUFLoader("/mnt/data/model/DeepseekV3-q4km-gguf") + self.gguf_loader = gguf_loader + assert device.lower() == "cpu", "KExpertsCPU can only be loaded on CPU" + self.n_routed_experts = n_routed_experts + self.out_device = out_device + self.backend = kwargs.get("backend", "llamafile") + + self.key = key + self.config = config + self.device = device + + self.call_count = 0 + self.flops_per_call = [] + self.times = [] + + self.tflops_fwd = [] + self.tflops_bwd = [] + + @classmethod + def _ensure_pinned_buffers(cls, num_tokens: int, hidden_size: int, num_experts_per_tok: int): + """ + Ensure pinned memory buffers are allocated with sufficient size. + Buffers are reused across calls and only reallocated if more space is needed. + """ + # Check if we need to allocate or expand buffers + if (cls._pinned_input_buf is None or + num_tokens > cls._pinned_buf_size or + hidden_size != cls._hidden_size or + num_experts_per_tok != cls._num_experts_per_tok): + + # Allocate with some extra capacity to reduce reallocations + new_size = max(num_tokens, cls._pinned_buf_size * 2) if cls._pinned_buf_size > 0 else num_tokens + new_size = max(new_size, 1024) # minimum 1024 tokens + + # Free old buffers + cls._pinned_input_buf = None + cls._pinned_output_buf = None + cls._pinned_expert_ids_buf = None + cls._pinned_weights_buf = None + cls._pinned_grad_out_buf = None + cls._pinned_grad_in_buf = None + + # Allocate new pinned buffers + cls._pinned_input_buf = torch.empty( + (new_size, hidden_size), dtype=torch.bfloat16, device="cpu", pin_memory=True + ) + cls._pinned_output_buf = torch.empty( + (new_size, hidden_size), dtype=torch.bfloat16, device="cpu", pin_memory=True + ) + cls._pinned_expert_ids_buf = torch.empty( + (new_size, num_experts_per_tok), dtype=torch.long, device="cpu", pin_memory=True + ) + cls._pinned_weights_buf = torch.empty( + (new_size, num_experts_per_tok), dtype=torch.float32, device="cpu", pin_memory=True + ) + cls._pinned_grad_out_buf = torch.empty( + (new_size, hidden_size), dtype=torch.bfloat16, device="cpu", pin_memory=True + ) + cls._pinned_grad_in_buf = torch.empty( + (new_size, hidden_size), dtype=torch.bfloat16, device="cpu", pin_memory=True + ) + + cls._pinned_buf_size = new_size + cls._hidden_size = hidden_size + cls._num_experts_per_tok = num_experts_per_tok + + print(f"[KSFTExpertsCPU] Allocated pinned memory buffers: " + f"size={new_size}, hidden={hidden_size}, k={num_experts_per_tok}") + + def load(self, w: dict | nn.Parameter | tuple | None = None, device:str|None = None, warmup:bool = False): + if device: + assert device.lower() == "cpu", "KSFTExpertsCPU can only be loaded on CPU, Parameter \"device\" can be cpu or None." + if w is None: w = self.load_weights()[self.key] + self.gate = w["gate"] + self.up = w["up"] + self.down = w["down"] + self.gate_type = w["gate_type"] + self.up_type = w["up_type"] + self.down_type = w["down_type"] + gate_ptr = ctypes.addressof( + ctypes.cast(self.gate.ctypes.data, ctypes.POINTER(ctypes.c_uint64)).contents + ) + up_ptr = ctypes.addressof( + ctypes.cast(self.up.ctypes.data, ctypes.POINTER(ctypes.c_uint64)).contents + ) + down_ptr = ctypes.addressof( + ctypes.cast(self.down.ctypes.data, ctypes.POINTER(ctypes.c_uint64)).contents + ) + #print(self.gate_type, self.up_type, self.down_type) + n_routed_experts = self.n_routed_experts + # n_routed_experts = len(self.orig_module) + self.cpu_infer = KSFTExpertsCPU.CPU_INFER + + model_dtype = torch.get_default_dtype() + if torch.xpu.is_available() and model_dtype == torch.float16: + hidden_type = 1 # fp16 + else: + hidden_type = 30 # bf16 + if self.backend == "llamafile": + # print("GO INTO LLAMAFILE!!") + moe_config = SFT_MOEConfig( + n_routed_experts, + self.config.num_experts_per_tok, + self.config.hidden_size, + self.config.moe_intermediate_size, + 64, + 10, + 1024, + gate_ptr, + up_ptr, + down_ptr, + self.gate_type, + self.up_type, + self.down_type, + hidden_type, # TODO: get from model.dtype + ) + self.moe = SFT_MOE(moe_config) + elif self.backend == "AMXBF16": + print("GO INTO AMXBF16!!") + from cpuinfer_ext.sft_moe import SFT_AMX_MOEConfig, SFT_AMXBF16_MOE + assert self.gate_type == GGMLQuantizationType.BF16 + assert self.up_type == GGMLQuantizationType.BF16 + assert self.down_type == GGMLQuantizationType.BF16 + moe_config = SFT_AMX_MOEConfig( + n_routed_experts, + self.config.num_experts_per_tok, + self.config.hidden_size, + self.config.moe_intermediate_size, + max(cuda_graphs) if isinstance(cuda_graphs, list) else Config().chunk_size, + gate_ptr, + up_ptr, + down_ptr, + ) + self.moe = SFT_AMXBF16_MOE(moe_config) + self.cpu_infer.submit(self.moe.load_weights()) + self.cpu_infer.sync() + elif self.backend == "AMXInt8": + print("GO INTO AMXInt8!!") + from cpuinfer_ext.sft_moe import SFT_AMX_MOEConfig, SFT_AMXInt8_MOE + assert self.gate_type == GGMLQuantizationType.BF16 + assert self.up_type == GGMLQuantizationType.BF16 + assert self.down_type == GGMLQuantizationType.BF16 + moe_config = SFT_AMX_MOEConfig( + n_routed_experts, + self.config.num_experts_per_tok, + self.config.hidden_size, + self.config.moe_intermediate_size, + max(cuda_graphs) if isinstance(cuda_graphs, list) else Config().chunk_size, + gate_ptr, + up_ptr, + down_ptr, + ) + self.moe = SFT_AMXInt8_MOE(moe_config) + self.cpu_infer.submit(self.moe.load_weights()) + self.cpu_infer.sync() + + # print(n_routed_experts, hidden_size, moe_intermediate_size) + num_experts_per_tok = self.config.num_experts_per_tok + if warmup: + self.cpu_infer.submit(self.moe.warm_up()) + self.cpu_infer.sync() + if self.out_device not in KSFTExpertsCPU.output_gpu_map: + KSFTExpertsCPU.output_gpu_map[self.out_device] = torch.zeros((self.config.hidden_size), device=self.out_device) + if KSFTExpertsCPU.input_tensor_cpu == None: + KSFTExpertsCPU.input_tensor_cpu = torch.zeros((self.config.hidden_size), device="cpu", pin_memory=True) + KSFTExpertsCPU.expert_ids_cpu = torch.zeros((num_experts_per_tok), device="cpu", dtype=torch.long, pin_memory=True) + KSFTExpertsCPU.weights_cpu = torch.zeros((num_experts_per_tok), device="cpu", dtype=torch.float32, pin_memory=True) + KSFTExpertsCPU.output_cpu = torch.zeros((self.config.hidden_size), device="cpu", pin_memory=True, dtype=torch.bfloat16) + + # Initialize pinned memory buffers for training (batch mode) + # Default size is 4096 tokens, will expand automatically if needed + default_max_tokens = 4096 + KSFTExpertsCPU._ensure_pinned_buffers( + default_max_tokens, + self.config.hidden_size, + num_experts_per_tok + ) + + self.gate = None + self.up = None + self.down = None + + def submit_for_one_decode(self, input_tensor, expert_ids, weights): + KSFTExpertsCPU.input_tensor_cpu.copy_(input_tensor, non_blocking=True) + KSFTExpertsCPU.expert_ids_cpu.copy_(expert_ids, non_blocking=True) + KSFTExpertsCPU.weights_cpu.copy_(weights, non_blocking=True) + self.cpu_infer.submit_with_cuda_stream(torch.cuda.current_stream(self.out_device).cuda_stream, self.moe.forward(1, expert_ids.size(0), KSFTExpertsCPU.expert_ids_cpu.data_ptr(), KSFTExpertsCPU.weights_cpu.data_ptr(), KSFTExpertsCPU.input_tensor_cpu.data_ptr(), KSFTExpertsCPU.output_cpu.data_ptr())) + + def sync_for_one_decode(self): + self.cpu_infer.sync_with_cuda_stream(torch.cuda.current_stream(self.out_device).cuda_stream) + KSFTExpertsCPU.output_gpu_map[self.out_device].copy_(KSFTExpertsCPU.output_cpu, non_blocking=True) + return KSFTExpertsCPU.output_gpu_map[self.out_device] + + @staticmethod + def forward(ctx, input_tensor, expert_ids, weights, cpu_infer, moe, out_device, layer_idx): + # print("Go into the forward") + + # generate, capture and run cuda graph + # torch.set_printoptions(threshold=float('inf')) + # print(expert_ids) + # expert_ids.cpu().numpy().tofile('debug_expert_ids.txt', sep='\n') + # print(expert_ids.size()) + # print(xx) + if input_tensor.size(0)==1 and torch.cuda.is_current_stream_capturing(): + # TODO: this branch is unreachable, but the shape of input_tensor([1,hidden_size]) and input_tensor_cpu([hidden_size]) is not compatible + #print("capturing experts") + wall_t0 = time.time() + KSFTExpertsCPU.input_tensor_cpu.copy_(input_tensor, non_blocking=True) + KSFTExpertsCPU.expert_ids_cpu.copy_(expert_ids, non_blocking=True) + KSFTExpertsCPU.weights_cpu.copy_(weights, non_blocking=True) + cpu_infer.submit_with_cuda_stream(torch.cuda.current_stream().cuda_stream, moe.forward(1, expert_ids.size(1), KSFTExpertsCPU.expert_ids_cpu.data_ptr(), KSFTExpertsCPU.weights_cpu.data_ptr(), KSFTExpertsCPU.input_tensor_cpu.data_ptr(), KSFTExpertsCPU.output_cpu.data_ptr())) + cpu_infer.sync_with_cuda_stream(torch.cuda.current_stream().cuda_stream) + t_fwd = time.time() - wall_t0 + KSFTExpertsCPU.output_gpu_map[out_device].copy_(KSFTExpertsCPU.output_cpu, non_blocking=True) + result = KSFTExpertsCPU.output_gpu_map[out_device] + # For backward compatibility, copy to CPU tensors + input_cpu = input_tensor.contiguous().cpu() + expert_ids_cpu = expert_ids.contiguous().cpu() + weights_cpu = weights.to(torch.float32).contiguous().cpu() + else: + num_tokens = input_tensor.size(0) + hidden_size = input_tensor.size(1) + num_experts_per_tok = expert_ids.size(1) + + # Ensure pinned buffers are large enough + KSFTExpertsCPU._ensure_pinned_buffers(num_tokens, hidden_size, num_experts_per_tok) + + # Use pinned memory buffers for efficient CPU-GPU transfer + input_buf = KSFTExpertsCPU._pinned_input_buf[:num_tokens] + output_buf = KSFTExpertsCPU._pinned_output_buf[:num_tokens] + expert_ids_buf = KSFTExpertsCPU._pinned_expert_ids_buf[:num_tokens] + weights_buf = KSFTExpertsCPU._pinned_weights_buf[:num_tokens] + + # Copy data to pinned memory (non_blocking for async transfer) + input_buf.copy_(input_tensor.to(torch.bfloat16), non_blocking=True) + expert_ids_buf.copy_(expert_ids, non_blocking=True) + weights_buf.copy_(weights.to(torch.float32), non_blocking=True) + + # Synchronize to ensure data is ready on CPU + if input_tensor.is_cuda: + torch.cuda.current_stream().synchronize() + + # Make contiguous views for CPU computation + input_cpu = input_buf.contiguous() + expert_ids_cpu = expert_ids_buf.contiguous() + weights_cpu = weights_buf.contiguous() + output_cpu = output_buf.contiguous() + + wall_t0 = time.time() + cpu_infer.submit( + moe.forward( + expert_ids_cpu.size(0), + expert_ids_cpu.size(1), + expert_ids_cpu.data_ptr(), + weights_cpu.data_ptr(), + input_cpu.data_ptr(), + output_cpu.data_ptr(), + ) + ) + cpu_infer.sync() + t_fwd = time.time() - wall_t0 + + # Copy result back to GPU using pinned memory (async) + result = torch.empty((num_tokens, hidden_size), dtype=input_tensor.dtype, device=out_device) + result.copy_(output_cpu, non_blocking=True) + + # Save CPU tensors for backward (already in pinned memory) + ctx.save_for_backward(input_cpu, expert_ids_cpu, weights_cpu) + ctx.cpu_infer = cpu_infer + ctx.moe = moe + ctx.out_device = out_device + ctx.layer_idx = layer_idx + + # ---------- FLOPs ---------- + qlen = expert_ids.size(0) + k = expert_ids.size(1) + + flops_fwd = 6 * qlen * k * H_FIXED * M_FIXED + tflops_f = flops_fwd / t_fwd / 1e12 + + ctx.saved_dims = (qlen, k) + ctx._time_fwd = t_fwd + # print(f"qlen ,k:{qlen}, {k}") + + # with open("test_V3_ESC.txt", "a", encoding="utf-8") as f: + # f.write(f"[KSFTExpertsCPU]Forward: {flops_fwd/1e9:.3f} GFLOPs {tflops_f:.2f} TFLOPS {t_fwd*1e3:.2f} ms\n") + + return result + + @staticmethod + def backward(ctx, output_grad): + # print("Go into the backward!!") + + # Pick back the middle results (already in pinned memory from forward) + input_tensor, expert_ids, weights = ctx.saved_tensors + + num_tokens = output_grad.size(0) + hidden_size = output_grad.size(1) + num_experts_per_tok = expert_ids.size(1) + + # Ensure pinned buffers are large enough (should already be from forward) + KSFTExpertsCPU._ensure_pinned_buffers(num_tokens, hidden_size, num_experts_per_tok) + + # Use pinned memory buffers for gradient transfer + grad_out_buf = KSFTExpertsCPU._pinned_grad_out_buf[:num_tokens] + grad_in_buf = KSFTExpertsCPU._pinned_grad_in_buf[:num_tokens] + + # Copy output_grad to pinned memory (async) + grad_out_buf.copy_(output_grad.to(torch.bfloat16), non_blocking=True) + + # Synchronize to ensure data is ready on CPU + if output_grad.is_cuda: + torch.cuda.current_stream().synchronize() + + # Make contiguous for CPU computation + output_grad_cpu = grad_out_buf.contiguous() + input_grad_cpu = grad_in_buf.contiguous() + + bw_start = time.time() + ctx.cpu_infer.submit( + ctx.moe.backward( + output_grad_cpu.size(0), # qlen + expert_ids.size(1), # k + expert_ids.data_ptr(), + weights.data_ptr(), + input_tensor.data_ptr(), + output_grad_cpu.data_ptr(), + input_grad_cpu.data_ptr(), + ) + ) + ctx.cpu_infer.sync() + + bw_end = time.time() + t_bw = bw_end - bw_start + + # Copy gradient back to GPU using pinned memory (async) + result_grad = torch.empty((num_tokens, hidden_size), dtype=output_grad.dtype, device=ctx.out_device) + result_grad.copy_(input_grad_cpu, non_blocking=True) + + # ---------- FLOPs ---------- + qlen, k = ctx.saved_dims + flops_bw = 10 * qlen * k * H_FIXED * M_FIXED + tflops_b = flops_bw / t_bw / 1e12 + # print(f"qlen:{qlen}, k:{k}") + + # with open("test_V3_ESC.txt", "a", encoding="utf-8") as f: + # f.write(f"[KSFTExpertsCPU]Backward: {flops_bw/1e9:.3f} GFLOPs {tflops_b:.2f} TFLOPS {t_bw*1e3:.2f} ms\n") + + return result_grad, None, None, None, None, None, None + + def unload(self): + return + + def load_weights(self, override_key: str | None = None, device: str = "cpu"): + # TODO: support Bias + res = {} + if override_key is not None: + keys = override_key + else: + keys = [self.key] + + gate = None + up = None + down = None + gate_type = None + up_type = None + down_type = None + + for key in keys: + if isinstance(self.gguf_loader, SafeTensorLoader): + res = self.gguf_loader.load_experts(key) + return {key: res} + elif self.gguf_loader.has_tensor(key + ".ffn_gate_exps.weight"): + gate = self.gguf_loader.get_mmap_tensor(key + ".ffn_gate_exps.weight") + up = self.gguf_loader.get_mmap_tensor(key + ".ffn_up_exps.weight") + down = self.gguf_loader.get_mmap_tensor(key + ".ffn_down_exps.weight") + # gate_type = self.gguf_loader.tensor_info[key + ".ffn_gate_exps.weight"]["ggml_type"] + # up_type = self.gguf_loader.tensor_info[key + ".ffn_up_exps.weight"]["ggml_type"] + # down_type = self.gguf_loader.tensor_info[key + ".ffn_down_exps.weight"]["ggml_type"] + gate_type = self.gguf_loader.get_ggml_type(key + ".ffn_gate_exps.weight") + up_type = self.gguf_loader.get_ggml_type(key + ".ffn_up_exps.weight") + down_type = self.gguf_loader.get_ggml_type(key + ".ffn_down_exps.weight") + + elif key + ".ffn_gate_exps.weight" in self.gguf_loader.tensor_info: + gate = self.gguf_loader.get_mmap_tensor(key + ".ffn_gate_exps.weight") + up = self.gguf_loader.get_mmap_tensor(key + ".ffn_up_exps.weight") + down = self.gguf_loader.get_mmap_tensor(key + ".ffn_down_exps.weight") + gate_type = self.gguf_loader.tensor_info[key + ".ffn_gate_exps.weight"]["ggml_type"] + up_type = self.gguf_loader.tensor_info[key + ".ffn_up_exps.weight"]["ggml_type"] + down_type = self.gguf_loader.tensor_info[key + ".ffn_down_exps.weight"]["ggml_type"] + elif key + ".ffn_down.0.weight" in self.gguf_loader.tensor_info: + # for supporting Mixtral-8x7B-Instuct + gate = [] + up = [] + down = [] + for i in range(8): + gate_it = self.gguf_loader.get_mmap_tensor(f"{key}.ffn_gate.{i}.weight") + up_it = self.gguf_loader.get_mmap_tensor(f"{key}.ffn_up.{i}.weight") + down_it = self.gguf_loader.get_mmap_tensor(f"{key}.ffn_down.{i}.weight") + gate.append(gate_it) + up.append(up_it) + down.append(down_it) + gate = np.stack(gate) + up = np.stack(up) + down = np.stack(down) + gate_type = self.gguf_loader.get_ggml_type(key + ".ffn_gate.0.weight") + up_type = self.gguf_loader.get_ggml_type(key + ".ffn_up.0.weight") + down_type = self.gguf_loader.get_ggml_type(key + ".ffn_down.0.weight") + else: + raise ValueError(f"Experts {key} not found in gguf_loader") + res = {key:{"gate": gate, "up": up, "down": down, "gate_type": gate_type, "up_type": up_type, "down_type": down_type}} + return res + +class KExpertsMarlin(KExpertsBase): + expert_num: int + loaded_experts_idx: list[int] + def __init__( + self, + key: str, + gguf_loader: GGUFLoader, + config: PretrainedConfig, + n_routed_experts: int, + orig_module: nn.Module = None, + device: str = "cuda", + **kwargs + ): + super().__init__(key, gguf_loader, config, orig_module, device, **kwargs) + self.expert_num = n_routed_experts + self.loaded_experts_idx = [] + self.act_fn = ACT2FN[config.hidden_act] + assert device.lower() != "cpu", "Marlin experts can only be loaded on GPU" + self.device = device + self.elements_per_tensor = config.moe_intermediate_size * config.hidden_size + + # create empty marlin experts according to the number of experts per token + # up + self.up_projs = [KLinearMarlin(key+ "." + "ffn_up_exps", gguf_loader, config, device=device) for i in range(self.expert_num)] + # gate + self.gate_projs = [KLinearMarlin(key+ "." + "ffn_gate_exps", gguf_loader, config, device=device) for i in range(self.expert_num)] + # down + self.down_projs = [KLinearMarlin(key+ "." + "ffn_down_exps", gguf_loader, config, device=device) for i in range(self.expert_num)] + + def load(self, w: dict | nn.Parameter | tuple | None = None, device: str | None = None, warmup: bool = False): + if device is None: device = self.device + assert device.lower() != "cpu", "Marlin experts can only be loaded on GPU" + if w is None: + w = self.load_weights() + load_by_experts = True + + if load_by_experts: + if isinstance(w, dict): + self.gate = w["gate"] + self.up = (w["up"]) + self.down = (w["down"]) + for i in tqdm(range(self.expert_num), desc=f"Dequanting and quanting for KExpertsMarlin {self.key}"): + up_weights = self.gguf_loader.load_expert_tensor(self.key + ".ffn_up_exps.weight", self.up, i, self.elements_per_tensor, device=self.device) + gate_weights = self.gguf_loader.load_expert_tensor(self.key + ".ffn_gate_exps.weight", self.gate, i, self.elements_per_tensor, device=self.device) + down_weights = self.gguf_loader.load_expert_tensor(self.key + ".ffn_down_exps.weight", self.down, i, self.elements_per_tensor, device=self.device) + + self.up_projs[i].load(nn.Parameter(up_weights), device=device) + self.gate_projs[i].load(nn.Parameter(gate_weights), device=device) + self.down_projs[i].load(nn.Parameter(down_weights), device=device) + self.loaded_experts_idx.append(i) + else: + if isinstance(w, dict): + self.gate = w["gate"] + self.up = (w["up"]) + self.down = (w["down"]) + for i in range(self.expert_num): + self.up_projs[i].load(nn.Parameter(self.up[i,...]), device=device) + self.gate_projs[i].load(nn.Parameter(self.gate[i,...]), device=device) + self.down_projs[i].load(nn.Parameter(self.down[i,...]), device=device) + self.loaded_experts_idx.append(i) + return + + def unload(self): + for i in self.loaded_experts_idx: + self.up_projs[i].unload() + self.gate_projs[i].unload() + self.down_projs[i].unload() + self.loaded_experts_idx = [] + + def load_weights(self, override_key: str | None = None): + res = {} + if override_key is not None: + keys = override_key + else: + keys = [self.key] + + gate = None + up = None + down = None + + for key in keys: + if self.gguf_loader.has_tensor(key + ".ffn_gate_exps.weight"): + gate = self.gguf_loader.get_mmap_tensor(key + ".ffn_gate_exps.weight") + up = self.gguf_loader.get_mmap_tensor(key + ".ffn_up_exps.weight") + down = self.gguf_loader.get_mmap_tensor(key + ".ffn_down_exps.weight") + res = {"gate": gate, "up": up, "down": down} + return res + + def forward(self, hidden_states_cpu: torch.Tensor, selected_experts_cpu: torch.Tensor, routing_weights_cpu: torch.Tensor) -> torch.Tensor: + org_dtype = hidden_states_cpu.dtype + org_device = hidden_states_cpu.device + hidden_states_cpu = hidden_states_cpu.to(self.device) + selected_experts_cpu = selected_experts_cpu.to(self.device) + routing_weights_cpu = routing_weights_cpu.to(self.device).to(org_dtype) + + batch_sequence_length, hidden_dim = hidden_states_cpu.size() + + final_hidden_states = torch.zeros( + (batch_sequence_length, hidden_dim), dtype=hidden_states_cpu.dtype, device=hidden_states_cpu.device + ) + # One hot encode the selected experts to create an expert mask + # this will be used to easily index which expert is going to be sollicitated + expert_mask = torch.nn.functional.one_hot(selected_experts_cpu, num_classes=self.expert_num).permute(2, 1, 0) + + # Loop over all available experts in the model and perform the computation on each expert + for expert_idx in range(self.expert_num): + if not expert_mask[expert_idx].any(): + continue + idx, top_x = torch.where(expert_mask[expert_idx]) + # Index the correct hidden states and compute the expert hidden state for + # the current expert. We need to make sure to multiply the output hidden + # states by `routing_weights` on the corresponding tokens (top-1 and top-2) + current_state = hidden_states_cpu[None, top_x].reshape(-1, hidden_dim) + G = self.gate_projs[expert_idx].forward(current_state) + A = self.act_fn(G) + U = self.up_projs[expert_idx].forward(current_state) + H = A * U # Element-wise multiplication + current_hidden_states = self.down_projs[expert_idx].forward(H) * routing_weights_cpu[top_x, idx, None] + # However `index_add_` only support torch tensors for indexing so we'll use + # the `top_x` tensor here. + final_hidden_states.index_add_(0, top_x, current_hidden_states) + + return final_hidden_states.to(dtype=org_dtype, device=org_device) + +# untested, CUDA OOM +class KExpertsTorch(KExpertsBase): + expert_num: int + loaded_experts_idx: list[int] + gate: torch.Tensor + up: torch.Tensor + down: torch.Tensor + def __init__( + self, + key: str, + gguf_loader: GGUFLoader, + config: PretrainedConfig, + n_routed_experts: int, + orig_module: nn.Module = None, + device: str = "cpu", + **kwargs + ): + super().__init__(key, gguf_loader, config, orig_module, device, **kwargs) + self.expert_num = n_routed_experts + # self.loaded_experts_idx = [] + self.act_fn = ACT2FN[config.hidden_act] + self.device = device + self.elements_per_tensor = config.moe_intermediate_size * config.hidden_size + self.gate = [None for _ in range(self.expert_num)] + self.up = [None for _ in range(self.expert_num)] + self.down = [None for _ in range(self.expert_num)] + self.dtype = torch.get_default_dtype() + + self.call_count = 0 + self.flops_per_call = [] + self.times = [] + self.expert_flops_details = [] + self.total_flops = 0 + + h = self.config.hidden_size + m = self.config.moe_intermediate_size + self.params_per_expert = 3 * h * m + self.total_params = self.expert_num * self.params_per_expert + + def load(self, w: dict | nn.Parameter | tuple | None = None, device: str | None = None, warmup: bool = False): + if device is None: device = self.device + if w is None: + w = self.load_weights() + load_by_experts = True + + if load_by_experts: + if isinstance(w, dict): + if isinstance(self.gguf_loader, SafeTensorLoader): + for i in tqdm(range(self.expert_num), desc=f"Loading experts(safetensors) for {self.key}"): + up_k = f"{self.key}.{i}.up_proj.weight" + gate_k = f"{self.key}.{i}.gate_proj.weight" + down_k = f"{self.key}.{i}.down_proj.weight" + + self.up[i] = self.gguf_loader.load_tensor(up_k, device=self.device).contiguous() + self.gate[i] = self.gguf_loader.load_tensor(gate_k, device=self.device).contiguous() + self.down[i] = self.gguf_loader.load_tensor(down_k, device=self.device).contiguous() + else: # GGUFLoader + for i in tqdm(range(self.expert_num), desc=f"Dequanting for KExpertsTorch {self.key}"): + up_weights = self.gguf_loader.load_expert_tensor(self.key + ".ffn_up_exps.weight", w["up"], i, self.elements_per_tensor, device=self.device) + gate_weights = self.gguf_loader.load_expert_tensor(self.key + ".ffn_gate_exps.weight", w["gate"], i, self.elements_per_tensor, device=self.device) + down_weights = self.gguf_loader.load_expert_tensor(self.key + ".ffn_down_exps.weight", w["down"], i, self.elements_per_tensor, device=self.device) + + self.up[i] = up_weights + self.gate[i] = gate_weights + self.down[i] = down_weights + else: + if isinstance(w, dict): + for i in range(self.expert_num): + self.gate[i] = w["gate"][i, ...].to(device=device, dtype=self.dtype) + self.up[i] = w["up"][i, ...].to(device=device, dtype=self.dtype) + self.down[i] = w["down"][i, ...].to(device=device, dtype=self.dtype) + + # self.up = torch.stack(self.up, dim=0) + # self.gate = torch.stack(self.gate, dim=0) + # self.down = torch.stack(self.down, dim=0) + self.up = nn.Parameter(torch.stack(self.up, dim=0)) + self.gate = nn.Parameter(torch.stack(self.gate, dim=0)) + self.down = nn.Parameter(torch.stack(self.down, dim=0)) + return + + def unload(self): + if self.gate is not None: + self.gate = None + self.up = None + self.down = None + + def load_weights(self, override_key: str | None = None): + res = {} + if override_key is not None: + keys = override_key + else: + keys = [self.key] + + gate = None + up = None + down = None + + for key in keys: + if isinstance(self.gguf_loader, SafeTensorLoader): + res = self.gguf_loader.load_experts(key) + return {key: res} + elif key + ".ffn_gate_exps.weight" in self.gguf_loader.tensor_info: + gate = self.gguf_loader.get_mmap_tensor(key + ".ffn_gate_exps.weight") + up = self.gguf_loader.get_mmap_tensor(key + ".ffn_up_exps.weight") + down = self.gguf_loader.get_mmap_tensor(key + ".ffn_down_exps.weight") + else: + import re + match = re.match(r'model\.layers\.(\d+)\.mlp\.experts(.*)', key) + if match: + layer_id = match.group(1) + suffix = match.group(2) + key = f"blk.{layer_id}{suffix}" + if key + ".ffn_gate_exps.weight" in self.gguf_loader.tensor_info: + gate = self.gguf_loader.get_mmap_tensor(key + ".ffn_gate_exps.weight") + up = self.gguf_loader.get_mmap_tensor(key + ".ffn_up_exps.weight") + down = self.gguf_loader.get_mmap_tensor(key + ".ffn_down_exps.weight") + res = {"gate": gate, "up": up, "down": down} + return res + + def forward(self, hidden_states_cpu: torch.Tensor, selected_experts_cpu: torch.Tensor, routing_weights_cpu: torch.Tensor) -> torch.Tensor: + start_time = time.time() + + org_device = hidden_states_cpu.device + hidden_states_cpu = hidden_states_cpu.to(self.device) + selected_experts_cpu = selected_experts_cpu.to(self.device) + routing_weights_cpu = routing_weights_cpu.to(self.device) + + batch_sequence_length, hidden_dim = hidden_states_cpu.size() + + final_hidden_states = torch.zeros( + (batch_sequence_length, hidden_dim), dtype=self.gate.dtype, device=hidden_states_cpu.device + ) + org_dtype = hidden_states_cpu.dtype + hidden_states_cpu = hidden_states_cpu.to(self.gate.dtype) + routing_weights_cpu = routing_weights_cpu.to(self.gate.dtype) + # One hot encode the selected experts to create an expert mask + # this will be used to easily index which expert is going to be sollicitated + expert_mask = torch.nn.functional.one_hot(selected_experts_cpu, num_classes=self.expert_num).permute(2, 1, 0) + + # Loop over all available experts in the model and perform the computation on each expert + for expert_idx in range(self.expert_num): + idx, top_x = torch.where(expert_mask[expert_idx]) + # Index the correct hidden states and compute the expert hidden state for + # the current expert. We need to make sure to multiply the output hidden + # states by `routing_weights` on the corresponding tokens (top-1 and top-2) + current_state = hidden_states_cpu[None, top_x].reshape(-1, hidden_dim) + G = current_state @ self.gate[expert_idx,...].T + A = self.act_fn(G) + U = current_state @ self.up[expert_idx,...].T + H = A * U # Element-wise multiplication + current_hidden_states = H @ self.down[expert_idx,...].T * routing_weights_cpu[top_x, idx, None] + # However `index_add_` only support torch tensors for indexing so we'll use + # the `top_x` tensor here. + final_hidden_states.index_add_(0, top_x, current_hidden_states) + + call_flops = 0 + expert_details = [] + + for expert_idx in range(self.expert_num): + idx, top_x = torch.where(expert_mask[expert_idx]) + t_e = len(top_x) + if t_e == 0: + expert_details.append({'gate':0, 'act':0, 'up':0, + 'element':0, 'down':0, 'routing':0}) + continue + + h = self.config.hidden_size + m = self.config.moe_intermediate_size + + flops_gate = 2 * t_e * h * m + flops_act = t_e * m + flops_up = 2 * t_e * h * m + flops_element = t_e * m + flops_down = 2 * t_e * m * h + flops_routing = t_e * h + + total_expert = sum([flops_gate, flops_act, flops_up, + flops_element, flops_down, flops_routing]) + call_flops += total_expert + + expert_details.append({ + 'gate': flops_gate, + 'act': flops_act, + 'up': flops_up, + 'element': flops_element, + 'down': flops_down, + 'routing': flops_routing + }) + + self.call_count += 1 + self.flops_per_call.append(call_flops) + self.total_flops += call_flops + self.expert_flops_details.append(expert_details) + self.times.append(time.time() - start_time) + + return final_hidden_states.to(dtype=org_dtype, device=org_device) + + # def forward(self, hidden_states_cpu: torch.Tensor, selected_experts_cpu: torch.Tensor, routing_weights_cpu: torch.Tensor) -> torch.Tensor: + # print("Enter the forward function!") + # current_call_start = time.perf_counter() + # if hasattr(self, 'last_call_end_time') and self.last_call_end_time is not None: + # inter_call_interval = current_call_start - self.last_call_end_time + # # print(f"\n[Forward Call Interval] Time since last forward call: {inter_call_interval:.6f} seconds") + # logging.info(f"\n[Forward Call Interval] Time since last forward call: {inter_call_interval:.6f} seconds") + # else: + # inter_call_interval = 0.0 + + # data_transfer_time = 0.0 + # tensor_init_time = 0.0 + # expert_mask_time = 0.0 + # expert_loop_total = 0.0 + # gate_time_total = 0.0 + # up_time_total = 0.0 + # elementwise_time_total = 0.0 + # down_time_total = 0.0 + # index_add_time_total = 0.0 + # cast_back_time = 0.0 + + # start = time.perf_counter() + # org_device = hidden_states_cpu.device + # hidden_states_cpu = hidden_states_cpu.to(self.device) + # selected_experts_cpu = selected_experts_cpu.to(self.device) + # routing_weights_cpu = routing_weights_cpu.to(self.device) + # data_transfer_time = time.perf_counter() - start + + # start = time.perf_counter() + # batch_sequence_length, hidden_dim = hidden_states_cpu.size() + # final_hidden_states = torch.zeros( + # (batch_sequence_length, hidden_dim), dtype=self.gate.dtype, device=hidden_states_cpu.device + # ) + # org_dtype = hidden_states_cpu.dtype + # hidden_states_cpu = hidden_states_cpu.to(self.gate.dtype) + # routing_weights_cpu = routing_weights_cpu.to(self.gate.dtype) + # tensor_init_time = time.perf_counter() - start + + # start = time.perf_counter() + # expert_mask = torch.nn.functional.one_hot(selected_experts_cpu, num_classes=self.expert_num).permute(2, 1, 0) + # expert_mask_time = time.perf_counter() - start + + # expert_loop_start = time.perf_counter() + # # for expert_idx in range(self.expert_num): + # for expert_idx in tqdm(range(self.expert_num), + # idx, top_x = torch.where(expert_mask[expert_idx]) + + # current_state = hidden_states_cpu[None, top_x].reshape(-1, hidden_dim) + + # gate_start = time.perf_counter() + # G = current_state @ self.gate[expert_idx,...].T + # A = self.act_fn(G) + # gate_time_total += time.perf_counter() - gate_start + + # up_start = time.perf_counter() + # U = current_state @ self.up[expert_idx,...].T + # up_time_total += time.perf_counter() - up_start + + # element_start = time.perf_counter() + # H = A * U # Element-wise multiplication + # elementwise_time_total += time.perf_counter() - element_start + + # down_start = time.perf_counter() + # current_hidden_states = H @ self.down[expert_idx,...].T * routing_weights_cpu[top_x, idx, None] + # down_time_total += time.perf_counter() - down_start + + # index_start = time.perf_counter() + # final_hidden_states.index_add_(0, top_x, current_hidden_states) + # index_add_time_total += time.perf_counter() - index_start + + # expert_loop_total = time.perf_counter() - expert_loop_start + # start = time.perf_counter() + # final_hidden_states = final_hidden_states.to(dtype=org_dtype, device=org_device) + # cast_back_time = time.perf_counter() - start + + # total_time = time.perf_counter() - current_call_start + # print(f""" + # [Timing Breakdown] + # Data Transfer: {data_transfer_time:.6f}s + # Tensor Initialization: {tensor_init_time:.6f}s + # Expert Mask Creation: {expert_mask_time:.6f}s + # Expert Loop Total: {expert_loop_total:.6f}s + # -> Gate Computations: {gate_time_total:.6f}s + # -> Up Projections: {up_time_total:.6f}s + # -> Elementwise Mult: {elementwise_time_total:.6f}s + # -> Down Projections: {down_time_total:.6f}s + # -> Index Add Ops: {index_add_time_total:.6f}s + # Cast Back to Original: {cast_back_time:.6f}s + # Total Forward Time: {total_time:.6f}s + # """) + # logging.info(f""" + # [Timing Breakdown] + # Data Transfer: {data_transfer_time:.6f}s + # Tensor Initialization: {tensor_init_time:.6f}s + # Expert Mask Creation: {expert_mask_time:.6f}s + # Expert Loop Total: {expert_loop_total:.6f}s + # -> Gate Computations: {gate_time_total:.6f}s + # -> Up Projections: {up_time_total:.6f}s + # -> Elementwise Mult: {elementwise_time_total:.6f}s + # -> Down Projections: {down_time_total:.6f}s + # -> Index Add Ops: {index_add_time_total:.6f}s + # Cast Back to Original: {cast_back_time:.6f}s + # Total Forward Time: {total_time:.6f}s + # """) + + # self.last_call_end_time = time.perf_counter() + + # return final_hidden_states + + +EXPERTS_MAP = { + "KExpertsCPU": KExpertsCPU, + "KSFTExpertsCPU": KSFTExpertsCPU, + "KExpertsTorch": KExpertsTorch, + "KExpertsMarlin": KExpertsMarlin, +} + +class KTransformersExperts(BaseInjectedModule, KExpertsBase): + def __init__(self, + key: str, + gguf_loader: GGUFLoader, + config: PretrainedConfig, + orig_module: nn.Module, + # device: str = "cuda", + prefill_device:str = "cuda", + prefill_op: str | None = "KExpertsTorch", + generate_device: str = "cpu", + generate_op: str | None = "KExpertsCPU", + **kwargs): + BaseInjectedModule.__init__(self, key, gguf_loader, config, orig_module, prefill_device, generate_device, **kwargs) + KExpertsBase.__init__(self, key, gguf_loader, config, orig_module, generate_device, **kwargs) + if generate_op is not None: + self.generate_experts = EXPERTS_MAP[generate_op](key, gguf_loader, config, len(orig_module), device=generate_device, **kwargs) + else: + self.generate_experts = None + if prefill_op is not None: + self.prefill_experts = EXPERTS_MAP[prefill_op](key, gguf_loader, config, len(orig_module), device=prefill_device, **kwargs) + else: + self.prefill_experts = None + self.gpu_mlp_type = prefill_op + self.cpu_mlp_type = generate_op + self.mode = InferenceState.UNLOAD + + def load(self, w: dict = None, mode: InferenceState = None, warmup: bool = True): + # TODO support w as input + if not mode: mode = InferenceState.GENERATE + if mode == InferenceState.GENERATE: + self.prefill_experts.unload() + self.generate_experts.load(w, warmup=warmup) + self.device = self.generate_experts.device + self.mode = mode + elif mode == InferenceState.PREFILL: + self.generate_experts.unload() + self.prefill_experts.load(w, warmup=warmup) + self.device = self.prefill_experts.device + self.mode = mode + elif mode == InferenceState.UNLOAD: + self.unload() + self.mode = mode + self.device = self.generate_experts.device + else: + raise ValueError("mode must be either InferenceState.GENERATE, InferenceState.PREFILL or InferenceState.UNLOAD") + + def unload(self): + if self.generate_experts is not None: + self.generate_experts.unload() + if self.prefill_experts is not None: + self.prefill_experts.unload() + self.device = self.generate_experts.device + + def forward(self, input_tensor, expert_ids, weights): + if self.mode == InferenceState.GENERATE: + assert self.generate_experts is not None, "generate_experts is None" + if type(self.generate_experts) == KSFTExpertsCPU: + layer_idx = int(re.search(r'\d+', self.key).group()) + return self.generate_experts.apply(input_tensor, expert_ids, weights, self.generate_experts.cpu_infer, self.generate_experts.moe, self.generate_experts.out_device, layer_idx) + else: + return self.generate_experts.forward(input_tensor, expert_ids, weights) + elif self.mode == InferenceState.PREFILL: + assert self.prefill_experts is not None, "prefill_experts is None" + return self.prefill_experts.forward(input_tensor, expert_ids, weights) + else: + raise ValueError("load or set_inference_mode before forward") + + def set_inference_mode(self, mode: InferenceState): + if mode == InferenceState.GENERATE: + self.load(mode=InferenceState.GENERATE, warmup=False) + elif mode == InferenceState.PREFILL: + self.load(mode=InferenceState.PREFILL, warmup=False) + elif mode == InferenceState.UNLOAD: + self.unload() + else: + raise ValueError("mode must be either InferenceState.GENERATE, InferenceState.PREFILL or InferenceState.UNLOAD") + + +from ktransformers.models.modeling_deepseek import DeepseekV2MoE +from ktransformers.models.modeling_deepseek_v3 import DeepseekV3MoE +from ktransformers.models.modeling_qwen2_moe import Qwen2MoeSparseMoeBlock +from ktransformers.models.modeling_qwen3_moe import Qwen3MoeSparseMoeBlock +from ktransformers.models.modeling_mixtral import MixtralSparseMoeBlock + + +class KQwen2MoeSparseMoeBlock(BaseInjectedModule, Qwen2MoeSparseMoeBlock): + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + """ """ + orig_shape = hidden_states.shape + batch_size, sequence_length, hidden_dim = hidden_states.shape + hidden_states = hidden_states.view(-1, hidden_dim) + # router_logits: (batch * sequence_length, n_experts) + router_logits = self.gate(hidden_states) + + routing_weights = F.softmax(router_logits, dim=1, dtype=torch.float) + routing_weights, selected_experts = torch.topk(routing_weights, self.top_k, dim=-1) + if self.norm_topk_prob: + routing_weights /= routing_weights.sum(dim=-1, keepdim=True) + # we cast back to the input dtype + routing_weights = routing_weights.to(hidden_states.dtype) + + if sequence_length == 1 and hasattr(self.experts.generate_experts, "submit_for_one_decode"): + self.experts.generate_experts.submit_for_one_decode(hidden_states[0], selected_experts[0], routing_weights[0]) + shared_expert_output = self.shared_expert(hidden_states) + shared_expert_output = F.sigmoid(self.shared_expert_gate(hidden_states)) * shared_expert_output + y = self.experts.generate_experts.sync_for_one_decode().unsqueeze(0) + y += shared_expert_output + y.resize_(*orig_shape) + return y, router_logits + + hidden_states_expert = hidden_states.to(self.experts.device) if isinstance(self.experts, KExpertsBase) else hidden_states.cpu() + selected_experts_expert = selected_experts.to(self.experts.device) if isinstance(self.experts, KExpertsBase) else selected_experts.cpu() + routing_weights_expert = routing_weights.to(self.experts.device) if isinstance(self.experts, KExpertsBase) else routing_weights.cpu() + + shared_expert_output = self.shared_expert(hidden_states) + shared_expert_output = ( + F.sigmoid(self.shared_expert_gate(hidden_states)) * shared_expert_output + ) + + if isinstance(self.experts, KExpertsBase): + y = ( + self.moe_kexperts( + hidden_states_expert, selected_experts_expert, routing_weights_expert + ) + .view(*orig_shape) + .to(device=hidden_states.device) + ) + elif hidden_states_expert.size(0) > 10: + y = self.moe_infer( + hidden_states_expert, selected_experts_expert, routing_weights_expert, orig_shape + ).to(device=hidden_states.device) + else: + y = self.moe_infer_simple( + hidden_states_expert, selected_experts_expert, routing_weights_expert + ).to(device=hidden_states.device) + y += shared_expert_output + y.resize_(*orig_shape) + return y, router_logits + + @maybe_no_grad() + def moe_kexperts(self, x: torch.Tensor, topk_ids: torch.Tensor, topk_weight: torch.Tensor) -> torch.Tensor: + outs = self.experts(x, topk_ids, topk_weight) + return outs + + @maybe_no_grad() + # TODO may bugs here + def moe_infer_simple(self, hidden_states_cpu: torch.Tensor, selected_experts_cpu: torch.Tensor, routing_weights_cpu: torch.Tensor) -> torch.Tensor: + ''' + hidden_states_cpu: [num_tokens, hidden_size] + topk_ids, topk_weight: [num_tokens, num_selected_experts] + ''' + outs = torch.zeros_like(hidden_states_cpu) + for token_idx in range(selected_experts_cpu.size(0)): + for expert_idx in range(selected_experts_cpu.size(1)): + expert = self.experts[selected_experts_cpu[token_idx, expert_idx]] + outs[token_idx] += expert.forward(hidden_states_cpu[token_idx]) * routing_weights_cpu[token_idx, expert_idx] + return outs + + @maybe_no_grad() + # TODO may bugs here + def moe_infer(self, hidden_states_cpu: torch.Tensor, selected_experts_cpu: torch.Tensor, routing_weights_cpu: torch.Tensor, orig_shape: tuple) -> torch.Tensor: + + batch_size, sequence_length, hidden_dim = orig_shape + + final_hidden_states = torch.zeros( + (batch_size * sequence_length, hidden_dim), dtype=hidden_states_cpu.dtype, device=hidden_states_cpu.device + ) + + # One hot encode the selected experts to create an expert mask + # this will be used to easily index which expert is going to be sollicitated + expert_mask = torch.nn.functional.one_hot(selected_experts_cpu, num_classes=self.num_experts).permute(2, 1, 0) + + # Loop over all available experts in the model and perform the computation on each expert + for expert_idx in range(self.num_experts): + expert_layer = self.experts[expert_idx] + idx, top_x = torch.where(expert_mask[expert_idx]) + + # Index the correct hidden states and compute the expert hidden state for + # the current expert. We need to make sure to multiply the output hidden + # states by `routing_weights` on the corresponding tokens (top-1 and top-2) + current_state = hidden_states_cpu[None, top_x].reshape(-1, hidden_dim) + current_hidden_states = expert_layer.forward(current_state) * routing_weights_cpu[top_x, idx, None] + + # However `index_add_` only support torch tensors for indexing so we'll use + # the `top_x` tensor here. + final_hidden_states.index_add_(0, top_x, current_hidden_states.to(hidden_states_cpu.dtype)) + + return final_hidden_states + +class KDeepseekV2MoE(BaseInjectedModule, DeepseekV2MoE): + def forward(self, hidden_states): + identity = hidden_states + orig_shape = hidden_states.shape + sequence_length = orig_shape[1] + topk_idx, topk_weight, aux_loss = self.gate(hidden_states) + hidden_states = hidden_states.view(-1, hidden_states.shape[-1]) + + if sequence_length == 1 and hasattr(self.experts.generate_experts, "submit_for_one_decode") and torch.cuda.is_available() and torch.cuda.is_current_stream_capturing(): + self.experts.generate_experts.submit_for_one_decode(hidden_states[0], topk_idx[0], topk_weight[0]) + if self.config.n_shared_experts is not None: + y_ = self.shared_experts(identity).squeeze(0) + y = self.experts.generate_experts.sync_for_one_decode().unsqueeze(0) + y += y_ + y.resize_(*orig_shape) + return y + + if self.config.n_shared_experts is not None: + y_ = self.shared_experts(identity).squeeze(0) + + if isinstance(self.experts, KExpertsBase): + y = self.moe_kexperts(hidden_states, topk_idx, topk_weight).view(*orig_shape).to(device=hidden_states.device) + elif hidden_states.size(0) > 10: + # TODO may bugs here + y = ( + self.moe_infer(hidden_states, topk_idx, topk_weight) + .view(*orig_shape) + .to(device=hidden_states.device) + ) + else: + # TODO may bugs here + y = ( + self.moe_infer_simple(hidden_states, topk_idx, topk_weight) + .view(*orig_shape) + .to(device=hidden_states.device) + ) + if self.config.n_shared_experts is not None: + y += y_ + return y + + @maybe_no_grad() + def moe_kexperts(self, x: torch.Tensor, topk_ids: torch.Tensor, topk_weight: torch.Tensor) -> torch.Tensor: + outs = self.experts(x, topk_ids, topk_weight) + return outs + + @maybe_no_grad() + # TODO may bugs here + def moe_infer_simple( + self, x: torch.Tensor, topk_ids: torch.Tensor, topk_weight: torch.Tensor + ) -> torch.Tensor: + """ + x: [num_tokens, hidden_size] + topk_ids, topk_weight: [num_tokens, num_selected_experts] + """ + outs = torch.zeros_like(x) + for token_idx in range(topk_ids.size(0)): + for expert_idx in range(topk_ids.size(1)): + expert = self.experts[topk_ids[token_idx, expert_idx]] + outs[token_idx] += ( + expert.forward(x[token_idx]) * topk_weight[token_idx, expert_idx] + ) + return outs + + @maybe_no_grad() + # TODO may bugs here + def moe_infer(self, x, topk_ids, topk_weight): + cnts = topk_ids.new_zeros((topk_ids.shape[0], len(self.experts))) + cnts.scatter_(1, topk_ids, 1) + tokens_per_expert = cnts.sum(dim=0) + idxs = topk_ids.view(-1).argsort() + sorted_tokens = x[idxs // topk_ids.shape[1]] + tokens_per_expert = tokens_per_expert.cpu().numpy() + + outputs = [] + start_idx = 0 + for i, num_tokens in enumerate(tokens_per_expert): + end_idx = start_idx + num_tokens + if num_tokens == 0: + continue + expert = self.experts[i + self.ep_rank * self.experts_per_rank] + tokens_for_this_expert = sorted_tokens[start_idx:end_idx] + expert_out = expert.forward(tokens_for_this_expert) + outputs.append(expert_out) + start_idx = end_idx + + outs = torch.cat(outputs, dim=0) if len(outputs) else sorted_tokens.new_empty(0) + + new_x = torch.empty_like(outs) + new_x[idxs] = outs + final_out = ( + new_x.view(*topk_ids.shape, -1) + .type(topk_weight.dtype) + .mul_(topk_weight.unsqueeze(dim=-1)) + .sum(dim=1) + .type(new_x.dtype) + ) + return final_out + +class KDeepseekV3MoE(BaseInjectedModule, DeepseekV3MoE): + + def forward(self, hidden_states): + identity = hidden_states + orig_shape = hidden_states.shape + sequence_length = orig_shape[1] + topk_idx, topk_weight = self.gate(hidden_states) + hidden_states = hidden_states.view(-1, hidden_states.shape[-1]) + + # only for generate phase + if sequence_length == 1 and hasattr(self.experts.generate_experts, "submit_for_one_decode") and torch.cuda.is_available() and torch.cuda.is_current_stream_capturing(): + self.experts.generate_experts.submit_for_one_decode(hidden_states[0], topk_idx[0], topk_weight[0]) + if self.config.n_shared_experts is not None: + y_ = self.shared_experts(identity).squeeze(0) + y = self.experts.generate_experts.sync_for_one_decode().unsqueeze(0) + y += y_ + y.resize_(*orig_shape) + return y + + if self.config.n_shared_experts is not None: + y_ = self.shared_experts(identity).squeeze(0) + + if isinstance(self.experts, KExpertsBase): + y = self.moe_kexperts(hidden_states, topk_idx, topk_weight).view(*orig_shape).to(device=hidden_states.device) + elif hidden_states.size(0) > 10: + # TODO may bugs here + y = ( + self.moe_infer(hidden_states, topk_idx, topk_weight) + .view(*orig_shape) + .to(device=hidden_states.device) + ) + else: + # TODO may bugs here + y = ( + self.moe_infer_simple(hidden_states, topk_idx, topk_weight) + .view(*orig_shape) + .to(device=hidden_states.device) + ) + if self.config.n_shared_experts is not None: + y += y_ + return y + + @maybe_no_grad() + def moe_kexperts(self, x: torch.Tensor, topk_ids: torch.Tensor, topk_weight: torch.Tensor) -> torch.Tensor: + outs = self.experts(x, topk_ids, topk_weight) + return outs + + @maybe_no_grad() + # TODO may bugs here + def moe_infer_simple( + self, x: torch.Tensor, topk_ids: torch.Tensor, topk_weight: torch.Tensor + ) -> torch.Tensor: + """ + x: [num_tokens, hidden_size] + topk_ids, topk_weight: [num_tokens, num_selected_experts] + """ + outs = torch.zeros_like(x) + for token_idx in range(topk_ids.size(0)): + for expert_idx in range(topk_ids.size(1)): + expert = self.experts[topk_ids[token_idx, expert_idx]] + outs[token_idx] += ( + expert.forward(x[token_idx]) * topk_weight[token_idx, expert_idx] + ) + return outs + + @maybe_no_grad() + # TODO may bugs here + def moe_infer(self, x, topk_ids, topk_weight): + cnts = topk_ids.new_zeros((topk_ids.shape[0], len(self.experts))) + cnts.scatter_(1, topk_ids, 1) + tokens_per_expert = cnts.sum(dim=0) + idxs = topk_ids.view(-1).argsort() + sorted_tokens = x[idxs // topk_ids.shape[1]] + tokens_per_expert = tokens_per_expert.cpu().numpy() + + outputs = [] + start_idx = 0 + for i, num_tokens in enumerate(tokens_per_expert): + end_idx = start_idx + num_tokens + if num_tokens == 0: + continue + expert = self.experts[i + self.ep_rank * self.experts_per_rank] + tokens_for_this_expert = sorted_tokens[start_idx:end_idx] + expert_out = expert.forward(tokens_for_this_expert) + outputs.append(expert_out) + start_idx = end_idx + + outs = torch.cat(outputs, dim=0) if len(outputs) else sorted_tokens.new_empty(0) + + new_x = torch.empty_like(outs) + new_x[idxs] = outs + final_out = ( + new_x.view(*topk_ids.shape, -1) + .type(topk_weight.dtype) + .mul_(topk_weight.unsqueeze(dim=-1)) + .sum(dim=1) + .type(new_x.dtype) + ) + return final_out + +class KMistralSparseMoEBlock(BaseInjectedModule, MixtralSparseMoeBlock): + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + """ """ + orig_shape = hidden_states.shape + batch_size, sequence_length, hidden_dim = hidden_states.shape + if self.training and self.jitter_noise > 0: + hidden_states *= torch.empty_like(hidden_states).uniform_(1.0 - self.jitter_noise, 1.0 + self.jitter_noise) + hidden_states = hidden_states.view(-1, hidden_dim) + # router_logits: (batch * sequence_length, n_experts) + router_logits = self.gate(hidden_states) + + routing_weights = F.softmax(router_logits, dim=1, dtype=torch.float) + routing_weights, selected_experts = torch.topk(routing_weights, self.top_k, dim=-1) + routing_weights /= routing_weights.sum(dim=-1, keepdim=True) + # we cast back to the input dtype + routing_weights = routing_weights.to(hidden_states.dtype) + + if sequence_length == 1 and hasattr(self.experts.generate_experts, "submit_for_one_decode"): + self.experts.generate_experts.submit_for_one_decode(hidden_states[0], selected_experts[0], routing_weights[0]) + y = self.experts.generate_experts.sync_for_one_decode().unsqueeze(0) + y.resize_(*orig_shape) + return y, router_logits + + hidden_states_expert = hidden_states.to(self.experts.device) if isinstance(self.experts, KExpertsBase) else hidden_states_expert.cpu() + selected_experts_expert = selected_experts.to(self.experts.device) if isinstance(self.experts, KExpertsBase) else selected_experts_expert.cpu() + routing_weights_expert = routing_weights.to(self.experts.device) if isinstance(self.experts, KExpertsBase) else routing_weights_expert.cpu() + + if isinstance(self.experts, KExpertsBase): + y = ( + self.moe_kexperts( + hidden_states_expert, selected_experts_expert, routing_weights_expert + ) + .view(*orig_shape) + .to(device=hidden_states.device) + ) + elif hidden_states_expert.size(0) > 10: + y = self.moe_infer( + hidden_states_expert, selected_experts_expert, routing_weights_expert, orig_shape + ).to(device=hidden_states.device) + else: + y = self.moe_infer_simple( + hidden_states_expert, selected_experts_expert, routing_weights_expert + ).to(device=hidden_states.device) + + y.resize_(*orig_shape) + return y, router_logits + + @maybe_no_grad() + def moe_kexperts(self, x: torch.Tensor, topk_ids: torch.Tensor, topk_weight: torch.Tensor) -> torch.Tensor: + outs = self.experts(x, topk_ids, topk_weight) + return outs + + @maybe_no_grad() + # TODO may bugs here + def moe_infer_simple(self, hidden_states_cpu: torch.Tensor, selected_experts_cpu: torch.Tensor, routing_weights_cpu: torch.Tensor) -> torch.Tensor: + ''' + hidden_states_cpu: [num_tokens, hidden_size] + topk_ids, topk_weight: [num_tokens, num_selected_experts] + ''' + outs = torch.zeros_like(hidden_states_cpu) + for token_idx in range(selected_experts_cpu.size(0)): + for expert_idx in range(selected_experts_cpu.size(1)): + expert = self.experts[selected_experts_cpu[token_idx, expert_idx]] + outs[token_idx] += expert.forward(hidden_states_cpu[token_idx]) * routing_weights_cpu[token_idx, expert_idx] + return outs + + @maybe_no_grad() + # TODO may bugs here + def moe_infer(self, hidden_states_cpu: torch.Tensor, selected_experts_cpu: torch.Tensor, routing_weights_cpu: torch.Tensor, orig_shape: tuple) -> torch.Tensor: + + batch_size, sequence_length, hidden_dim = orig_shape + + final_hidden_states = torch.zeros( + (batch_size * sequence_length, hidden_dim), dtype=hidden_states_cpu.dtype, device=hidden_states_cpu.device + ) + + # One hot encode the selected experts to create an expert mask + # this will be used to easily index which expert is going to be sollicitated + expert_mask = torch.nn.functional.one_hot(selected_experts_cpu, num_classes=self.num_experts).permute(2, 1, 0) + + # Loop over all available experts in the model and perform the computation on each expert + for expert_idx in range(self.num_experts): + expert_layer = self.experts[expert_idx] + idx, top_x = torch.where(expert_mask[expert_idx]) + + # Index the correct hidden states and compute the expert hidden state for + # the current expert. We need to make sure to multiply the output hidden + # states by `routing_weights` on the corresponding tokens (top-1 and top-2) + current_state = hidden_states_cpu[None, top_x].reshape(-1, hidden_dim) + current_hidden_states = expert_layer.forward(current_state) * routing_weights_cpu[top_x, idx, None] + + # However `index_add_` only support torch tensors for indexing so we'll use + # the `top_x` tensor here. + final_hidden_states.index_add_(0, top_x, current_hidden_states.to(hidden_states_cpu.dtype)) + + return final_hidden_states + +class KDeepseekV3MoEV2(BaseInjectedModule, DeepseekV3MoE): + def forward(self, hidden_states, bsz_tensor, cuda_graph_idx=0): + identity = hidden_states + orig_shape = hidden_states.shape + sequence_length = orig_shape[1] + topk_idx, topk_weight = self.gate(hidden_states) + hidden_states = hidden_states.view(-1, hidden_states.shape[-1]) + + + # only for generate phase + if hasattr(self.experts.generate_experts, "submit_for_one_decode") and torch.cuda.is_available() and torch.cuda.is_current_stream_capturing(): # TODO: this branch cause jit bug + self.experts.generate_experts.submit_for_one_decode(hidden_states, topk_idx, topk_weight, bsz_tensor, cuda_graph_idx) + if self.config.n_shared_experts is not None: + y_ = self.shared_experts(identity, bsz_tensor).squeeze(0) + y = self.experts.generate_experts.sync_for_one_decode(cuda_graph_idx).unsqueeze(0) + y += y_ + y.resize_(*orig_shape) + return y + + if self.config.n_shared_experts is not None: + y_ = self.shared_experts(identity, bsz_tensor).squeeze(0) + + if isinstance(self.experts, KExpertsBase): + y = self.moe_on_cpuinfer(hidden_states, topk_idx, topk_weight, bsz_tensor, cuda_graph_idx).view(*orig_shape).to(device=hidden_states.device) + elif hidden_states.size(0) > 10: + # TODO may bugs here + y = ( + self.moe_infer(hidden_states, topk_idx, topk_weight) + .view(*orig_shape) + .to(device=hidden_states.device) + ) + else: + # TODO may bugs here + y = ( + self.moe_infer_simple(hidden_states, topk_idx, topk_weight) + .view(*orig_shape) + .to(device=hidden_states.device) + ) + if self.config.n_shared_experts is not None: + y += y_ + return y + + @maybe_no_grad() + def moe_on_cpuinfer(self, x: torch.Tensor, topk_ids: torch.Tensor, topk_weight: torch.Tensor, bsz_tensor, cuda_graph_idx=0) -> torch.Tensor: + outs = torch.empty_like(x) + outs = self.experts(x, topk_ids, topk_weight, bsz_tensor, cuda_graph_idx) + return outs + + @maybe_no_grad() + # TODO may bugs here + def moe_infer_simple( + self, x: torch.Tensor, topk_ids: torch.Tensor, topk_weight: torch.Tensor + ) -> torch.Tensor: + """ + x: [num_tokens, hidden_size] + topk_ids, topk_weight: [num_tokens, num_selected_experts] + """ + outs = torch.zeros_like(x) + for token_idx in range(topk_ids.size(0)): + for expert_idx in range(topk_ids.size(1)): + expert = self.experts[topk_ids[token_idx, expert_idx]] + outs[token_idx] += ( + expert.forward(x[token_idx]) * topk_weight[token_idx, expert_idx] + ) + return outs + + @maybe_no_grad() + # TODO may bugs here + def moe_infer(self, x, topk_ids, topk_weight): + cnts = topk_ids.new_zeros((topk_ids.shape[0], len(self.experts))) + cnts.scatter_(1, topk_ids, 1) + tokens_per_expert = cnts.sum(dim=0) + idxs = topk_ids.view(-1).argsort() + sorted_tokens = x[idxs // topk_ids.shape[1]] + tokens_per_expert = tokens_per_expert.cpu().numpy() + + outputs = [] + start_idx = 0 + for i, num_tokens in enumerate(tokens_per_expert): + end_idx = start_idx + num_tokens + if num_tokens == 0: + continue + expert = self.experts[i + self.ep_rank * self.experts_per_rank] + tokens_for_this_expert = sorted_tokens[start_idx:end_idx] + expert_out = expert.forward(tokens_for_this_expert) + outputs.append(expert_out) + start_idx = end_idx + + outs = torch.cat(outputs, dim=0) if len(outputs) else sorted_tokens.new_empty(0) + + new_x = torch.empty_like(outs) + new_x[idxs] = outs + final_out = ( + new_x.view(*topk_ids.shape, -1) + .type(topk_weight.dtype) + .mul_(topk_weight.unsqueeze(dim=-1)) + .sum(dim=1) + .type(new_x.dtype) + ) + return final_out + +class KTransformersExpertsV2(BaseInjectedModule, KExpertsBase): + def __init__(self, + key: str, + gguf_loader: GGUFLoader, + config: PretrainedConfig, + orig_module: nn.Module, + # device: str = "cuda", + prefill_device:str = "cuda", + prefill_op: str | None = "KExpertsTorch", + generate_device: str = "cpu", + generate_op: str | None = "KExpertsCPU", + **kwargs): + BaseInjectedModule.__init__(self, key, gguf_loader, config, orig_module, generate_device, **kwargs) + KExpertsBase.__init__(self, key, gguf_loader, config, orig_module, generate_device, **kwargs) + if generate_op is not None: + self.generate_experts = EXPERTS_MAP[generate_op](key, gguf_loader, config, len(orig_module), device=generate_device, **kwargs) + else: + self.generate_experts = None + if prefill_op is not None: + self.prefill_experts = EXPERTS_MAP[prefill_op](key, gguf_loader, config, len(orig_module), device=prefill_device, **kwargs) + else: + self.prefill_experts = None + self.gpu_mlp_type = prefill_op + self.cpu_mlp_type = generate_op + self.mode = InferenceState.UNLOAD + + def load(self, w: dict = None, mode: InferenceState = None, warmup: bool = True): + # TODO support w as input + if not mode: mode = InferenceState.GENERATE + if mode == InferenceState.GENERATE: + self.prefill_experts.unload() + self.generate_experts.load(w, warmup=warmup) + self.device = self.generate_experts.device + self.mode = mode + elif mode == InferenceState.PREFILL: + self.generate_experts.unload() + self.prefill_experts.load(w, warmup=warmup) + self.device = self.prefill_experts.device + self.mode = mode + elif mode == InferenceState.UNLOAD: + self.unload() + self.mode = mode + self.device = self.generate_experts.device + else: + raise ValueError("mode must be either InferenceState.GENERATE, InferenceState.PREFILL or InferenceState.UNLOAD") + + def unload(self): + if self.generate_experts is not None: + self.generate_experts.unload() + if self.prefill_experts is not None: + self.prefill_experts.unload() + self.device = self.generate_experts.device + + def forward(self, input_tensor, expert_ids, weights, bsz_tensor, cuda_graph_idx=0): + if self.mode == InferenceState.GENERATE: + assert self.generate_experts is not None, "generate_experts is None" + return self.generate_experts.forward(input_tensor, expert_ids, weights, bsz_tensor, cuda_graph_idx) + elif self.mode == InferenceState.PREFILL: + assert self.prefill_experts is not None, "prefill_experts is None" + return self.prefill_experts.forward(input_tensor, expert_ids, weights, bsz_tensor, cuda_graph_idx) + else: + raise ValueError("load or set_inference_mode before forward") + + def set_inference_mode(self, mode: InferenceState): + if mode == InferenceState.GENERATE: + self.load(mode=InferenceState.GENERATE, warmup=False) + elif mode == InferenceState.PREFILL: + self.load(mode=InferenceState.PREFILL, warmup=False) + elif mode == InferenceState.UNLOAD: + self.unload() + else: + raise ValueError("mode must be either InferenceState.GENERATE, InferenceState.PREFILL or InferenceState.UNLOAD") + +class KQwen2MoeSparseMoeBlockV2(BaseInjectedModule, Qwen2MoeSparseMoeBlock): + def forward(self, hidden_states, bsz_tensor, cuda_graph_idx=0): + + orig_shape = hidden_states.shape + sequence_length = orig_shape[1] + + hidden_states = hidden_states.view(-1, hidden_states.shape[-1]) + + router_logits = self.gate(hidden_states, bsz_tensor) + + routing_weights = F.softmax(router_logits, dim=1, dtype=torch.float) + routing_weights, selected_experts = torch.topk(routing_weights, self.top_k, dim=-1) + if self.norm_topk_prob: + routing_weights /= routing_weights.sum(dim=-1, keepdim=True) + # we cast back to the input dtype + routing_weights = routing_weights.to(hidden_states.dtype) + + # only for generate phase + if hasattr(self.experts.generate_experts, "submit_for_one_decode") and torch.cuda.is_available() and torch.cuda.is_current_stream_capturing(): # TODO: this branch cause jit bug + self.experts.generate_experts.submit_for_one_decode(hidden_states, selected_experts, routing_weights, bsz_tensor, cuda_graph_idx) + y_ = self.shared_expert(hidden_states, bsz_tensor).squeeze(0) + y_ = F.sigmoid(self.shared_expert_gate(hidden_states)) * y_ + + y = self.experts.generate_experts.sync_for_one_decode(cuda_graph_idx).unsqueeze(0) + + y += y_ + y.resize_(*orig_shape) + return y + + y_ = self.shared_expert(hidden_states, bsz_tensor).squeeze(0) + y_ = ( + F.sigmoid(self.shared_expert_gate(hidden_states)) * y_ + ) + + + if isinstance(self.experts, KExpertsBase): + y = self.moe_on_cpuinfer(hidden_states, selected_experts, routing_weights, bsz_tensor, cuda_graph_idx).view(*orig_shape).to(device=hidden_states.device) + elif hidden_states.size(0) > 10: + # TODO may bugs here + y = ( + self.moe_infer(hidden_states, selected_experts, routing_weights) + .view(*orig_shape) + .to(device=hidden_states.device) + ) + else: + # TODO may bugs here + y = ( + self.moe_infer_simple(hidden_states, selected_experts, routing_weights) + .view(*orig_shape) + .to(device=hidden_states.device) + ) + y += y_ + return y + + @maybe_no_grad() + def moe_on_cpuinfer(self, x: torch.Tensor, topk_ids: torch.Tensor, topk_weight: torch.Tensor, bsz_tensor, cuda_graph_idx=0) -> torch.Tensor: + outs = torch.empty_like(x) + outs = self.experts(x, topk_ids, topk_weight, bsz_tensor, cuda_graph_idx) + return outs + + @maybe_no_grad() + # TODO may bugs here + def moe_infer_simple( + self, x: torch.Tensor, topk_ids: torch.Tensor, topk_weight: torch.Tensor + ) -> torch.Tensor: + """ + x: [num_tokens, hidden_size] + topk_ids, topk_weight: [num_tokens, num_selected_experts] + """ + outs = torch.zeros_like(x) + for token_idx in range(topk_ids.size(0)): + for expert_idx in range(topk_ids.size(1)): + expert = self.experts[topk_ids[token_idx, expert_idx]] + outs[token_idx] += ( + expert.forward(x[token_idx]) * topk_weight[token_idx, expert_idx] + ) + return outs + + @maybe_no_grad() + # TODO may bugs here + def moe_infer(self, x, topk_ids, topk_weight): + cnts = topk_ids.new_zeros((topk_ids.shape[0], len(self.experts))) + cnts.scatter_(1, topk_ids, 1) + tokens_per_expert = cnts.sum(dim=0) + idxs = topk_ids.view(-1).argsort() + sorted_tokens = x[idxs // topk_ids.shape[1]] + tokens_per_expert = tokens_per_expert.cpu().numpy() + + outputs = [] + start_idx = 0 + for i, num_tokens in enumerate(tokens_per_expert): + end_idx = start_idx + num_tokens + if num_tokens == 0: + continue + expert = self.experts[i + self.ep_rank * self.experts_per_rank] + tokens_for_this_expert = sorted_tokens[start_idx:end_idx] + expert_out = expert.forward(tokens_for_this_expert) + outputs.append(expert_out) + start_idx = end_idx + + outs = torch.cat(outputs, dim=0) if len(outputs) else sorted_tokens.new_empty(0) + + new_x = torch.empty_like(outs) + new_x[idxs] = outs + final_out = ( + new_x.view(*topk_ids.shape, -1) + .type(topk_weight.dtype) + .mul_(topk_weight.unsqueeze(dim=-1)) + .sum(dim=1) + .type(new_x.dtype) + ) + return final_out + +class KQwen3MoeSparseMoeBlockV2(BaseInjectedModule, Qwen3MoeSparseMoeBlock): + def forward(self, hidden_states, bsz_tensor=None, cuda_graph_idx=0): + + orig_shape = hidden_states.shape + sequence_length = orig_shape[1] + + hidden_states = hidden_states.view(-1, hidden_states.shape[-1]) + + if bsz_tensor is None: + router_logits = self.gate(hidden_states) + else: + router_logits = self.gate(hidden_states, bsz_tensor) + + if router_logits.device.type == "xpu": + from ipex_llm.transformers.models.common import moe_softmax_topk + selected_experts, routing_weights = moe_softmax_topk( + router_logits.half(), self.top_k, self.norm_topk_prob + ) + else: + routing_weights = torch.nn.functional.softmax(router_logits, dim=1, dtype=torch.float) + routing_weights, selected_experts = torch.topk(routing_weights, self.top_k, dim=-1) + if self.norm_topk_prob: + routing_weights /= routing_weights.sum(dim=-1, keepdim=True) + # we cast back to the input dtype + routing_weights = routing_weights.to(hidden_states.dtype) + + # only for generate phase + if hasattr(self.experts.generate_experts, "submit_for_one_decode") and torch.cuda.is_available() and torch.cuda.is_current_stream_capturing(): # TODO: this branch cause jit bug + self.experts.generate_experts.submit_for_one_decode(hidden_states, selected_experts, routing_weights, bsz_tensor, cuda_graph_idx) + # y_ = self.shared_expert(hidden_states, bsz_tensor).squeeze(0) + # y_ = F.sigmoid(self.shared_expert_gate(hidden_states)) * y_ + + y = self.experts.generate_experts.sync_for_one_decode(cuda_graph_idx).unsqueeze(0) + + # y += y_ + y.resize_(*orig_shape) + return y + + # y_ = self.shared_expert(hidden_states, bsz_tensor).squeeze(0) + # y_ = ( + # F.sigmoid(self.shared_expert_gate(hidden_states)) * y_ + # ) + + + if isinstance(self.experts, KExpertsBase): + y = self.moe_on_cpuinfer(hidden_states, selected_experts, routing_weights, bsz_tensor, cuda_graph_idx).view(*orig_shape).to(device=hidden_states.device) + elif hidden_states.size(0) > 10: + # TODO may bugs here + y = ( + self.moe_infer(hidden_states, selected_experts, routing_weights) + .view(*orig_shape) + .to(device=hidden_states.device) + ) + else: + # TODO may bugs here + y = ( + self.moe_infer_simple(hidden_states, selected_experts, routing_weights) + .view(*orig_shape) + .to(device=hidden_states.device) + ) + # y += y_ + return y + + @maybe_no_grad() + def moe_on_cpuinfer(self, x: torch.Tensor, topk_ids: torch.Tensor, topk_weight: torch.Tensor, bsz_tensor, cuda_graph_idx=0) -> torch.Tensor: + outs = torch.empty_like(x) + outs = self.experts(x, topk_ids, topk_weight, bsz_tensor, cuda_graph_idx) + return outs + + @maybe_no_grad() + # TODO may bugs here + def moe_infer_simple( + self, x: torch.Tensor, topk_ids: torch.Tensor, topk_weight: torch.Tensor + ) -> torch.Tensor: + """ + x: [num_tokens, hidden_size] + topk_ids, topk_weight: [num_tokens, num_selected_experts] + """ + outs = torch.zeros_like(x) + for token_idx in range(topk_ids.size(0)): + for expert_idx in range(topk_ids.size(1)): + expert = self.experts[topk_ids[token_idx, expert_idx]] + outs[token_idx] += ( + expert.forward(x[token_idx]) * topk_weight[token_idx, expert_idx] + ) + return outs + + @maybe_no_grad() + # TODO may bugs here + def moe_infer(self, x, topk_ids, topk_weight): + cnts = topk_ids.new_zeros((topk_ids.shape[0], len(self.experts))) + cnts.scatter_(1, topk_ids, 1) + tokens_per_expert = cnts.sum(dim=0) + idxs = topk_ids.view(-1).argsort() + sorted_tokens = x[idxs // topk_ids.shape[1]] + tokens_per_expert = tokens_per_expert.cpu().numpy() + + outputs = [] + start_idx = 0 + for i, num_tokens in enumerate(tokens_per_expert): + end_idx = start_idx + num_tokens + if num_tokens == 0: + continue + expert = self.experts[i + self.ep_rank * self.experts_per_rank] + tokens_for_this_expert = sorted_tokens[start_idx:end_idx] + expert_out = expert.forward(tokens_for_this_expert) + outputs.append(expert_out) + start_idx = end_idx + + outs = torch.cat(outputs, dim=0) if len(outputs) else sorted_tokens.new_empty(0) + + new_x = torch.empty_like(outs) + new_x[idxs] = outs + final_out = ( + new_x.view(*topk_ids.shape, -1) + .type(topk_weight.dtype) + .mul_(topk_weight.unsqueeze(dim=-1)) + .sum(dim=1) + .type(new_x.dtype) + ) + return final_out + + +class KQwen3MoeSparseMoeBlock(BaseInjectedModule, Qwen3MoeSparseMoeBlock): + def forward(self, hidden_states): + + orig_shape = hidden_states.shape + sequence_length = orig_shape[1] + + hidden_states = hidden_states.view(-1, hidden_states.shape[-1]) + + router_logits = self.gate(hidden_states) + + if router_logits.device.type == "xpu": + from ipex_llm.transformers.models.common import moe_softmax_topk + selected_experts, routing_weights = moe_softmax_topk( + router_logits.half(), self.top_k, self.norm_topk_prob + ) + else: + routing_weights = torch.nn.functional.softmax(router_logits, dim=1, dtype=torch.float) + routing_weights, selected_experts = torch.topk(routing_weights, self.top_k, dim=-1) + if self.norm_topk_prob: + routing_weights /= routing_weights.sum(dim=-1, keepdim=True) + # we cast back to the input dtype + routing_weights = routing_weights.to(hidden_states.dtype) + + # only for generate phase + if sequence_length == 1 and hasattr(self.experts.generate_experts, + "submit_for_one_decode") and torch.cuda.is_available() and torch.cuda.is_current_stream_capturing(): # TODO: this branch cause jit bug + self.experts.generate_experts.submit_for_one_decode(hidden_states[0], selected_experts[0], + routing_weights[0]) + # y_ = self.shared_expert(hidden_states).squeeze(0) + # y_ = F.sigmoid(self.shared_expert_gate(hidden_states)) * y_ + + y = self.experts.generate_experts.sync_for_one_decode().unsqueeze(0) + + # y += y_ + y.resize_(*orig_shape) + return y + + # y_ = self.shared_expert(hidden_states).squeeze(0) + # y_ = ( + # F.sigmoid(self.shared_expert_gate(hidden_states)) * y_ + # ) + + if isinstance(self.experts, KExpertsBase): + y = self.moe_kexperts(hidden_states, selected_experts, routing_weights).view(*orig_shape).to( + device=hidden_states.device) + elif hidden_states.size(0) > 10: + # TODO may bugs here + y = ( + self.moe_infer(hidden_states, selected_experts, routing_weights) + .view(*orig_shape) + .to(device=hidden_states.device) + ) + else: + # TODO may bugs here + y = ( + self.moe_infer_simple(hidden_states, selected_experts, routing_weights) + .view(*orig_shape) + .to(device=hidden_states.device) + ) + # y += y_ + return y + + @maybe_no_grad() + def moe_kexperts(self, x: torch.Tensor, topk_ids: torch.Tensor, topk_weight: torch.Tensor) -> torch.Tensor: + outs = self.experts(x, topk_ids, topk_weight) + return outs + + @maybe_no_grad() + # TODO may bugs here + def moe_infer_simple( + self, x: torch.Tensor, topk_ids: torch.Tensor, topk_weight: torch.Tensor + ) -> torch.Tensor: + """ + x: [num_tokens, hidden_size] + topk_ids, topk_weight: [num_tokens, num_selected_experts] + """ + outs = torch.zeros_like(x) + for token_idx in range(topk_ids.size(0)): + for expert_idx in range(topk_ids.size(1)): + expert = self.experts[topk_ids[token_idx, expert_idx]] + outs[token_idx] += ( + expert.forward(x[token_idx]) * topk_weight[token_idx, expert_idx] + ) + return outs + + @maybe_no_grad() + # TODO may bugs here + def moe_infer(self, x, topk_ids, topk_weight): + cnts = topk_ids.new_zeros((topk_ids.shape[0], len(self.experts))) + cnts.scatter_(1, topk_ids, 1) + tokens_per_expert = cnts.sum(dim=0) + idxs = topk_ids.view(-1).argsort() + sorted_tokens = x[idxs // topk_ids.shape[1]] + tokens_per_expert = tokens_per_expert.cpu().numpy() + + outputs = [] + start_idx = 0 + for i, num_tokens in enumerate(tokens_per_expert): + end_idx = start_idx + num_tokens + if num_tokens == 0: + continue + expert = self.experts[i + self.ep_rank * self.experts_per_rank] + tokens_for_this_expert = sorted_tokens[start_idx:end_idx] + expert_out = expert.forward(tokens_for_this_expert) + outputs.append(expert_out) + start_idx = end_idx + + outs = torch.cat(outputs, dim=0) if len(outputs) else sorted_tokens.new_empty(0) + + new_x = torch.empty_like(outs) + new_x[idxs] = outs + final_out = ( + new_x.view(*topk_ids.shape, -1) + .type(topk_weight.dtype) + .mul_(topk_weight.unsqueeze(dim=-1)) + .sum(dim=1) + .type(new_x.dtype) + ) + return final_out diff --git a/archive/kt-sft/ktransformers/operators/flashinfer_batch_prefill_wrapper.py b/archive/kt-sft/ktransformers/operators/flashinfer_batch_prefill_wrapper.py new file mode 100644 index 0000000..287affb --- /dev/null +++ b/archive/kt-sft/ktransformers/operators/flashinfer_batch_prefill_wrapper.py @@ -0,0 +1,324 @@ +import torch +import flashinfer +import gc +try: + from flash_attn import flash_attn_with_kvcache + print("found flash_attn") + +except ImportError: + print("flash_attn not found, flashinfer unit test needed it. If you are using balance serve, ignore this.") + +from typing import Union, Optional + +def setup_seed(seed): + torch.manual_seed(seed) + torch.cuda.manual_seed_all(seed) + +setup_seed(998244353) + +torch.set_grad_enabled(False) +torch.set_default_dtype(torch.bfloat16) +global_dtype=torch.bfloat16 +global_device=torch.device("cuda",0) +torch.cuda.set_device(0) +torch.backends.cudnn.enabled =True +torch.backends.cudnn.benchmark = True + +class flashInferAttn(): + + float_workspace_buffer = None + def __init__(self, + max_batch_token, + max_batch_size, + max_pages, + device = "cuda:0", + kv_layout: str = "NHD", + use_cuda_graph: bool = False, + ) -> None: + self.device = device + self.max_batch_token = max_batch_token + self.kv_layout = kv_layout + self.use_cuda_graph = use_cuda_graph + if flashInferAttn.float_workspace_buffer is None: + flashInferAttn.float_workspace_buffer = torch.empty(max_batch_token * 1024 * 1024, dtype=torch.uint8, device=device) + self.qo_indptr_buf = torch.empty((max_batch_size+1,), dtype=torch.int32, device=device) + self.paged_kv_indptr_buf = torch.empty((max_batch_size+1,), dtype=torch.int32, device=device) + self.paged_kv_indices_buf = torch.empty((max_pages,), dtype=torch.int32, device=device) + self.paged_kv_last_page_len_buf = torch.empty((max_batch_size,), dtype=torch.int32, device=device) + self.batch_size_tensor_buf = torch.empty((1,), dtype=torch.int32, device=device) + self.num_tokens_tensor_buf = torch.empty((1,), dtype=torch.uint32, device=device) + + # TODO: custom mask + self.custom_mask_buf = None + self.qk_indptr_buf = None + self.warpper = flashinfer.BatchPrefillWithPagedKVCacheWrapper( + flashInferAttn.float_workspace_buffer, + self.kv_layout, + use_cuda_graph=self.use_cuda_graph, + qo_indptr_buf=self.qo_indptr_buf, + paged_kv_indptr_buf=self.paged_kv_indptr_buf, + paged_kv_indices_buf=self.paged_kv_indices_buf, + paged_kv_last_page_len_buf=self.paged_kv_last_page_len_buf, + backend = "fa2", + ) + + def plan(self, + qo_indptr: torch.Tensor, + paged_kv_indptr: torch.Tensor, + paged_kv_indices: torch.Tensor, + paged_kv_last_page_len: torch.Tensor, + batch_size_tensor: torch.Tensor, + num_tokens_tensor: torch.Tensor, + num_qo_heads: int, + num_kv_heads: int, + head_dim: int, + page_size: int, + causal: bool = True, + pos_encoding_mode: str = "NONE", + q_data_type: Union[str, torch.dtype] = torch.bfloat16, + kv_data_type: Optional[Union[str, torch.dtype]] = None): + + self.batch_size_tensor_buf.copy_(batch_size_tensor, non_blocking=True) + self.num_tokens_tensor_buf.copy_(num_tokens_tensor, non_blocking=True) + self.page_size = page_size + self.warpper.plan( + qo_indptr, + paged_kv_indptr, + paged_kv_indices, + paged_kv_last_page_len, + num_qo_heads, + num_kv_heads, + head_dim, + page_size, + causal = causal, + pos_encoding_mode = pos_encoding_mode, + q_data_type = q_data_type, + kv_data_type = kv_data_type + ) + + def calc_batch_indices(self, ragged_size = None): + if self.use_cuda_graph: + self.batch_indices, self.positions = flashinfer.get_batch_indices_positions( + self.qo_indptr_buf, flashinfer.get_seq_lens(self.paged_kv_indptr_buf, self.paged_kv_last_page_len_buf, self.page_size), self.batch_size_tensor_buf, self.max_batch_token) + else: + self.batch_indices, self.positions = flashinfer.get_batch_indices_positions( + self.warpper._qo_indptr_buf, flashinfer.get_seq_lens(self.warpper._paged_kv_indptr_buf, self.warpper._paged_kv_last_page_len_buf, self.page_size), self.batch_size_tensor_buf, ragged_size) + + def forward(self, q, k_cache, v_cache, k, v): + if self.use_cuda_graph: + flashinfer.page.append_paged_kv_cache(k, v, self.batch_indices, self.positions, (k_cache, v_cache), self.paged_kv_indices_buf, self.paged_kv_indptr_buf, self.paged_kv_last_page_len_buf, self.num_tokens_tensor_buf) + return self.warpper.run(q, (k_cache, v_cache)) + else: + flashinfer.page.append_paged_kv_cache(k, v, self.batch_indices, self.positions, (k_cache, v_cache), self.warpper._paged_kv_indices_buf, self.warpper._paged_kv_indptr_buf, self.warpper._paged_kv_last_page_len_buf, self.num_tokens_tensor_buf) + return self.warpper.run(q, (k_cache, v_cache)) + + +def testCudaGraph(): + + # use max batch to create buffer + batch_decode = 8 + prefill_chunk = 48 + past_kv_0 = 4090 + past_kv_1 = 4096 + raged_size = prefill_chunk + batch_decode + num_key_value_heads = 8 + head_dim = 128 + num_attention_heads = 64 + page_size = 256 + num_pages_per_seq = (past_kv_1 + page_size - 1) // page_size + total_num_pages = (num_pages_per_seq + 1) * (batch_decode + 1) + prefill_chunk // page_size + attn = flashInferAttn(raged_size, batch_decode+1, total_num_pages, use_cuda_graph=True) + + batch_size_tensor = torch.tensor([batch_decode + 1], device=global_device, dtype=torch.int32) + + k_caches = [] + v_caches = [] + ks = [] + vs = [] + qs = [] + for layer_idx in range(3): + k_caches.append(torch.randn(total_num_pages, page_size, num_key_value_heads, head_dim, device=global_device, dtype=torch.bfloat16)) + v_caches.append(torch.randn(total_num_pages, page_size, num_key_value_heads, head_dim, device=global_device, dtype=torch.bfloat16)) + ks.append(torch.randn(raged_size, num_key_value_heads, head_dim, device=global_device, dtype=torch.bfloat16)) + vs.append(torch.randn(raged_size, num_key_value_heads, head_dim, device=global_device, dtype=torch.bfloat16)) + qs.append(torch.randn(raged_size, num_attention_heads, head_dim, device=global_device, dtype=torch.bfloat16)) + + # warmup and capture small batch + past_kv_0 = 250 + past_kv_1 = 256 + num_pages_per_seq = (past_kv_1 + page_size - 1) // page_size + total_num_pages = (num_pages_per_seq + 1) * (batch_decode + 1) + prefill_chunk // page_size + q_indptr = torch.empty((batch_decode + 2,), dtype=torch.int32, device=global_device) + q_indptr[0] = 0 + q_indptr[1:] = torch.arange(prefill_chunk, prefill_chunk + batch_decode + 1, device=global_device, dtype=torch.int32) + kv_indptr = torch.arange(0, batch_decode + 2, device=global_device, dtype=torch.int32) * num_pages_per_seq + kv_indices = torch.arange(0, total_num_pages, device=global_device, dtype=torch.int32) + kv_last_page_len = torch.empty((batch_decode + 1,), dtype=torch.int32, device=global_device) + kv_last_page_len[:1+batch_decode//2] = int((past_kv_0 - 1) % page_size + 1) + kv_last_page_len[1+batch_decode//2:] = int((past_kv_1 - 1) % page_size + 1) + + print(q_indptr) + print(kv_indptr) + print(kv_indices) + print(kv_last_page_len) + attn.plan(q_indptr, + kv_indptr, + kv_indices, + kv_last_page_len, + batch_size_tensor, + num_attention_heads, + num_key_value_heads, + head_dim, + page_size, + causal = True, + pos_encoding_mode="NONE", + q_data_type=torch.bfloat16) + + attn.calc_batch_indices(raged_size) + for layer_idx in range(3): + attn.forward(qs[layer_idx], k_caches[layer_idx], v_caches[layer_idx], ks[layer_idx], vs[layer_idx]) + torch.cuda.synchronize() + + outs = [] + g = torch.cuda.CUDAGraph() + with torch.cuda.graph(g): + for layer_idx in range(3): + outs.append(attn.forward(qs[layer_idx], k_caches[layer_idx], v_caches[layer_idx], ks[layer_idx], vs[layer_idx])) + g.replay() + + kv_last_page_len[:1+batch_decode//2] = int(past_kv_0) + kv_last_page_len[1+batch_decode//2:] = int(past_kv_1) + for layer_idx in range(3): + for i in range(batch_decode + 1): + + qi = qs[layer_idx][q_indptr[i] : q_indptr[i + 1]] + o_ref_i = flash_attn_with_kvcache( + qi.unsqueeze(0), + k_caches[layer_idx], + v_caches[layer_idx], + causal=True, + block_table=kv_indices[kv_indptr[i]:kv_indptr[i+1]].unsqueeze(0), + cache_seqlens=torch.tensor([past_kv_0 if i < 1+batch_decode//2 else past_kv_1], device=global_device, dtype=torch.int32) + ) + o_i = outs[layer_idx][q_indptr[i] : q_indptr[i + 1]] + print(layer_idx, i) + torch.testing.assert_close(o_i.unsqueeze(0), o_ref_i, rtol=5e-3, atol=5e-3) + + # run another batch size use capture cuda graph + past_kv_0 = 4090 + past_kv_1 = 4096 + prefill_chunk = 24 + batch_decode = 4 + num_pages_per_seq = (past_kv_1 + page_size - 1) // page_size + total_num_pages = (num_pages_per_seq + 1) * (batch_decode + 1) + prefill_chunk // page_size + batch_size_tensor = torch.tensor([batch_decode + 1], device=global_device, dtype=torch.int32) + num_tokens_tensor = torch.tensor([batch_decode + prefill_chunk], device=global_device, dtype=torch.int32) + + q_indptr = torch.empty((batch_decode + 2,), dtype=torch.int32, device=global_device) + q_indptr[0] = 0 + q_indptr[1:] = torch.arange(prefill_chunk, prefill_chunk + batch_decode + 1, device=global_device, dtype=torch.int32) + kv_indptr = torch.arange(0, batch_decode + 2, device=global_device, dtype=torch.int32) * num_pages_per_seq + kv_indices = torch.arange(0, total_num_pages, device=global_device, dtype=torch.int32) + kv_last_page_len = torch.empty((batch_decode + 1,), dtype=torch.int32, device=global_device) + kv_last_page_len[:1+batch_decode//2] = int((past_kv_0 - 1) % page_size + 1) + kv_last_page_len[1+batch_decode//2:] = int((past_kv_1 - 1) % page_size + 1) + attn.plan(q_indptr, + kv_indptr, + kv_indices, + kv_last_page_len, + batch_size_tensor, + num_attention_heads, + num_key_value_heads, + head_dim, + page_size, + causal = True, + pos_encoding_mode="NONE", + q_data_type=torch.bfloat16) + attn.calc_batch_indices(raged_size) + g.replay() + + kv_last_page_len[:1+batch_decode//2] = int(past_kv_0) + kv_last_page_len[1+batch_decode//2:] = int(past_kv_1) + for layer_idx in range(3): + for i in range(batch_decode + 1): + + qi = qs[layer_idx][q_indptr[i] : q_indptr[i + 1]] + o_ref_i = flash_attn_with_kvcache( + qi.unsqueeze(0), + k_caches[layer_idx], + v_caches[layer_idx], + causal=True, + block_table=kv_indices[kv_indptr[i]:kv_indptr[i+1]].unsqueeze(0), + cache_seqlens=torch.tensor([past_kv_0 if i < 1+batch_decode//2 else past_kv_1], device=global_device, dtype=torch.int32) + ) + o_i = outs[layer_idx][q_indptr[i] : q_indptr[i + 1]] + print(layer_idx, i) + torch.testing.assert_close(o_i.unsqueeze(0), o_ref_i, rtol=5e-3, atol=5e-3) + + + +def testAttentionFlashInfer( + ): + batch_decode = 32 + prefill_chunk = 64 + past_kv_0 = 510 + past_kv_1 = 512 + raged_size = prefill_chunk + batch_decode + num_key_value_heads = 8 + head_dim = 128 + num_attention_heads = 64 + cases = 1 + page_size = 32 + num_pages_per_seq = (past_kv_1 + page_size - 1) // page_size + total_num_pages = (num_pages_per_seq + 1) * (batch_decode + 1) + prefill_chunk // page_size + workspace_buffer = torch.empty(128 * 1024 * 1024, dtype=torch.uint8, device="cuda:0") + qs = [] + kvs = [] + q_indptrs = [] + kv_indptrs = [] + kv_indicess = [] + kv_last_page_lens = [] + wrappers = [] + for case_id in range(cases): + kvs.append(torch.randn(total_num_pages, 2, page_size, num_key_value_heads, head_dim, device=global_device, dtype=torch.bfloat16)) + qs.append(torch.randn(raged_size, num_attention_heads, head_dim, device=global_device, dtype=torch.bfloat16)) + q_indptr = torch.empty((batch_decode + 2,), dtype=torch.int32, device=global_device) + q_indptr[0] = 0 + q_indptr[1:] = torch.arange(prefill_chunk, prefill_chunk + batch_decode + 1, device=global_device, dtype=torch.int32) + q_indptrs.append(q_indptr) + kv_indptrs.append(torch.arange(0, batch_decode + 2, device=global_device, dtype=torch.int32) * num_pages_per_seq) + kv_indicess.append(torch.arange(0, total_num_pages, device=global_device, dtype=torch.int32)) + kv_last_page_len = torch.empty((batch_decode + 1,), dtype=torch.int32, device=global_device) + kv_last_page_len[:1+batch_decode//2] = int((past_kv_0 - 1) % page_size + 1) + kv_last_page_len[1+batch_decode//2:] = int((past_kv_1 - 1) % page_size + 1) + kv_last_page_lens.append(kv_last_page_len) + wrappers.append(flashinfer.BatchPrefillWithPagedKVCacheWrapper( + workspace_buffer, + "NHD", + use_cuda_graph=True, + qo_indptr_buf=q_indptrs[case_id], + paged_kv_indptr_buf=kv_indptrs[case_id], + paged_kv_indices_buf=kv_indicess[case_id], + paged_kv_last_page_len_buf=kv_last_page_lens[case_id], + )) + wrappers[case_id].plan( + q_indptrs[case_id], + kv_indptrs[case_id], + kv_indicess[case_id], + kv_last_page_lens[case_id], + num_attention_heads, + num_key_value_heads, + head_dim, + page_size, + causal = True, + pos_encoding_mode="ROPE_LLAMA", + q_data_type=torch.bfloat16 + ) + + def custom_forward(case_id): + out = wrappers[case_id].run(qs[case_id], kvs[case_id]) + + custom_forward(0) + +# testCudaGraph() +# pass \ No newline at end of file diff --git a/archive/kt-sft/ktransformers/operators/flashinfer_wrapper.py b/archive/kt-sft/ktransformers/operators/flashinfer_wrapper.py new file mode 100644 index 0000000..81fd75e --- /dev/null +++ b/archive/kt-sft/ktransformers/operators/flashinfer_wrapper.py @@ -0,0 +1,481 @@ +''' +Description : flashinfer MLA wrapper +Author : Boxin Zhang +Version : 0.2.3 +''' +import torch +import os +from ktransformers.operators.triton_attention import decode_attention_fwd_grouped + +flashinfer_enabled = False + +try: + import flashinfer + flashinfer_enabled = True + print("found flashinfer") + +except ImportError: + print("flashinfer not found, use triton for linux") + +import math + +def attention_ref_torch( + batch_size, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + causal: bool, + sm_scale: float, +) -> torch.Tensor: + qo_len = q.shape[0] // batch_size + kv_len = k.shape[0] // batch_size + num_qo_heads = q.shape[1] + head_dim_qk = q.shape[2] + head_dim_vo = v.shape[2] + logits = ( + torch.einsum( + "bmhd,bnhd->bhmn", + q.view(batch_size, qo_len, num_qo_heads, head_dim_qk).float(), + k.view(batch_size, kv_len, num_qo_heads, head_dim_qk).float(), + ) + * sm_scale + ) + + #print("attn weights", logits) + + if causal: + mask = ( + torch.arange(kv_len - qo_len, kv_len).unsqueeze(1) + >= torch.arange(0, kv_len).unsqueeze(0) + ).to(q.device) + else: + mask = torch.ones(qo_len, kv_len).to(q.device) + + logits = logits.masked_fill(mask.unsqueeze(0).unsqueeze(0) == 0, float("-inf")) + lse_ref = torch.logsumexp(logits, -1).transpose(-1, -2) + p = torch.softmax(logits, dim=-1) + o_ref = ( + torch.einsum( + "bhmn,bnhd->bmhd", + p, + v.view(batch_size, kv_len, num_qo_heads, head_dim_vo).float(), + ) + .contiguous() + .view(batch_size * qo_len, num_qo_heads, head_dim_vo) + .to(q) + ) + + return o_ref, lse_ref * math.log2(math.e) + +class MLAWrapper(): + def __init__(self, + max_batch_size, + max_pages, + use_cuda_graph = True, + device = "cuda", + ): + self.float_workspace_buffer = torch.empty(128*1024*1024, dtype=torch.int8, device=device) + self.max_batch_size = max_batch_size + self.max_pages = max_pages + if use_cuda_graph: + if self.max_batch_size == 1: + self.qo_indptr_buf = torch.arange(0, max_batch_size+1, dtype=torch.int32, device=device) + self.kv_indptr_buf = torch.tensor([0, max_pages], dtype=torch.int32, device=device) + self.kv_indices_buf = torch.arange(0, max_pages, dtype=torch.int32, device=device) + else: + self.qo_indptr_buf = torch.empty(max_batch_size+1, dtype=torch.int32, device=device) + self.kv_indptr_buf = torch.empty(max_batch_size+1, dtype=torch.int32, device=device) + self.kv_indices_buf = torch.empty(max_pages, dtype=torch.int32, device=device) + self.batch_size_tensor_buf = torch.tensor([self.max_batch_size], dtype=torch.int32, device=device) + self.kv_len_arr_buf = torch.empty(max_batch_size, dtype=torch.int32, device=device) + else: + self.qo_indptr_buf = None + self.kv_indptr_buf = None + self.kv_indices_buf = None + self.kv_len_arr_buf = None + self.wrapper = flashinfer.mla.BatchMLAPagedAttentionWrapper( + self.float_workspace_buffer, + use_cuda_graph=use_cuda_graph, + qo_indptr=self.qo_indptr_buf, + kv_indptr=self.kv_indptr_buf, + kv_indices=self.kv_indices_buf, + kv_len_arr=self.kv_len_arr_buf, + bsz_tensor=self.batch_size_tensor_buf, + backend = "fa2", + ) + self.need_plan = True + + + def plan(self, + qo_indptr, + kv_indptr, + kv_indices, + kv_len_arr, + bsz_tensor, + num_heads, + head_dim_ckv, + head_dim_kpe, + page_size, + sm_scale, + q_data_type, + kv_data_type, + ): + if qo_indptr is None: + assert self.max_batch_size == 1 + qo_indptr = self.qo_indptr_buf + if kv_indptr is None: + assert self.max_batch_size == 1 + kv_indptr = self.kv_indptr_buf + if kv_indices is None: + assert self.max_batch_size == 1 + kv_indices = self.kv_indices_buf + if bsz_tensor is None: + assert self.max_batch_size == 1 + bsz_tensor = self.batch_size_tensor_buf + + self.wrapper.plan( + qo_indptr, + kv_indptr, + kv_indices, + kv_len_arr, + num_heads, + head_dim_ckv, + head_dim_kpe, + page_size, + True, # causal + sm_scale, + q_data_type, + kv_data_type, + bsz_tensor + ) + + def run(self, q_nope, q_pe, ckv, k_pe, return_lse = False): + return self.wrapper.run(q_nope, q_pe, ckv, k_pe, return_lse = return_lse) + +class MLAWrapperSingleton(): + wrappers:dict = {} + + @classmethod + def get_instance(cls, device, *args, **kwargs)->MLAWrapper: + if device not in cls.wrappers: + cls.make_instance(device, *args, **kwargs) + return cls.wrappers[device] + + @classmethod + def make_instance(cls, device, *args, **kwargs): + cls.wrappers[device] = MLAWrapper(*args, **kwargs, device=device) + + @classmethod + def plan_all(cls, qo_indptr, + kv_indptr, + kv_indices, + kv_len_arr, + bsz_tensor, + num_heads, + head_dim_ckv, + head_dim_kpe, + page_size, + sm_scale, + q_data_type, + kv_data_type,): + for device, wrapper in cls.wrappers.items(): + kv_len_arr_cur_device = kv_len_arr.to(device) + wrapper.plan(qo_indptr, + kv_indptr, + kv_indices, + kv_len_arr_cur_device, + bsz_tensor, + num_heads, + head_dim_ckv, + head_dim_kpe, + page_size, + sm_scale, + q_data_type, + kv_data_type,) + wrapper.need_plan = False + + @classmethod + def need_plan_all(cls): + for device, wrapper in cls.wrappers.items(): + wrapper.need_plan = True + + @classmethod + def reset_buffer(cls): + for device, wrapper in cls.wrappers.items(): + wrapper.qo_indptr_buf[1] = 1 # assert max_batch_size=1 here. + + @classmethod + def update_buffer(cls, max_pages): + for device, wrapper in cls.wrappers.items(): + wrapper.kv_indptr_buf[1] = max_pages # assert max_batch_size=1 here. + wrapper.kv_indices_buf = torch.arange(0, max_pages, dtype=torch.int32, device=device) + wrapper.wrapper._kv_indices_buf = wrapper.kv_indices_buf + +def checksame(): + flashinfer_folder = "./flashinfer_output" + flashinfer_folder = "./kv_cache_flashinfer" + triton_folder = "./triton_output" + triton_folder = "./kv_cache_triton" + + max_layer_id = 1 + max_forward_id = 2 + + for forward_id in range(0, 19): + print("forward_id", forward_id) + for layer_id in range(max_layer_id): + print(layer_id) + #file_name = f"layer_{layer_id}_forward_{forward_id}_attn_output.pt" + #file_name = f"layer_{layer_id}_forward_{forward_id}_q_pe.pt" + file_name = f"layer_{layer_id}.pt" + + flashinfer_path = os.path.join(flashinfer_folder, file_name) + triton_path = os.path.join(triton_folder, file_name) + + if not os.path.exists(triton_path): + print(f"{file_name} not exist in {triton_folder}") + continue + if not os.path.exists(flashinfer_path): + print(f"{file_name} not exist in {flashinfer_folder}") + continue + + + flashinfer_tensor = torch.load(flashinfer_path)[1:2, :62]# + triton_tensor = torch.load(triton_path)[1:2, :62]#.squeeze(1)# + try: + torch.testing.assert_close(flashinfer_tensor, triton_tensor, rtol=1e-9, atol=1e-9) + except AssertionError as e: + print(e) + +if __name__ == "__main__": + + #checksame() + #exit(0) + + max_batch_size = 2 + max_batch_tokens = 256 + max_pages = 128 + page_size = 64 + num_heads = 128 + + # warm-up + kv_len = 4023 + q_len = 1 + q_nope_buf = torch.randn((max_batch_tokens, num_heads, 512), dtype=torch.bfloat16, device="cuda") + q_pe_buf = torch.randn((max_batch_tokens, num_heads, 64), dtype=torch.bfloat16, device="cuda") + kv_buf = torch.randn((max_pages, page_size, 576), dtype=torch.bfloat16, device="cuda") + ckv, k_pe = torch.split(kv_buf, [512, 64], dim=-1) + + + wrapper = MLAWrapperSingleton.get_instance( + "cuda", + max_batch_size, + max_pages, + ) + + used_pages = (kv_len + page_size - 1)// page_size + kv_len_arr = torch.tensor([kv_len], dtype=torch.int32, device="cuda") + qo_indptr = torch.tensor([0, q_len], dtype=torch.int32, device="cuda") + kv_indptr = torch.tensor([0, used_pages], dtype=torch.int32, device="cuda") + kv_indices = torch.empty(max_pages, dtype=torch.int32, device="cuda") + kv_indices[:used_pages] = torch.arange(0, used_pages, dtype=torch.int32, device="cuda") + bsz_tensor = torch.tensor([1], dtype=torch.int32, device="cuda") + wrapper.plan( + qo_indptr, + kv_indptr, + kv_indices, + kv_len_arr, + bsz_tensor, + 128, + 512, + 64, + page_size, + 192 ** (-0.5), + torch.bfloat16, + torch.bfloat16, + ) + + attn_output = wrapper.run(q_nope_buf[:q_len], q_pe_buf[:q_len], ckv, k_pe) + print(attn_output.shape) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + attn_output = wrapper.run(q_nope_buf, q_pe_buf, ckv, k_pe) + graph.replay() + + q = torch.cat([q_nope_buf, q_pe_buf], dim=-1) + k = ( + torch.cat([ckv, k_pe], dim=-1) + .view(-1, 1, 512 + 64) + .repeat_interleave(num_heads, dim=1) + ) + v = ckv.view(-1, 1, 512).repeat_interleave(num_heads, dim=1) + attn_ref, lse_ref = attention_ref_torch( + 1, + q[:q_len], + k[:kv_len], + v[:kv_len], + True, + 192 ** (-0.5) + ) + torch.testing.assert_close(attn_output[:q_len], attn_ref, rtol=5e-3, atol=5e-3) + # warm-up finished + + kv_len = 512 + q_len = 128 + pages = max_pages + used_pages = (kv_len + page_size - 1)// page_size + q_nope = torch.randn((q_len*2, num_heads, 512), dtype=torch.bfloat16, device="cuda") + q_nope[q_len:] = q_nope[:q_len] + q_pe = torch.randn((q_len*2, num_heads, 64), dtype=torch.bfloat16, device="cuda") + q_pe[q_len:] = q_pe[:q_len] + kv_cache = torch.randn((max_pages, page_size, 576), dtype=torch.bfloat16, device="cuda") + kv_cache[used_pages:2*used_pages] = kv_cache[:used_pages] + ckv, k_pe = torch.split(kv_cache, [512, 64], dim=-1) + + kv_len_arr = torch.tensor([kv_len, kv_len], dtype=torch.int32, device="cuda") + qo_indptr = torch.tensor([0, q_len, q_len*2], dtype=torch.int32, device="cuda") + kv_indptr = torch.tensor([0, used_pages, used_pages*2], dtype=torch.int32, device="cuda") + kv_indices = torch.empty(max_pages, dtype=torch.int32, device="cuda") + kv_indices[:2*used_pages] = torch.arange(0, 2*used_pages, dtype=torch.int32, device="cuda") + bsz_tensor = torch.tensor([2], dtype=torch.int32, device="cuda") + wrapper.plan( + qo_indptr, + kv_indptr, + kv_indices, + kv_len_arr, + bsz_tensor, + 128, + 512, + 64, + page_size, + 192 ** (-0.5), + torch.bfloat16, + torch.bfloat16, + ) + + q_nope_buf.copy_(q_nope) + q_pe_buf.copy_(q_pe) + kv_buf[:pages].copy_(kv_cache) + + torch.cuda.synchronize() + graph.replay() + torch.cuda.synchronize() + + # ref_torch + q = torch.cat([q_nope, q_pe], dim=-1) + k = ( + torch.cat([ckv, k_pe], dim=-1) + .view(-1, 1, 512 + 64) + .repeat_interleave(num_heads, dim=1) + ) + v = ckv.view(-1, 1, 512).repeat_interleave(num_heads, dim=1) + attn_ref, lse_ref = attention_ref_torch( + max_batch_size, + q, + k[:2*kv_len], + v[:2*kv_len], + True, + 192 ** (-0.5) + ) + + torch.testing.assert_close(attn_ref[:q_len], attn_ref[q_len:q_len*2], rtol=1e-9, atol=1e-9) + torch.testing.assert_close(attn_output[:q_len], attn_output[q_len:q_len*2], rtol=1e-9, atol=1e-9) + torch.testing.assert_close(attn_output[:q_len], attn_ref[:q_len], rtol=5e-3, atol=5e-3) + torch.testing.assert_close(attn_output[q_len:q_len*2], attn_ref[q_len:q_len*2], rtol=5e-3, atol=5e-3) + #torch.testing.assert_close(attn_output[:q_len], attn_output[q_len:q_len*2], rtol=1e-9, atol=1e-9) + #torch.testing.assert_close(attn_output, attn_ref, rtol=5e-3, atol=5e-3) + + exit(0) + + for forward_id in range(0, 1): + print("forward_id", forward_id) + for layer_id in range(1): + print(layer_id) + flashinfer_folder = "./kv_cache_flashinfer" + forward_id = 17 + layer_id = 0 + file_name = f"layer_{layer_id}.pt" + kv_cache_path = os.path.join(flashinfer_folder, file_name) + flashinfer_folder = "./flashinfer_output" + + q_len = 1 + kv_len = 126 + file_name = f"layer_{layer_id}_forward_{forward_id}_q_nope.pt" + q_nope = torch.load(os.path.join(flashinfer_folder, file_name)).view(q_len,128,512).to(device="cuda") + file_name = f"layer_{layer_id}_forward_{forward_id}_q_pe.pt" + q_pe = torch.load(os.path.join(flashinfer_folder, file_name)).view(q_len,128,64).to(device="cuda") + q = torch.cat([q_nope, q_pe], dim=-1) + kv_cache = torch.load(kv_cache_path).to(device="cuda") + pages, page_size, _, head_dim = kv_cache.shape + kv_cache = kv_cache.view(pages, page_size, head_dim) + ckv, k_pe = torch.split(kv_cache, [512, 64], dim=-1) + + kv_len_arr = torch.tensor([kv_len], dtype=torch.int32, device="cuda") + qo_indptr = torch.tensor([0, q_len], dtype=torch.int32, device="cuda") + wrapper.plan( + None, + None, + None, + kv_len_arr, + 128, + 512, + 64, + page_size, + 192 ** (-0.5), + torch.bfloat16, + torch.bfloat16, + ) + + q_nope_buf.copy_(q_nope) + q_pe_buf.copy_(q_pe) + kv_buf[:pages].copy_(kv_cache) + + torch.cuda.synchronize() + graph.replay() + torch.cuda.synchronize() + + # ref_torch + k = ( + torch.cat([ckv, k_pe], dim=-1) + .view(-1, 1, 512 + 64) + .repeat_interleave(num_heads, dim=1) + ) + v = ckv.view(-1, 1, 512).repeat_interleave(num_heads, dim=1) + attn_ref, lse_ref = attention_ref_torch( + max_batch_size, + q, + k[:kv_len], + v[:kv_len], + False, + 192 ** (-0.5) + ) + torch.testing.assert_close(attn_output, attn_ref, rtol=1e-3, atol=1e-3) + + # ref_triton + attn_logits = torch.empty( + ( + max_batch_size, + num_heads, + 4, #num_kv_splits # follow vLLM, fix it TODO + 512 + 1, + ), + dtype=torch.float32, + device = "cuda" + ) + + triton_ref = torch.zeros_like(q_nope) + page_table = torch.arange(max_pages, dtype=torch.int32, device="cuda") + ckv_with_pe = torch.cat([ckv, k_pe], dim=-1).contiguous().view(pages, page_size, 1, 576) + ckv = ckv.view(pages, page_size, 1, 512) + decode_attention_fwd_grouped(q, ckv_with_pe, ckv, triton_ref, + page_table, + kv_len_arr, attn_logits, + 4, #num_kv_splits # follow vLLM, fix it TODO + 192 ** (-0.5), + page_size) + + torch.testing.assert_close(attn_output, triton_ref, rtol=1e-3, atol=1e-3) + + #file_name = f"./flashinfer_output/layer_{layer_id}_forward_{forward_id}_attn_output.pt" + #ktrans_output = torch.load(file_name) + #torch.testing.assert_close(attn_output, ktrans_output.squeeze(1), rtol=1e-3, atol=1e-3) + print("test past") \ No newline at end of file diff --git a/archive/kt-sft/ktransformers/operators/gate.py b/archive/kt-sft/ktransformers/operators/gate.py new file mode 100644 index 0000000..f5f96c1 --- /dev/null +++ b/archive/kt-sft/ktransformers/operators/gate.py @@ -0,0 +1,215 @@ +from typing import Optional +from torch import nn +import torch +import torch.nn.functional as F +import os +from ktransformers.operators.base_operator import BaseInjectedModule +from ktransformers.operators.base_operator import BaseInjectedModule +from ktransformers.operators.linear import KTransformersLinear +from ktransformers.util.custom_loader import GGUFLoader, ModelLoader, SafeTensorLoader +from transformers.configuration_utils import PretrainedConfig +from abc import ABC, abstractmethod + + +# class Base(BaseInjectedModule, ABC): +class KMoEGateBase(ABC): + def __init__(self, + key: str, + gguf_loader: GGUFLoader, + config: PretrainedConfig, + orig_module: nn.Module, + device: str = "cuda", + **kwargs): + # super().__init__(key, gguf_loader, config, orig_module, device, **kwargs) + super().__init__() + self.key = key + self.gguf_loader = gguf_loader + self.config = config + self.device = device + self.orig_module = orig_module + + @abstractmethod + def forward(self, input_tensor, expert_ids, weights): + pass + + @abstractmethod + def load(self, w: dict | nn.Parameter | tuple | None = None, device: str = "cpu", warmup: bool = False): + pass + + @abstractmethod + def unload(): + pass + + def load_weights(self, override_key: str | None = None, device: str = "cpu"): + res = {} + if override_key is not None: + keys = override_key + else: + keys = [self.key] + + gate = None + up = None + down = None + gate_type = None + up_type = None + down_type = None + + for key in keys: + # key = ".".join(key.split(".")[:-1]) + if isinstance(self.gguf_loader, SafeTensorLoader): + res = self.gguf_loader.load_gate(key, device=device) + elif self.gguf_loader.has_tensor(key+".weight"): + # targets = [".ffn_gate_inp.weight", ".exp_probs_b.bias"] + targets = [".weight", ".e_score_correction_bias"] + tensors = self.load_multi(key, targets, device=device) + weight = tensors[".weight"] + e_score_correction_bias = tensors[".e_score_correction_bias"] + # weight_type = self.gguf_loader.tensor_info[key + ".weight"]["ggml_type"] + res = {"weight": weight, "e_score_correction_bias": e_score_correction_bias} + else: + raise ValueError(f"Experts {key} not found in gguf_loader") + + return res + + def load_multi(self, key: str, keys: list[str], device: str = "cpu"): + tensors = {} + for k in keys: + tensors[k] = self.gguf_loader.load_gguf_tensor(key + k, device=device) + return tensors + + +class KMoEGate(BaseInjectedModule, KMoEGateBase): + def __init__( + self, + key: str, + gguf_loader: GGUFLoader, + config: PretrainedConfig, + orig_module: nn.Module = None, + generate_device: str = "cuda", + prefill_device: str = "cuda", + **kwargs, + ): + BaseInjectedModule.__init__(self, key, gguf_loader, config, orig_module, prefill_device, generate_device, **kwargs) + KMoEGateBase.__init__(self, key, gguf_loader, config, orig_module, generate_device, **kwargs) + self.generate_device = generate_device + self.prefill_device = prefill_device + + def forward(self, hidden_states) -> torch.Tensor: + return self.orig_module.forward(hidden_states) + + def load(self, w: dict | nn.Parameter | tuple | None = None, device: str|None = None): + if device is None: device = self.device + if w is None: w = self.load_weights(device=device) + + if isinstance(w, dict): + self.orig_module.weight = nn.Parameter(w["weight"]) + self.orig_module.e_score_correction_bias = nn.Parameter(w["e_score_correction_bias"]) + else: + raise ValueError("Invalid weight type") + self.orig_module.weight = nn.Parameter(self.orig_module.weight.to(device)) + self.orig_module.e_score_correction_bias = nn.Parameter(self.orig_module.e_score_correction_bias.to(device)) + + def unload(self): + if self.weight is not None: + self.weight = None + if self.e_score_correction_bias is not None: + self.e_score_correction_bias = None + + +class KMoEGateQwen2Moe(BaseInjectedModule, KMoEGateBase): + def __init__( + self, + key: str, + gguf_loader: GGUFLoader, + config: PretrainedConfig, + orig_module: nn.Module = None, + generate_device: str = "cuda", + generate_op: str| None = "KLinearMarlin", + prefill_device: str = "cuda", + prefill_op: str| None = "KLinearMarlin", + use_quant: bool = False, + **kwargs, + ): + BaseInjectedModule.__init__(self, key, gguf_loader, config, orig_module, prefill_device, generate_device, **kwargs) + KMoEGateBase.__init__(self, key, gguf_loader, config, orig_module, generate_device, **kwargs) + self.generate_device = generate_device + self.prefill_device = prefill_device + self.generate_op = generate_op + self.prefill_op = prefill_op + self.is_windows = os.name == 'nt' + self.use_quant = use_quant + if not self.is_windows and use_quant: + self.gate_linear = nn.Linear(self.gating_dim, self.n_routed_experts, device=generate_device) + self.gate_linear = KTransformersLinear(key + ".ffn_gate_inp", + gguf_loader, config, self.gate_linear, #orig_module + generate_device, generate_op, prefill_device, prefill_op) + else: + self.gate_linear = None + + def forward(self, hidden_states) -> torch.Tensor: + if self.is_windows: + return self.orig_module.forward(hidden_states) + + bsz, seq_len, h = hidden_states.shape + ### compute gating score + hidden_states = hidden_states.view(-1, h) + if self.use_quant: + logits = self.gate_linear.forward(logits) + else: + logits = F.linear( + hidden_states.type(torch.float32), self.weight.type(torch.float32), None + ) + + return grouped_topk(hidden_states, logits, + self.top_k, self.norm_topk_prob, + self.n_group, self.topk_group) + + def load(self, w: dict | nn.Parameter | tuple | None = None, device: str|None = None): + if device is None: device = self.device + if w is None: w = self.load_weights(device=device) + + if isinstance(w, dict): + self.orig_module.weight = nn.Parameter(w["weight"]) + self.orig_module.e_score_correction_bias = nn.Parameter(w["e_score_correction_bias"]) + else: + raise ValueError("Invalid weight type") + self.orig_module.weight = nn.Parameter(self.orig_module.weight.to(device)) + self.orig_module.e_score_correction_bias = nn.Parameter(self.orig_module.e_score_correction_bias.to(device)) + if not self.is_windows and self.use_quant: + self.gate_linear.load(self.orig_module.weight) + + def unload(self): + if self.weight is not None: + self.weight = None + if self.e_score_correction_bias is not None: + self.e_score_correction_bias = None + + +class KMoEGateIPEXLLM(KMoEGate): + def __init__( + self, + key: str, + gguf_loader: GGUFLoader, + config: PretrainedConfig, + orig_module: nn.Module = None, + generate_device: str = "xpu", + prefill_device: str = "xpu", + **kwargs, + ): + BaseInjectedModule.__init__(self, key, gguf_loader, config, orig_module, prefill_device, generate_device, **kwargs) + KMoEGate.__init__(self, key, gguf_loader, config, orig_module, generate_device, **kwargs) + self.generate_device = generate_device + self.prefill_device = prefill_device + + def forward(self, hidden_states) -> torch.Tensor: + x = hidden_states.view(-1, hidden_states.size(-1)) + logits = torch.nn.functional.linear( + x.type(torch.float32), self.orig_module.weight.type(torch.float32), None + ) + scores = logits.sigmoid() + + from ipex_llm.transformers.models.common import moe_group_topk + topk_idx, topk_weight = moe_group_topk(scores, self.orig_module.e_score_correction_bias, + self.n_group, self.topk_group, self.top_k, + self.norm_topk_prob, self.routed_scaling_factor) + return topk_idx, topk_weight.to(x.dtype) \ No newline at end of file diff --git a/archive/kt-sft/ktransformers/operators/layernorm.py b/archive/kt-sft/ktransformers/operators/layernorm.py new file mode 100644 index 0000000..796592c --- /dev/null +++ b/archive/kt-sft/ktransformers/operators/layernorm.py @@ -0,0 +1,225 @@ +''' +Date: 2024-11-13 15:05:52 +LastEditors: Xie Weiyu ervinxie@qq.com +LastEditTime: 2024-11-25 08:59:19 +''' +""" +Copyright 2023-2024 SGLang Team +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +""" + +"""Fused operators for normalization layers.""" + +import logging +from typing import Optional, Tuple, Union +from transformers import PretrainedConfig +import torch +import torch.nn as nn +from ktransformers.models.modeling_deepseek_v3 import DeepseekV3RMSNorm +from ktransformers.models.modeling_qwen2_moe import Qwen2MoeRMSNorm +from ktransformers.models.modeling_qwen3_moe import Qwen3MoeRMSNorm +from ktransformers.operators.base_operator import BaseInjectedModule +from ktransformers.util.custom_loader import GGUFLoader +if not torch.xpu.is_available(): + from flashinfer.norm import ( + fused_add_rmsnorm, + rmsnorm, + ) + + +logger = logging.getLogger(__name__) + + +class RMSNorm(DeepseekV3RMSNorm, BaseInjectedModule): + def __init__(self, + key: str, + gguf_loader : GGUFLoader, + config: PretrainedConfig, + orig_module: nn.Module, + prefill_device: str = "cuda", + generate_device: str = "cuda", + **kwargs): + BaseInjectedModule.__init__(self, key, gguf_loader, config, orig_module, prefill_device, **kwargs) + self.orig_module.__init__(orig_module.hidden_size, + orig_module.variance_epsilon) + + def forward( + self, + x: torch.Tensor, + batch_size_tensor: torch.Tensor = None, + residual: Optional[torch.Tensor] = None, + ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]: + #return self.forward_native(x, residual) + if batch_size_tensor is None: + return self.forward_native(x) + if residual is not None: + fused_add_rmsnorm(x, residual, self.weight.data, batch_size_tensor, self.variance_epsilon) + #residual = x + residual + #out = rmsnorm(residual, self.weight.data, batch_size_tensor, self.variance_epsilon) + return x, residual + # print(x.shape, self.weight.data.shape, self.variance_epsilon, x.dtype, self.weight.data.dtype, x.device, self.weight.device, x.is_contiguous(), self.weight.data.is_contiguous()) + out = rmsnorm(x, self.weight.data, batch_size_tensor,self.variance_epsilon) + return out + + def forward_native( + self, hidden_states + ): + input_dtype = hidden_states.dtype + hidden_states = hidden_states.to(torch.float32) + variance = hidden_states.pow(2).mean(-1, keepdim=True) + hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon) + return self.weight * hidden_states.to(input_dtype) + + +class KQwen2MoeRMSNorm(Qwen2MoeRMSNorm, BaseInjectedModule): + def __init__(self, + key: str, + gguf_loader : GGUFLoader, + config: PretrainedConfig, + orig_module: nn.Module, + prefill_device: str = "cuda", + generate_device: str = "cuda", + **kwargs): + BaseInjectedModule.__init__(self, key, gguf_loader, config, orig_module, prefill_device, **kwargs) + self.orig_module.__init__(config.hidden_size, + orig_module.variance_epsilon) + + def forward( + self, + x: torch.Tensor, + batch_size_tensor: torch.Tensor = None, + residual: Optional[torch.Tensor] = None, + ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]: + #return self.forward_native(x, residual) + if batch_size_tensor is None: + return self.forward_native(x) + if residual is not None: + fused_add_rmsnorm(x, residual, self.weight.data, batch_size_tensor, self.variance_epsilon) + #residual = x + residual + #out = rmsnorm(residual, self.weight.data, batch_size_tensor, self.variance_epsilon) + return x, residual + # print(x.shape, self.weight.data.shape, self.variance_epsilon, x.dtype, self.weight.data.dtype, x.device, self.weight.device, x.is_contiguous(), self.weight.data.is_contiguous()) + out = rmsnorm(x, self.weight.data, batch_size_tensor,self.variance_epsilon) + return out + + def forward_native( + self, hidden_states + ): + input_dtype = hidden_states.dtype + hidden_states = hidden_states.to(torch.float32) + variance = hidden_states.pow(2).mean(-1, keepdim=True) + hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon) + return self.weight * hidden_states.to(input_dtype) + + +class KQwen3MoeRMSNorm(Qwen3MoeRMSNorm, BaseInjectedModule): + def __init__(self, + key: str, + gguf_loader : GGUFLoader, + config: PretrainedConfig, + orig_module: nn.Module, + prefill_device: str = "cuda", + generate_device: str = "cuda", + **kwargs): + BaseInjectedModule.__init__(self, key, gguf_loader, config, orig_module, prefill_device, **kwargs) + self.orig_module.__init__(orig_module.hidden_size, + orig_module.variance_epsilon) + + def forward( + self, + x: torch.Tensor, + batch_size_tensor: torch.Tensor = None, + residual: Optional[torch.Tensor] = None, + ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]: + #return self.forward_native(x, residual) + bsz, hidden_size = x.shape + x = x.view(-1, self.orig_module.hidden_size) + if batch_size_tensor is None: + return self.forward_native(x) + if residual is not None: + fused_add_rmsnorm(x, residual, self.weight.data, batch_size_tensor, self.variance_epsilon) + #residual = x + residual + #out = rmsnorm(residual, self.weight.data, batch_size_tensor, self.variance_epsilon) + return x, residual + # print(x.shape, self.weight.data.shape, self.variance_epsilon, x.dtype, self.weight.data.dtype, x.device, self.weight.device, x.is_contiguous(), self.weight.data.is_contiguous()) + out = rmsnorm(x, self.weight.data, batch_size_tensor,self.variance_epsilon) + out = out.view(bsz, hidden_size) + return out + + def forward_native( + self, hidden_states + ): + input_dtype = hidden_states.dtype + hidden_states = hidden_states.to(torch.float32) + variance = hidden_states.pow(2).mean(-1, keepdim=True) + hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon) + return self.weight * hidden_states.to(input_dtype) + +class DeepseekV3RMSNormTorch(DeepseekV3RMSNorm, BaseInjectedModule): + def __init__(self, + key: str, + gguf_loader : GGUFLoader, + config: PretrainedConfig, + orig_module: nn.Module, + prefill_device: str = "cuda", + generate_device: str = "cuda", + **kwargs): + BaseInjectedModule.__init__(self, key, gguf_loader, config, orig_module, prefill_device, **kwargs) + self.orig_module.__init__(orig_module.hidden_size, + orig_module.variance_epsilon) + + def forward( + self, + x, + batch_size_tensor: torch.Tensor = None, + residual: Optional[torch.Tensor] = None, + )-> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]: + if residual is not None: + x = x + residual + residual = x + # range batch_size_tensor for x + input_dtype = x.dtype + x = x.to(torch.float32) + variance = x.pow(2).mean(-1, keepdim=True) + x = x * torch.rsqrt(variance + self.variance_epsilon) + if residual is not None: + return self.weight * x.to(input_dtype), residual + return self.weight * x.to(input_dtype) + + +class KDeepseekRMSNormIPEXLLM(DeepseekV3RMSNorm, BaseInjectedModule): + def __init__(self, + key: str, + gguf_loader : GGUFLoader, + config: PretrainedConfig, + orig_module: nn.Module, + prefill_device: str = "xpu", + generate_device: str = "xpu", + **kwargs): + BaseInjectedModule.__init__(self, key, gguf_loader, config, orig_module, prefill_device, **kwargs) + self.orig_module.__init__(orig_module.weight.shape[0], + orig_module.variance_epsilon) + self.eps = orig_module.variance_epsilon + + def forward(self, x: torch.Tensor) -> torch.Tensor: + from ipex_llm.transformers.models.common import rms_norm_forward + if x.dtype not in [torch.float32, torch.float16]: + output = rms_norm_forward(self, x.float()) + else: + output = rms_norm_forward(self, x) + return output.to(x.dtype) + + def load(self): + BaseInjectedModule.load(self) + if self.weight.dtype not in [torch.float32, torch.float16]: + self.weight = self.weight.float() \ No newline at end of file diff --git a/archive/kt-sft/ktransformers/operators/linear.py b/archive/kt-sft/ktransformers/operators/linear.py new file mode 100644 index 0000000..d617bcd --- /dev/null +++ b/archive/kt-sft/ktransformers/operators/linear.py @@ -0,0 +1,981 @@ +#!/usr/bin/env python +# coding=utf-8 +''' +Description : +Author : Azure-Tang, Boxin Zhang +Date : 2024-07-25 11:25:24 +Version : 0.1.0 +LastEditors : Azure +LastEditTime : 2024-08-29 09:11:16 +Copyright (c) 2024 by KVCache.AI, All Rights Reserved. +''' + + +import ctypes +import time +import torch +from torch import Tensor, nn +if not torch.xpu.is_available(): + import KTransformersOps + import vLLMMarlin +from ktransformers.util.custom_loader import GGUFLoader, SafeTensorLoader +from ktransformers.util.inference_state import InferenceState +if not torch.xpu.is_available(): + from ktransformers.ktransformers_ext.operators.custom_marlin.quantize.utils.marlin_utils import ( + MarlinWorkspace, + marlin_quantize, + GPTQ_MARLIN_MIN_THREAD_N, + GPTQ_MARLIN_MIN_THREAD_K, + GPTQ_MARLIN_MAX_PARALLEL, + vllm_marlin_quantize + ) +from ktransformers.operators.base_operator import BaseInjectedModule +from transformers.configuration_utils import PretrainedConfig +from ktransformers.ktransformers_ext.triton.fp8gemm import fp8_gemm, act_quant, weight_dequant +from ktransformers.util.globals import GLOBAL_CONFIG +from abc import ABC, abstractmethod +import sys, os +sys.path.append(os.path.join(os.path.dirname(__file__), "..", "ktransformers_ext", "build")) +sys.path.append(os.path.join(os.path.dirname(__file__), "..", "ktransformers_ext", "build", "Release")) +sys.path.append(os.path.join(os.path.dirname(__file__), "..", "ktransformers_ext", "build", "Debug")) +import cpuinfer_ext +from ktransformers.operators.cpuinfer import CPUInfer +from ktransformers.server.config.config import Config +from typing import Dict, Tuple, Optional, Union +import numpy as np + +#class KLinearBase(BaseInjectedModule, ABC): +class KLinearBase(nn.Module, ABC): + def __init__( + self, + key: str, + gguf_loader: GGUFLoader, + config: PretrainedConfig, + orig_module: nn.Module = None, + device: str = "cuda", + **kwargs, + ): + # super().__init__(key, gguf_loader, config, orig_module, device, **kwargs) + super().__init__() + self.key = key + self.gguf_loader = gguf_loader + self.device = device + self.config = config + + self.has_bias = False + self.dtype = torch.get_default_dtype() + if orig_module is not None: + self.in_features = orig_module.in_features + self.out_features = orig_module.out_features + else: + shape = self.gguf_loader.tensor_info[key + ".weight"]["shape"] + if len(shape) == 1: + print("Warning: orig_module is not set, but has in_features or out_features equals to 1, can't get in_features and out_features from GGUF") + self.in_features = self.gguf_loader.tensor_info[key + ".weight"]["shape"][0] + self.out_features = self.gguf_loader.tensor_info[key + ".weight"]["shape"][1] + + self.loaded = False # for lm_head pre-load, TODO: use new way to do lm_head pre-load when layer wise prefill. + + @abstractmethod + def forward(self, x: torch.Tensor) -> torch.Tensor: + pass + + def load_weight(self, override_key: str | None = None, device: str | None = None): + if override_key is not None: + keys = override_key + else: + keys = [self.key] + + for key in keys: + if isinstance(self.gguf_loader, SafeTensorLoader): + # using safetensor_loader + tensor = self.gguf_loader.load_tensor(key+'.weight') + if self.gguf_loader.has_tensor(key+'.weight_scale_inv'): + weight_scale_inv = self.gguf_loader.load_tensor(key+'.weight_scale_inv') + return nn.Parameter(tensor), nn.Parameter(weight_scale_inv) + return nn.Parameter(tensor) + + elif self.gguf_loader.has_tensor(key + ".weight") or "kv_b_proj" in key: + if key + ".bias" in self.gguf_loader.tensor_file_map: + tensors = self.load_multi(key, ["weight", "bias"], device=device) + tensor = tensors["weight"] + bias = tensors["bias"] + # self.qtype = GGML_TYPE_QTYPE_MAP[tensorinfo[key + ".weight"]["ggml_type"]] + # print(torch.isinf(tensor).any(), torch.isinf(bias).any()) + return nn.Parameter(tensor), nn.Parameter(bias) + elif "kv_b_proj" in key and not self.gguf_loader.has_tensor(key + ".weight"): + attn_k_b_tensors = self.load_multi(key.replace("self_attn.kv_b_proj", "attn_k_b"), ["weight"], device=device) + attn_k_b = attn_k_b_tensors["weight"] + del attn_k_b_tensors + attn_k_b = attn_k_b.transpose(1, 2).contiguous() + attn_v_b_tensors = self.load_multi(key.replace("self_attn.kv_b_proj", "attn_v_b"), ["weight"], device=device) + attn_v_b = attn_v_b_tensors["weight"] + del attn_v_b_tensors + kv_b_proj = torch.cat((attn_k_b, attn_v_b), dim=1) + kv_b_proj = kv_b_proj.contiguous() if kv_b_proj.ndim == 2 else kv_b_proj.flatten(0, 1).contiguous() + del attn_k_b + del attn_v_b + return nn.Parameter(kv_b_proj) + else: + tensors = self.load_multi(key, ["weight"], device=device) + tensor = tensors["weight"] + # self.qtype = GGML_TYPE_QTYPE_MAP[tensorinfo[key + ".weight"]["ggml_type"]] + return nn.Parameter(tensor) + else: + raise FileNotFoundError(f"Weight file not found for key {key}") + + def load_multi(self, key: str, keys: list[str], device: str = "cpu"): + tensors = {} + for k in keys: + tensors[k] = self.gguf_loader.load_gguf_tensor(key + "." + k, device=device) + return tensors + + @abstractmethod + def load(self, w: dict | nn.Parameter | tuple | None = None, device: str|None = "cuda"): + pass + + @abstractmethod + def unload(self): + pass + + +class KLinearTorch(KLinearBase): + def __init__( + self, + key: str, + gguf_loader: GGUFLoader, + config: PretrainedConfig, + orig_module: nn.Module = None, + device: str = "cuda", + **kwargs, + ): + super().__init__(key, gguf_loader, config, orig_module, device, **kwargs) + self.has_bias = False + self.dtype = torch.get_default_dtype() + self.weight = None + self.has_bias = False + + def forward(self, x: torch.Tensor, bsz_tensor: torch.Tensor=None, **kwargs) -> torch.Tensor: + dtype = x.dtype + out_device = x.device + + if (not x.requires_grad) and GLOBAL_CONFIG._config["mod"] == "sft": + x = x.requires_grad_(True) + # TODO: support CUDA Graph when using cpu, but CPUInfer is recommended. + x = x.to(device=self.device, dtype=self.dtype) + x = x @ self.weight + if self.has_bias: + x = x + self.bias + x = x.to(dtype=dtype, device=out_device) + return x + + def load(self, w: dict | nn.Parameter | tuple | None = None, device: str|None = None): + if self.loaded: return + if device is None: device = self.device + if w is None: w = self.load_weight(device=device) + # else: self.out_features = w.shape[0], self.in_features = w.shape[1] + + if isinstance(w, nn.Parameter): + try: + self.weight = w.to(dtype=self.dtype).view(self.out_features, self.in_features).T + except: + self.weight = w.to(dtype=self.dtype).T + self.has_bias = False + elif isinstance(w, tuple): + try: + self.weight = w[0].to(dtype=self.dtype).view(self.out_features, self.in_features).T + except: + self.weight = w[0].to(dtype=self.dtype).T + self.bias = w[1].to(dtype=self.dtype) + self.has_bias = True + else: + raise ValueError("Invalid weight type") + # self.linear = self.linear.to(device) + self.weight = self.weight.to(device) + if self.has_bias: + self.bias = self.bias.to(device) + self.loaded = True + + def unload(self): + if self.weight is not None: + self.weight = None + if self.has_bias: + self.bias = None + +class KLinearQ8(KLinearBase): + def __init__( + self, + key: str, + gguf_loader: GGUFLoader, + config: PretrainedConfig, + orig_module: nn.Module = None, + device: str = "cuda", + **kwargs, + ): + super().__init__(key, gguf_loader, config, orig_module, device, **kwargs) + self.has_bias = False + self.compute_dtype = torch.float32 + self.weight = None + self.weight_scale = None + self.weight_zero_point = None + self.bias = None + self.loaded = False + + def forward(self, x: torch.Tensor, bsz_tensor: torch.Tensor=None) -> torch.Tensor: + orig_dtype = x.dtype + out_device = x.device + + x = x.to(device=self.device, dtype=self.compute_dtype) + + + weight_dequant = self._dequantize_weight(self.weight, self.weight_scale, bits=8) + out = x @ weight_dequant.T + + if self.has_bias: + out = out + self.bias + + return out.to(dtype=orig_dtype, device=out_device) + + def _dequantize_weight(self, q_matrix, scales, bits=8): + """ + Dequantize a low-precision matrix back to floating-point + + Args: + q_matrix (torch.Tensor): Quantized int matrix + scales (torch.Tensor): Scale factors for each column + bits (int): Quantization bits used (8 or 4) + + Returns: + torch.Tensor: Dequantized floating-point matrix + """ + # Ensure inputs are torch tensors + if not isinstance(q_matrix, torch.Tensor): + q_matrix = torch.tensor(q_matrix, dtype=torch.int8) + if not isinstance(scales, torch.Tensor): + scales = torch.tensor(scales, dtype=torch.float32) + + # Convert to correct dtype if needed + if q_matrix.dtype != torch.int8: + q_matrix = q_matrix.to(torch.int8) + if scales.dtype != torch.float32: + scales = scales.to(torch.float32) + + # For Q4, ensure the values stay within 4-bit range + if bits == 4: + q_matrix = torch.clamp(q_matrix, -7, 7) + rows, cols = q_matrix.shape + dequant_matrix = q_matrix.to(torch.float32) + scales_broadcast = scales.view(1, cols) + # Apply dequantization to all columns at once using matrix multiplication + dequant_matrix = dequant_matrix * scales_broadcast + + return dequant_matrix + + + def _quantize_weight(self, matrix, bits=8): + """ + Quantize a floating-point matrix to lower precision (Q8 or Q4) + + Args: + matrix (torch.Tensor): Input matrix in floating-point format + bits (int): Quantization bits, either 8 or 4 + + Returns: + tuple: (quantized int matrix, scale factors for each column) + """ + if not isinstance(matrix, torch.Tensor): + matrix = torch.tensor(matrix, dtype=torch.float32) + + # Convert to float32 if needed + if matrix.dtype != torch.float32: + matrix = matrix.to(torch.float32) + + # Get matrix shape + rows, cols = matrix.shape + + # Determine quantization parameters based on bits + if bits == 8: + max_int = 127 + qtype = torch.int8 + elif bits == 4: + max_int = 7 + qtype = torch.int8 # We'll still use int8 storage but limit to 4-bit range, wait for native support + else: + raise ValueError("Quantization bits must be either 8 or 4") + + scales = torch.zeros(cols, dtype=torch.float32, device=matrix.device) + + # Calculate max absolute value for each column + max_abs_vals, _ = torch.max(torch.abs(matrix), dim=0) + + # Handle zero columns (avoid division by zero) + zero_cols = max_abs_vals == 0 + max_abs_vals[zero_cols] = 1.0 + + # Calculate scale factors for all columns at once + scales = max_abs_vals / max_int + + # Prepare the scales for broadcasting [1, cols] + scales_broadcast = scales.view(1, cols) + + # Apply quantization to the entire matrix at once + q_matrix = torch.round(matrix / scales_broadcast).to(qtype) + + # For Q4, clamp values to ensure they stay within 4-bit range + if bits == 4: + q_matrix = torch.clamp(q_matrix, -max_int, max_int) + + return q_matrix, scales + + def load(self, w: Union[Dict, nn.Parameter, Tuple, None] = None, device: Optional[str] = None): + if self.loaded: return + if device is None: device = self.device + if w is None: w = self.load_weight(device=device) + + if isinstance(w, nn.Parameter): + try: + weight = w.to(dtype=self.compute_dtype).view(self.out_features, self.in_features) + except: + weight = w.to(dtype=self.compute_dtype) + self.has_bias = False + elif isinstance(w, tuple): + try: + weight = w[0].to(dtype=self.compute_dtype).view(self.out_features, self.in_features) + except: + weight = w[0].to(dtype=self.compute_dtype) + self.bias = w[1].to(dtype=self.compute_dtype).to(device) + self.has_bias = True + else: + raise ValueError("Invalid weight type") + + self.weight, self.weight_scale = self._quantize_weight(weight, bits=8) + + self.weight = self.weight.to(device) + self.weight_scale = self.weight_scale.to(device) + + if self.has_bias: + self.bias = self.bias.to(device) + + self.loaded = True + + def unload(self): + self.weight = None + self.weight_scale = None + self.weight_zero_point = None + self._orig_weight = None + + if self.has_bias: + self.bias = None + + self.loaded = False + + +class KLinearFP8(KLinearBase): + # this kernel requires special handling for weight + # Please load the weight file downloaded from KVCache.AI + has_bias: bool + weight: torch.Tensor + bias: torch.Tensor + def __init__( + self, + key: str, + gguf_loader: GGUFLoader, + config: PretrainedConfig, + orig_module: nn.Module = None, + device: str = "cuda", + block_size: int = 128, + **kwargs, + ): + super().__init__(key, gguf_loader, config, orig_module, device, **kwargs) + self.has_bias = False + self.dtype = torch.get_default_dtype() + self.block_size = block_size + + def forward(self, x: torch.Tensor, bsz_tensor: torch.Tensor) -> torch.Tensor: + x = x.to(self.device) + orig_dtype = x.dtype + x_quantized, scale_x = act_quant(x, self.block_size) + y = fp8_gemm(x_quantized, scale_x, self.weight, self.weight_scale_inv) + return y.to(dtype=orig_dtype) + + def load(self, w: dict | nn.Parameter | tuple | None = None, device: str|None = None): + if device is None: device = self.device + if w is None: + w = self.load_weight(device=device) + ### TODO fit weight_inv format + if isinstance(w, tuple): + self.weight = w[0].to(device) + self.weight_scale_inv = w[1].to(device) + self.has_bias = False + else: + raise ValueError("Invalid weight type") + self.weight = self.weight.to(device) + if self.has_bias: + self.bias = self.bias.to(device) + + def unload(self): + if self.weight is not None: + self.weight = None + if self.has_bias: + self.bias = None + +# TODO: merge two marlin class + +class VLinearMarlin(KLinearBase): + marlin_q_w: torch.Tensor + marlin_s: torch.Tensor + g_idx: torch.Tensor + sort_indices: torch.Tensor + has_bias: bool + def __init__( + self, + key: str, + gguf_loader: GGUFLoader, + config: PretrainedConfig, + orig_module: nn.Module = None, + device: str = "cuda", + num_bits: int = 4, # 4-bit/8-bit is supported + group_size: int = 64, # -1, 32, 64, 128 + act_order: bool = False, + is_k_full=True, + **kwargs, + ): + assert device.lower() != "cpu", "Marlin quantized linear only supports GPU device" + super().__init__(key, gguf_loader, config, orig_module, device, **kwargs) + self.num_bits = num_bits + self.group_size = group_size + self.act_order = act_order + self.is_k_full = is_k_full + self.padding = False + self.orin_in_features = self.in_features + self.orin_out_features = self.out_features + if self.in_features%GPTQ_MARLIN_MIN_THREAD_K!=0 or self.out_features%GPTQ_MARLIN_MIN_THREAD_K!=0: + #print(f"warning!, in_features={in_features} or out_features={out_features} is undivisible by GPTQ_MARLIN_MIN_THREAD_K={GPTQ_MARLIN_MIN_THREAD_K} and GPTQ_MARLIN_MIN_THREAD_N={GPTQ_MARLIN_MIN_THREAD_N}, padding") + self.padding = True + self.in_features = (self.in_features+GPTQ_MARLIN_MIN_THREAD_K-1)//GPTQ_MARLIN_MIN_THREAD_K*GPTQ_MARLIN_MIN_THREAD_K + self.out_features = (self.out_features+GPTQ_MARLIN_MIN_THREAD_N-1)//GPTQ_MARLIN_MIN_THREAD_N*GPTQ_MARLIN_MIN_THREAD_N + #print(f"After padding: in_features={in_features}, out_features={out_features}") + + self.k = self.in_features + self.n = self.out_features + + def load(self, w: dict | nn.Parameter | tuple | None = None, device: str|None = None): + if self.loaded: return + if device is None: device = self.device + assert device.lower() != "cpu", "Marlin quantized linear only supports GPU device" + + #if self.in_features * self.out_features: + if w is None: + w = self.load_weight(device=device) + + if isinstance(w, nn.Parameter): + # pad weight + weight = w.view(self.orin_out_features, self.orin_in_features).T + self.has_bias = False + elif isinstance(w, tuple): + w = list(w) + weight = w[0].view(self.orin_out_features, self.orin_in_features).T + self.bias = w[1].view(self.orin_out_features) + self.bias = w[1] + self.has_bias = True + else: + raise ValueError("Invalid weight type") + weight = weight.to(device) + if self.has_bias: + self.bias = self.bias.to(device) + + if self.padding: + padded_weight = torch.zeros(self.in_features, self.out_features, device=self.device) + padded_weight[:self.orin_in_features, :self.orin_out_features] = weight + weight = padded_weight + + # Pack Marlin linear + marlin_q_w, marlin_s, g_idx, sort_indices, _ = marlin_quantize( + weight, self.num_bits, self.group_size, self.act_order + ) + self.workspace = MarlinWorkspace( + self.out_features, GPTQ_MARLIN_MIN_THREAD_N, GPTQ_MARLIN_MAX_PARALLEL,self.device + ) + self.weight = marlin_q_w + self.marlin_q_w = marlin_q_w + self.marlin_s = marlin_s + self.g_idx = g_idx + self.sort_indices = sort_indices + self.k = weight.shape[0] + self.n = weight.shape[1] + # self.shape_buffer = torch.tensor([60], dtype=torch.int32, device=self.device) + self.loaded = True + + + def forward(self, x: torch.Tensor, bsz_tensor: torch.Tensor = None) -> torch.Tensor: + if bsz_tensor is None: + bsz_tensor = torch.tensor([x.shape[0]], dtype=torch.int32, device=self.device) + + + # Only support input x as BF16 and FP16 + x = x.to(self.device) + orig_shape = list(x.shape) + orig_dtype = x.dtype + x = x.reshape(-1, orig_shape[-1]) + marlin_s = self.marlin_s.to(x.dtype) + sms = -1 + + # padding x.shape[0] to avoid CUDA illegal memory access error + x, orig_size_m = self._pad_input(x) + + x = vLLMMarlin.gptq_marlin_gemm( + x, + self.marlin_q_w, + marlin_s, + self.g_idx, + self.sort_indices, + self.workspace.scratch, + self.num_bits, + bsz_tensor, + x.shape[0], + self.n, + x.shape[-1], + sms, + self.is_k_full, + ) + + x = x[:orig_size_m] + + if self.has_bias: + x = x + self.bias + orig_shape[-1] = self.n + return x.reshape(orig_shape).to(orig_dtype) + + def unload(self): + + if self.has_bias: + self.bias = None + self.marlin_q_w = None + self.marlin_s = None + self.g_idx = None + self.sort_indices = None + self.workspace = None + + def _pad_input(self, x): + + size_m = x.shape[0] + size_k = x.shape[1] + + # size_m and align value depends on VLinearMarlin implementation + if size_m > 1024: + align = 1024 + elif size_m > 64: + align = 64 + else: + align = 1 + + padded_size_m = ((size_m + align - 1) // align) * align + + if padded_size_m > size_m: + pad_len = padded_size_m - size_m + pad_tensor = torch.zeros((pad_len, size_k), dtype=x.dtype, device=x.device) + x = torch.cat([x, pad_tensor], dim = 0).contiguous() + return x, size_m + +class KLinearMarlin(KLinearBase): + marlin_q_w: torch.Tensor + marlin_s: torch.Tensor + g_idx: torch.Tensor + sort_indices: torch.Tensor + has_bias: bool + def __init__( + self, + key: str, + gguf_loader: GGUFLoader, + config: PretrainedConfig, + orig_module: nn.Module = None, + device: str = "cuda", + num_bits: int = 4, # 4-bit/8-bit is supported + group_size: int = 64, # -1, 32, 64, 128 + act_order: bool = False, + is_k_full=True, + **kwargs, + ): + assert device.lower() != "cpu", "Marlin quantized linear only supports GPU device" + super().__init__(key, gguf_loader, config, orig_module, device, **kwargs) + self.num_bits = num_bits + self.group_size = group_size + self.act_order = act_order + self.is_k_full = is_k_full + self.padding = False + self.orin_in_features = self.in_features + self.orin_out_features = self.out_features + if self.in_features%GPTQ_MARLIN_MIN_THREAD_K!=0 or self.out_features%GPTQ_MARLIN_MIN_THREAD_K!=0: + #print(f"warning!, in_features={in_features} or out_features={out_features} is undivisible by GPTQ_MARLIN_MIN_THREAD_K={GPTQ_MARLIN_MIN_THREAD_K} and GPTQ_MARLIN_MIN_THREAD_N={GPTQ_MARLIN_MIN_THREAD_N}, padding") + self.padding = True + self.in_features = (self.in_features+GPTQ_MARLIN_MIN_THREAD_K-1)//GPTQ_MARLIN_MIN_THREAD_K*GPTQ_MARLIN_MIN_THREAD_K + self.out_features = (self.out_features+GPTQ_MARLIN_MIN_THREAD_N-1)//GPTQ_MARLIN_MIN_THREAD_N*GPTQ_MARLIN_MIN_THREAD_N + #print(f"After padding: in_features={in_features}, out_features={out_features}") + + self.k = self.in_features + self.n = self.out_features + + def load(self, w: dict | nn.Parameter | tuple | None = None, device: str|None = None): + if self.loaded: return + if device is None: device = self.device + assert device.lower() != "cpu", "Marlin quantized linear only supports GPU device" + + #if self.in_features * self.out_features: + if w is None: + w = self.load_weight(device=device) + + if isinstance(w, nn.Parameter): + # pad weight + weight = w.view(self.orin_out_features, self.orin_in_features).T + self.has_bias = False + elif isinstance(w, tuple): + w = list(w) + weight = w[0].view(self.orin_out_features, self.orin_in_features).T + self.bias = w[1].view(self.orin_out_features) + self.bias = w[1] + self.has_bias = True + else: + raise ValueError("Invalid weight type") + weight = weight.to(device) + if self.has_bias: + self.bias = self.bias.to(device) + + if self.padding: + padded_weight = torch.zeros(self.in_features, self.out_features, device=self.device) + padded_weight[:self.orin_in_features, :self.orin_out_features] = weight + weight = padded_weight + + # Pack Marlin linear + marlin_q_w, marlin_s, g_idx, sort_indices, _ = marlin_quantize( + weight, self.num_bits, self.group_size, self.act_order + ) + self.workspace = MarlinWorkspace( + self.out_features, GPTQ_MARLIN_MIN_THREAD_N, GPTQ_MARLIN_MAX_PARALLEL,self.device + ) + self.weight = marlin_q_w # modeling_xxx.py may use linear.weight + self.marlin_q_w = marlin_q_w + self.marlin_s = marlin_s + self.g_idx = g_idx + self.sort_indices = sort_indices + self.k = weight.shape[0] + self.n = weight.shape[1] + self.loaded = True + + def forward(self, x: torch.Tensor, bsz_tensor: torch.Tensor=None, **kwargs) -> torch.Tensor: + # Only support input x as BF16 and FP16 + x = x.to(self.device) + orig_shape = list(x.shape) + orig_dtype = x.dtype + x = x.reshape(-1, orig_shape[-1]) + x = x.reshape(-1, x.shape[-1]) + if self.padding: + padding_input=torch.empty(x.shape[0], self.in_features, device=x.device, dtype=x.dtype) + padding_input[:,:self.orin_in_features] = x + x = padding_input + marlin_s = self.marlin_s.to(x.dtype) + x = KTransformersOps.gptq_marlin_gemm( + x, + self.marlin_q_w, + marlin_s, + self.g_idx, + self.sort_indices, + self.workspace.scratch, + self.num_bits, + x.shape[0], + self.n, + x.shape[-1], + self.is_k_full, + ) + if self.padding: + x = x[:,:self.orin_out_features] + orig_shape[-1] = self.orin_out_features + else: + orig_shape[-1] = self.out_features + if self.has_bias: + x = x + self.bias + return x.reshape(orig_shape).to(orig_dtype) + + def unload(self): + + if self.has_bias: + self.bias = None + self.marlin_q_w = None + self.marlin_s = None + self.g_idx = None + self.sort_indices = None + self.workspace = None + +class KLinearCPUInfer(KLinearBase): + CPU_INFER = None + def __init__( + self, + key: str, + gguf_loader: GGUFLoader, + config: PretrainedConfig, + orig_module: nn.Module = None, + device: str = "cpu", + out_device: str = "cuda", # this device mean which device the output should on. TODO: support cpu. + stride = 16, + group_max_len = 1024, + **kwargs, + ): + super().__init__(key, gguf_loader, config, orig_module, device, **kwargs) + if KLinearCPUInfer.CPU_INFER is None: + KLinearCPUInfer.CPU_INFER = CPUInfer(Config().cpu_infer) + self.has_bias = False + self.dtype = torch.get_default_dtype() + self.w = None + self.has_bias = False + self.stride = stride + self.group_max_len = group_max_len + self.out_device = out_device + + def forward(self, x: torch.Tensor, bsz_tensor: torch.Tensor = None) -> torch.Tensor: + origin_shape = x.shape # [batch_size, q_len, hidden_size] + if origin_shape[1] == 1 and torch.cuda.is_current_stream_capturing(): + out_device = x.device + self.input_tensor_cpu.copy_(x, non_blocking=True) + qlen = origin_shape[1] + KLinearCPUInfer.CPU_INFER.submit_with_cuda_stream( + torch.cuda.current_stream().cuda_stream, + self.linear.forward( + qlen, + self.input_tensor_cpu.data_ptr(), + self.output_cpu.data_ptr() + ) + ) + KLinearCPUInfer.CPU_INFER.sync_with_cuda_stream(torch.cuda.current_stream().cuda_stream) + self.output_gpu.copy_(self.output_cpu, non_blocking=True) + if self.has_bias: + self.output_gpu += self.bias + return self.output_gpu + else: + dtype = x.dtype + out_device = x.device + x = x.to(device=self.device) + qlen = origin_shape[1] + output_shape = (*origin_shape[:-1], self.out_features) + output = torch.empty(output_shape, device=x.device, dtype=x.dtype) + KLinearCPUInfer.CPU_INFER.submit( + self.linear.forward( + qlen, + x.data_ptr(), + output.data_ptr() + ) + ) + KLinearCPUInfer.CPU_INFER.sync() + if self.has_bias: + output = output + self.bias + output = output.to(dtype=dtype, device=out_device) + return output + + def load(self, w: dict | nn.Parameter | tuple | None = None, device: str|None = None, warmup:bool = True): + print(f"loading {self.key} to {self.device} using CPUInfer") + if device is None: device = self.device + self.load_weights(w=w, device=device) + if self.bias is not None: + self.has_bias = True + self.bias = self.bias.to(device) + + weight_ptr = ctypes.addressof( + ctypes.cast(self.weight.ctypes.data, ctypes.POINTER(ctypes.c_uint64)).contents + ) + config = cpuinfer_ext.linear.LinearConfig(self.in_features, self.out_features, self.stride, self.group_max_len, weight_ptr, self.weight_type, 30) + self.linear = cpuinfer_ext.linear.Linear(config) + + if warmup: + KLinearCPUInfer.CPU_INFER.submit(self.linear.warm_up()) + KLinearCPUInfer.CPU_INFER.sync() + self.input_tensor_cpu = torch.zeros((1, 1, self.in_features), device="cpu", pin_memory=True) + self.output_cpu = torch.zeros((1, 1, self.out_features), device="cpu", pin_memory=True, dtype=torch.bfloat16) + self.output_gpu = torch.zeros((1, 1, self.out_features), device=self.out_device) + + def load_weights(self, w: dict | nn.Parameter | tuple | None = None, device: str = "cpu"): + if self.gguf_loader.has_tensor(self.key + ".weight"): + if self.key + ".bias" in self.gguf_loader.tensor_file_map: + self.weight = self.gguf_loader.get_mmap_tensor(self.key + ".weight") + self.weight_type = self.gguf_loader.tensor_info[self.key + ".weight"]["ggml_type"] + self.bias = self.gguf_loader.load_gguf_tensor(self.key + ".bias", device=device) + else: + self.weight = self.gguf_loader.get_mmap_tensor(self.key + ".weight") + self.weight_type = self.gguf_loader.tensor_info[self.key + ".weight"]["ggml_type"] + self.bias = None + else: + raise ValueError(f"Linear {self.key} not found in gguf_loader") + + def unload(self): + if self.w is not None: + self.w = None + if self.has_bias: + self.bias = None + +class KLinearIPEXLLM(KLinearBase): + def __init__( + self, + key: str, + gguf_loader: GGUFLoader, + config: PretrainedConfig, + orig_module: nn.Module = None, + device: str = "xpu", + precision: str = "sym_int4", + **kwargs, + ): + super().__init__(key, gguf_loader, config, orig_module, device, **kwargs) + self.has_bias = False + self.dtype = torch.get_default_dtype() + self.weight = None + self.has_bias = False + self.precision = precision + self.qtype = None + + def forward(self, x: torch.Tensor, bsz_tensor: torch.Tensor = None) -> torch.Tensor: + dtype = x.dtype + out_device = x.device + from ipex_llm.transformers.models.common import linear_forward + x = linear_forward(x.half(), self.weight, self.qtype, self.out_features) + + if self.has_bias: + x = x + self.bias + x = x.to(dtype=dtype, device=out_device) + return x + + def load(self, w: dict | nn.Parameter | tuple | None = None, device: str|None = None): + if self.loaded: return + if device is None: device = self.device + assert device.lower()[:3] == "xpu", "IPEX-LLM quantized linear only supports XPU device" + if w is None: w = self.load_weight(device=device) + + if isinstance(w, nn.Parameter): + try: + weight = w.to(dtype=self.dtype).view(self.out_features, self.in_features).T + except: + weight = w.to(dtype=self.dtype).T + self.has_bias = False + elif isinstance(w, tuple): + try: + weight = w[0].to(dtype=self.dtype).view(self.out_features, self.in_features).T + except: + weight = w[0].to(dtype=self.dtype).T + self.bias = w[1].to(dtype=self.dtype) + self.has_bias = True + else: + raise ValueError("Invalid weight type") + weight = weight.to("cpu").float().transpose(0, 1).contiguous() + + if self.has_bias: + self.bias = self.bias.to(device) + + # quantize linear weight + from ipex_llm.transformers.models.common import quantize_linear + paramsLowBit, qtype = quantize_linear(weight, self.in_features, self.precision) + self.weight = paramsLowBit.to(device) + self.qtype = qtype + self.loaded = True + + def unload(self): + if self.weight is not None: + self.weight = None + if self.has_bias: + self.bias = None + +LINEAR_MAP = { + "KLinearMarlin": KLinearMarlin, + "KLinearTorch": KLinearTorch, + "KLinearCPUInfer": KLinearCPUInfer, + "VLinearMarlin": VLinearMarlin, + "KLinearFP8": KLinearFP8, + "KLinearQ8": KLinearQ8, + "KLinearIPEXLLM": KLinearIPEXLLM, +} + +class KTransformersLinear(BaseInjectedModule, KLinearBase): + def __init__( + self, + key: str, + gguf_loader: GGUFLoader, + config: PretrainedConfig, + orig_module: nn.Module, + generate_device: str = "cuda", + generate_op: str| None = "KLinearMarlin", + prefill_device: str = "cuda", + prefill_op: str| None = "KLinearTorch", + **kwargs, + ): + BaseInjectedModule.__init__(self, key, gguf_loader, config, orig_module, prefill_device, generate_device, **kwargs) + KLinearBase.__init__(self, key, gguf_loader, config, orig_module, generate_device, **kwargs) + # build all the linear operators + if prefill_op is not None: + assert prefill_op in LINEAR_MAP, f"linear_type {prefill_op} not supported" + self.prefill_linear = LINEAR_MAP[prefill_op](key, gguf_loader, config, orig_module, prefill_device, **kwargs) + else: + self.prefill_linear = None + + if generate_op is not None: + assert generate_op in LINEAR_MAP, f"linear_type {generate_op} not supported" + self.generate_linear = LINEAR_MAP[generate_op](key, gguf_loader, config, orig_module, generate_device, **kwargs) + else: + self.generate_linear = None + self.mode = InferenceState.UNLOAD + + def forward(self, x, bsz_tensor=None): + # linear_fwd_st = time.time() + if self.mode == InferenceState.PREFILL: + assert self.prefill_linear is not None, "cpu linear is not initialized" + y = self.prefill_linear.forward(x, bsz_tensor) + else: + assert self.generate_linear is not None, "gpu linear is not initialized" + # TODO: A violence way to solve the weight=None, for Lora inference Test, need modify it later + try: + y = self.generate_linear.forward(x, bsz_tensor) + except TypeError as e: + Warning("A Dange way to avoid the none weight, Need to check it later in KTransformersLinear forward!!") + self.generate_linear.weight = self.orig_module.generate_linear.weight + self.weight = self.orig_module.generate_linear.weight + y = self.generate_linear.forward(x, bsz_tensor) + + # linear_fwd_end = time.time() + # print(f"[KTLinear] Forward time: {linear_fwd_end-linear_fwd_st}") + return y + + def load(self, w: dict | nn.Parameter | tuple | None = None, mode: InferenceState = InferenceState.GENERATE): + if not mode: + mode = InferenceState.GENERATE + # load to device + if mode == InferenceState.PREFILL: + self.generate_linear.unload() + self.prefill_linear.load(w=w) + self.device = self.prefill_linear.device + self.weight = self.prefill_linear.weight # modeling_xxx.py may use linear.weight + elif mode == InferenceState.GENERATE: + self.prefill_linear.unload() + self.generate_linear.load(w=w) + self.device = self.generate_linear.device + self.weight = self.generate_linear.weight # modeling_xxx.py may use linear.weight + elif mode == InferenceState.UNLOAD: + self.prefill_linear.unload() + self.generate_linear.unload() + self.device = "cpu" + else: + raise ValueError("mode must be either InferenceState.GENERATE, InferenceState.PREFILL or InferenceState.UNLOAD") + self.mode = mode + + def unload(self): + if self.prefill_linear is not None: + self.prefill_linear.unload() + if self.generate_linear is not None: + self.generate_linear.unload() + self.device = self.generate_linear.device + + def set_inference_mode(self, mode: InferenceState): + if not mode: + mode = InferenceState.GENERATE + if mode == InferenceState.GENERATE: + self.load(mode=InferenceState.GENERATE) + elif mode == InferenceState.PREFILL: + self.load(mode=InferenceState.PREFILL) + elif mode == InferenceState.UNLOAD: + self.unload() + else: + raise ValueError("mode must be either InferenceState.GENERATE, InferenceState.PREFILL or InferenceState.UNLOAD") + + diff --git a/archive/kt-sft/ktransformers/operators/mlp.py b/archive/kt-sft/ktransformers/operators/mlp.py new file mode 100644 index 0000000..77d7d05 --- /dev/null +++ b/archive/kt-sft/ktransformers/operators/mlp.py @@ -0,0 +1,37 @@ + +from ktransformers.operators.base_operator import BaseInjectedModule +from ktransformers.util.custom_loader import GGUFLoader +from transformers import PretrainedConfig +import torch.nn as nn +from ktransformers.models.modeling_deepseek_v3 import DeepseekV3MLP +from ktransformers.models.modeling_qwen2_moe import Qwen2MoeMLP +class kDeepseekV3MLP(DeepseekV3MLP, BaseInjectedModule): + def __init__(self, + key: str, + gguf_loader : GGUFLoader, + config: PretrainedConfig, + orig_module: nn.Module, + prefill_device: str = "cuda", + generate_device: str = "cuda", + **kwargs): + BaseInjectedModule.__init__(self, key, gguf_loader, config, orig_module, prefill_device, **kwargs) + self.orig_module.__init__(orig_module.config, + orig_module.hidden_size, orig_module.intermediate_size) + def forward(self, x, bsz_tensor): + down_proj = self.down_proj(self.act_fn(self.gate_proj(x, bsz_tensor)) * self.up_proj(x, bsz_tensor), bsz_tensor) + return down_proj +class KQwen2MoeMLP(Qwen2MoeMLP, BaseInjectedModule): + def __init__(self, + key: str, + gguf_loader : GGUFLoader, + config: PretrainedConfig, + orig_module: nn.Module, + prefill_device: str = "cuda", + generate_device: str = "cuda", + **kwargs): + BaseInjectedModule.__init__(self, key, gguf_loader, config, orig_module, prefill_device, **kwargs) + self.orig_module.__init__(orig_module.config, + orig_module.intermediate_size) + def forward(self, x, bsz_tensor): + down_proj = self.down_proj(self.act_fn(self.gate_proj(x, bsz_tensor)) * self.up_proj(x, bsz_tensor), bsz_tensor) + return down_proj \ No newline at end of file diff --git a/archive/kt-sft/ktransformers/operators/models.py b/archive/kt-sft/ktransformers/operators/models.py new file mode 100644 index 0000000..941f0fb --- /dev/null +++ b/archive/kt-sft/ktransformers/operators/models.py @@ -0,0 +1,1756 @@ +#!/usr/bin/env python +# coding=utf-8 +""" +Description : +Author : Azure-Tang +Date : 2024-07-25 11:25:24 +Version : 1.0.0 +LastEditors : Azure +LastEditTime : 2024-08-27 07:29:04 +Copyright (c) 2024 by KVCache.AI, All Rights Reserved. +""" + +import inspect +import math +from typing import List, Optional, Tuple, Union +import time +import torch +import torch.nn.functional as F +import torch.utils.checkpoint +from torch import nn +from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss +from ktransformers.operators.dynamic_attention import DynamicScaledDotProductAttention +from ktransformers.server.config.config import Config +import os +import yaml +from transformers.activations import ACT2FN +from transformers.cache_utils import Cache, DynamicCache, StaticCache +from transformers.modeling_attn_mask_utils import ( + AttentionMaskConverter, +) +from transformers.modeling_outputs import ( + MoeCausalLMOutputWithPast, + MoeModelOutputWithPast, + SequenceClassifierOutputWithPast, + TokenClassifierOutput, +) +from transformers.modeling_utils import PreTrainedModel +from transformers.utils import ( + add_start_docstrings, + add_start_docstrings_to_model_forward, + is_flash_attn_2_available, + is_flash_attn_greater_or_equal_2_10, + logging, + replace_return_docstrings, +) +from ktransformers.models.modeling_qwen2_moe import ( + Qwen2MoeSparseMoeBlock, + Qwen2MoeMLP, + Qwen2MoeDecoderLayer, + Qwen2MoeRotaryEmbedding, +) + +from ktransformers.models.modeling_qwen3_moe import ( + Qwen3MoeSparseMoeBlock, + Qwen3MoeMLP, + Qwen3MoeDecoderLayer, +) + +from ktransformers.models.modeling_deepseek import ( + BaseModelOutputWithPast, + DeepseekV2DecoderLayer, + DeepseekV2MoE, +) +from ktransformers.util.vendors import device_manager, get_device, to_device, GPUVendor +from transformers.models.qwen2_moe.configuration_qwen2_moe import Qwen2MoeConfig +from ktransformers.models.configuration_llama import LlamaConfig +from ktransformers.operators.base_operator import BaseInjectedModule +from ktransformers.util.inference_state import InferenceState +from ktransformers.util.utils import get_compute_capability +from ktransformers.util.custom_loader import GGUFLoader +from transformers.configuration_utils import PretrainedConfig +from ktransformers.models.modeling_llama import ( + LlamaDecoderLayer, + LlamaRMSNorm, + LlamaRotaryEmbedding, +) + +if is_flash_attn_2_available(): + from flash_attn import flash_attn_func, flash_attn_varlen_func + from flash_attn.bert_padding import index_first_axis, pad_input, unpad_input # noqa + + _flash_supports_window_size = "window_size" in list( + inspect.signature(flash_attn_func).parameters + ) + +logger = logging.get_logger(__name__) + +_CHECKPOINT_FOR_DOC = "Qwen/Qwen1.5-MoE-A2.7B" +_CONFIG_FOR_DOC = "Qwen2MoeConfig" + +QWEN2MOE_START_DOCSTRING = r""" + This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the + library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads + etc.) + + This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass. + Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage + and behavior. + + Parameters: + config ([`Qwen2MoeConfig`]): + Model configuration class with all the parameters of the model. Initializing with a config file does not + load the weights associated with the model, only the configuration. Check out the + [`~PreTrainedModel.from_pretrained`] method to load the model weights. +""" + +QWEN2MOE_INPUTS_DOCSTRING = r""" + Args: + input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): + Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide + it. + + Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and + [`PreTrainedTokenizer.__call__`] for details. + + [What are input IDs?](../glossary#input-ids) + attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*): + Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`: + + - 1 for tokens that are **not masked**, + - 0 for tokens that are **masked**. + + [What are attention masks?](../glossary#attention-mask) + + Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and + [`PreTrainedTokenizer.__call__`] for details. + + If `past_key_values` is used, optionally only the last `decoder_input_ids` have to be input (see + `past_key_values`). + + If you want to change padding behavior, you should read [`modeling_opt._prepare_decoder_attention_mask`] + and modify to your needs. See diagram 1 in [the paper](https://arxiv.org/abs/1910.13461) for more + information on the default strategy. + + - 1 indicates the head is **not masked**, + - 0 indicates the head is **masked**. + position_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0, + config.n_positions - 1]`. + + [What are position IDs?](../glossary#position-ids) + past_key_values (`Cache` or `tuple(tuple(torch.FloatTensor))`, *optional*): + Pre-computed hidden-states (key and values in the self-attention blocks and in the cross-attention + blocks) that can be used to speed up sequential decoding. This typically consists in the `past_key_values` + returned by the model at a previous stage of decoding, when `use_cache=True` or `config.use_cache=True`. + + Two formats are allowed: + - a [`~cache_utils.Cache`] instance; + - Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of + shape `(batch_size, num_heads, sequence_length, embed_size_per_head)`). This is also known as the legacy + cache format. + + The model will output the same cache format that is fed as input. If no `past_key_values` are passed, the + legacy cache format will be returned. + + If `past_key_values` are used, the user can optionally input only the last `input_ids` (those that don't + have their past key value states given to this model) of shape `(batch_size, 1)` instead of all `input_ids` + of shape `(batch_size, sequence_length)`. + inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*): + Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This + is useful if you want more control over how to convert `input_ids` indices into associated vectors than the + model's internal embedding lookup matrix. + use_cache (`bool`, *optional*): + If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see + `past_key_values`). + output_attentions (`bool`, *optional*): + Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned + tensors for more detail. + output_hidden_states (`bool`, *optional*): + Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for + more detail. + output_router_logits (`bool`, *optional*): + Whether or not to return the logits of all the routers. They are useful for computing the router loss, and + should not be returned during inference. + return_dict (`bool`, *optional*): + Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple. + cache_position (`torch.LongTensor` of shape `(sequence_length)`, *optional*): + Indices depicting the position of the input sequence tokens in the sequence. Contrarily to `position_ids`, + this tensor is not affected by padding. It is used to update the cache in the correct position and to infer + the complete sequence length. +""" + + +@add_start_docstrings( + "The bare Qwen2MoE Model outputting raw hidden-states without any specific head on top.", + QWEN2MOE_START_DOCSTRING, +) +class KQwen2MoeModel(BaseInjectedModule): + """ + Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`Qwen2MoeDecoderLayer`] + + Args: + config: Qwen2MoeConfig + """ + + def __init__( + self, + key: str, + gguf_loader: GGUFLoader, + config: PretrainedConfig, + orig_module: nn.Module, + device: str = "cuda", + per_layer_prefill_intput_threshold: int = 30000, # if None, no per-layer prefill + transfer_map: dict = None, + **kwargs, + ): + BaseInjectedModule.__init__( + self, key, gguf_loader, config, orig_module, device, **kwargs + ) + self.per_layer_prefill_intput_threshold = per_layer_prefill_intput_threshold + self.transfer_map = transfer_map + self.stream_device_map = dict() + + @add_start_docstrings_to_model_forward(QWEN2MOE_INPUTS_DOCSTRING) + def forward( + self, + input_ids: torch.LongTensor = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[List[torch.FloatTensor]] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + output_router_logits: Optional[bool] = None, + return_dict: Optional[bool] = None, + cache_position: Optional[torch.LongTensor] = None, + per_layer_prefill_intput_threshold: ( + int | None + ) = None, # if None or 0, close per-layer prefill + ) -> Union[Tuple, MoeModelOutputWithPast]: + # print(f'Total length of input_ids: {input_ids.size(1)}, {input_ids.size()}') + + if per_layer_prefill_intput_threshold is None: + per_layer_prefill_intput_threshold = self.per_layer_prefill_intput_threshold + per_layer_prefill_flag = False + seq_lenth = ( + inputs_embeds.size(1) if inputs_embeds is not None else input_ids.size(1) + ) + if ( + per_layer_prefill_intput_threshold + and per_layer_prefill_intput_threshold < seq_lenth + ): + per_layer_prefill_flag = True + for layer in self.layers: + self.load_layer_to(layer, InferenceState.UNLOAD) + else: + pass + output_attentions = ( + output_attentions + if output_attentions is not None + else self.config.output_attentions + ) + output_router_logits = ( + output_router_logits + if output_router_logits is not None + else self.config.output_router_logits + ) + output_hidden_states = ( + output_hidden_states + if output_hidden_states is not None + else self.config.output_hidden_states + ) + use_cache = use_cache if use_cache is not None else self.config.use_cache + + return_dict = ( + return_dict if return_dict is not None else self.config.use_return_dict + ) + + if (input_ids is None) ^ (inputs_embeds is not None): + raise ValueError( + "You cannot specify both input_ids and inputs_embeds at the same time, and must specify either one" + ) + + if self.gradient_checkpointing and self.training: + if use_cache: + logger.warning_once( + "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..." + ) + use_cache = False + + use_legacy_cache = False + if use_cache and not isinstance(past_key_values, Cache): + use_legacy_cache = True + past_key_values = DynamicCache.from_legacy_cache(past_key_values) + logger.warning_once( + "We detected that you are passing `past_key_values` as a tuple and this is deprecated and will be removed in v4.43. " + "Please use an appropriate `Cache` class (https://huggingface.co/docs/transformers/v4.41.3/en/internal/generation_utils#transformers.Cache)" + ) + + if inputs_embeds is None: + input_ids = input_ids.to("cpu") + inputs_embeds = self.embed_tokens(input_ids) + inputs_embeds = inputs_embeds.to("cuda") + + if cache_position is None: + past_seen_tokens = ( + past_key_values.get_seq_length() if past_key_values is not None else 0 + ) + cache_position = torch.arange( + past_seen_tokens, + past_seen_tokens + inputs_embeds.shape[1], + device=inputs_embeds.device, + ) + if position_ids is None: + position_ids = cache_position.unsqueeze(0) + + causal_mask = self._update_causal_mask( + attention_mask, + inputs_embeds, + cache_position, + past_key_values, + output_attentions, + ) + + hidden_states = inputs_embeds + + # create position embeddings to be shared across the decoder layers + if torch.xpu.is_available() and inputs_embeds.device.type == "xpu": + position_embeddings = self.rotary_emb(hidden_states, position_ids) + else: + position_embeddings = None + + # decoder layers + all_hidden_states = () if output_hidden_states else None + all_self_attns = () if output_attentions else None + all_router_logits = () if output_router_logits else None + next_decoder_cache = None + + for i, decoder_layer in enumerate(self.layers): + if self.transfer_map is not None and i in self.transfer_map: + prev_stream = torch.cuda.current_stream() + cur_device = self.transfer_map[i] + if cur_device not in self.stream_device_map: + self.stream_device_map[cur_device] = torch.cuda.Stream(cur_device) + torch.cuda.set_device(cur_device) + self.stream_device_map[cur_device].wait_stream(prev_stream) + torch.cuda.set_stream(self.stream_device_map[cur_device]) + hidden_states = hidden_states.to( + self.transfer_map[i], non_blocking=True + ) + causal_mask = ( + causal_mask.to(self.transfer_map[i], non_blocking=True) + if causal_mask is not None + else None + ) + position_ids = ( + position_ids.to(self.transfer_map[i], non_blocking=True) + if position_ids is not None + else None + ) + cache_position = ( + cache_position.to(self.transfer_map[i], non_blocking=True) + if cache_position is not None + else None + ) + + if output_hidden_states: + all_hidden_states += (hidden_states,) + + if self.gradient_checkpointing and self.training: + layer_outputs = self._gradient_checkpointing_func( + decoder_layer.__call__, + hidden_states, + causal_mask, + position_ids, + past_key_values, + output_attentions, + output_router_logits, + use_cache, + cache_position, + ) + else: + if per_layer_prefill_flag: + # print(f"to gpu") + self.load_layer_to(decoder_layer, InferenceState.PREFILL) + torch.cuda.empty_cache() + layer_outputs = decoder_layer( + hidden_states, + attention_mask=causal_mask, + position_ids=position_ids, + past_key_value=past_key_values, + output_attentions=output_attentions, + output_router_logits=output_router_logits, + use_cache=use_cache, + cache_position=cache_position, + position_embeddings=position_embeddings, + ) + if per_layer_prefill_flag: + # print(f"to cpu") + self.load_layer_to(decoder_layer, InferenceState.UNLOAD) + torch.cuda.empty_cache() + hidden_states = layer_outputs[0] + + if use_cache and len(layer_outputs) > 1: + next_decoder_cache = layer_outputs[2 if output_attentions else 1] + else: + next_decoder_cache = None + + if output_attentions: + all_self_attns += (layer_outputs[1],) + + if output_router_logits and layer_outputs[-1] is not None: + all_router_logits += (layer_outputs[-1],) + + hidden_states = self.norm(hidden_states) + + if per_layer_prefill_flag: + per_layer_prefill_flag = False + for layer in self.layers: + self.load_layer_to(layer, InferenceState.GENERATE) + if output_hidden_states: + all_hidden_states += (hidden_states,) + + next_cache = None + if use_cache: + if next_decoder_cache is not None: + next_cache = ( + next_decoder_cache.to_legacy_cache() + if use_legacy_cache + else next_decoder_cache + ) + else: + next_cache = past_key_values + + if not return_dict: + return tuple( + v + for v in [ + hidden_states, + next_cache, + all_hidden_states, + all_self_attns, + all_router_logits, + ] + if v is not None + ) + return MoeModelOutputWithPast( + last_hidden_state=hidden_states, + past_key_values=next_cache, + hidden_states=all_hidden_states, + attentions=all_self_attns, + router_logits=all_router_logits, + ) + + def load_layer_to(self, layer: Qwen2MoeDecoderLayer, target: InferenceState): + assert isinstance( + layer, Qwen2MoeDecoderLayer + ), "module should be nn.ModuleList of decoder layers" + + # TODO Support restore to original device, not only cuda + device = "cpu" if target == InferenceState.UNLOAD else "cuda" + + # attn + layer.self_attn.q_proj.set_inference_mode(target) + layer.self_attn.k_proj.set_inference_mode(target) + layer.self_attn.v_proj.set_inference_mode(target) + layer.self_attn.o_proj.set_inference_mode(target) + layer.self_attn.rotary_emb = layer.self_attn.rotary_emb.to(device) + + # mlp + if isinstance(layer.mlp, Qwen2MoeSparseMoeBlock): + layer.mlp.gate.set_inference_mode(target) + layer.mlp.experts.set_inference_mode(target) + layer.mlp.shared_expert.gate_proj.set_inference_mode(target) + layer.mlp.shared_expert.up_proj.set_inference_mode(target) + layer.mlp.shared_expert.down_proj.set_inference_mode(target) + layer.mlp.shared_expert.act_fn.to(device) + layer.mlp.shared_expert_gate.to(device) + else: + layer.mlp.gate_proj.set_inference_mode(target) + layer.mlp.up_proj.set_inference_mode(target) + layer.mlp.down_proj.set_inference_mode(target) + layer.mlp.act_fn.to(device) + # layer norm + layer.input_layernorm.to(device) + layer.post_attention_layernorm.to(device) + + +DeepseekV2_INPUTS_DOCSTRING = r""" + Args: + input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): + Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide + it. + + Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and + [`PreTrainedTokenizer.__call__`] for details. + + [What are input IDs?](../glossary#input-ids) + attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*): + Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`: + + - 1 for tokens that are **not masked**, + - 0 for tokens that are **masked**. + + [What are attention masks?](../glossary#attention-mask) + + Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and + [`PreTrainedTokenizer.__call__`] for details. + + If `past_key_values` is used, optionally only the last `input_ids` have to be input (see + `past_key_values`). + + If you want to change padding behavior, you should read [`modeling_opt._prepare_decoder_attention_mask`] + and modify to your needs. See diagram 1 in [the paper](https://arxiv.org/abs/1910.13461) for more + information on the default strategy. + + - 1 indicates the head is **not masked**, + - 0 indicates the head is **masked**. + position_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0, + config.n_positions - 1]`. + + [What are position IDs?](../glossary#position-ids) + past_key_values (`Cache` or `tuple(tuple(torch.FloatTensor))`, *optional*): + Pre-computed hidden-states (key and values in the self-attention blocks and in the cross-attention + blocks) that can be used to speed up sequential decoding. This typically consists in the `past_key_values` + returned by the model at a previous stage of decoding, when `use_cache=True` or `config.use_cache=True`. + + Two formats are allowed: + - a [`~cache_utils.Cache`] instance; + - Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of + shape `(batch_size, num_heads, sequence_length, embed_size_per_head)`). This is also known as the legacy + cache format. + + The model will output the same cache format that is fed as input. If no `past_key_values` are passed, the + legacy cache format will be returned. + + If `past_key_values` are used, the user can optionally input only the last `input_ids` (those that don't + have their past key value states given to this model) of shape `(batch_size, 1)` instead of all `input_ids` + of shape `(batch_size, sequence_length)`. + inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*): + Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This + is useful if you want more control over how to convert `input_ids` indices into associated vectors than the + model's internal embedding lookup matrix. + use_cache (`bool`, *optional*): + If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see + `past_key_values`). + output_attentions (`bool`, *optional*): + Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned + tensors for more detail. + output_hidden_states (`bool`, *optional*): + Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for + more detail. + return_dict (`bool`, *optional*): + Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple. +""" + + +class KDeepseekV2Model(BaseInjectedModule): + """ + Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`DeepseekV2DecoderLayer`] + + Args: + config: DeepseekV2Config + """ + + def __init__( + self, + key: str, + gguf_loader: GGUFLoader, + config: PretrainedConfig, + orig_module: nn.Module, + device: str = "cuda", + per_layer_prefill_intput_threshold: int = 30000, # if None, no per-layer prefill + transfer_map: dict = None, + **kwargs, + ): + BaseInjectedModule.__init__( + self, key, gguf_loader, config, orig_module, device, **kwargs + ) + self.per_layer_prefill_intput_threshold = per_layer_prefill_intput_threshold + self.transfer_map = transfer_map + self.stream_device_map = dict() + + @add_start_docstrings_to_model_forward(DeepseekV2_INPUTS_DOCSTRING) + def forward( + self, + input_ids: torch.LongTensor = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[List[torch.FloatTensor]] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + cache_position: Optional[torch.LongTensor] = None, + per_layer_prefill_intput_threshold: ( + int | None + ) = None, # if None, no per-layer prefill + ) -> Union[Tuple, BaseModelOutputWithPast]: + self.gradient_checkpointing = False + if per_layer_prefill_intput_threshold is None: + per_layer_prefill_intput_threshold = self.per_layer_prefill_intput_threshold + per_layer_prefill_flag = False + seq_lenth = ( + inputs_embeds.size(1) if inputs_embeds is not None else input_ids.size(1) + ) + if ( + per_layer_prefill_intput_threshold + and per_layer_prefill_intput_threshold < seq_lenth + ): + per_layer_prefill_flag = True + for layer in self.layers: + self.load_layer_to(layer, InferenceState.UNLOAD) + torch.cuda.empty_cache() + else: + pass + output_attentions = ( + output_attentions + if output_attentions is not None + else self.config.output_attentions + ) + output_hidden_states = ( + output_hidden_states + if output_hidden_states is not None + else self.config.output_hidden_states + ) + use_cache = use_cache if use_cache is not None else self.config.use_cache + + return_dict = ( + return_dict if return_dict is not None else self.config.use_return_dict + ) + + # retrieve input_ids and inputs_embeds + if input_ids is not None and inputs_embeds is not None: + raise ValueError( + "You cannot specify both input_ids and inputs_embeds at the same time" + ) + elif input_ids is not None: + batch_size, seq_length = input_ids.shape[:2] + elif inputs_embeds is not None: + batch_size, seq_length = inputs_embeds.shape[:2] + else: + raise ValueError("You have to specify either input_ids or inputs_embeds") + + if self.gradient_checkpointing and self.training: + if use_cache: + logger.warning_once( + "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`transformers." + ) + use_cache = False + + past_key_values_length = 0 + if use_cache: + use_legacy_cache = not isinstance(past_key_values, Cache) + if use_legacy_cache: + past_key_values = DynamicCache.from_legacy_cache(past_key_values) + past_key_values_length = past_key_values.get_usable_length(seq_length) + + if inputs_embeds is None: + org_device = input_ids.device + # TODO move to embed_tokens's device, not hard code to cpu + # input_ids = input_ids.to("cpu") + input_ids = input_ids.to(self.embed_tokens.weight.device) + inputs_embeds = self.embed_tokens(input_ids).to(org_device) + input_ids = input_ids.to(org_device) + + if cache_position is None: + past_seen_tokens = ( + past_key_values.get_seq_length() if past_key_values is not None else 0 + ) + cache_position = torch.arange( + past_seen_tokens, + past_seen_tokens + inputs_embeds.shape[1], + device=inputs_embeds.device, + ) + + if position_ids is None: + position_ids = cache_position.unsqueeze(0) + + if inputs_embeds.device.type == "xpu" and position_ids is not None: + cos, sin = self.layers[0].self_attn.rotary_emb(inputs_embeds, + position_ids) + position_embeddings = (cos, sin) + else: + position_embeddings = None + + if per_layer_prefill_flag: + causal_mask = None + else: + if (os.name == 'nt' + or get_compute_capability() < 8 + or (self.transfer_map is not None and 'cpu' in self.transfer_map.values()) + or device_manager.gpu_vendor != GPUVendor.NVIDIA): + # print("for Windows or GPU before ampere, use forward_windows") + # only use mask in forward windows or can't flash attn + causal_mask = self._update_causal_mask( + attention_mask, inputs_embeds, cache_position, past_key_values, output_attentions + ) + else: + causal_mask = None + + # embed positions + hidden_states = inputs_embeds + if per_layer_prefill_flag: + print(f"Total length of input_ids: {hidden_states.size(1)}") + + # decoder layers + all_hidden_states = () if output_hidden_states else None + all_self_attns = () if output_attentions else None + next_decoder_cache = None + + t_gpu = 0 + t_cpu = 0 + t_f = 0 + + for i, decoder_layer in enumerate(self.layers): + # print(f"@@@@@@@@@@@@@@@@@layer {i}@@@@@@@@@@@@@@@@@@@@ \n") + if self.transfer_map is not None and i in self.transfer_map: + prev_stream = torch.cuda.current_stream() + cur_device = self.transfer_map[i] + if cur_device not in self.stream_device_map and cur_device.lower() != "cpu": + self.stream_device_map[cur_device] = torch.cuda.Stream(cur_device) + if cur_device.lower() != "cpu": + torch.cuda.set_device(cur_device) + self.stream_device_map[cur_device].wait_stream(prev_stream) + torch.cuda.set_stream(self.stream_device_map[cur_device]) + hidden_states = hidden_states.to( + self.transfer_map[i], non_blocking=True + ) + causal_mask = ( + causal_mask.to(self.transfer_map[i], non_blocking=True) + if causal_mask is not None + else None + ) + position_ids = ( + position_ids.to(self.transfer_map[i], non_blocking=True) + if position_ids is not None + else None + ) + cache_position = ( + cache_position.to(self.transfer_map[i], non_blocking=True) + if cache_position is not None + else None + ) + + if output_hidden_states: + all_hidden_states += (hidden_states,) + + if self.gradient_checkpointing and self.training: + layer_outputs = self._gradient_checkpointing_func( + decoder_layer.__call__, + hidden_states, + causal_mask, + position_ids, + past_key_values, + output_attentions, + use_cache, + cache_position, + ) + else: + t3 = time.time() + if per_layer_prefill_flag: + # print(f"to gpu") + self.load_layer_to(decoder_layer, InferenceState.PREFILL) + torch.cuda.empty_cache() + t4 = time.time() + # with open("log.txt", "a") as f: + # f.write(f"@@@@@@@@@@@@@@@@@layer {i}@@@@@@@@@@@@@@@@@@@@ \n") + layer_outputs = decoder_layer( + hidden_states, + attention_mask=causal_mask, + position_ids=position_ids, + past_key_value=past_key_values, + output_attentions=output_attentions, + use_cache=use_cache, + cache_position=cache_position, + position_embeddings=position_embeddings, + ) + t5 = time.time() + if per_layer_prefill_flag: + # print(f"to cpu") + self.load_layer_to(decoder_layer, InferenceState.UNLOAD) + torch.cuda.empty_cache() + t6 = time.time() + t_gpu += t4 - t3 + t_cpu += t6 - t5 + t_f += t5 - t4 + + hidden_states = layer_outputs[0] + + # @@@@@@@ TODO open this notes, tmp close to fit deepseekv3 + if use_cache: + next_decoder_cache = layer_outputs[2 if output_attentions else 1] + + if output_attentions: + all_self_attns += (layer_outputs[1],) + + hidden_states = self.norm(hidden_states) + # with open("log.txt", "a") as f: + # f.write(f"@@@After layers\n") + # f.write(f"hidden_states={hidden_states}\n") + # f.write(f"hidden_states.shape={hidden_states.shape}\n") + + if per_layer_prefill_flag: + t6 = time.time() + # print(f"restore") + per_layer_prefill_flag = False + for layer in self.layers: + self.load_layer_to(layer, InferenceState.GENERATE) + torch.cuda.empty_cache() + t7 = time.time() + + print( + f"total time: {t7-t3}, \n layer num{len(self.layers)}, gpu time: {t_gpu}, cpu time: {t_cpu}, forward time: {t_f}, restore time: {t7-t6}" + ) + + # add hidden states from the last decoder layer + if output_hidden_states: + all_hidden_states += (hidden_states,) + + next_cache = None + if use_cache: + next_cache = ( + next_decoder_cache.to_legacy_cache() + if use_legacy_cache + else next_decoder_cache + ) + if not return_dict: + return tuple( + v + for v in [hidden_states, next_cache, all_hidden_states, all_self_attns] + if v is not None + ) + return BaseModelOutputWithPast( + last_hidden_state=hidden_states, + past_key_values=next_cache, + hidden_states=all_hidden_states, + attentions=all_self_attns, + ) + + def load_layer_to(self, layer: DeepseekV2DecoderLayer, target: InferenceState): + assert isinstance( + layer, DeepseekV2DecoderLayer + ), "module should be nn.ModuleList of decoder layers" + + # TODO Support restore to original device, not only cuda + device = "cpu" if target == InferenceState.UNLOAD else "cuda" + + # TODO Support DFS to auto use {to, set_inference_mode} according to the module type + + # attn + layer.self_attn.to(device) # + + # mlp + if isinstance(layer.mlp, DeepseekV2MoE): + layer.mlp.gate.to(device) + layer.mlp.experts.set_inference_mode(target) + layer.mlp.shared_experts.gate_proj.set_inference_mode(target) + layer.mlp.shared_experts.up_proj.set_inference_mode(target) + layer.mlp.shared_experts.down_proj.set_inference_mode(target) + layer.mlp.shared_experts.act_fn.to(device) + # layer.mlp.shared_expert_gate.to(device) + else: + layer.mlp.gate_proj.set_inference_mode(target) + layer.mlp.up_proj.set_inference_mode(target) + layer.mlp.down_proj.set_inference_mode(target) + layer.mlp.act_fn.to(device) + # layer norm + layer.input_layernorm.to(device) + layer.post_attention_layernorm.to(device) + + +LLAMA_START_DOCSTRING = r""" + This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the + library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads + etc.) + + This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass. + Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage + and behavior. + + Parameters: + config ([`LlamaConfig`]): + Model configuration class with all the parameters of the model. Initializing with a config file does not + load the weights associated with the model, only the configuration. Check out the + [`~PreTrainedModel.from_pretrained`] method to load the model weights. +""" + +LLAMA_INPUTS_DOCSTRING = r""" + Args: + input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): + Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide + it. + + Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and + [`PreTrainedTokenizer.__call__`] for details. + + [What are input IDs?](../glossary#input-ids) + attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*): + Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`: + + - 1 for tokens that are **not masked**, + - 0 for tokens that are **masked**. + + [What are attention masks?](../glossary#attention-mask) + + Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and + [`PreTrainedTokenizer.__call__`] for details. + + If `past_key_values` is used, optionally only the last `input_ids` have to be input (see + `past_key_values`). + + If you want to change padding behavior, you should read [`modeling_opt._prepare_decoder_attention_mask`] + and modify to your needs. See diagram 1 in [the paper](https://arxiv.org/abs/1910.13461) for more + information on the default strategy. + + - 1 indicates the head is **not masked**, + - 0 indicates the head is **masked**. + position_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0, + config.n_positions - 1]`. + + [What are position IDs?](../glossary#position-ids) + past_key_values (`Cache` or `tuple(tuple(torch.FloatTensor))`, *optional*): + Pre-computed hidden-states (key and values in the self-attention blocks and in the cross-attention + blocks) that can be used to speed up sequential decoding. This typically consists in the `past_key_values` + returned by the model at a previous stage of decoding, when `use_cache=True` or `config.use_cache=True`. + + Two formats are allowed: + - a [`~cache_utils.Cache`] instance; + - Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of + shape `(batch_size, num_heads, sequence_length, embed_size_per_head)`). This is also known as the legacy + cache format. + + The model will output the same cache format that is fed as input. If no `past_key_values` are passed, the + legacy cache format will be returned. + + If `past_key_values` are used, the user can optionally input only the last `input_ids` (those that don't + have their past key value states given to this model) of shape `(batch_size, 1)` instead of all `input_ids` + of shape `(batch_size, sequence_length)`. + inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*): + Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This + is useful if you want more control over how to convert `input_ids` indices into associated vectors than the + model's internal embedding lookup matrix. + use_cache (`bool`, *optional*): + If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see + `past_key_values`). + output_attentions (`bool`, *optional*): + Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned + tensors for more detail. + output_hidden_states (`bool`, *optional*): + Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for + more detail. + return_dict (`bool`, *optional*): + Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple. + cache_position (`torch.LongTensor` of shape `(sequence_length)`, *optional*): + Indices depicting the position of the input sequence tokens in the sequence. Contrarily to `position_ids`, + this tensor is not affected by padding. It is used to update the cache in the correct position and to infer + the complete sequence length. +""" + + +@add_start_docstrings( + "The bare LLaMA Model outputting raw hidden-states without any specific head on top.", + LLAMA_START_DOCSTRING, +) +class LlamaPreTrainedModel(PreTrainedModel): + config_class = LlamaConfig + base_model_prefix = "model" + supports_gradient_checkpointing = True + _no_split_modules = ["LlamaDecoderLayer"] + _skip_keys_device_placement = ["past_key_values"] + _supports_flash_attn_2 = True + _supports_sdpa = True + _supports_cache_class = True + _supports_quantized_cache = True + _supports_static_cache = True + + def _init_weights(self, module): + std = self.config.initializer_range + if isinstance(module, nn.Linear): + module.weight.data.normal_(mean=0.0, std=std) + if module.bias is not None: + module.bias.data.zero_() + elif isinstance(module, nn.Embedding): + module.weight.data.normal_(mean=0.0, std=std) + if module.padding_idx is not None: + module.weight.data[module.padding_idx].zero_() + + +class KLlamaModel(BaseInjectedModule): + """ + Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`LlamaDecoderLayer`] + + Args: + config: LlamaConfig + """ + + dynamic_sdpa = None + + def __init__( + self, + key: str, + gguf_loader: GGUFLoader, + config: PretrainedConfig, + orig_module: nn.Module, + device: str = "cuda", + per_layer_prefill_intput_threshold: int = 30000, # if None, no per-layer prefill + transfer_map: dict = None, + **kwargs, + ): + + BaseInjectedModule.__init__( + self, key, gguf_loader, config, orig_module, device, **kwargs + ) + self.per_layer_prefill_intput_threshold = per_layer_prefill_intput_threshold + self.transfer_map = transfer_map + self.stream_device_map = dict() + user_path: str = os.path.expanduser('~') + localstore_path: str = os.path.join(user_path,'.ktransformers') + config_path: str = os.path.join(localstore_path,Config.CONFIG_FILE_NAME) + with open(config_path,"r") as file: + config_yaml = yaml.safe_load(file.read()) + self.long_context_config = config_yaml.get("long_context") + self.ext_config = config_yaml.get("ext") + + KLlamaModel.dynamic_sdpa = DynamicScaledDotProductAttention( + max_seq_len=self.long_context_config["max_seq_len"], + block_size=self.long_context_config["block_size"], + config=config, + device=torch.device("cuda"), + local_windows_len=self.long_context_config["local_windows_len"], + topk=self.long_context_config["second_select_num"], + threads_num=self.ext_config["cpu_infer"], + anchor_type=self.long_context_config["anchor_type"], + kv_type=self.long_context_config["kv_type"], + dense_layer_num=self.long_context_config["dense_layer_num"], + anchor_num=self.long_context_config["anchor_num"], + preselect_block=self.long_context_config["preselect_block"], + block_selection_mode=self.long_context_config["head_select_mode"], + preselect_block_count=self.long_context_config["preselect_block_count"], + layer_step=self.long_context_config["layer_step"], + token_step=self.long_context_config["token_step"], + prefill_chunk_size=self.long_context_config["chunk_size"], + use_attn_sparsity=False, + ) + + def get_input_embeddings(self): + return self.embed_tokens + + def set_input_embeddings(self, value): + self.embed_tokens = value + + @add_start_docstrings_to_model_forward(LLAMA_INPUTS_DOCSTRING) + def forward( + self, + input_ids: torch.LongTensor = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[Union[Cache, List[torch.FloatTensor]]] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + cache_position: Optional[torch.LongTensor] = None, + ) -> Union[Tuple, BaseModelOutputWithPast]: + output_attentions = ( + output_attentions + if output_attentions is not None + else self.config.output_attentions + ) + output_hidden_states = ( + output_hidden_states + if output_hidden_states is not None + else self.config.output_hidden_states + ) + use_cache = use_cache if use_cache is not None else self.config.use_cache + return_dict = ( + return_dict if return_dict is not None else self.config.use_return_dict + ) + + if (input_ids is None) ^ (inputs_embeds is not None): + raise ValueError( + "You cannot specify both input_ids and inputs_embeds at the same time, and must specify either one" + ) + + if self.gradient_checkpointing and self.training and use_cache: + logger.warning_once( + "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`." + ) + use_cache = False + + return_legacy_cache = False + if ( + use_cache and not isinstance(past_key_values, Cache) and not self.training + ): # kept for BC (non `Cache` `past_key_values` inputs) + return_legacy_cache = True + past_key_values = DynamicCache.from_legacy_cache(past_key_values) + logger.warning_once( + "We detected that you are passing `past_key_values` as a tuple and this is deprecated and will be removed in v4.43. " + "Please use an appropriate `Cache` class (https://huggingface.co/docs/transformers/v4.41.3/en/internal/generation_utils#transformers.Cache)" + ) + + if cache_position is None: + past_seen_tokens = ( + past_key_values.get_seq_length() if past_key_values is not None else 0 + ) + cache_position = torch.arange( + past_seen_tokens, + past_seen_tokens + inputs_embeds.shape[1], + device="cuda", + ) + if position_ids is None: + position_ids = cache_position.unsqueeze(0) + + causal_mask = None + chunck_size = self.long_context_config["chunk_size"] + cur_idx = 0 + if inputs_embeds is None: + inputs_embeds = self.embed_tokens(input_ids.to("cpu")) + q_len = cache_position.size(0) + + # generate + if q_len == 1: + x = inputs_embeds[:, -1:, :] + position_ids = position_ids[:, -1:] + return self.forward_chunk( + x, + causal_mask, + position_ids, + past_key_values, + output_attentions, + use_cache, + cache_position, + output_hidden_states, + return_dict, + ) + elif q_len <= chunck_size: + inputs_embeds = inputs_embeds.to('cuda') + output = self.forward_chunk( + inputs_embeds, + causal_mask, + position_ids, + past_key_values, + output_attentions, + use_cache, + cache_position, + output_hidden_states, + return_dict, + ) + KLlamaModel.dynamic_sdpa.calc_anchor(cache_position[-1] + 1) + KLlamaModel.dynamic_sdpa.clear_importance(cache_position[-1] + 1) + return output + cur_idx = 0 + assert ( + output_attentions == False + ), "output_attentions is not supported when using chunked attention" + attn_output = None + # prefill + KLlamaModel.dynamic_sdpa.remaining_length = q_len + while cur_idx < q_len: + print(f'current prefill length: {cur_idx}') + chunk_mask = None + if inputs_embeds.device.type == 'cpu': + tmp_inputs_embeds = inputs_embeds[:, cur_idx : min(cur_idx + chunck_size, q_len)].to("cuda") + else: + tmp_inputs_embeds = inputs_embeds[:, cur_idx : min(cur_idx + chunck_size, q_len)] + output_with_past = self.forward_chunk( + tmp_inputs_embeds, + chunk_mask, + position_ids[:, cur_idx : min(cur_idx + chunck_size, q_len)], + past_key_values, + output_attentions, + use_cache, + cache_position[cur_idx : min(cur_idx + chunck_size, q_len)], + ) + cur_output = output_with_past.last_hidden_state + KLlamaModel.dynamic_sdpa.remaining_length -= ( + min(cur_idx + chunck_size, q_len) - cur_idx + ) + cur_idx += chunck_size + # if attn_output is None: + attn_output = cur_output + # else: + # attn_output = torch.cat((attn_output, cur_output), dim=-2) + + KLlamaModel.dynamic_sdpa.calc_anchor(cache_position[-1] + 1) + KLlamaModel.dynamic_sdpa.clear_importance(cache_position[-1] + 1) + return BaseModelOutputWithPast(last_hidden_state=attn_output) + + def forward_chunk( + self, + inputs_embeds, + causal_mask, + position_ids, + past_key_values, + output_attentions, + use_cache, + cache_position, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + ): + + output_hidden_states = ( + output_hidden_states + if output_hidden_states is not None + else self.config.output_hidden_states + ) + return_legacy_cache = False + if use_cache and not isinstance( + past_key_values, Cache + ): # kept for BC (non `Cache` `past_key_values` inputs) + return_legacy_cache = True + past_key_values = DynamicCache.from_legacy_cache(past_key_values) + return_dict = ( + return_dict if return_dict is not None else self.config.use_return_dict + ) + + hidden_states = inputs_embeds + + # create position embeddings to be shared across the decoder layers + position_embeddings = self.rotary_emb(hidden_states, position_ids) + + # decoder layers + all_hidden_states = () if output_hidden_states else None + all_self_attns = () if output_attentions else None + next_decoder_cache = None + # decoder layers + all_hidden_states = () if output_hidden_states else None + all_self_attns = () if output_attentions else None + next_decoder_cache = None + + for decoder_layer in self.layers: + if output_hidden_states: + all_hidden_states += (hidden_states,) + + if self.gradient_checkpointing and self.training: + layer_outputs = self._gradient_checkpointing_func( + decoder_layer.__call__, + hidden_states, + causal_mask, + position_ids, + past_key_values, + output_attentions, + use_cache, + cache_position, + position_embeddings, + ) + else: + layer_outputs = decoder_layer( + hidden_states, + attention_mask=causal_mask, + position_ids=position_ids, + past_key_value=past_key_values, + output_attentions=output_attentions, + use_cache=use_cache, + cache_position=cache_position, + position_embeddings=position_embeddings, + ) + + hidden_states = layer_outputs[0] + + if use_cache: + next_decoder_cache = layer_outputs[2 if output_attentions else 1] + + if output_attentions: + all_self_attns += (layer_outputs[1],) + + hidden_states = self.norm(hidden_states) + + # add hidden states from the last decoder layer + if output_hidden_states: + all_hidden_states += (hidden_states,) + + next_cache = next_decoder_cache if use_cache else None + if return_legacy_cache: + next_cache = next_cache.to_legacy_cache() + + if not return_dict: + return tuple( + v + for v in [hidden_states, next_cache, all_hidden_states, all_self_attns] + if v is not None + ) + return BaseModelOutputWithPast( + last_hidden_state=hidden_states, + past_key_values=next_cache, + hidden_states=all_hidden_states, + attentions=all_self_attns, + ) + + def _update_causal_mask( + self, + attention_mask: torch.Tensor, + input_tensor: torch.Tensor, + cache_position: torch.Tensor, + past_key_values: Cache, + output_attentions: bool, + ): + # TODO: As of torch==2.2.0, the `attention_mask` passed to the model in `generate` is 2D and of dynamic length even when the static + # KV cache is used. This is an issue for torch.compile which then recaptures cudagraphs at each decode steps due to the dynamic shapes. + # (`recording cudagraph tree for symint key 13`, etc.), which is VERY slow. A workaround is `@torch.compiler.disable`, but this prevents using + # `fullgraph=True`. See more context in https://github.com/huggingface/transformers/pull/29114 + + if self.config._attn_implementation == "flash_attention_2": + if attention_mask is not None and 0.0 in attention_mask: + return attention_mask + return None + + # For SDPA, when possible, we will rely on its `is_causal` argument instead of its `attn_mask` argument, in + # order to dispatch on Flash Attention 2. This feature is not compatible with static cache, as SDPA will fail + # to infer the attention mask. + past_seen_tokens = ( + past_key_values.get_seq_length() if past_key_values is not None else 0 + ) + using_static_cache = isinstance(past_key_values, StaticCache) + + # When output attentions is True, sdpa implementation's forward method calls the eager implementation's forward + if ( + self.config._attn_implementation == "sdpa" + and not using_static_cache + and not output_attentions + ): + if AttentionMaskConverter._ignore_causal_mask_sdpa( + attention_mask, + inputs_embeds=input_tensor, + past_key_values_length=past_seen_tokens, + is_training=self.training, + ): + return None + + dtype, device = input_tensor.dtype, input_tensor.device + min_dtype = torch.finfo(dtype).min + sequence_length = input_tensor.shape[1] + if using_static_cache: + target_length = past_key_values.get_max_length() + else: + target_length = ( + attention_mask.shape[-1] + if isinstance(attention_mask, torch.Tensor) + else past_seen_tokens + sequence_length + 1 + ) + + if attention_mask is not None and attention_mask.dim() == 4: + # in this case we assume that the mask comes already in inverted form and requires no inversion or slicing + if attention_mask.max() != 0: + raise ValueError( + "Custom 4D attention mask should be passed in inverted form with max==0`" + ) + causal_mask = attention_mask + else: + causal_mask = torch.full( + (sequence_length, target_length), + fill_value=min_dtype, + dtype=dtype, + device=device, + ) + if sequence_length != 1: + causal_mask = torch.triu(causal_mask, diagonal=1) + causal_mask *= torch.arange( + target_length, device=device + ) > cache_position.reshape(-1, 1) + causal_mask = causal_mask[None, None, :, :].expand( + input_tensor.shape[0], 1, -1, -1 + ) + if attention_mask is not None: + causal_mask = ( + causal_mask.clone() + ) # copy to contiguous memory for in-place edit + mask_length = attention_mask.shape[-1] + padding_mask = ( + causal_mask[:, :, :, :mask_length] + + attention_mask[:, None, None, :] + ) + padding_mask = padding_mask == 0 + causal_mask[:, :, :, :mask_length] = causal_mask[ + :, :, :, :mask_length + ].masked_fill(padding_mask, min_dtype) + if ( + self.config._attn_implementation == "sdpa" + and attention_mask is not None + and attention_mask.device.type == "cuda" + and not output_attentions + ): + # Attend to all tokens in fully masked rows in the causal_mask, for example the relevant first rows when + # using left padding. This is required by F.scaled_dot_product_attention memory-efficient attention path. + # Details: https://github.com/pytorch/pytorch/issues/110213 + causal_mask = AttentionMaskConverter._unmask_unattended( + causal_mask, min_dtype + ) + + return causal_mask + + + +QWEN3MOE_INPUTS_DOCSTRING = r""" + Args: + input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): + Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide + it. + + Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and + [`PreTrainedTokenizer.__call__`] for details. + + [What are input IDs?](../glossary#input-ids) + attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*): + Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`: + + - 1 for tokens that are **not masked**, + - 0 for tokens that are **masked**. + + [What are attention masks?](../glossary#attention-mask) + + Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and + [`PreTrainedTokenizer.__call__`] for details. + + If `past_key_values` is used, optionally only the last `decoder_input_ids` have to be input (see + `past_key_values`). + + If you want to change padding behavior, you should read [`modeling_opt._prepare_decoder_attention_mask`] + and modify to your needs. See diagram 1 in [the paper](https://arxiv.org/abs/1910.13461) for more + information on the default strategy. + + - 1 indicates the head is **not masked**, + - 0 indicates the head is **masked**. + position_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0, + config.n_positions - 1]`. + + [What are position IDs?](../glossary#position-ids) + past_key_values (`Cache` or `tuple(tuple(torch.FloatTensor))`, *optional*): + Pre-computed hidden-states (key and values in the self-attention blocks and in the cross-attention + blocks) that can be used to speed up sequential decoding. This typically consists in the `past_key_values` + returned by the model at a previous stage of decoding, when `use_cache=True` or `config.use_cache=True`. + + Two formats are allowed: + - a [`~cache_utils.Cache`] instance; + - Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of + shape `(batch_size, num_heads, sequence_length, embed_size_per_head)`). This is also known as the legacy + cache format. + + The model will output the same cache format that is fed as input. If no `past_key_values` are passed, the + legacy cache format will be returned. + + If `past_key_values` are used, the user can optionally input only the last `input_ids` (those that don't + have their past key value states given to this model) of shape `(batch_size, 1)` instead of all `input_ids` + of shape `(batch_size, sequence_length)`. + inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*): + Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This + is useful if you want more control over how to convert `input_ids` indices into associated vectors than the + model's internal embedding lookup matrix. + use_cache (`bool`, *optional*): + If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see + `past_key_values`). + output_attentions (`bool`, *optional*): + Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned + tensors for more detail. + output_hidden_states (`bool`, *optional*): + Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for + more detail. + output_router_logits (`bool`, *optional*): + Whether or not to return the logits of all the routers. They are useful for computing the router loss, and + should not be returned during inference. + return_dict (`bool`, *optional*): + Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple. + cache_position (`torch.LongTensor` of shape `(sequence_length)`, *optional*): + Indices depicting the position of the input sequence tokens in the sequence. Contrarily to `position_ids`, + this tensor is not affected by padding. It is used to update the cache in the correct position and to infer + the complete sequence length. +""" + +class KQwen3MoeModel(BaseInjectedModule): + """ + Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`Qwen3MoeDecoderLayer`] + + Args: + config: Qwen3MoeConfig + """ + + def __init__( + self, + key: str, + gguf_loader: GGUFLoader, + config: PretrainedConfig, + orig_module: nn.Module, + device: str = "cuda", + per_layer_prefill_intput_threshold: int = 30000, # if None, no per-layer prefill + transfer_map: dict = None, + **kwargs, + ): + BaseInjectedModule.__init__( + self, key, gguf_loader, config, orig_module, device, **kwargs + ) + self.per_layer_prefill_intput_threshold = per_layer_prefill_intput_threshold + self.transfer_map = transfer_map + self.stream_device_map = dict() + self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads) + self.max_position_embeddings = config.max_position_embeddings + self.rope_theta = config.rope_theta + self.rotary_emb = Qwen2MoeRotaryEmbedding( + self.head_dim, + max_position_embeddings=self.max_position_embeddings, + base=self.rope_theta, + ) + + @add_start_docstrings_to_model_forward(QWEN3MOE_INPUTS_DOCSTRING) + def forward( + self, + input_ids: torch.LongTensor = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[List[torch.FloatTensor]] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + output_router_logits: Optional[bool] = None, + return_dict: Optional[bool] = None, + cache_position: Optional[torch.LongTensor] = None, + per_layer_prefill_intput_threshold: ( + int | None + ) = None, # if None or 0, close per-layer prefill + ) -> Union[Tuple, MoeModelOutputWithPast]: + # print(f'Total length of input_ids: {input_ids.size(1)}, {input_ids.size()}') + + if per_layer_prefill_intput_threshold is None: + per_layer_prefill_intput_threshold = self.per_layer_prefill_intput_threshold + per_layer_prefill_flag = False + seq_lenth = ( + inputs_embeds.size(1) if inputs_embeds is not None else input_ids.size(1) + ) + if ( + per_layer_prefill_intput_threshold + and per_layer_prefill_intput_threshold < seq_lenth + ): + per_layer_prefill_flag = True + for layer in self.layers: + self.load_layer_to(layer, InferenceState.UNLOAD) + else: + pass + output_attentions = ( + output_attentions + if output_attentions is not None + else self.config.output_attentions + ) + output_router_logits = ( + output_router_logits + if output_router_logits is not None + else self.config.output_router_logits + ) + output_hidden_states = ( + output_hidden_states + if output_hidden_states is not None + else self.config.output_hidden_states + ) + use_cache = use_cache if use_cache is not None else self.config.use_cache + + return_dict = ( + return_dict if return_dict is not None else self.config.use_return_dict + ) + + if (input_ids is None) ^ (inputs_embeds is not None): + raise ValueError( + "You cannot specify both input_ids and inputs_embeds at the same time, and must specify either one" + ) + + if self.gradient_checkpointing and self.training: + if use_cache: + logger.warning_once( + "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..." + ) + use_cache = False + + if use_cache and past_key_values is None: + past_key_values = DynamicCache() + # use_legacy_cache = False + # if use_cache and not isinstance(past_key_values, Cache): + # use_legacy_cache = True + # past_key_values = DynamicCache.from_legacy_cache(past_key_values) + # logger.warning_once( + # "We detected that you are passing `past_key_values` as a tuple and this is deprecated and will be removed in v4.43. " + # "Please use an appropriate `Cache` class (https://huggingface.co/docs/transformers/v4.41.3/en/internal/generation_utils#transformers.Cache)" + # ) + + if inputs_embeds is None: + input_ids = input_ids.to("cpu") + inputs_embeds = self.embed_tokens(input_ids) + inputs_embeds = inputs_embeds.to("cuda") + + if cache_position is None: + past_seen_tokens = ( + past_key_values.get_seq_length() if past_key_values is not None else 0 + ) + cache_position = torch.arange( + past_seen_tokens, + past_seen_tokens + inputs_embeds.shape[1], + device=inputs_embeds.device, + ) + if position_ids is None: + position_ids = cache_position.unsqueeze(0) + + causal_mask = self._update_causal_mask( + attention_mask, + inputs_embeds, + cache_position, + past_key_values, + output_attentions, + ) + + hidden_states = inputs_embeds + + # position_embeddings = self.rotary_emb(hidden_states, position_ids) + # decoder layers + all_hidden_states = () if output_hidden_states else None + all_self_attns = () if output_attentions else None + all_router_logits = () if output_router_logits else None + # next_decoder_cache = None + + for i, decoder_layer in enumerate(self.layers): + # if self.transfer_map is not None and i in self.transfer_map: + # prev_stream = torch.cuda.current_stream() + # cur_device = self.transfer_map[i] + # if cur_device not in self.stream_device_map: + # self.stream_device_map[cur_device] = torch.cuda.Stream(cur_device) + # torch.cuda.set_device(cur_device) + # self.stream_device_map[cur_device].wait_stream(prev_stream) + # torch.cuda.set_stream(self.stream_device_map[cur_device]) + # hidden_states = hidden_states.to( + # self.transfer_map[i], non_blocking=True + # ) + # causal_mask = ( + # causal_mask.to(self.transfer_map[i], non_blocking=True) + # if causal_mask is not None + # else None + # ) + # position_ids = ( + # position_ids.to(self.transfer_map[i], non_blocking=True) + # if position_ids is not None + # else None + # ) + # cache_position = ( + # cache_position.to(self.transfer_map[i], non_blocking=True) + # if cache_position is not None + # else None + # ) + + if output_hidden_states: + all_hidden_states += (hidden_states,) + + if self.gradient_checkpointing and self.training: + layer_outputs = self._gradient_checkpointing_func( + decoder_layer.__call__, + hidden_states, + causal_mask, + position_ids, + past_key_values, + output_attentions, + output_router_logits, + use_cache, + cache_position, + # position_embeddings, + ) + else: + if per_layer_prefill_flag: + # print(f"to gpu") + self.load_layer_to(decoder_layer, InferenceState.PREFILL) + torch.cuda.empty_cache() + layer_outputs = decoder_layer( + hidden_states, + attention_mask=causal_mask, + position_ids=position_ids, + past_key_value=past_key_values, + output_attentions=output_attentions, + output_router_logits=output_router_logits, + use_cache=use_cache, + cache_position=cache_position, + # position_embeddings=position_embeddings, + ) + if per_layer_prefill_flag: + # print(f"to cpu") + self.load_layer_to(decoder_layer, InferenceState.UNLOAD) + torch.cuda.empty_cache() + hidden_states = layer_outputs[0] + # use_cache=False + # if use_cache: + # next_decoder_cache = layer_outputs[2 if output_attentions else 1] + + if output_attentions: + all_self_attns += (layer_outputs[1],) + + if output_router_logits and layer_outputs[-1] is not None: + all_router_logits += (layer_outputs[-1],) + + hidden_states = self.norm(hidden_states) + + if per_layer_prefill_flag: + per_layer_prefill_flag = False + for layer in self.layers: + self.load_layer_to(layer, InferenceState.GENERATE) + if output_hidden_states: + all_hidden_states += (hidden_states,) + + # next_cache = None + # if use_cache: + # next_cache = ( + # next_decoder_cache.to_legacy_cache() + # if use_legacy_cache + # else next_decoder_cache + # ) + + if not return_dict: + return tuple( + v + for v in [ + hidden_states, + past_key_values, + all_hidden_states, + all_self_attns, + all_router_logits, + ] + if v is not None + ) + return MoeModelOutputWithPast( + last_hidden_state=hidden_states, + past_key_values=past_key_values, + hidden_states=all_hidden_states, + attentions=all_self_attns, + router_logits=all_router_logits, + ) + + def load_layer_to(self, layer: Qwen3MoeDecoderLayer, target: InferenceState): + assert isinstance( + layer, Qwen3MoeDecoderLayer + ), "module should be nn.ModuleList of decoder layers" + + # TODO Support restore to original device, not only cuda + device = "cpu" if target == InferenceState.UNLOAD else "cuda" + + # attn + layer.self_attn.q_proj.set_inference_mode(target) + layer.self_attn.k_proj.set_inference_mode(target) + layer.self_attn.v_proj.set_inference_mode(target) + layer.self_attn.o_proj.set_inference_mode(target) + layer.self_attn.rotary_emb = layer.self_attn.rotary_emb.to(device) + + # mlp + if isinstance(layer.mlp, Qwen3MoeSparseMoeBlock): + layer.mlp.gate.set_inference_mode(target) + layer.mlp.experts.set_inference_mode(target) + layer.mlp.shared_expert.gate_proj.set_inference_mode(target) + layer.mlp.shared_expert.up_proj.set_inference_mode(target) + layer.mlp.shared_expert.down_proj.set_inference_mode(target) + layer.mlp.shared_expert.act_fn.to(device) + layer.mlp.shared_expert_gate.to(device) + else: + layer.mlp.gate_proj.set_inference_mode(target) + layer.mlp.up_proj.set_inference_mode(target) + layer.mlp.down_proj.set_inference_mode(target) + layer.mlp.act_fn.to(device) + # layer norm + layer.input_layernorm.to(device) + layer.post_attention_layernorm.to(device) diff --git a/archive/kt-sft/ktransformers/operators/triton_attention.py b/archive/kt-sft/ktransformers/operators/triton_attention.py new file mode 100644 index 0000000..aafdea0 --- /dev/null +++ b/archive/kt-sft/ktransformers/operators/triton_attention.py @@ -0,0 +1,385 @@ +# Adapted from +# https://github.com/sgl-project/sglang/blob/9f635ea50de920aa507f486daafba26a5b837574/python/sglang/srt/layers/attention/triton_ops/decode_attention.py +# which was originally adapted from +# https://github.com/ModelTC/lightllm/blob/96353e868a840db4d103138caf15ed9dbea8c186/lightllm/models/deepseek2/triton_kernel/gqa_flash_decoding_stage1.py +# https://github.com/ModelTC/lightllm/blob/96353e868a840db4d103138caf15ed9dbea8c186/lightllm/models/deepseek2/triton_kernel/gqa_flash_decoding_stage2.py + +import triton +import triton.language as tl +from ktransformers.util.vendors import device_manager, get_device, to_device, GPUVendor +@triton.jit +def tanh(x): + # Tanh is just a scaled sigmoid + return 2 * tl.sigmoid(2 * x) - 1 + +@triton.jit +def _fwd_grouped_kernel_stage1( + Q, + K_Buffer, + V_Buffer, + sm_scale, + Req_to_tokens, + B_Seqlen, + Att_Out, + stride_req_to_tokens_b, + stride_qbs, + stride_qh, + stride_buf_kbs, + stride_buf_kh, + stride_buf_vbs, + stride_buf_vh, + stride_mid_ob, + stride_mid_oh, + stride_mid_os, + kv_group_num: tl.constexpr, + q_head_num: tl.constexpr, + BLOCK_DMODEL: tl.constexpr, + BLOCK_DPE: tl.constexpr, + BLOCK_DV: tl.constexpr, + BLOCK_N: tl.constexpr, + BLOCK_H: tl.constexpr, + NUM_KV_SPLITS: tl.constexpr, + PAGE_SIZE: tl.constexpr, + logit_cap: tl.constexpr, + Lk: tl.constexpr, + Lv: tl.constexpr, +): + cur_batch = tl.program_id(0) + cur_head_id = tl.program_id(1) + cur_kv_head = cur_head_id // tl.cdiv(kv_group_num, BLOCK_H) + split_kv_id = tl.program_id(2) + + if kv_group_num > BLOCK_H: + VALID_BLOCK_H: tl.constexpr = BLOCK_H + else: + VALID_BLOCK_H: tl.constexpr = kv_group_num + cur_head = cur_head_id * VALID_BLOCK_H + tl.arange(0, BLOCK_H) + mask_h = cur_head < (cur_head_id + 1) * VALID_BLOCK_H + mask_h = mask_h & (cur_head < q_head_num) + + offs_d = tl.arange(0, BLOCK_DMODEL) + offs_dv = tl.arange(0, BLOCK_DV) + mask_d = offs_d < Lk + mask_dv = offs_dv < Lv + cur_batch_seq_len = tl.load(B_Seqlen + cur_batch) + cur_batch_req_idx = cur_batch + + offs_q = cur_batch * stride_qbs + cur_head[:, None] * stride_qh + offs_d[ + None, :] + q = tl.load(Q + offs_q, + mask=(mask_h[:, None]) & (mask_d[None, :]), + other=0.0) + + if BLOCK_DPE > 0: + offs_dpe = BLOCK_DMODEL + tl.arange(0, BLOCK_DPE) + mask_dpe = offs_dpe < Lk + off_qpe = (cur_batch * stride_qbs + cur_head[:, None] * stride_qh + + offs_dpe[None, :]) + qpe = tl.load(Q + off_qpe, + mask=(mask_h[:, None]) & (mask_dpe[None, :]), + other=0.0) + + kv_len_per_split = tl.cdiv(cur_batch_seq_len, NUM_KV_SPLITS) + split_kv_start = kv_len_per_split * split_kv_id + split_kv_end = tl.minimum(split_kv_start + kv_len_per_split, + cur_batch_seq_len) + + e_max = tl.zeros([BLOCK_H], dtype=tl.float32) - float("inf") + e_sum = tl.zeros([BLOCK_H], dtype=tl.float32) + acc = tl.zeros([BLOCK_H, BLOCK_DV], dtype=tl.float32) + + if split_kv_end > split_kv_start: + for start_n in range(split_kv_start, split_kv_end, BLOCK_N): + offs_n = start_n + tl.arange(0, BLOCK_N) + kv_page_number = tl.load( + Req_to_tokens + stride_req_to_tokens_b * cur_batch_req_idx + + offs_n // PAGE_SIZE, + mask=offs_n < split_kv_end, + other=0, + ) + kv_loc = kv_page_number * PAGE_SIZE + offs_n % PAGE_SIZE + offs_buf_k = (kv_loc[None, :] * stride_buf_kbs + + cur_kv_head * stride_buf_kh + offs_d[:, None]) + k = tl.load( + K_Buffer + offs_buf_k, + mask=(offs_n[None, :] < split_kv_end) & (mask_d[:, None]), + other=0.0, + ) + qk = tl.dot(q, k.to(q.dtype)) + + if BLOCK_DPE > 0: + offs_buf_kpe = (kv_loc[None, :] * stride_buf_kbs + + cur_kv_head * stride_buf_kh + + offs_dpe[:, None]) + kpe = tl.load( + K_Buffer + offs_buf_kpe, + mask=(offs_n[None, :] < split_kv_end) & + (mask_dpe[:, None]), + other=0.0, + ) + qk += tl.dot(qpe, kpe.to(qpe.dtype)) + qk *= sm_scale + + if logit_cap > 0: + qk = logit_cap * tanh(qk / logit_cap) + + qk = tl.where(mask_h[:, None] & (offs_n[None, :] < split_kv_end), + qk, float("-inf")) + + offs_buf_v = (kv_loc[:, None] * stride_buf_vbs + + cur_kv_head * stride_buf_vh + offs_dv[None, :]) + v = tl.load( + V_Buffer + offs_buf_v, + mask=(offs_n[:, None] < split_kv_end) & (mask_dv[None, :]), + other=0.0, + ) + + n_e_max = tl.maximum(tl.max(qk, 1), e_max) + re_scale = tl.exp(e_max - n_e_max) + p = tl.exp(qk - n_e_max[:, None]) + acc *= re_scale[:, None] + acc += tl.dot(p.to(v.dtype), v) + + e_sum = e_sum * re_scale + tl.sum(p, 1) + e_max = n_e_max + + offs_mid_o = (cur_batch * stride_mid_ob + + cur_head[:, None] * stride_mid_oh + + split_kv_id * stride_mid_os + offs_dv[None, :]) + + tl.store( + Att_Out + offs_mid_o, + acc / e_sum[:, None], + mask=(mask_h[:, None]) & (mask_dv[None, :]), + ) + + offs_mid_o_1 = (cur_batch * stride_mid_ob + cur_head * stride_mid_oh + + split_kv_id * stride_mid_os + Lv) + + tl.store( + Att_Out + offs_mid_o_1, + e_max + tl.log(e_sum), + mask=mask_h, + ) + +def _decode_grouped_att_m_fwd( + q, + k_buffer, + v_buffer, + att_out, + Req_to_tokens, + B_Seqlen, + num_kv_splits, + sm_scale, + page_size, + logit_cap, +): + BLOCK = 32 + Lk = k_buffer.shape[-1] + Lv = v_buffer.shape[-1] + + # [TODO] work around shmem limit on MI3xx + + # TODO: support hip + if device_manager.gpu_vendor == GPUVendor.AMD and Lk >= 576: + BLOCK = 16 + + if Lk == 576: + BLOCK_DMODEL = 512 + BLOCK_DPE = 64 + elif Lk == 288: + BLOCK_DMODEL = 256 + BLOCK_DPE = 32 + else: + BLOCK_DMODEL = triton.next_power_of_2(Lk) + BLOCK_DPE = 0 + BLOCK_DV = triton.next_power_of_2(Lv) + + batch, head_num = q.shape[0], q.shape[1] + kv_group_num = q.shape[1] // k_buffer.shape[-2] + + BLOCK_H = 16 + NUM_KV_SPLITS = num_kv_splits + grid = ( + batch, + triton.cdiv(head_num, min(BLOCK_H, kv_group_num)), + NUM_KV_SPLITS, + ) + + extra_kargs = {} + # TODO: support hip + """ + if is_hip_: + # https://rocm.docs.amd.com/en/docs-6.2.0/how-to/llm-fine-tuning-optimization/optimizing-triton-kernel.html + # https://github.com/triton-lang/triton/blob/main/third_party/amd/backend/compiler.py + extra_kargs = { + "waves_per_eu": 4, + "matrix_instr_nonkdim": 16, + "kpack": 2 + } + """ + + _fwd_grouped_kernel_stage1[grid]( + q, + k_buffer, + v_buffer, + sm_scale, + Req_to_tokens, + B_Seqlen, + att_out, + Req_to_tokens.stride(0), + q.stride(0), + q.stride(1), + k_buffer.stride(-3), # Assume (..., PAGE_SIZE, NUM_HEADS, HEAD_DIM) + k_buffer.stride(-2), # Assume (..., PAGE_SIZE, NUM_HEADS, HEAD_DIM) + v_buffer.stride(-3), # Assume (..., PAGE_SIZE, NUM_HEADS, HEAD_DIM) + v_buffer.stride(-2), # Assume (..., PAGE_SIZE, NUM_HEADS, HEAD_DIM) + att_out.stride(0), + att_out.stride(1), + att_out.stride(2), + kv_group_num=kv_group_num, + q_head_num=head_num, + BLOCK_DMODEL=BLOCK_DMODEL, + BLOCK_DPE=BLOCK_DPE, + BLOCK_DV=BLOCK_DV, + BLOCK_N=BLOCK, + BLOCK_H=BLOCK_H, + NUM_KV_SPLITS=NUM_KV_SPLITS, + PAGE_SIZE=page_size, + logit_cap=logit_cap, + num_warps=4, + num_stages=2, + Lk=Lk, + Lv=Lv, + **extra_kargs, + ) + +@triton.jit +def _fwd_kernel_stage2( + Mid_O, + o, + B_Seqlen, + stride_mid_ob, + stride_mid_oh, + stride_mid_os, + stride_obs, + stride_oh, + NUM_KV_SPLITS: tl.constexpr, + BLOCK_DV: tl.constexpr, + Lv: tl.constexpr, +): + cur_batch = tl.program_id(0) + cur_head = tl.program_id(1) + + cur_batch_seq_len = tl.load(B_Seqlen + cur_batch) + + offs_d = tl.arange(0, BLOCK_DV) + mask_d = offs_d < Lv + + e_sum = 0.0 + e_max = -float("inf") + acc = tl.zeros([BLOCK_DV], dtype=tl.float32) + + offs_v = cur_batch * stride_mid_ob + cur_head * stride_mid_oh + offs_d + offs_logic = cur_batch * stride_mid_ob + cur_head * stride_mid_oh + Lv + + for split_kv_id in range(0, NUM_KV_SPLITS): + kv_len_per_split = tl.cdiv(cur_batch_seq_len, NUM_KV_SPLITS) + split_kv_start = kv_len_per_split * split_kv_id + split_kv_end = tl.minimum(split_kv_start + kv_len_per_split, + cur_batch_seq_len) + + if split_kv_end > split_kv_start: + tv = tl.load(Mid_O + offs_v + split_kv_id * stride_mid_os, + mask=mask_d, + other=0.0) + tlogic = tl.load(Mid_O + offs_logic + split_kv_id * stride_mid_os) + n_e_max = tl.maximum(tlogic, e_max) + + old_scale = tl.exp(e_max - n_e_max) + acc *= old_scale + exp_logic = tl.exp(tlogic - n_e_max) + acc += exp_logic * tv + + e_sum = e_sum * old_scale + exp_logic + e_max = n_e_max + + tl.store( + o + cur_batch * stride_obs + cur_head * stride_oh + offs_d, + acc / e_sum, + mask=mask_d, + ) + +def _decode_softmax_reducev_fwd( + logits, + q, + o, + v_buffer, + b_seq_len, + num_kv_splits, +): + batch, head_num = q.shape[0], q.shape[1] + Lv = v_buffer.shape[-1] + BLOCK_DV = triton.next_power_of_2(Lv) + + NUM_KV_SPLITS = num_kv_splits + + extra_kargs = {} + # TODO: support hip + """ + if is_hip_: + # https://rocm.docs.amd.com/en/docs-6.2.0/how-to/llm-fine-tuning-optimization/optimizing-triton-kernel.html + # https://github.com/triton-lang/triton/blob/main/third_party/amd/backend/compiler.py + extra_kargs = { + "waves_per_eu": 4, + "matrix_instr_nonkdim": 16, + "kpack": 2 + } + """ + + grid = (batch, head_num) + _fwd_kernel_stage2[grid]( + logits, + o, + b_seq_len, + logits.stride(0), + logits.stride(1), + logits.stride(2), + o.stride(0), + o.stride(1), + NUM_KV_SPLITS=NUM_KV_SPLITS, + BLOCK_DV=BLOCK_DV, + Lv=Lv, + num_warps=4, + num_stages=2, + **extra_kargs, + ) + +def decode_attention_fwd_grouped( + q, + k_buffer, + v_buffer, + o, + req_to_token, + b_seq_len, + attn_logits, + num_kv_splits, + sm_scale, + page_size, + logit_cap=0.0, +): + _decode_grouped_att_m_fwd( + q, + k_buffer, + v_buffer, + attn_logits, + req_to_token, + b_seq_len, + num_kv_splits, + sm_scale, + page_size, + logit_cap, + ) + + _decode_softmax_reducev_fwd(attn_logits, q, o, v_buffer, b_seq_len, + num_kv_splits) diff --git a/archive/kt-sft/ktransformers/operators/triton_attention_prefill.py b/archive/kt-sft/ktransformers/operators/triton_attention_prefill.py new file mode 100644 index 0000000..a807ef3 --- /dev/null +++ b/archive/kt-sft/ktransformers/operators/triton_attention_prefill.py @@ -0,0 +1,206 @@ + +# Adapted from +# https://github.com/sgl-project/sglang/blob/9f635ea50de920aa507f486daafba26a5b837574/python/sglang/srt/layers/attention/triton_ops/prefill_attention.py +# which was originally adapted from +# https://github.com/ModelTC/lightllm/blob/f2a54f0912293f683bf1d1695fd12c4098a5bf82/lightllm/models/llama/triton_kernel/context_flashattention_nopad.py#L1 + +""" +Memory-efficient attention for prefill. +It supporst page size = 1. +""" + +# Adapted from +# https://github.com/ModelTC/lightllm/blob/f2a54f0912293f683bf1d1695fd12c4098a5bf82/lightllm/models/llama/triton_kernel/context_flashattention_nopad.py#L1 +import torch +import triton +import triton.language as tl + +is_cuda_available = torch.cuda.is_available() +if is_cuda_available: + CUDA_CAPABILITY = torch.cuda.get_device_capability() + + +@triton.jit +def _fwd_kernel( + Q, + K, + V, + sm_scale, + B_Start_Loc, + B_Seqlen, + Out, + stride_qbs, + stride_qh, + stride_kbs, + stride_kh, + stride_vbs, + stride_vh, + stride_obs, + stride_oh, + kv_group_num: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_DMODEL: tl.constexpr, + BLOCK_N: tl.constexpr, + IS_CAUSAL: tl.constexpr, + Lk: tl.constexpr, +): + cur_batch = tl.program_id(0) + cur_head = tl.program_id(1) + start_m = tl.program_id(2) + + cur_kv_head = cur_head // kv_group_num + + cur_batch_seq_len = tl.load(B_Seqlen + cur_batch) + cur_batch_in_all_start_index = tl.load(B_Start_Loc + cur_batch) + + block_start_loc = BLOCK_M * start_m + + # initialize offsets + offs_n = tl.arange(0, BLOCK_N) + offs_d = tl.arange(0, BLOCK_DMODEL) + offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M) + off_q = ( + (cur_batch_in_all_start_index + offs_m[:, None]) * stride_qbs + + cur_head * stride_qh + + offs_d[None, :] + ) + off_k = offs_n[None, :] * stride_kbs + cur_kv_head * stride_kh + offs_d[:, None] + off_v = offs_n[:, None] * stride_vbs + cur_kv_head * stride_vh + offs_d[None, :] + + mask_d = offs_d < Lk + + q = tl.load( + Q + off_q, + mask=(offs_m[:, None] < cur_batch_seq_len) & (mask_d[None, :]), + other=0.0, + ) + + k_ptrs = K + off_k + v_ptrs = V + off_v + + # initialize pointer to m and l + m_i = tl.zeros([BLOCK_M], dtype=tl.float32) - float("inf") + l_i = tl.zeros([BLOCK_M], dtype=tl.float32) + acc = tl.zeros([BLOCK_M, BLOCK_DMODEL], dtype=tl.float32) + + block_mask = tl.where(block_start_loc < cur_batch_seq_len, 1, 0) + + end_n = ( + cur_batch_seq_len + if not IS_CAUSAL + else tl.minimum((start_m + 1) * BLOCK_M, cur_batch_seq_len) + ) + for start_n in range(0, block_mask * end_n, BLOCK_N): + start_n = tl.multiple_of(start_n, BLOCK_N) + # -- compute qk ---- + k = tl.load( + k_ptrs + (cur_batch_in_all_start_index + start_n) * stride_kbs, + mask=((start_n + offs_n[None, :]) < cur_batch_seq_len) & (mask_d[:, None]), + other=0.0, + ) + # mask = tl.load(mask_ptrs + start_n, mask=start_n + offs_n < cur_batch_end_loc, other=0.0) + + qk = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32) + qk += tl.dot(q, k) + qk *= sm_scale + + if IS_CAUSAL: + qk += tl.where( + (start_n + offs_n[None, :] < cur_batch_seq_len) + & (offs_m[:, None] >= (start_n + offs_n[None, :])), + 0, + float("-inf"), + ) + else: + qk += tl.where( + (start_n + offs_n[None, :]) < cur_batch_seq_len, 0, float("-inf") + ) + + # -- compute m_ij, p, l_ij + m_ij = tl.max(qk, 1) + p = tl.exp(qk - m_ij[:, None]) + l_ij = tl.sum(p, 1) + # -- update m_i and l_i + m_i_new = tl.maximum(m_i, m_ij) + alpha = tl.exp(m_i - m_i_new) + beta = tl.exp(m_ij - m_i_new) + l_i_new = alpha * l_i + beta * l_ij + # -- update output accumulator -- + # scale p + p_scale = beta / l_i_new + p = p * p_scale[:, None] + # scale acc + acc_scale = l_i / l_i_new * alpha + acc = acc * acc_scale[:, None] + # update acc + v = tl.load( + v_ptrs + (cur_batch_in_all_start_index + start_n) * stride_vbs, + mask=((start_n + offs_n[:, None]) < cur_batch_seq_len) & (mask_d[None, :]), + other=0.0, + ) + + p = p.to(v.dtype) + acc += tl.dot(p, v) + # update m_i and l_i + l_i = l_i_new + m_i = m_i_new + # initialize pointers to output + off_o = ( + (cur_batch_in_all_start_index + offs_m[:, None]) * stride_obs + + cur_head * stride_oh + + offs_d[None, :] + ) + out_ptrs = Out + off_o + tl.store( + out_ptrs, acc, mask=(offs_m[:, None] < cur_batch_seq_len) & (mask_d[None, :]) + ) + + +def context_attention_fwd( + q, k, v, o, b_start_loc, b_seq_len, max_input_len, is_causal=True +): + """ + q, k, v: [b * s, head, head_dim] + b_start_loc: [b] + b_seq_len: [b] + out: [b * s, head, head_dim] + """ + if is_cuda_available and CUDA_CAPABILITY[0] > 8: + BLOCK = 128 + else: + BLOCK = 64 + + Lq, Lk, Lv = q.shape[-1], k.shape[-1], v.shape[-1] + + sm_scale = 1.0 / (Lq**0.5) + batch, head = b_seq_len.shape[0], q.shape[1] + kv_group_num = q.shape[1] // k.shape[1] + + grid = (batch, head, triton.cdiv(max_input_len, BLOCK)) + num_warps = 4 if Lk <= 64 else 8 + + _fwd_kernel[grid]( + q, + k, + v, + sm_scale, + b_start_loc, + b_seq_len, + o, + q.stride(0), + q.stride(1), + k.stride(0), + k.stride(1), + v.stride(0), + v.stride(1), + o.stride(0), + o.stride(1), + kv_group_num=kv_group_num, + BLOCK_M=BLOCK, + BLOCK_DMODEL=triton.next_power_of_2(Lk), + BLOCK_N=BLOCK, + IS_CAUSAL=is_causal, + num_warps=num_warps, + num_stages=1, + Lk=Lk, + ) \ No newline at end of file diff --git a/archive/kt-sft/ktransformers/optimize/optimize.py b/archive/kt-sft/ktransformers/optimize/optimize.py new file mode 100644 index 0000000..bbe08c8 --- /dev/null +++ b/archive/kt-sft/ktransformers/optimize/optimize.py @@ -0,0 +1,137 @@ +''' +Description : +Author : Boxin Zhang, Azure-Tang +Version : 0.1.0 +Copyright (c) 2024 by KVCache.AI, All Rights Reserved. +''' +from typing import Mapping, List +import torch +import yaml +import re +from torch import nn +from transformers import AutoConfig +from transformers.configuration_utils import PretrainedConfig +# from operators import BaseInjectedModule +from ktransformers.util.custom_loader import GGUFLoader, ModelLoaderFactory +from ktransformers.util.utils import set_module, load_weights +import itertools +import copy + +def inject(module, local_optimization_dict, model_config:AutoConfig ,gguf_loader:GGUFLoader, prefix=''): + for name, child in module._modules.items(): + if child is not None: + child_prefix = prefix + name + if child_prefix in local_optimization_dict: + inject_module_meta=local_optimization_dict[child_prefix] + if inject_module_meta["class"] != "default": + import_path = inject_module_meta["class"].split(".") + import_module_name = ".".join(import_path[:-1]) + gguf_loader.tensor_device_map[inject_module_meta["key"]] = inject_module_meta["kwargs"] if "kwargs" in inject_module_meta else dict() + import_class_name = import_path[-1] + module_cls=getattr(__import__(import_module_name, fromlist=[""]), import_class_name) + print(f"Injecting {child_prefix} as", import_module_name, ".", import_class_name) + inject_module=module_cls(key = inject_module_meta["key"], gguf_loader = gguf_loader, config = model_config, orig_module=child, **inject_module_meta["kwargs"]) + set_module(module, name, inject_module) + elif inject_module_meta["class"] == "default": + print(f"Injecting {child_prefix} as default") + gguf_loader.tensor_device_map[inject_module_meta["key"]] = inject_module_meta["kwargs"] if "kwargs" in inject_module_meta else dict() + else: + raise Exception("inject_module_meta[\"class\"] must be \"default\" or a class path") + child_prefix += "." + child_optimization_dict = {k: v for k, v in local_optimization_dict.items() if k.startswith(child_prefix)} + inject(child, child_optimization_dict, model_config, gguf_loader, child_prefix) + +def del_meta(module:nn.Module): + #print("default loading weights", prefix) + persistent_buffers = {k: v for k, v in module._buffers.items() if k not in module._non_persistent_buffers_set} + local_name_params = itertools.chain(module._parameters.items(), persistent_buffers.items()) + local_state = {k: v for k, v in local_name_params if v is not None} + for name, param in local_state.items(): + if param.device == "meta" or param.device == torch.device("meta"): + module.__delattr__(name) + for name, child in module._modules.items(): + del_meta(child) + +def gen_optimize_config(module: nn.Module, out_data: Mapping, rule_list: List, prefix: str="", default_device: str = "cuda:0"): + module_name = prefix[:-1] + # translated_name = translate_name_to_gguf(prefix)[:-1] + #print("gen_optimize_config", prefix, module_name, translated_name) + recursive = True + for rule in rule_list: + match_meta = rule["match"] + if "class" not in match_meta and "name" not in match_meta: + raise Exception("match must have at least one of \"class\" and \"name\"") + if "class" in match_meta: + import_path = match_meta["class"].split(".") + import_module_name = ".".join(import_path[:-1]) + import_class_name = import_path[-1] + module_cls=getattr(__import__(import_module_name, fromlist=[""]), import_class_name) + if not isinstance(module, module_cls): + continue + if "name" in match_meta: + if re.search(match_meta["name"], module_name) is None: + continue + if "replace" not in rule: + raise Exception("replace must be in rule") + if "replace" in rule: + replace_meta = rule["replace"] + if module_name not in out_data: + out_data[module_name]={"key": module_name, + "class": replace_meta["class"] if "class" in replace_meta else "default", + # "device": replace_meta["device"] if "device" in replace_meta else default_device, + "kwargs": copy.deepcopy(replace_meta["kwargs"]) if "kwargs" in replace_meta else dict()} + else: + if out_data[module_name]["class"] == "default": + out_data[module_name]["class"] = replace_meta["class"] if "class" in replace_meta else "default" + out_data[module_name]["kwargs"].update(copy.deepcopy(replace_meta["kwargs"]) if "kwargs" in replace_meta else dict()) + if "recursive" in rule: + recursive = bool(rule["recursive"]) + break + + if module_name not in out_data: + out_data[module_name]= { + "class": "default", + "key": module_name, + "kwargs": {"generate_device": default_device, + "prefill_device": default_device} + } + + #print(out_data[module_name]) + #input() + + if recursive: + for name, child in module._modules.items(): + if child is not None: + child_prefix = prefix + name + "." + gen_optimize_config(child, out_data, rule_list, child_prefix, default_device = default_device) + + +def translate_model_config(model_config: PretrainedConfig): + # for supporting some special model + if model_config.model_type == "mixtral": + model_config.moe_intermediate_size = model_config.intermediate_size + + return model_config + + +def optimize_and_load_gguf(module: nn.Module, rule_file: str, gguf_path: str, model_config: PretrainedConfig, default_device: str = "cuda:0"): + with open(rule_file, 'r', encoding='utf-8') as f: + rule_list = yaml.load(f.read(), Loader=yaml.FullLoader) + + optimize_config = dict() + gen_optimize_config(module, optimize_config, rule_list, default_device = default_device) + + model_config = translate_model_config(model_config) + + weights_loader = ModelLoaderFactory.create_loader(gguf_path) + with torch.device("meta"): + inject(module, optimize_config, model_config, weights_loader) + # pre load lm_head because its big inter result + load_weights(module.lm_head, weights_loader, "lm_head.", device=default_device) + load_weights(module, weights_loader, device=default_device) + module.gguf_loader = weights_loader + del_meta(module) + if torch.cuda.is_available(): + torch.cuda.empty_cache() + elif torch.xpu.is_available(): + torch.xpu.empty_cache() diff --git a/archive/kt-sft/ktransformers/optimize/optimize_rules/DeepSeek-V2-Chat-multi-gpu-4.yaml b/archive/kt-sft/ktransformers/optimize/optimize_rules/DeepSeek-V2-Chat-multi-gpu-4.yaml new file mode 100644 index 0000000..66a420a --- /dev/null +++ b/archive/kt-sft/ktransformers/optimize/optimize_rules/DeepSeek-V2-Chat-multi-gpu-4.yaml @@ -0,0 +1,240 @@ +- match: + name: "^model.embed_tokens" + replace: + class: "default" + kwargs: + generate_device: "cpu" + prefill_device: "cpu" + +- match: + name: "^model\\.layers\\.([0-9]|[1][0-4])\\." + class: ktransformers.models.modeling_deepseek.DeepseekV2YarnRotaryEmbedding + replace: + class: ktransformers.operators.RoPE.YarnRotaryEmbedding + kwargs: + generate_device: "cuda:0" + prefill_device: "cuda:0" +- match: + name: "^model\\.layers\\.([2][0-9]|[1][5-9])\\." + class: ktransformers.models.modeling_deepseek.DeepseekV2YarnRotaryEmbedding + replace: + class: ktransformers.operators.RoPE.YarnRotaryEmbedding + kwargs: + generate_device: "cuda:1" + prefill_device: "cuda:1" +- match: + name: "^model\\.layers\\.([3][0-9]|[4][0-4])\\." + class: ktransformers.models.modeling_deepseek.DeepseekV2YarnRotaryEmbedding + replace: + class: ktransformers.operators.RoPE.YarnRotaryEmbedding + kwargs: + generate_device: "cuda:2" + prefill_device: "cuda:2" +- match: + name: "^model\\.layers\\.([5][0-9]|[4][5-9])\\." + class: ktransformers.models.modeling_deepseek.DeepseekV2YarnRotaryEmbedding + replace: + class: ktransformers.operators.RoPE.YarnRotaryEmbedding + kwargs: + generate_device: "cuda:3" + prefill_device: "cuda:3" + +- match: + name: "^model\\.layers\\.([0-9]|[1][0-4])\\.(?!self_attn\\.kv_b_proj).*$" # regular expression + class: torch.nn.Linear # only match modules matching name and class simultaneously + replace: + class: ktransformers.operators.linear.KTransformersLinear # optimized Kernel on quantized data types + kwargs: + generate_device: "cuda:0" + prefill_device: "cuda:0" + generate_op: "KLinearMarlin" + prefill_op: "KLinearTorch" +- match: + name: "^model\\.layers\\.([2][0-9]|[1][5-9])\\.(?!self_attn\\.kv_b_proj).*$" # regular expression + class: torch.nn.Linear # only match modules matching name and class simultaneously + replace: + class: ktransformers.operators.linear.KTransformersLinear # optimized Kernel on quantized data types + kwargs: + generate_device: "cuda:1" + prefill_device: "cuda:1" + generate_op: "KLinearMarlin" + prefill_op: "KLinearTorch" +- match: + name: "^model\\.layers\\.([3][0-9]|[4][0-4])\\.(?!self_attn\\.kv_b_proj).*$" # regular expression + class: torch.nn.Linear # only match modules matching name and class simultaneously + replace: + class: ktransformers.operators.linear.KTransformersLinear # optimized Kernel on quantized data types + kwargs: + generate_device: "cuda:2" + prefill_device: "cuda:2" + generate_op: "KLinearMarlin" + prefill_op: "KLinearTorch" +- match: + name: "^model\\.layers\\.([5][0-9]|[4][5-9])\\.(?!self_attn\\.kv_b_proj).*$" # regular expression + class: torch.nn.Linear # only match modules matching name and class simultaneously + replace: + class: ktransformers.operators.linear.KTransformersLinear # optimized Kernel on quantized data types + kwargs: + generate_device: "cuda:3" + prefill_device: "cuda:3" + generate_op: "KLinearMarlin" + prefill_op: "KLinearTorch" + +- match: + name: "^model\\.layers\\.([0-9]|[1][0-4])\\.mlp$" + class: ktransformers.models.modeling_deepseek.DeepseekV2MoE + replace: + class: ktransformers.operators.experts.KDeepseekV2MoE # mlp module with custom forward function + kwargs: + generate_device: "cuda:0" + prefill_device: "cuda:0" +- match: + name: "^model\\.layers\\.([2][0-9]|[1][5-9])\\.mlp$" + class: ktransformers.models.modeling_deepseek.DeepseekV2MoE + replace: + class: ktransformers.operators.experts.KDeepseekV2MoE # mlp module with custom forward function + kwargs: + generate_device: "cuda:1" + prefill_device: "cuda:1" +- match: + name: "^model\\.layers\\.([3][0-9]|[4][0-4])\\.mlp$" + class: ktransformers.models.modeling_deepseek.DeepseekV2MoE + replace: + class: ktransformers.operators.experts.KDeepseekV2MoE # mlp module with custom forward function + kwargs: + generate_device: "cuda:2" + prefill_device: "cuda:2" +- match: + name: "^model\\.layers\\.([5][0-9]|[4][5-9])\\.mlp$" + class: ktransformers.models.modeling_deepseek.DeepseekV2MoE + replace: + class: ktransformers.operators.experts.KDeepseekV2MoE # mlp module with custom forward function + kwargs: + generate_device: "cuda:3" + prefill_device: "cuda:3" + +- match: + name: "^model\\.layers\\.([0-9]|[1][0-4])\\.mlp\\.experts$" + replace: + class: ktransformers.operators.experts.KTransformersExperts # custom MoE Kernel with expert paralleism + kwargs: + prefill_device: "cuda:0" + prefill_op: "KExpertsTorch" + generate_device: "cpu" + generate_op: "KExpertsCPU" + out_device: "cuda:0" + recursive: False # don't recursively inject submodules of this module +- match: + name: "^model\\.layers\\.([2][0-9]|[1][5-9])\\.mlp\\.experts$" + replace: + class: ktransformers.operators.experts.KTransformersExperts # custom MoE Kernel with expert paralleism + kwargs: + prefill_device: "cuda:1" + prefill_op: "KExpertsTorch" + generate_device: "cpu" + generate_op: "KExpertsCPU" + out_device: "cuda:1" + recursive: False # don't recursively inject submodules of this module +- match: + name: "^model\\.layers\\.([3][0-9]|[4][0-4])\\.mlp\\.experts$" + replace: + class: ktransformers.operators.experts.KTransformersExperts # custom MoE Kernel with expert paralleism + kwargs: + prefill_device: "cuda:2" + prefill_op: "KExpertsTorch" + generate_device: "cpu" + generate_op: "KExpertsCPU" + out_device: "cuda:2" + recursive: False # don't recursively inject submodules of this module +- match: + name: "^model\\.layers\\.([5][0-9]|[4][5-9])\\.mlp\\.experts$" + replace: + class: ktransformers.operators.experts.KTransformersExperts # custom MoE Kernel with expert paralleism + kwargs: + prefill_device: "cuda:3" + prefill_op: "KExpertsTorch" + generate_device: "cpu" + generate_op: "KExpertsCPU" + out_device: "cuda:3" + recursive: False # don't recursively inject submodules of this module + +- match: + name: "^model\\.layers\\.([0-9]|[1][0-4])\\.self_attn$" + replace: + class: ktransformers.operators.attention.KDeepseekV2Attention # optimized MLA implementation + kwargs: + generate_device: "cuda:0" + prefill_device: "cuda:0" +- match: + name: "^model\\.layers\\.([2][0-9]|[1][5-9])\\.self_attn$" + replace: + class: ktransformers.operators.attention.KDeepseekV2Attention # optimized MLA implementation + kwargs: + generate_device: "cuda:1" + prefill_device: "cuda:1" +- match: + name: "^model\\.layers\\.([3][0-9]|[4][0-4])\\.self_attn$" + replace: + class: ktransformers.operators.attention.KDeepseekV2Attention # optimized MLA implementation + kwargs: + generate_device: "cuda:2" + prefill_device: "cuda:2" +- match: + name: "^model\\.layers\\.([5][0-9]|[4][5-9])\\.self_attn$" + replace: + class: ktransformers.operators.attention.KDeepseekV2Attention # optimized MLA implementation + kwargs: + generate_device: "cuda:3" + prefill_device: "cuda:3" + +- match: + name: "^model$" + replace: + class: "ktransformers.operators.models.KDeepseekV2Model" + kwargs: + per_layer_prefill_intput_threshold: 0 # 0 is close layer wise prefill + transfer_map: + 15: "cuda:1" + 30: "cuda:2" + 45: "cuda:3" + +- match: + name: "^model\\.layers\\.([0-9]|[1][0-4])\\." + replace: + class: "default" + kwargs: + generate_device: "cuda:0" + prefill_device: "cuda:0" +- match: + name: "(^model\\.layers\\.([2][0-9]|[1][5-9])\\.)" + replace: + class: "default" + kwargs: + generate_device: "cuda:1" + prefill_device: "cuda:1" +- match: + name: "(^model\\.layers\\.([3][0-9]|[4][0-4])\\.)" + replace: + class: "default" + kwargs: + generate_device: "cuda:2" + prefill_device: "cuda:2" + +- match: + name: "^lm_head" + class: torch.nn.Linear + replace: + class: ktransformers.operators.linear.KTransformersLinear + kwargs: + generate_device: "cuda:3" + prefill_device: "cuda:3" + generate_op: "KLinearMarlin" + prefill_op: "KLinearTorch" + +- match: + name: "(^model\\.layers\\.([5][0-9]|[4][5-9])\\.)|(^model.norm)" + replace: + class: "default" + kwargs: + generate_device: "cuda:3" + prefill_device: "cuda:3" \ No newline at end of file diff --git a/archive/kt-sft/ktransformers/optimize/optimize_rules/DeepSeek-V2-Chat-multi-gpu.yaml b/archive/kt-sft/ktransformers/optimize/optimize_rules/DeepSeek-V2-Chat-multi-gpu.yaml new file mode 100644 index 0000000..f409376 --- /dev/null +++ b/archive/kt-sft/ktransformers/optimize/optimize_rules/DeepSeek-V2-Chat-multi-gpu.yaml @@ -0,0 +1,137 @@ +- match: + name: "^model.embed_tokens" + replace: + class: "default" + kwargs: + generate_device: "cpu" + prefill_device: "cpu" + +- match: + name: "^model\\.layers\\.(0|[1-9]|[12][0-9])\\." + class: ktransformers.models.modeling_deepseek.DeepseekV2YarnRotaryEmbedding + replace: + class: ktransformers.operators.RoPE.YarnRotaryEmbedding + kwargs: + generate_device: "cuda:0" + prefill_device: "cuda:0" +- match: + name: "^model\\.layers\\.([345][0-9])\\." + class: ktransformers.models.modeling_deepseek.DeepseekV2YarnRotaryEmbedding + replace: + class: ktransformers.operators.RoPE.YarnRotaryEmbedding + kwargs: + generate_device: "cuda:1" + prefill_device: "cuda:1" + +- match: + name: "^model\\.layers\\.(0|[1-9]|[12][0-9])\\.(?!self_attn\\.kv_b_proj).*$" # regular expression + class: torch.nn.Linear # only match modules matching name and class simultaneously + replace: + class: ktransformers.operators.linear.KTransformersLinear # optimized Kernel on quantized data types + kwargs: + generate_device: "cuda:0" + prefill_device: "cuda:0" + generate_op: "KLinearMarlin" + prefill_op: "KLinearTorch" + +- match: + name: "^model\\.layers\\.([345][0-9])\\.(?!self_attn\\.kv_b_proj).*$" # regular expression + class: torch.nn.Linear # only match modules matching name and class simultaneously + replace: + class: ktransformers.operators.linear.KTransformersLinear # optimized Kernel on quantized data types + kwargs: + generate_device: "cuda:1" + prefill_device: "cuda:1" + generate_op: "KLinearMarlin" + prefill_op: "KLinearTorch" + +- match: + name: "^model\\.layers\\.(0|[1-9]|[12][0-9])\\.mlp$" + class: ktransformers.models.modeling_deepseek.DeepseekV2MoE + replace: + class: ktransformers.operators.experts.KDeepseekV2MoE # mlp module with custom forward function + kwargs: + generate_device: "cuda:0" + prefill_device: "cuda:0" +- match: + name: "^model\\.layers\\.([345][0-9])\\.mlp$" + class: ktransformers.models.modeling_deepseek.DeepseekV2MoE + replace: + class: ktransformers.operators.experts.KDeepseekV2MoE # mlp module with custom forward function + kwargs: + generate_device: "cuda:1" + prefill_device: "cuda:1" + +- match: + name: "^model\\.layers\\.(0|[1-9]|[12][0-9])\\.mlp\\.experts$" + replace: + class: ktransformers.operators.experts.KTransformersExperts # custom MoE Kernel with expert paralleism + kwargs: + prefill_device: "cuda:0" + prefill_op: "KExpertsTorch" + generate_device: "cpu" + generate_op: "KExpertsCPU" + out_device: "cuda:0" + recursive: False # don't recursively inject submodules of this module + +- match: + name: "^model\\.layers\\.([345][0-9])\\.mlp\\.experts$" + replace: + class: ktransformers.operators.experts.KTransformersExperts # custom MoE Kernel with expert paralleism + kwargs: + prefill_device: "cuda:1" + prefill_op: "KExpertsTorch" + generate_device: "cpu" + generate_op: "KExpertsCPU" + out_device: "cuda:1" + recursive: False # don't recursively inject submodules of this module + +- match: + name: "^model\\.layers\\.(0|[1-9]|[12][0-9])\\.self_attn$" + replace: + class: ktransformers.operators.attention.KDeepseekV2Attention # optimized MLA implementation + kwargs: + generate_device: "cuda:0" + prefill_device: "cuda:0" +- match: + name: "^model\\.layers\\.([345][0-9])\\.self_attn$" + replace: + class: ktransformers.operators.attention.KDeepseekV2Attention # optimized MLA implementation + kwargs: + generate_device: "cuda:1" + prefill_device: "cuda:1" +- match: + name: "^model$" + replace: + class: "ktransformers.operators.models.KDeepseekV2Model" + kwargs: + per_layer_prefill_intput_threshold: 0 # 0 is close layer wise prefill + transfer_map: + 30: "cuda:1" + +- match: + name: "^model\\.layers\\.(0|[1-9]|[12][0-9])\\." + replace: + class: "default" + kwargs: + generate_device: "cuda:0" + prefill_device: "cuda:0" + +- match: + name: "^lm_head" + class: torch.nn.Linear + replace: + class: ktransformers.operators.linear.KTransformersLinear + kwargs: + generate_device: "cuda:1" + prefill_device: "cuda:1" + generate_op: "KLinearMarlin" + prefill_op: "KLinearTorch" + +- match: + name: "(^model\\.layers\\.([345][0-9])\\.)|(model.norm)" + replace: + class: "default" + kwargs: + generate_device: "cuda:1" + prefill_device: "cuda:1" diff --git a/archive/kt-sft/ktransformers/optimize/optimize_rules/DeepSeek-V2-Chat-sft-amx.yaml b/archive/kt-sft/ktransformers/optimize/optimize_rules/DeepSeek-V2-Chat-sft-amx.yaml new file mode 100644 index 0000000..7e6e340 --- /dev/null +++ b/archive/kt-sft/ktransformers/optimize/optimize_rules/DeepSeek-V2-Chat-sft-amx.yaml @@ -0,0 +1,69 @@ +- match: + class: ktransformers.models.modeling_deepseek.DeepseekV2YarnRotaryEmbedding + replace: + class: ktransformers.operators.RoPE.YarnRotaryEmbedding + kwargs: + generate_device: "cuda" + prefill_device: "cuda" +- match: + name: "^model\\.layers\\.(?!.*self_attn\\.kv_b_proj).*$" # regular expression + class: torch.nn.Linear # only match modules matching name and class simultaneously + replace: + class: ktransformers.operators.linear.KTransformersLinear # optimized Kernel on quantized data types + kwargs: + generate_device: "cuda" + prefill_device: "cuda" + generate_op: "KLinearTorch" + prefill_op: "KLinearTorch" + +- match: + name: "^lm_head" + class: torch.nn.Linear + replace: + class: ktransformers.operators.linear.KTransformersLinear + kwargs: + generate_device: "cuda" + prefill_device: "cuda" + generate_op: "KLinearTorch" + prefill_op: "KLinearTorch" + +- match: + name: "^model\\.layers\\..*\\.mlp$" + class: ktransformers.models.modeling_deepseek.DeepseekV2MoE + replace: + class: ktransformers.operators.experts.KDeepseekV2MoE # mlp module with custom forward function + kwargs: + generate_device: "cuda" + prefill_device: "cuda" +- match: + name: "^model\\.layers\\..*\\.mlp\\.experts$" + replace: + class: ktransformers.operators.experts.KTransformersExperts # custom MoE Kernel with expert paralleism + kwargs: + prefill_device: "cpu" + prefill_op: "KExpertsTorch" + generate_device: "cpu" + generate_op: "KSFTExpertsCPU" + out_device: "cuda" + backend: "AMXInt8" # or "AMXBF16" or "llamafile" (default) + recursive: False # don't recursively inject submodules of this module +- match: + name: "^model\\.layers\\..*\\.self_attn$" + replace: + class: ktransformers.operators.attention.KDeepseekV2Attention # optimized MLA implementation + kwargs: + generate_device: "cuda" + prefill_device: "cuda" +- match: + name: "^model$" + replace: + class: "ktransformers.operators.models.KDeepseekV2Model" + kwargs: + per_layer_prefill_intput_threshold: 0 # 0 is close layer wise prefill +- match: + name: "^model.embed_tokens" + replace: + class: "default" + kwargs: + generate_device: "cpu" + prefill_device: "cpu" \ No newline at end of file diff --git a/archive/kt-sft/ktransformers/optimize/optimize_rules/DeepSeek-V2-Chat.yaml b/archive/kt-sft/ktransformers/optimize/optimize_rules/DeepSeek-V2-Chat.yaml new file mode 100644 index 0000000..7f3e44e --- /dev/null +++ b/archive/kt-sft/ktransformers/optimize/optimize_rules/DeepSeek-V2-Chat.yaml @@ -0,0 +1,68 @@ +- match: + class: ktransformers.models.modeling_deepseek.DeepseekV2YarnRotaryEmbedding + replace: + class: ktransformers.operators.RoPE.YarnRotaryEmbedding + kwargs: + generate_device: "cuda" + prefill_device: "cuda" +- match: + name: "^model\\.layers\\.(?!.*self_attn\\.kv_b_proj).*$" # regular expression + class: torch.nn.Linear # only match modules matching name and class simultaneously + replace: + class: ktransformers.operators.linear.KTransformersLinear # optimized Kernel on quantized data types + kwargs: + generate_device: "cuda" + prefill_device: "cuda" + generate_op: "KLinearMarlin" + prefill_op: "KLinearTorch" + +- match: + name: "^lm_head" + class: torch.nn.Linear + replace: + class: ktransformers.operators.linear.KTransformersLinear + kwargs: + generate_device: "cuda" + prefill_device: "cuda" + generate_op: "KLinearMarlin" + prefill_op: "KLinearTorch" + +- match: + name: "^model\\.layers\\..*\\.mlp$" + class: ktransformers.models.modeling_deepseek.DeepseekV2MoE + replace: + class: ktransformers.operators.experts.KDeepseekV2MoE # mlp module with custom forward function + kwargs: + generate_device: "cuda" + prefill_device: "cuda" +- match: + name: "^model\\.layers\\..*\\.mlp\\.experts$" + replace: + class: ktransformers.operators.experts.KTransformersExperts # custom MoE Kernel with expert paralleism + kwargs: + prefill_device: "cuda" + prefill_op: "KExpertsTorch" + generate_device: "cpu" + generate_op: "KExpertsCPU" + out_device: "cuda" + recursive: False # don't recursively inject submodules of this module +- match: + name: "^model\\.layers\\..*\\.self_attn$" + replace: + class: ktransformers.operators.attention.KDeepseekV2Attention # optimized MLA implementation + kwargs: + generate_device: "cuda" + prefill_device: "cuda" +- match: + name: "^model$" + replace: + class: "ktransformers.operators.models.KDeepseekV2Model" + kwargs: + per_layer_prefill_intput_threshold: 0 # 0 is close layer wise prefill +- match: + name: "^model.embed_tokens" + replace: + class: "default" + kwargs: + generate_device: "cpu" + prefill_device: "cpu" \ No newline at end of file diff --git a/archive/kt-sft/ktransformers/optimize/optimize_rules/DeepSeek-V2-Lite-Chat-multi-gpu.yaml b/archive/kt-sft/ktransformers/optimize/optimize_rules/DeepSeek-V2-Lite-Chat-multi-gpu.yaml new file mode 100644 index 0000000..75b79f0 --- /dev/null +++ b/archive/kt-sft/ktransformers/optimize/optimize_rules/DeepSeek-V2-Lite-Chat-multi-gpu.yaml @@ -0,0 +1,137 @@ +- match: + name: "^model.embed_tokens" + replace: + class: "default" + kwargs: + generate_device: "cpu" + prefill_device: "cpu" + +- match: + name: "^model\\.layers\\.(0|[1-9])\\." + class: ktransformers.models.modeling_deepseek.DeepseekV2YarnRotaryEmbedding + replace: + class: ktransformers.operators.RoPE.YarnRotaryEmbedding + kwargs: + generate_device: "cuda:0" + prefill_device: "cuda:0" +- match: + name: "^model\\.layers\\.([12][0-9])\\." + class: ktransformers.models.modeling_deepseek.DeepseekV2YarnRotaryEmbedding + replace: + class: ktransformers.operators.RoPE.YarnRotaryEmbedding + kwargs: + generate_device: "cuda:1" + prefill_device: "cuda:1" + +- match: + name: "^model\\.layers\\.(0|[1-9])\\.(?!self_attn).*$" # regular expression + class: torch.nn.Linear # only match modules matching name and class simultaneously + replace: + class: ktransformers.operators.linear.KTransformersLinear # optimized Kernel on quantized data types + kwargs: + generate_device: "cuda:0" + prefill_device: "cuda:0" + generate_op: "KLinearMarlin" + prefill_op: "KLinearTorch" + +- match: + name: "^model\\.layers\\.([12][0-9])\\.(?!self_attn).*$" # regular expression + class: torch.nn.Linear # only match modules matching name and class simultaneously + replace: + class: ktransformers.operators.linear.KTransformersLinear # optimized Kernel on quantized data types + kwargs: + generate_device: "cuda:1" + prefill_device: "cuda:1" + generate_op: "KLinearMarlin" + prefill_op: "KLinearTorch" + +- match: + name: "^model\\.layers\\.(0|[1-9])\\.mlp$" + class: ktransformers.models.modeling_deepseek.DeepseekV2MoE + replace: + class: ktransformers.operators.experts.KDeepseekV2MoE # mlp module with custom forward function + kwargs: + generate_device: "cuda:0" + prefill_device: "cuda:0" +- match: + name: "^model\\.layers\\.([12][0-9])\\.mlp$" + class: ktransformers.models.modeling_deepseek.DeepseekV2MoE + replace: + class: ktransformers.operators.experts.KDeepseekV2MoE # mlp module with custom forward function + kwargs: + generate_device: "cuda:1" + prefill_device: "cuda:1" + +- match: + name: "^model\\.layers\\.(0|[1-9])\\.mlp\\.experts$" + replace: + class: ktransformers.operators.experts.KTransformersExperts # custom MoE Kernel with expert paralleism + kwargs: + prefill_device: "cuda:0" + prefill_op: "KExpertsTorch" + generate_device: "cpu" + generate_op: "KExpertsCPU" + out_device: "cuda:0" + recursive: False # don't recursively inject submodules of this module + +- match: + name: "^model\\.layers\\.([12][0-9])\\.mlp\\.experts$" + replace: + class: ktransformers.operators.experts.KTransformersExperts # custom MoE Kernel with expert paralleism + kwargs: + prefill_device: "cuda:1" + prefill_op: "KExpertsTorch" + generate_device: "cpu" + generate_op: "KExpertsCPU" + out_device: "cuda:1" + recursive: False # don't recursively inject submodules of this module + +- match: + name: "^model\\.layers\\.(0|[1-9])\\.self_attn$" + replace: + class: ktransformers.operators.attention.KDeepseekV2Attention # optimized MLA implementation + kwargs: + generate_device: "cuda:0" + prefill_device: "cuda:0" +- match: + name: "^model\\.layers\\.([12][0-9])\\.self_attn$" + replace: + class: ktransformers.operators.attention.KDeepseekV2Attention # optimized MLA implementation + kwargs: + generate_device: "cuda:1" + prefill_device: "cuda:1" +- match: + name: "^model$" + replace: + class: "ktransformers.operators.models.KDeepseekV2Model" + kwargs: + per_layer_prefill_intput_threshold: 0 # 0 is close layer wise prefill + transfer_map: + 10: "cuda:1" + +- match: + name: "^model\\.layers\\.(0|[1-9])\\." + replace: + class: "default" + kwargs: + generate_device: "cuda:0" + prefill_device: "cuda:0" + +- match: + name: "^lm_head" + class: torch.nn.Linear + replace: + class: ktransformers.operators.linear.KTransformersLinear + kwargs: + generate_device: "cuda:1" + prefill_device: "cuda:1" + generate_op: "KLinearMarlin" + prefill_op: "KLinearTorch" + +- match: + name: "(^model\\.layers\\.([12][0-9])\\.)|(model.norm)" + replace: + class: "default" + kwargs: + generate_device: "cuda:1" + prefill_device: "cuda:1" \ No newline at end of file diff --git a/archive/kt-sft/ktransformers/optimize/optimize_rules/DeepSeek-V2-Lite-Chat-sft-amx-multi-gpu.yaml b/archive/kt-sft/ktransformers/optimize/optimize_rules/DeepSeek-V2-Lite-Chat-sft-amx-multi-gpu.yaml new file mode 100644 index 0000000..54f97e5 --- /dev/null +++ b/archive/kt-sft/ktransformers/optimize/optimize_rules/DeepSeek-V2-Lite-Chat-sft-amx-multi-gpu.yaml @@ -0,0 +1,139 @@ +- match: + name: "^model.embed_tokens" + replace: + class: "default" + kwargs: + generate_device: "cpu" + prefill_device: "cpu" + +- match: + name: "^model\\.layers\\.(0|[1-9])\\." + class: ktransformers.models.modeling_deepseek.DeepseekV2YarnRotaryEmbedding + replace: + class: ktransformers.operators.RoPE.YarnRotaryEmbedding + kwargs: + generate_device: "cuda:0" + prefill_device: "cuda:0" +- match: + name: "^model\\.layers\\.([12][0-9])\\." + class: ktransformers.models.modeling_deepseek.DeepseekV2YarnRotaryEmbedding + replace: + class: ktransformers.operators.RoPE.YarnRotaryEmbedding + kwargs: + generate_device: "cuda:1" + prefill_device: "cuda:1" + +- match: + name: "^model\\.layers\\.(0|[1-9])\\.(?!.*self_attn\\.kv_b_proj).*$" # regular expression + class: torch.nn.Linear # only match modules matching name and class simultaneously + replace: + class: ktransformers.operators.linear.KTransformersLinear # optimized Kernel on quantized data types + kwargs: + generate_device: "cuda:0" + prefill_device: "cuda:0" + generate_op: "KLinearTorch" + prefill_op: "KLinearTorch" + +- match: + name: "^model\\.layers\\.([12][0-9])\\.(?!.*self_attn\\.kv_b_proj).*$" # regular expression + class: torch.nn.Linear # only match modules matching name and class simultaneously + replace: + class: ktransformers.operators.linear.KTransformersLinear # optimized Kernel on quantized data types + kwargs: + generate_device: "cuda:1" + prefill_device: "cuda:1" + generate_op: "KLinearTorch" + prefill_op: "KLinearTorch" + +- match: + name: "^model\\.layers\\.(0|[1-9])\\.mlp$" + class: ktransformers.models.modeling_deepseek.DeepseekV2MoE + replace: + class: ktransformers.operators.experts.KDeepseekV2MoE # mlp module with custom forward function + kwargs: + generate_device: "cuda:0" + prefill_device: "cuda:0" +- match: + name: "^model\\.layers\\.([12][0-9])\\.mlp$" + class: ktransformers.models.modeling_deepseek.DeepseekV2MoE + replace: + class: ktransformers.operators.experts.KDeepseekV2MoE # mlp module with custom forward function + kwargs: + generate_device: "cuda:1" + prefill_device: "cuda:1" + +- match: + name: "^model\\.layers\\.(0|[1-9])\\.mlp\\.experts$" + replace: + class: ktransformers.operators.experts.KTransformersExperts # custom MoE Kernel with expert paralleism + kwargs: + prefill_device: "cuda:0" + prefill_op: "KExpertsTorch" + generate_device: "cpu" + generate_op: "KSFTExpertsCPU" + out_device: "cuda:0" + backend: "AMXInt8" # or "AMXBF16" or "llamafile" (default) + recursive: False # don't recursively inject submodules of this module + +- match: + name: "^model\\.layers\\.([12][0-9])\\.mlp\\.experts$" + replace: + class: ktransformers.operators.experts.KTransformersExperts # custom MoE Kernel with expert paralleism + kwargs: + prefill_device: "cuda:1" + prefill_op: "KExpertsTorch" + generate_device: "cpu" + generate_op: "KSFTExpertsCPU" + out_device: "cuda:1" + backend: "AMXInt8" # or "AMXBF16" or "llamafile" (default) + recursive: False # don't recursively inject submodules of this module + +- match: + name: "^model\\.layers\\.(0|[1-9])\\.self_attn$" + replace: + class: ktransformers.operators.attention.KDeepseekV2Attention # optimized MLA implementation + kwargs: + generate_device: "cuda:0" + prefill_device: "cuda:0" +- match: + name: "^model\\.layers\\.([12][0-9])\\.self_attn$" + replace: + class: ktransformers.operators.attention.KDeepseekV2Attention # optimized MLA implementation + kwargs: + generate_device: "cuda:1" + prefill_device: "cuda:1" +- match: + name: "^model$" + replace: + class: "ktransformers.operators.models.KDeepseekV2Model" + kwargs: + per_layer_prefill_intput_threshold: 0 # 0 is close layer wise prefill + transfer_map: + 10: "cuda:1" + +- match: + name: "^model\\.layers\\.(0|[1-9])\\." + replace: + class: "default" + kwargs: + generate_device: "cuda:0" + prefill_device: "cuda:0" + +- match: + name: "^lm_head" + class: torch.nn.Linear + replace: + class: ktransformers.operators.linear.KTransformersLinear + kwargs: + generate_device: "cuda:1" + prefill_device: "cuda:1" + generate_op: "KLinearTorch" + prefill_op: "KLinearTorch" + +- match: + name: "(^model\\.layers\\.([12][0-9])\\.)|(model.norm)" + replace: + class: "default" + kwargs: + generate_device: "cuda:1" + prefill_device: "cuda:1" diff --git a/archive/kt-sft/ktransformers/optimize/optimize_rules/DeepSeek-V2-Lite-Chat-sft-amx.yaml b/archive/kt-sft/ktransformers/optimize/optimize_rules/DeepSeek-V2-Lite-Chat-sft-amx.yaml new file mode 100644 index 0000000..7e6e340 --- /dev/null +++ b/archive/kt-sft/ktransformers/optimize/optimize_rules/DeepSeek-V2-Lite-Chat-sft-amx.yaml @@ -0,0 +1,69 @@ +- match: + class: ktransformers.models.modeling_deepseek.DeepseekV2YarnRotaryEmbedding + replace: + class: ktransformers.operators.RoPE.YarnRotaryEmbedding + kwargs: + generate_device: "cuda" + prefill_device: "cuda" +- match: + name: "^model\\.layers\\.(?!.*self_attn\\.kv_b_proj).*$" # regular expression + class: torch.nn.Linear # only match modules matching name and class simultaneously + replace: + class: ktransformers.operators.linear.KTransformersLinear # optimized Kernel on quantized data types + kwargs: + generate_device: "cuda" + prefill_device: "cuda" + generate_op: "KLinearTorch" + prefill_op: "KLinearTorch" + +- match: + name: "^lm_head" + class: torch.nn.Linear + replace: + class: ktransformers.operators.linear.KTransformersLinear + kwargs: + generate_device: "cuda" + prefill_device: "cuda" + generate_op: "KLinearTorch" + prefill_op: "KLinearTorch" + +- match: + name: "^model\\.layers\\..*\\.mlp$" + class: ktransformers.models.modeling_deepseek.DeepseekV2MoE + replace: + class: ktransformers.operators.experts.KDeepseekV2MoE # mlp module with custom forward function + kwargs: + generate_device: "cuda" + prefill_device: "cuda" +- match: + name: "^model\\.layers\\..*\\.mlp\\.experts$" + replace: + class: ktransformers.operators.experts.KTransformersExperts # custom MoE Kernel with expert paralleism + kwargs: + prefill_device: "cpu" + prefill_op: "KExpertsTorch" + generate_device: "cpu" + generate_op: "KSFTExpertsCPU" + out_device: "cuda" + backend: "AMXInt8" # or "AMXBF16" or "llamafile" (default) + recursive: False # don't recursively inject submodules of this module +- match: + name: "^model\\.layers\\..*\\.self_attn$" + replace: + class: ktransformers.operators.attention.KDeepseekV2Attention # optimized MLA implementation + kwargs: + generate_device: "cuda" + prefill_device: "cuda" +- match: + name: "^model$" + replace: + class: "ktransformers.operators.models.KDeepseekV2Model" + kwargs: + per_layer_prefill_intput_threshold: 0 # 0 is close layer wise prefill +- match: + name: "^model.embed_tokens" + replace: + class: "default" + kwargs: + generate_device: "cpu" + prefill_device: "cpu" \ No newline at end of file diff --git a/archive/kt-sft/ktransformers/optimize/optimize_rules/DeepSeek-V2-Lite-Chat-sft.yaml b/archive/kt-sft/ktransformers/optimize/optimize_rules/DeepSeek-V2-Lite-Chat-sft.yaml new file mode 100644 index 0000000..eebc24b --- /dev/null +++ b/archive/kt-sft/ktransformers/optimize/optimize_rules/DeepSeek-V2-Lite-Chat-sft.yaml @@ -0,0 +1,68 @@ +- match: + class: ktransformers.models.modeling_deepseek.DeepseekV2YarnRotaryEmbedding + replace: + class: ktransformers.operators.RoPE.YarnRotaryEmbedding + kwargs: + generate_device: "cuda" + prefill_device: "cuda" +- match: + name: "^model\\.layers\\.(?!.*self_attn\\.kv_b_proj).*$" # regular expression + class: torch.nn.Linear # only match modules matching name and class simultaneously + replace: + class: ktransformers.operators.linear.KTransformersLinear # optimized Kernel on quantized data types + kwargs: + generate_device: "cuda" + prefill_device: "cuda" + generate_op: "KLinearTorch" + prefill_op: "KLinearTorch" + +- match: + name: "^lm_head" + class: torch.nn.Linear + replace: + class: ktransformers.operators.linear.KTransformersLinear + kwargs: + generate_device: "cuda" + prefill_device: "cuda" + generate_op: "KLinearTorch" + prefill_op: "KLinearTorch" + +- match: + name: "^model\\.layers\\..*\\.mlp$" + class: ktransformers.models.modeling_deepseek.DeepseekV2MoE + replace: + class: ktransformers.operators.experts.KDeepseekV2MoE # mlp module with custom forward function + kwargs: + generate_device: "cuda" + prefill_device: "cuda" +- match: + name: "^model\\.layers\\..*\\.mlp\\.experts$" + replace: + class: ktransformers.operators.experts.KTransformersExperts # custom MoE Kernel with expert paralleism + kwargs: + prefill_device: "cpu" + prefill_op: "KExpertsTorch" + generate_device: "cpu" + generate_op: "KSFTExpertsCPU" + out_device: "cuda" + recursive: False # don't recursively inject submodules of this module +- match: + name: "^model\\.layers\\..*\\.self_attn$" + replace: + class: ktransformers.operators.attention.KDeepseekV2Attention # optimized MLA implementation + kwargs: + generate_device: "cuda" + prefill_device: "cuda" +- match: + name: "^model$" + replace: + class: "ktransformers.operators.models.KDeepseekV2Model" + kwargs: + per_layer_prefill_intput_threshold: 0 # 0 is close layer wise prefill +- match: + name: "^model.embed_tokens" + replace: + class: "default" + kwargs: + generate_device: "cpu" + prefill_device: "cpu" \ No newline at end of file diff --git a/archive/kt-sft/ktransformers/optimize/optimize_rules/DeepSeek-V2-Lite-Chat-use-adapter.yaml b/archive/kt-sft/ktransformers/optimize/optimize_rules/DeepSeek-V2-Lite-Chat-use-adapter.yaml new file mode 100644 index 0000000..85a3aeb --- /dev/null +++ b/archive/kt-sft/ktransformers/optimize/optimize_rules/DeepSeek-V2-Lite-Chat-use-adapter.yaml @@ -0,0 +1,68 @@ +- match: + class: ktransformers.models.modeling_deepseek.DeepseekV2YarnRotaryEmbedding + replace: + class: ktransformers.operators.RoPE.YarnRotaryEmbedding + kwargs: + generate_device: "cuda" + prefill_device: "cuda" +- match: + name: "^model\\.layers\\.(?!.*self_attn\\.kv_b_proj).*$" # regular expression + class: torch.nn.Linear # only match modules matching name and class simultaneously + replace: + class: ktransformers.operators.linear.KTransformersLinear # optimized Kernel on quantized data types + kwargs: + generate_device: "cuda" + prefill_device: "cuda" + generate_op: "KLinearTorch" + prefill_op: "KLinearTorch" + +- match: + name: "^lm_head" + class: torch.nn.Linear + replace: + class: ktransformers.operators.linear.KTransformersLinear + kwargs: + generate_device: "cuda" + prefill_device: "cuda" + generate_op: "KLinearTorch" + prefill_op: "KLinearTorch" + +- match: + name: "^model\\.layers\\..*\\.mlp$" + class: ktransformers.models.modeling_deepseek.DeepseekV2MoE + replace: + class: ktransformers.operators.experts.KDeepseekV2MoE # mlp module with custom forward function + kwargs: + generate_device: "cuda" + prefill_device: "cuda" +- match: + name: "^model\\.layers\\..*\\.mlp\\.experts$" + replace: + class: ktransformers.operators.experts.KTransformersExperts # custom MoE Kernel with expert paralleism + kwargs: + prefill_device: "cuda" + prefill_op: "KExpertsTorch" + generate_device: "cpu" + generate_op: "KExpertsCPU" + out_device: "cuda" + recursive: False # don't recursively inject submodules of this module +- match: + name: "^model\\.layers\\..*\\.self_attn$" + replace: + class: ktransformers.operators.attention.KDeepseekV2Attention # optimized MLA implementation + kwargs: + generate_device: "cuda" + prefill_device: "cuda" +- match: + name: "^model$" + replace: + class: "ktransformers.operators.models.KDeepseekV2Model" + kwargs: + per_layer_prefill_intput_threshold: 0 # 0 is close layer wise prefill +- match: + name: "^model.embed_tokens" + replace: + class: "default" + kwargs: + generate_device: "cpu" + prefill_device: "cpu" \ No newline at end of file diff --git a/archive/kt-sft/ktransformers/optimize/optimize_rules/DeepSeek-V2-Lite-Chat.yaml b/archive/kt-sft/ktransformers/optimize/optimize_rules/DeepSeek-V2-Lite-Chat.yaml new file mode 100644 index 0000000..7f3e44e --- /dev/null +++ b/archive/kt-sft/ktransformers/optimize/optimize_rules/DeepSeek-V2-Lite-Chat.yaml @@ -0,0 +1,68 @@ +- match: + class: ktransformers.models.modeling_deepseek.DeepseekV2YarnRotaryEmbedding + replace: + class: ktransformers.operators.RoPE.YarnRotaryEmbedding + kwargs: + generate_device: "cuda" + prefill_device: "cuda" +- match: + name: "^model\\.layers\\.(?!.*self_attn\\.kv_b_proj).*$" # regular expression + class: torch.nn.Linear # only match modules matching name and class simultaneously + replace: + class: ktransformers.operators.linear.KTransformersLinear # optimized Kernel on quantized data types + kwargs: + generate_device: "cuda" + prefill_device: "cuda" + generate_op: "KLinearMarlin" + prefill_op: "KLinearTorch" + +- match: + name: "^lm_head" + class: torch.nn.Linear + replace: + class: ktransformers.operators.linear.KTransformersLinear + kwargs: + generate_device: "cuda" + prefill_device: "cuda" + generate_op: "KLinearMarlin" + prefill_op: "KLinearTorch" + +- match: + name: "^model\\.layers\\..*\\.mlp$" + class: ktransformers.models.modeling_deepseek.DeepseekV2MoE + replace: + class: ktransformers.operators.experts.KDeepseekV2MoE # mlp module with custom forward function + kwargs: + generate_device: "cuda" + prefill_device: "cuda" +- match: + name: "^model\\.layers\\..*\\.mlp\\.experts$" + replace: + class: ktransformers.operators.experts.KTransformersExperts # custom MoE Kernel with expert paralleism + kwargs: + prefill_device: "cuda" + prefill_op: "KExpertsTorch" + generate_device: "cpu" + generate_op: "KExpertsCPU" + out_device: "cuda" + recursive: False # don't recursively inject submodules of this module +- match: + name: "^model\\.layers\\..*\\.self_attn$" + replace: + class: ktransformers.operators.attention.KDeepseekV2Attention # optimized MLA implementation + kwargs: + generate_device: "cuda" + prefill_device: "cuda" +- match: + name: "^model$" + replace: + class: "ktransformers.operators.models.KDeepseekV2Model" + kwargs: + per_layer_prefill_intput_threshold: 0 # 0 is close layer wise prefill +- match: + name: "^model.embed_tokens" + replace: + class: "default" + kwargs: + generate_device: "cpu" + prefill_device: "cpu" \ No newline at end of file diff --git a/archive/kt-sft/ktransformers/optimize/optimize_rules/DeepSeek-V3-Chat-amx.yaml b/archive/kt-sft/ktransformers/optimize/optimize_rules/DeepSeek-V3-Chat-amx.yaml new file mode 100644 index 0000000..724e1a4 --- /dev/null +++ b/archive/kt-sft/ktransformers/optimize/optimize_rules/DeepSeek-V3-Chat-amx.yaml @@ -0,0 +1,77 @@ +- match: + class: ktransformers.models.modeling_deepseek_v3.DeepseekV3RotaryEmbedding + replace: + class: ktransformers.operators.RoPE.YarnRotaryEmbeddingV3 + kwargs: + generate_device: "cuda" + prefill_device: "cuda" + +- match: + name: "^lm_head$" # regular expression + class: torch.nn.Linear # only match modules matching name and class simultaneously + replace: + class: ktransformers.operators.linear.KTransformersLinear # optimized Kernel on quantized data types + kwargs: + generate_device: "cuda" + prefill_device: "cuda" + generate_op: "KLinearMarlin" + prefill_op: "KLinearTorch" + +- match: + name: "^model\\.layers\\.(?!.*self_attn\\.kv_b_proj).*$" # regular expression + class: torch.nn.Linear # only match modules matching name and class simultaneously + replace: + class: ktransformers.operators.linear.KTransformersLinear # optimized Kernel on quantized data types + kwargs: + generate_device: "cuda" + prefill_device: "cuda" + generate_op: "KLinearMarlin" + prefill_op: "KLinearTorch" +- match: + name: "^model\\.layers\\..*\\.mlp$" + class: ktransformers.models.modeling_deepseek_v3.DeepseekV3MoE + replace: + class: ktransformers.operators.experts.KDeepseekV3MoE # mlp module with custom forward function + kwargs: + generate_device: "cuda" + prefill_device: "cuda" +- match: + class: ktransformers.models.modeling_deepseek_v3.MoEGate + replace: + class: ktransformers.operators.gate.KMoEGate + kwargs: + generate_device: "cuda:0" + prefill_device: "cuda:0" +- match: + name: "^model\\.layers\\..*\\.mlp\\.experts$" + replace: + class: ktransformers.operators.experts.KTransformersExperts # custom MoE Kernel with expert paralleism + kwargs: + prefill_device: "cuda" + prefill_op: "KExpertsTorch" + generate_device: "cpu" + generate_op: "KExpertsCPU" + out_device: "cuda" + backend: "AMXInt8" # or "AMXBF16" or "llamafile" (default) + recursive: False # don't recursively inject submodules of this module +- match: + name: "^model\\.layers\\..*\\.self_attn$" + replace: + class: ktransformers.operators.attention.KDeepseekV2Attention # optimized MLA implementation + kwargs: + generate_device: "cuda" + prefill_device: "cuda" + absorb_for_prefill: False # change this to True to enable long context(prefill may slower). +- match: + name: "^model$" + replace: + class: "ktransformers.operators.models.KDeepseekV2Model" + kwargs: + per_layer_prefill_intput_threshold: 0 # 0 is close layer wise prefill +- match: + name: "^model.embed_tokens" + replace: + class: "default" + kwargs: + generate_device: "cpu" + prefill_device: "cpu" \ No newline at end of file diff --git a/archive/kt-sft/ktransformers/optimize/optimize_rules/DeepSeek-V3-Chat-fp8-linear-ggml-experts-serve-amx.yaml b/archive/kt-sft/ktransformers/optimize/optimize_rules/DeepSeek-V3-Chat-fp8-linear-ggml-experts-serve-amx.yaml new file mode 100644 index 0000000..670f6d5 --- /dev/null +++ b/archive/kt-sft/ktransformers/optimize/optimize_rules/DeepSeek-V3-Chat-fp8-linear-ggml-experts-serve-amx.yaml @@ -0,0 +1,91 @@ +- match: + class: ktransformers.models.modeling_deepseek_v3.DeepseekV3RotaryEmbedding + replace: + class: ktransformers.operators.RoPE.YarnRotaryEmbeddingV3 + kwargs: + generate_device: "cuda" + prefill_device: "cuda" +- match: + name: "^model\\.layers\\.(?!.*self_attn\\.kv_b_proj).*$" # regular expression + class: torch.nn.Linear # only match modules matching name and class simultaneously + replace: + class: ktransformers.operators.linear.KTransformersLinear # optimized Kernel on quantized data types + kwargs: + generate_device: "cuda" + prefill_device: "cuda" + generate_op: "KLinearFP8" + prefill_op: "KLinearTorch" +- match: + name: "^model\\.layers\\..*\\.mlp$" + class: ktransformers.models.modeling_deepseek_v3.DeepseekV3MoE + replace: + class: ktransformers.operators.experts.KDeepseekV3MoEV2 # mlp module with custom forward function + kwargs: + generate_device: "cuda" + prefill_device: "cuda" +- match: + class: ktransformers.models.modeling_deepseek_v3.MoEGate + replace: + class: ktransformers.operators.gate.KMoEGate + kwargs: + generate_device: "cuda:0" + prefill_device: "cuda:0" +- match: + name: "^model\\.layers\\..*\\.mlp\\.experts$" + replace: + class: ktransformers.operators.experts.KTransformersExpertsV2 # custom MoE Kernel with expert paralleism + kwargs: + prefill_device: "cuda" + prefill_op: "KExpertsTorch" + generate_device: "cpu" + generate_op: "KExpertsCPU" + out_device: "cuda" + backend: "llamafile" + recursive: False # don't recursively inject submodules of this module +- match: + name: "^model\\.layers\\..*\\.self_attn$" + replace: + class: ktransformers.operators.balance_serve_attention.flashinfer_attn # optimized MLA implementation + kwargs: + generate_device: "cuda" + prefill_device: "cuda" +- match: + name: "^model$" + replace: + class: "ktransformers.operators.models.KDeepseekV2Model" + kwargs: + per_layer_prefill_intput_threshold: 0 # 0 is close layer wise prefill +- match: + name: "^model.embed_tokens" + replace: + class: "default" + kwargs: + generate_device: "cpu" + prefill_device: "cpu" + +- match: + class: ktransformers.models.modeling_deepseek_v3.DeepseekV3RMSNorm + replace: + class: ktransformers.operators.layernorm.RMSNorm + kwargs: + generate_device: "cuda" + prefill_device: "cuda" + +- match: + class: ktransformers.models.modeling_deepseek_v3.DeepseekV3MLP + replace: + class: ktransformers.operators.mlp.kDeepseekV3MLP + kwargs: + generate_device: "cuda" + prefill_device: "cuda" + +- match: + name: "^lm_head$" # regular expression + class: torch.nn.Linear # only match modules matching name and class simultaneously + replace: + class: ktransformers.operators.linear.KTransformersLinear # optimized Kernel on quantized data types + kwargs: + generate_device: "cuda" + prefill_device: "cuda" + generate_op: "VLinearMarlin" + prefill_op: "KLinearTorch" \ No newline at end of file diff --git a/archive/kt-sft/ktransformers/optimize/optimize_rules/DeepSeek-V3-Chat-fp8-linear-ggml-experts-serve.yaml b/archive/kt-sft/ktransformers/optimize/optimize_rules/DeepSeek-V3-Chat-fp8-linear-ggml-experts-serve.yaml new file mode 100644 index 0000000..56320bf --- /dev/null +++ b/archive/kt-sft/ktransformers/optimize/optimize_rules/DeepSeek-V3-Chat-fp8-linear-ggml-experts-serve.yaml @@ -0,0 +1,90 @@ +- match: + class: ktransformers.models.modeling_deepseek_v3.DeepseekV3RotaryEmbedding + replace: + class: ktransformers.operators.RoPE.YarnRotaryEmbeddingV3 + kwargs: + generate_device: "cuda" + prefill_device: "cuda" +- match: + name: "^model\\.layers\\.(?!.*self_attn\\.kv_b_proj).*$" # regular expression + class: torch.nn.Linear # only match modules matching name and class simultaneously + replace: + class: ktransformers.operators.linear.KTransformersLinear # optimized Kernel on quantized data types + kwargs: + generate_device: "cuda" + prefill_device: "cuda" + generate_op: "KLinearFP8" + prefill_op: "KLinearTorch" +- match: + name: "^model\\.layers\\..*\\.mlp$" + class: ktransformers.models.modeling_deepseek_v3.DeepseekV3MoE + replace: + class: ktransformers.operators.experts.KDeepseekV3MoEV2 # mlp module with custom forward function + kwargs: + generate_device: "cuda" + prefill_device: "cuda" +- match: + class: ktransformers.models.modeling_deepseek_v3.MoEGate + replace: + class: ktransformers.operators.gate.KMoEGate + kwargs: + generate_device: "cuda:0" + prefill_device: "cuda:0" +- match: + name: "^model\\.layers\\..*\\.mlp\\.experts$" + replace: + class: ktransformers.operators.experts.KTransformersExpertsV2 # custom MoE Kernel with expert paralleism + kwargs: + prefill_device: "cuda" + prefill_op: "KExpertsTorch" + generate_device: "cpu" + generate_op: "KExpertsCPU" + out_device: "cuda" + recursive: False # don't recursively inject submodules of this module +- match: + name: "^model\\.layers\\..*\\.self_attn$" + replace: + class: ktransformers.operators.balance_serve_attention.flashinfer_attn # optimized MLA implementation + kwargs: + generate_device: "cuda" + prefill_device: "cuda" +- match: + name: "^model$" + replace: + class: "ktransformers.operators.models.KDeepseekV2Model" + kwargs: + per_layer_prefill_intput_threshold: 0 # 0 is close layer wise prefill +- match: + name: "^model.embed_tokens" + replace: + class: "default" + kwargs: + generate_device: "cpu" + prefill_device: "cpu" + +- match: + class: ktransformers.models.modeling_deepseek_v3.DeepseekV3RMSNorm + replace: + class: ktransformers.operators.layernorm.RMSNorm + kwargs: + generate_device: "cuda" + prefill_device: "cuda" + +- match: + class: ktransformers.models.modeling_deepseek_v3.DeepseekV3MLP + replace: + class: ktransformers.operators.mlp.kDeepseekV3MLP + kwargs: + generate_device: "cuda" + prefill_device: "cuda" + +- match: + name: "^lm_head$" # regular expression + class: torch.nn.Linear # only match modules matching name and class simultaneously + replace: + class: ktransformers.operators.linear.KTransformersLinear # optimized Kernel on quantized data types + kwargs: + generate_device: "cuda" + prefill_device: "cuda" + generate_op: "VLinearMarlin" + prefill_op: "KLinearTorch" \ No newline at end of file diff --git a/archive/kt-sft/ktransformers/optimize/optimize_rules/DeepSeek-V3-Chat-fp8-linear-ggml-experts.yaml b/archive/kt-sft/ktransformers/optimize/optimize_rules/DeepSeek-V3-Chat-fp8-linear-ggml-experts.yaml new file mode 100644 index 0000000..25f021e --- /dev/null +++ b/archive/kt-sft/ktransformers/optimize/optimize_rules/DeepSeek-V3-Chat-fp8-linear-ggml-experts.yaml @@ -0,0 +1,63 @@ +- match: + class: ktransformers.models.modeling_deepseek_v3.DeepseekV3RotaryEmbedding + replace: + class: ktransformers.operators.RoPE.YarnRotaryEmbeddingV3 + kwargs: + generate_device: "cuda" + prefill_device: "cuda" +- match: + name: "^model\\.layers\\.(?!.*self_attn\\.kv_b_proj).*$" # regular expression + class: torch.nn.Linear # only match modules matching name and class simultaneously + replace: + class: ktransformers.operators.linear.KTransformersLinear # optimized Kernel on quantized data types + kwargs: + generate_device: "cuda" + prefill_device: "cuda" + generate_op: "KLinearFP8" + prefill_op: "KLinearTorch" +- match: + name: "^model\\.layers\\..*\\.mlp$" + class: ktransformers.models.modeling_deepseek_v3.DeepseekV3MoE + replace: + class: ktransformers.operators.experts.KDeepseekV3MoE # mlp module with custom forward function + kwargs: + generate_device: "cuda" + prefill_device: "cuda" +- match: + class: ktransformers.models.modeling_deepseek_v3.MoEGate + replace: + class: ktransformers.operators.gate.KMoEGate + kwargs: + generate_device: "cuda:0" + prefill_device: "cuda:0" +- match: + name: "^model\\.layers\\..*\\.mlp\\.experts$" + replace: + class: ktransformers.operators.experts.KTransformersExperts # custom MoE Kernel with expert paralleism + kwargs: + prefill_device: "cuda" + prefill_op: "KExpertsTorch" + generate_device: "cpu" + generate_op: "KExpertsCPU" + out_device: "cuda" + recursive: False # don't recursively inject submodules of this module +- match: + name: "^model\\.layers\\..*\\.self_attn$" + replace: + class: ktransformers.operators.attention.KDeepseekV2Attention # optimized MLA implementation + kwargs: + generate_device: "cuda" + prefill_device: "cuda" +- match: + name: "^model$" + replace: + class: "ktransformers.operators.models.KDeepseekV2Model" + kwargs: + per_layer_prefill_intput_threshold: 0 # 0 is close layer wise prefill +- match: + name: "^model.embed_tokens" + replace: + class: "default" + kwargs: + generate_device: "cpu" + prefill_device: "cpu" \ No newline at end of file diff --git a/archive/kt-sft/ktransformers/optimize/optimize_rules/DeepSeek-V3-Chat-multi-gpu-4.yaml b/archive/kt-sft/ktransformers/optimize/optimize_rules/DeepSeek-V3-Chat-multi-gpu-4.yaml new file mode 100644 index 0000000..ea75b30 --- /dev/null +++ b/archive/kt-sft/ktransformers/optimize/optimize_rules/DeepSeek-V3-Chat-multi-gpu-4.yaml @@ -0,0 +1,388 @@ +- match: + name: "^model.embed_tokens" + replace: + class: "default" + kwargs: + generate_device: "cpu" + prefill_device: "cpu" + +# === Rotary Embedding Replacement === + +# GPU 0: layers 0–14 +- match: + name: "^model\\.layers\\.([0-9]|1[0-4])\\." + class: ktransformers.models.modeling_deepseek_v3.DeepseekV3RotaryEmbedding + replace: + class: ktransformers.operators.RoPE.YarnRotaryEmbeddingV3 + kwargs: + generate_device: "cuda:0" + prefill_device: "cuda:0" + +# GPU 1: layers 15–29 +- match: + name: "^model\\.layers\\.(1[5-9]|2[0-9])\\." + class: ktransformers.models.modeling_deepseek_v3.DeepseekV3RotaryEmbedding + replace: + class: ktransformers.operators.RoPE.YarnRotaryEmbeddingV3 + kwargs: + generate_device: "cuda:1" + prefill_device: "cuda:1" + +# GPU 2: layers 30–44 +- match: + name: "^model\\.layers\\.(3[0-9]|4[0-4])\\." + class: ktransformers.models.modeling_deepseek_v3.DeepseekV3RotaryEmbedding + replace: + class: ktransformers.operators.RoPE.YarnRotaryEmbeddingV3 + kwargs: + generate_device: "cuda:2" + prefill_device: "cuda:2" + +# GPU 3: layers 45–60 +- match: + name: "^model\\.layers\\.(4[5-9]|5[0-9]|60)\\." + class: ktransformers.models.modeling_deepseek_v3.DeepseekV3RotaryEmbedding + replace: + class: ktransformers.operators.RoPE.YarnRotaryEmbeddingV3 + kwargs: + generate_device: "cuda:3" + prefill_device: "cuda:3" + +# === Linear Layers Replacement (excluding self_attn.kv_b_proj) === + +# GPU 0: layers 0–14 +- match: + name: "^model\\.layers\\.([0-9]|1[0-4])\\.(?!self_attn\\.kv_b_proj).*$" + class: torch.nn.Linear + replace: + class: ktransformers.operators.linear.KTransformersLinear + kwargs: + generate_device: "cuda:0" + prefill_device: "cuda:0" + generate_op: "KLinearMarlin" + prefill_op: "KLinearTorch" + +# GPU 1: layers 15–29 +- match: + name: "^model\\.layers\\.(1[5-9]|2[0-9])\\.(?!self_attn\\.kv_b_proj).*$" + class: torch.nn.Linear + replace: + class: ktransformers.operators.linear.KTransformersLinear + kwargs: + generate_device: "cuda:1" + prefill_device: "cuda:1" + generate_op: "KLinearMarlin" + prefill_op: "KLinearTorch" + +# GPU 2: layers 30–44 +- match: + name: "^model\\.layers\\.(3[0-9]|4[0-4])\\.(?!self_attn\\.kv_b_proj).*$" + class: torch.nn.Linear + replace: + class: ktransformers.operators.linear.KTransformersLinear + kwargs: + generate_device: "cuda:2" + prefill_device: "cuda:2" + generate_op: "KLinearMarlin" + prefill_op: "KLinearTorch" + +# GPU 3: layers 45–60 +- match: + name: "^model\\.layers\\.(4[5-9]|5[0-9]|60)\\.(?!self_attn\\.kv_b_proj).*$" + class: torch.nn.Linear + replace: + class: ktransformers.operators.linear.KTransformersLinear + kwargs: + generate_device: "cuda:3" + prefill_device: "cuda:3" + generate_op: "KLinearMarlin" + prefill_op: "KLinearTorch" + +# === MLP (MoE) Replacement === + +# GPU 0: layers 0–14 +- match: + name: "^model\\.layers\\.([0-9]|1[0-4])\\.mlp$" + class: ktransformers.models.modeling_deepseek_v3.DeepseekV3MoE + replace: + class: ktransformers.operators.experts.KDeepseekV3MoE + kwargs: + generate_device: "cuda:0" + prefill_device: "cuda:0" + +# GPU 1: layers 15–29 +- match: + name: "^model\\.layers\\.(1[5-9]|2[0-9])\\.mlp$" + class: ktransformers.models.modeling_deepseek_v3.DeepseekV3MoE + replace: + class: ktransformers.operators.experts.KDeepseekV3MoE + kwargs: + generate_device: "cuda:1" + prefill_device: "cuda:1" + +# GPU 2: layers 30–44 +- match: + name: "^model\\.layers\\.(3[0-9]|4[0-4])\\.mlp$" + class: ktransformers.models.modeling_deepseek_v3.DeepseekV3MoE + replace: + class: ktransformers.operators.experts.KDeepseekV3MoE + kwargs: + generate_device: "cuda:2" + prefill_device: "cuda:2" + +# GPU 3: layers 45–60 +- match: + name: "^model\\.layers\\.(4[5-9]|5[0-9]|60)\\.mlp$" + class: ktransformers.models.modeling_deepseek_v3.DeepseekV3MoE + replace: + class: ktransformers.operators.experts.KDeepseekV3MoE + kwargs: + generate_device: "cuda:3" + prefill_device: "cuda:3" + +# === MLP Gate Replacement === + +# GPU 0: layers 0–14 +- match: + name: "^model\\.layers\\.([0-9]|1[0-4])\\.mlp\\.gate$" + class: ktransformers.models.modeling_deepseek_v3.MoEGate + replace: + class: ktransformers.operators.gate.KMoEGate + kwargs: + generate_device: "cuda:0" + prefill_device: "cuda:0" + +# GPU 1: layers 15–29 +- match: + name: "^model\\.layers\\.(1[5-9]|2[0-9])\\.mlp\\.gate$" + class: ktransformers.models.modeling_deepseek_v3.MoEGate + replace: + class: ktransformers.operators.gate.KMoEGate + kwargs: + generate_device: "cuda:1" + prefill_device: "cuda:1" + +# GPU 2: layers 30–44 +- match: + name: "^model\\.layers\\.(3[0-9]|4[0-4])\\.mlp\\.gate$" + class: ktransformers.models.modeling_deepseek_v3.MoEGate + replace: + class: ktransformers.operators.gate.KMoEGate + kwargs: + generate_device: "cuda:2" + prefill_device: "cuda:2" + +# GPU 3: layers 45–60 +- match: + name: "^model\\.layers\\.(4[5-9]|5[0-9]|60)\\.mlp\\.gate$" + class: ktransformers.models.modeling_deepseek_v3.MoEGate + replace: + class: ktransformers.operators.gate.KMoEGate + kwargs: + generate_device: "cuda:3" + prefill_device: "cuda:3" + +# === MLP Experts Replacement === +# replace with marlin expert. Open and modify layer-num as needed. +# Each layer of malin experts takes about 6GB of GPU memory. +# !!!Do remember 'close' cuda graph if you are using marlin expert.!!! +# !!!KExpertsTorch is untested, we don't have enough VRAM.!!! + +# GPU 0: layers 3–4 +# - match: +# name: "^model\\.layers\\.([3-4])\\.mlp\\.experts$" +# replace: +# class: ktransformers.operators.experts.KTransformersExperts +# kwargs: +# generate_device: "cuda:0" +# generate_op: "KExpertsMarlin" +# recursive: False + +# # GPU 1: layers 15–17 +# - match: +# name: "^model\\.layers\\.(1[5-7])\\.mlp\\.experts$" +# replace: +# class: ktransformers.operators.experts.KTransformersExperts +# kwargs: +# generate_device: "cuda:1" +# generate_op: "KExpertsMarlin" +# recursive: False + +# # GPU 2: layers 30–32 +# - match: +# name: "^model\\.layers\\.(3[0-2])\\.mlp\\.experts$" +# replace: +# class: ktransformers.operators.experts.KTransformersExperts +# kwargs: +# generate_device: "cuda:2" +# generate_op: "KExpertsMarlin" +# recursive: False + +# # GPU 3: layers 45–46 +# - match: +# name: "^model\\.layers\\.(4[5-6])\\.mlp\\.experts$" +# replace: +# class: ktransformers.operators.experts.KTransformersExperts +# kwargs: +# generate_device: "cuda:3" +# generate_op: "KExpertsMarlin" +# recursive: False + + +# === MLP Experts Replacement === + +# GPU 0: layers 0–14 +- match: + name: "^model\\.layers\\.([0-9]|1[0-4])\\.mlp\\.experts$" + replace: + class: ktransformers.operators.experts.KTransformersExperts + kwargs: + prefill_device: "cuda:0" + prefill_op: "KExpertsTorch" + generate_device: "cpu" + generate_op: "KExpertsCPU" + out_device: "cuda:0" + recursive: False + +# GPU 1: layers 15–29 +- match: + name: "^model\\.layers\\.(1[5-9]|2[0-9])\\.mlp\\.experts$" + replace: + class: ktransformers.operators.experts.KTransformersExperts + kwargs: + prefill_device: "cuda:1" + prefill_op: "KExpertsTorch" + generate_device: "cpu" + generate_op: "KExpertsCPU" + out_device: "cuda:1" + recursive: False + +# GPU 2: layers 30–44 +- match: + name: "^model\\.layers\\.(3[0-9]|4[0-4])\\.mlp\\.experts$" + replace: + class: ktransformers.operators.experts.KTransformersExperts + kwargs: + prefill_device: "cuda:2" + prefill_op: "KExpertsTorch" + generate_device: "cpu" + generate_op: "KExpertsCPU" + out_device: "cuda:2" + recursive: False + +# GPU 3: layers 45–60 +- match: + name: "^model\\.layers\\.(4[5-9]|5[0-9]|60)\\.mlp\\.experts$" + replace: + class: ktransformers.operators.experts.KTransformersExperts + kwargs: + prefill_device: "cuda:3" + prefill_op: "KExpertsTorch" + generate_device: "cpu" + generate_op: "KExpertsCPU" + out_device: "cuda:3" + recursive: False + +# === Self-Attention Replacement === + +# GPU 0: layers 0–14 +- match: + name: "^model\\.layers\\.([0-9]|1[0-4])\\.self_attn$" + replace: + class: ktransformers.operators.attention.KDeepseekV2Attention + kwargs: + generate_device: "cuda:0" + prefill_device: "cuda:0" + absorb_for_prefill: False + +# GPU 1: layers 15–29 +- match: + name: "^model\\.layers\\.(1[5-9]|2[0-9])\\.self_attn$" + replace: + class: ktransformers.operators.attention.KDeepseekV2Attention + kwargs: + generate_device: "cuda:1" + prefill_device: "cuda:1" + absorb_for_prefill: False + +# GPU 2: layers 30–44 +- match: + name: "^model\\.layers\\.(3[0-9]|4[0-4])\\.self_attn$" + replace: + class: ktransformers.operators.attention.KDeepseekV2Attention + kwargs: + generate_device: "cuda:2" + prefill_device: "cuda:2" + absorb_for_prefill: False + +# GPU 3: layers 45–60 +- match: + name: "^model\\.layers\\.(4[5-9]|5[0-9]|60)\\.self_attn$" + replace: + class: ktransformers.operators.attention.KDeepseekV2Attention + kwargs: + generate_device: "cuda:3" + prefill_device: "cuda:3" + absorb_for_prefill: False + +# === Overall Model Replacement with Transfer Map === + +- match: + name: "^model$" + replace: + class: "ktransformers.operators.models.KDeepseekV2Model" + kwargs: + per_layer_prefill_intput_threshold: 0 # 0 means close layer‐wise prefill + transfer_map: + 15: "cuda:1" # Layers 15+ on GPU 1 + 30: "cuda:2" # Layers 30+ on GPU 2 + 45: "cuda:3" # Layers 45+ on GPU 3 + +# === Default Catch-All for Other Modules === + +# GPU 0: layers 0–14 +- match: + name: "^model\\.layers\\.([0-9]|1[0-4])\\." + replace: + class: "default" + kwargs: + generate_device: "cuda:0" + prefill_device: "cuda:0" + +# GPU 1: layers 15–29 +- match: + name: "^model\\.layers\\.(1[5-9]|2[0-9])\\." + replace: + class: "default" + kwargs: + generate_device: "cuda:1" + prefill_device: "cuda:1" + +# GPU 2: layers 30–44 +- match: + name: "^model\\.layers\\.(3[0-9]|4[0-4])\\." + replace: + class: "default" + kwargs: + generate_device: "cuda:2" + prefill_device: "cuda:2" + +- match: + name: "^lm_head" + class: torch.nn.Linear + replace: + class: ktransformers.operators.linear.KTransformersLinear + kwargs: + generate_device: "cuda:3" + prefill_device: "cuda:3" + generate_op: "KLinearMarlin" + prefill_op: "KLinearTorch" + +# For final modules (model.norm), ensure they are on GPU 3 (as in your original config) +- match: + name: "(^model\\.layers\\.(4[5-9]|5[0-9]|60)\\.)|(^model\\.norm)" + replace: + class: "default" + kwargs: + generate_device: "cuda:3" + prefill_device: "cuda:3" diff --git a/archive/kt-sft/ktransformers/optimize/optimize_rules/DeepSeek-V3-Chat-multi-gpu-8.yaml b/archive/kt-sft/ktransformers/optimize/optimize_rules/DeepSeek-V3-Chat-multi-gpu-8.yaml new file mode 100644 index 0000000..b00d2b4 --- /dev/null +++ b/archive/kt-sft/ktransformers/optimize/optimize_rules/DeepSeek-V3-Chat-multi-gpu-8.yaml @@ -0,0 +1,734 @@ +- match: + name: "^model.embed_tokens" + replace: + class: "default" + kwargs: + generate_device: "cpu" + prefill_device: "cpu" + +# === Rotary Embedding Replacement === + +# GPU 0: layers 0–7 +- match: + name: "^model\\.layers\\.([0-7])\\." + class: ktransformers.models.modeling_deepseek_v3.DeepseekV3RotaryEmbedding + replace: + class: ktransformers.operators.RoPE.YarnRotaryEmbeddingV3 + kwargs: + generate_device: "cuda:0" + prefill_device: "cuda:0" + +# GPU 1: layers 8–15 +- match: + name: "^model\\.layers\\.(8|9|1[0-5])\\." + class: ktransformers.models.modeling_deepseek_v3.DeepseekV3RotaryEmbedding + replace: + class: ktransformers.operators.RoPE.YarnRotaryEmbeddingV3 + kwargs: + generate_device: "cuda:1" + prefill_device: "cuda:1" + +# GPU 2: layers 16–23 +- match: + name: "^model\\.layers\\.(1[6-9]|2[0-3])\\." + class: ktransformers.models.modeling_deepseek_v3.DeepseekV3RotaryEmbedding + replace: + class: ktransformers.operators.RoPE.YarnRotaryEmbeddingV3 + kwargs: + generate_device: "cuda:2" + prefill_device: "cuda:2" + +# GPU 3: layers 24–31 +- match: + name: "^model\\.layers\\.(2[4-9]|3[0-1])\\." + class: ktransformers.models.modeling_deepseek_v3.DeepseekV3RotaryEmbedding + replace: + class: ktransformers.operators.RoPE.YarnRotaryEmbeddingV3 + kwargs: + generate_device: "cuda:3" + prefill_device: "cuda:3" + +# GPU 4: layers 32–39 +- match: + name: "^model\\.layers\\.([3][2-9])\\." + class: ktransformers.models.modeling_deepseek_v3.DeepseekV3RotaryEmbedding + replace: + class: ktransformers.operators.RoPE.YarnRotaryEmbeddingV3 + kwargs: + generate_device: "cuda:4" + prefill_device: "cuda:4" + +# GPU 5: layers 40–47 +- match: + name: "^model\\.layers\\.(4[0-7])\\." + class: ktransformers.models.modeling_deepseek_v3.DeepseekV3RotaryEmbedding + replace: + class: ktransformers.operators.RoPE.YarnRotaryEmbeddingV3 + kwargs: + generate_device: "cuda:5" + prefill_device: "cuda:5" + +# GPU 6: layers 48–55 +- match: + name: "^model\\.layers\\.(4[8-9]|5[0-5])\\." + class: ktransformers.models.modeling_deepseek_v3.DeepseekV3RotaryEmbedding + replace: + class: ktransformers.operators.RoPE.YarnRotaryEmbeddingV3 + kwargs: + generate_device: "cuda:6" + prefill_device: "cuda:6" + +# GPU 7: layers 56–60 +- match: + name: "^model\\.layers\\.(5[6-9]|60)\\." + class: ktransformers.models.modeling_deepseek_v3.DeepseekV3RotaryEmbedding + replace: + class: ktransformers.operators.RoPE.YarnRotaryEmbeddingV3 + kwargs: + generate_device: "cuda:7" + prefill_device: "cuda:7" + + +# === Linear Layers Replacement (excluding self_attn.kv_b_proj) === + +# GPU 0: layers 0–7 +- match: + name: "^model\\.layers\\.([0-7])\\.(?!self_attn\\.kv_b_proj).*$" + class: torch.nn.Linear + replace: + class: ktransformers.operators.linear.KTransformersLinear + kwargs: + generate_device: "cuda:0" + prefill_device: "cuda:0" + generate_op: "KLinearMarlin" + prefill_op: "KLinearTorch" + +# GPU 1: layers 8–15 +- match: + name: "^model\\.layers\\.(8|9|1[0-5])\\.(?!self_attn\\.kv_b_proj).*$" + class: torch.nn.Linear + replace: + class: ktransformers.operators.linear.KTransformersLinear + kwargs: + generate_device: "cuda:1" + prefill_device: "cuda:1" + generate_op: "KLinearMarlin" + prefill_op: "KLinearTorch" + +# GPU 2: layers 16–23 +- match: + name: "^model\\.layers\\.(1[6-9]|2[0-3])\\.(?!self_attn\\.kv_b_proj).*$" + class: torch.nn.Linear + replace: + class: ktransformers.operators.linear.KTransformersLinear + kwargs: + generate_device: "cuda:2" + prefill_device: "cuda:2" + generate_op: "KLinearMarlin" + prefill_op: "KLinearTorch" + +# GPU 3: layers 24–31 +- match: + name: "^model\\.layers\\.(2[4-9]|3[0-1])\\.(?!self_attn\\.kv_b_proj).*$" + class: torch.nn.Linear + replace: + class: ktransformers.operators.linear.KTransformersLinear + kwargs: + generate_device: "cuda:3" + prefill_device: "cuda:3" + generate_op: "KLinearMarlin" + prefill_op: "KLinearTorch" + +# GPU 4: layers 32–39 +- match: + name: "^model\\.layers\\.(3[2-9])\\.(?!self_attn\\.kv_b_proj).*$" + class: torch.nn.Linear + replace: + class: ktransformers.operators.linear.KTransformersLinear + kwargs: + generate_device: "cuda:4" + prefill_device: "cuda:4" + generate_op: "KLinearMarlin" + prefill_op: "KLinearTorch" + +# GPU 5: layers 40–47 +- match: + name: "^model\\.layers\\.(4[0-7])\\.(?!self_attn\\.kv_b_proj).*$" + class: torch.nn.Linear + replace: + class: ktransformers.operators.linear.KTransformersLinear + kwargs: + generate_device: "cuda:5" + prefill_device: "cuda:5" + generate_op: "KLinearMarlin" + prefill_op: "KLinearTorch" + +# GPU 6: layers 48–55 +- match: + name: "^model\\.layers\\.(4[8-9]|5[0-5])\\.(?!self_attn\\.kv_b_proj).*$" + class: torch.nn.Linear + replace: + class: ktransformers.operators.linear.KTransformersLinear + kwargs: + generate_device: "cuda:6" + prefill_device: "cuda:6" + generate_op: "KLinearMarlin" + prefill_op: "KLinearTorch" + +# GPU 7: layers 56–63 +- match: + name: "^model\\.layers\\.(5[6-9]|60)\\.(?!self_attn\\.kv_b_proj).*$" + class: torch.nn.Linear + replace: + class: ktransformers.operators.linear.KTransformersLinear + kwargs: + generate_device: "cuda:7" + prefill_device: "cuda:7" + generate_op: "KLinearMarlin" + prefill_op: "KLinearTorch" + + + +# === MLP (MoE) Replacement === + +# GPU 0: layers 0–7 +- match: + name: "^model\\.layers\\.([0-7])\\.mlp$" + class: ktransformers.models.modeling_deepseek_v3.DeepseekV3MoE + replace: + class: ktransformers.operators.experts.KDeepseekV3MoE + kwargs: + generate_device: "cuda:0" + prefill_device: "cuda:0" + +# GPU 1: layers 8–15 +- match: + name: "^model\\.layers\\.(8|9|1[0-5])\\.mlp$" + class: ktransformers.models.modeling_deepseek_v3.DeepseekV3MoE + replace: + class: ktransformers.operators.experts.KDeepseekV3MoE + kwargs: + generate_device: "cuda:1" + prefill_device: "cuda:1" + +# GPU 2: layers 16–23 +- match: + name: "^model\\.layers\\.(1[6-9]|2[0-3])\\.mlp$" + class: ktransformers.models.modeling_deepseek_v3.DeepseekV3MoE + replace: + class: ktransformers.operators.experts.KDeepseekV3MoE + kwargs: + generate_device: "cuda:2" + prefill_device: "cuda:2" + +# GPU 3: layers 24–31 +- match: + name: "^model\\.layers\\.(2[4-9]|3[0-1])\\.mlp$" + class: ktransformers.models.modeling_deepseek_v3.DeepseekV3MoE + replace: + class: ktransformers.operators.experts.KDeepseekV3MoE + kwargs: + generate_device: "cuda:3" + prefill_device: "cuda:3" + +# GPU 4: layers 32–39 +- match: + name: "^model\\.layers\\.(3[2-9])\\.mlp$" + class: ktransformers.models.modeling_deepseek_v3.DeepseekV3MoE + replace: + class: ktransformers.operators.experts.KDeepseekV3MoE + kwargs: + generate_device: "cuda:4" + prefill_device: "cuda:4" + +# GPU 5: layers 40–47 +- match: + name: "^model\\.layers\\.(4[0-7])\\.mlp$" + class: ktransformers.models.modeling_deepseek_v3.DeepseekV3MoE + replace: + class: ktransformers.operators.experts.KDeepseekV3MoE + kwargs: + generate_device: "cuda:5" + prefill_device: "cuda:5" + +# GPU 6: layers 48–55 +- match: + name: "^model\\.layers\\.(4[8-9]|5[0-5])\\.mlp$" + class: ktransformers.models.modeling_deepseek_v3.DeepseekV3MoE + replace: + class: ktransformers.operators.experts.KDeepseekV3MoE + kwargs: + generate_device: "cuda:6" + prefill_device: "cuda:6" + +# GPU 7: layers 56–60 +- match: + name: "^model\\.layers\\.(5[6-9]|60)\\.mlp$" + class: ktransformers.models.modeling_deepseek_v3.DeepseekV3MoE + replace: + class: ktransformers.operators.experts.KDeepseekV3MoE + kwargs: + generate_device: "cuda:7" + prefill_device: "cuda:7" + +# === MLP Gate Replacement === + +# GPU 0: layers 0–7 +- match: + name: "^model\\.layers\\.([0-7])\\.mlp\\.gate$" + class: ktransformers.models.modeling_deepseek_v3.MoEGate + replace: + class: ktransformers.operators.gate.KMoEGate + kwargs: + generate_device: "cuda:0" + prefill_device: "cuda:0" + +# GPU 1: layers 8–15 +- match: + name: "^model\\.layers\\.(8|9|1[0-5])\\.mlp\\.gate$" + class: ktransformers.models.modeling_deepseek_v3.MoEGate + replace: + class: ktransformers.operators.gate.KMoEGate + kwargs: + generate_device: "cuda:1" + prefill_device: "cuda:1" + +# GPU 2: layers 16–23 +- match: + name: "^model\\.layers\\.(1[6-9]|2[0-3])\\.mlp\\.gate$" + class: ktransformers.models.modeling_deepseek_v3.MoEGate + replace: + class: ktransformers.operators.gate.KMoEGate + kwargs: + generate_device: "cuda:2" + prefill_device: "cuda:2" + +# GPU 3: layers 24–31 +- match: + name: "^model\\.layers\\.(2[4-9]|3[0-1])\\.mlp\\.gate$" + class: ktransformers.models.modeling_deepseek_v3.MoEGate + replace: + class: ktransformers.operators.gate.KMoEGate + kwargs: + generate_device: "cuda:3" + prefill_device: "cuda:3" + +# GPU 4: layers 32–39 +- match: + name: "^model\\.layers\\.(3[2-9])\\.mlp\\.gate$" + class: ktransformers.models.modeling_deepseek_v3.MoEGate + replace: + class: ktransformers.operators.gate.KMoEGate + kwargs: + generate_device: "cuda:4" + prefill_device: "cuda:4" + +# GPU 5: layers 40–47 +- match: + name: "^model\\.layers\\.(4[0-7])\\.mlp\\.gate$" + class: ktransformers.models.modeling_deepseek_v3.MoEGate + replace: + class: ktransformers.operators.gate.KMoEGate + kwargs: + generate_device: "cuda:5" + prefill_device: "cuda:5" + +# GPU 6: layers 48–55 +- match: + name: "^model\\.layers\\.(4[8-9]|5[0-5])\\.mlp\\.gate$" + class: ktransformers.models.modeling_deepseek_v3.MoEGate + replace: + class: ktransformers.operators.gate.KMoEGate + kwargs: + generate_device: "cuda:6" + prefill_device: "cuda:6" + +# GPU 7: layers 56–60 +- match: + name: "^model\\.layers\\.(5[6-9]|60)\\.mlp\\.gate$" + class: ktransformers.models.modeling_deepseek_v3.MoEGate + replace: + class: ktransformers.operators.gate.KMoEGate + kwargs: + generate_device: "cuda:7" + prefill_device: "cuda:7" + + +# === MLP Experts Replacement === +# replace with marlin expert. Open and modify layer-num as needed. +# Each layer of malin experts takes about 6GB of GPU memory. +# !!!Do remember 'close' cuda graph if you are using marlin expert.!!! +# !!!Loading marlin expert will take signifcant time.!!! + +# GPU 0: layers 0–7 +# - match: +# name: "^model\\.layers\\.([0-7])\\.mlp\\.experts$" # inject experts in layer 0~4 as marlin expert +# replace: +# class: ktransformers.operators.experts.KTransformersExperts +# kwargs: +# generate_device: "cuda:0" +# generate_op: "KExpertsMarlin" +# recursive: False + +# # GPU 1: layers 8–15 +# - match: +# name: "^model\\.layers\\.([8-9]|1[0-5)\\.mlp\\.experts$" # inject experts in layer 30~31 as marlin expert +# replace: +# class: ktransformers.operators.experts.KTransformersExperts +# kwargs: +# generate_device: "cuda:1" +# generate_op: "KExpertsMarlin" +# recursive: False + +# # GPU 2: layers 16–23 +# - match: +# name: "^model\\.layers\\.(1[6-9]|2[0-3])\\.mlp\\.experts$" # inject experts in layer 0~4 as marlin expert +# replace: +# class: ktransformers.operators.experts.KTransformersExperts +# kwargs: +# generate_device: "cuda:0" +# generate_op: "KExpertsMarlin" +# recursive: False + +# # GPU 3: layers 24–31 +# - match: +# name: "^model\\.layers\\.(2[4-9]|3[0-1])\\.mlp\\.experts$" # inject experts in layer 30~31 as marlin expert +# replace: +# class: ktransformers.operators.experts.KTransformersExperts +# kwargs: +# generate_device: "cuda:1" +# generate_op: "KExpertsMarlin" +# recursive: False + +# # GPU 4: layers 32–39 +# - match: +# name: "^model\\.layers\\.(3[2-9])\\.mlp\\.experts$" # inject experts in layer 0~4 as marlin expert +# replace: +# class: ktransformers.operators.experts.KTransformersExperts +# kwargs: +# generate_device: "cuda:0" +# generate_op: "KExpertsMarlin" +# recursive: False + +# # GPU 5: layers 40–47 +# - match: +# name: "^model\\.layers\\.(4[0-7])\\.mlp\\.experts$" # inject experts in layer 30~31 as marlin expert +# replace: +# class: ktransformers.operators.experts.KTransformersExperts +# kwargs: +# generate_device: "cuda:1" +# generate_op: "KExpertsMarlin" +# recursive: False + +# # GPU 6: layers 48–55 +# - match: +# name: "^model\\.layers\\.(4[8-9]|5[0-5])\\.mlp\\.experts$" # inject experts in layer 0~4 as marlin expert +# replace: +# class: ktransformers.operators.experts.KTransformersExperts +# kwargs: +# generate_device: "cuda:0" +# generate_op: "KExpertsMarlin" +# recursive: False + +# # GPU 7: layers 56–60 +# - match: +# name: "^model\\.layers\\.(5[6-9]|60)\\.mlp\\.experts$" # inject experts in layer 30~31 as marlin expert +# replace: +# class: ktransformers.operators.experts.KTransformersExperts +# kwargs: +# generate_device: "cuda:1" +# generate_op: "KExpertsMarlin" +# recursive: False + + +# === MLP Experts Replacement === + +# GPU 0: layers 0–7 +- match: + name: "^model\\.layers\\.([0-7])\\.mlp\\.experts$" + replace: + class: ktransformers.operators.experts.KTransformersExperts + kwargs: + prefill_device: "cuda:0" + prefill_op: "KExpertsTorch" + generate_device: "cpu" + generate_op: "KExpertsCPU" + out_device: "cuda:0" + recursive: False + +# GPU 1: layers 8–15 +- match: + name: "^model\\.layers\\.(8|9|1[0-5])\\.mlp\\.experts$" + replace: + class: ktransformers.operators.experts.KTransformersExperts + kwargs: + prefill_device: "cuda:1" + prefill_op: "KExpertsTorch" + generate_device: "cpu" + generate_op: "KExpertsCPU" + out_device: "cuda:1" + recursive: False + +# GPU 2: layers 16–23 +- match: + name: "^model\\.layers\\.(1[6-9]|2[0-3])\\.mlp\\.experts$" + replace: + class: ktransformers.operators.experts.KTransformersExperts + kwargs: + prefill_device: "cuda:2" + prefill_op: "KExpertsTorch" + generate_device: "cpu" + generate_op: "KExpertsCPU" + out_device: "cuda:2" + recursive: False + +# GPU 3: layers 24–31 +- match: + name: "^model\\.layers\\.(2[4-9]|3[0-1])\\.mlp\\.experts$" + replace: + class: ktransformers.operators.experts.KTransformersExperts + kwargs: + prefill_device: "cuda:3" + prefill_op: "KExpertsTorch" + generate_device: "cpu" + generate_op: "KExpertsCPU" + out_device: "cuda:3" + recursive: False + +# GPU 4: layers 32–39 +- match: + name: "^model\\.layers\\.(3[2-9])\\.mlp\\.experts$" + replace: + class: ktransformers.operators.experts.KTransformersExperts + kwargs: + prefill_device: "cuda:4" + prefill_op: "KExpertsTorch" + generate_device: "cpu" + generate_op: "KExpertsCPU" + out_device: "cuda:4" + recursive: False + +# GPU 5: layers 40–47 +- match: + name: "^model\\.layers\\.(4[0-7])\\.mlp\\.experts$" + replace: + class: ktransformers.operators.experts.KTransformersExperts + kwargs: + prefill_device: "cuda:5" + prefill_op: "KExpertsTorch" + generate_device: "cpu" + generate_op: "KExpertsCPU" + out_device: "cuda:5" + recursive: False + +# GPU 6: layers 48–55 +- match: + name: "^model\\.layers\\.(4[8-9]|5[0-5])\\.mlp\\.experts$" + replace: + class: ktransformers.operators.experts.KTransformersExperts + kwargs: + prefill_device: "cuda:6" + prefill_op: "KExpertsTorch" + generate_device: "cpu" + generate_op: "KExpertsCPU" + out_device: "cuda:6" + recursive: False + +# GPU 7: layers 56–60 +- match: + name: "^model\\.layers\\.(5[6-9]|60)\\.mlp\\.experts$" + replace: + class: ktransformers.operators.experts.KTransformersExperts + kwargs: + prefill_device: "cuda:7" + prefill_op: "KExpertsTorch" + generate_device: "cpu" + generate_op: "KExpertsCPU" + out_device: "cuda:7" + recursive: False + + +# === Self-Attention Replacement === + +# GPU 0: layers 0–7 +- match: + name: "^model\\.layers\\.([0-7])\\.self_attn$" + replace: + class: ktransformers.operators.attention.KDeepseekV2Attention + kwargs: + generate_device: "cuda:0" + prefill_device: "cuda:0" + +# GPU 1: layers 8–15 +- match: + name: "^model\\.layers\\.(8|9|1[0-5])\\.self_attn$" + replace: + class: ktransformers.operators.attention.KDeepseekV2Attention + kwargs: + generate_device: "cuda:1" + prefill_device: "cuda:1" + +# GPU 2: layers 16–23 +- match: + name: "^model\\.layers\\.(1[6-9]|2[0-3])\\.self_attn$" + replace: + class: ktransformers.operators.attention.KDeepseekV2Attention + kwargs: + generate_device: "cuda:2" + prefill_device: "cuda:2" + +# GPU 3: layers 24–31 +- match: + name: "^model\\.layers\\.(2[4-9]|3[0-1])\\.self_attn$" + replace: + class: ktransformers.operators.attention.KDeepseekV2Attention + kwargs: + generate_device: "cuda:3" + prefill_device: "cuda:3" + +# GPU 4: layers 32–39 +- match: + name: "^model\\.layers\\.(3[2-9])\\.self_attn$" + replace: + class: ktransformers.operators.attention.KDeepseekV2Attention + kwargs: + generate_device: "cuda:4" + prefill_device: "cuda:4" + +# GPU 5: layers 40–47 +- match: + name: "^model\\.layers\\.(4[0-7])\\.self_attn$" + replace: + class: ktransformers.operators.attention.KDeepseekV2Attention + kwargs: + generate_device: "cuda:5" + prefill_device: "cuda:5" + +# GPU 6: layers 48–55 +- match: + name: "^model\\.layers\\.(4[8-9]|5[0-5])\\.self_attn$" + replace: + class: ktransformers.operators.attention.KDeepseekV2Attention + kwargs: + generate_device: "cuda:6" + prefill_device: "cuda:6" + +# GPU 7: layers 56–60 +- match: + name: "^model\\.layers\\.(5[6-9]|60)\\.self_attn$" + replace: + class: ktransformers.operators.attention.KDeepseekV2Attention + kwargs: + generate_device: "cuda:7" + prefill_device: "cuda:7" + +# === Overall Model Replacement with Transfer Map === + +- match: + name: "^model$" + replace: + class: "ktransformers.operators.models.KDeepseekV2Model" + kwargs: + per_layer_prefill_intput_threshold: 0 # 0 means close layer‐wise prefill + transfer_map: + 8: "cuda:1" + 16: "cuda:2" + 24: "cuda:3" + 32: "cuda:4" + 40: "cuda:5" + 48: "cuda:6" + 56: "cuda:7" + +# === Default Catch-All for Other Modules === + +# GPU 0: layers 0–7 +- match: + name: "^model\\.layers\\.([0-7])\\." + replace: + class: "default" + kwargs: + generate_device: "cuda:0" + prefill_device: "cuda:0" + +# GPU 1: layers 8–15 +- match: + name: "^model\\.layers\\.(8|9|1[0-5])\\." + replace: + class: "default" + kwargs: + generate_device: "cuda:1" + prefill_device: "cuda:1" + +# GPU 2: layers 16–23 +- match: + name: "^model\\.layers\\.(1[6-9]|2[0-3])\\." + replace: + class: "default" + kwargs: + generate_device: "cuda:2" + prefill_device: "cuda:2" + +# GPU 3: layers 24–31 +- match: + name: "^model\\.layers\\.(2[4-9]|3[0-1])\\." + replace: + class: "default" + kwargs: + generate_device: "cuda:3" + prefill_device: "cuda:3" + +# GPU 4: layers 32–39 +- match: + name: "^model\\.layers\\.(3[2-9])\\." + replace: + class: "default" + kwargs: + generate_device: "cuda:4" + prefill_device: "cuda:4" + +# GPU 5: layers 40–47 +- match: + name: "^model\\.layers\\.(4[0-7])\\." + replace: + class: "default" + kwargs: + generate_device: "cuda:5" + prefill_device: "cuda:5" + +# GPU 6: layers 48–55 +- match: + name: "^model\\.layers\\.(4[8-9]|5[0-5])\\." + replace: + class: "default" + kwargs: + generate_device: "cuda:6" + prefill_device: "cuda:6" + +# GPU 7: layers 56–63 +- match: + name: "^model\\.layers\\.(5[6-9]|60)\\." + replace: + class: "default" + kwargs: + generate_device: "cuda:7" + prefill_device: "cuda:7" + +- match: + name: "^lm_head" + class: torch.nn.Linear + replace: + class: ktransformers.operators.linear.KTransformersLinear + kwargs: + generate_device: "cuda:7" + prefill_device: "cuda:7" + generate_op: "KLinearMarlin" + prefill_op: "KLinearTorch" + +# For final modules (model.norm), ensure they are on GPU 7 (as in your original config) +- match: + name: "(^model\\.layers\\.(4[5-9]|5[0-9]|60)\\.)|(^model\\.norm)" + replace: + class: "default" + kwargs: + generate_device: "cuda:7" + prefill_device: "cuda:7" \ No newline at end of file diff --git a/archive/kt-sft/ktransformers/optimize/optimize_rules/DeepSeek-V3-Chat-multi-gpu-fp8-linear-ggml-experts.yaml b/archive/kt-sft/ktransformers/optimize/optimize_rules/DeepSeek-V3-Chat-multi-gpu-fp8-linear-ggml-experts.yaml new file mode 100644 index 0000000..fa8c03d --- /dev/null +++ b/archive/kt-sft/ktransformers/optimize/optimize_rules/DeepSeek-V3-Chat-multi-gpu-fp8-linear-ggml-experts.yaml @@ -0,0 +1,157 @@ +- match: + name: "^model.embed_tokens" + replace: + class: "default" + kwargs: + generate_device: "cpu" + prefill_device: "cpu" + +- match: + name: "^model\\.layers\\.(0|[1-9]|[12][0-9])\\." + class: ktransformers.models.modeling_deepseek_v3.DeepseekV3RotaryEmbedding + replace: + class: ktransformers.operators.RoPE.YarnRotaryEmbeddingV3 + kwargs: + generate_device: "cuda:0" + prefill_device: "cuda:0" +- match: + name: "^model\\.layers\\.([3456][0-9])\\." + class: ktransformers.models.modeling_deepseek_v3.DeepseekV3RotaryEmbedding + replace: + class: ktransformers.operators.RoPE.YarnRotaryEmbeddingV3 + kwargs: + generate_device: "cuda:1" + prefill_device: "cuda:1" + +- match: + name: "^model\\.layers\\.(0|[1-9]|[12][0-9])\\.(?!self_attn\\.kv_b_proj).*$" # regular expression + class: torch.nn.Linear # only match modules matching name and class simultaneously + replace: + class: ktransformers.operators.linear.KTransformersLinear # optimized Kernel on quantized data types + kwargs: + generate_device: "cuda:0" + prefill_device: "cuda:0" + generate_op: "KLinearFP8" + prefill_op: "KLinearTorch" + +- match: + name: "^model\\.layers\\.([3456][0-9])\\.(?!self_attn\\.kv_b_proj).*$" # regular expression + class: torch.nn.Linear # only match modules matching name and class simultaneously + replace: + class: ktransformers.operators.linear.KTransformersLinear # optimized Kernel on quantized data types + kwargs: + generate_device: "cuda:1" + prefill_device: "cuda:1" + generate_op: "KLinearFP8" + prefill_op: "KLinearTorch" + +- match: + name: "^model\\.layers\\.(0|[1-9]|[12][0-9])\\.mlp$" + class: ktransformers.models.modeling_deepseek_v3.DeepseekV3MoE + replace: + class: ktransformers.operators.experts.KDeepseekV3MoE # mlp module with custom forward function + kwargs: + generate_device: "cuda:0" + prefill_device: "cuda:0" +- match: + name: "^model\\.layers\\.([3456][0-9])\\.mlp$" + class: ktransformers.models.modeling_deepseek_v3.DeepseekV3MoE + replace: + class: ktransformers.operators.experts.KDeepseekV3MoE # mlp module with custom forward function + kwargs: + generate_device: "cuda:1" + prefill_device: "cuda:1" + +- match: + name: "^model\\.layers\\.(0|[1-9]|[12][0-9])\\.mlp\\.gate$" + class: ktransformers.models.modeling_deepseek_v3.MoEGate + replace: + class: ktransformers.operators.gate.KMoEGate + kwargs: + generate_device: "cuda:0" + prefill_device: "cuda:0" +- match: + name: "^model\\.layers\\.([3456][0-9])\\.mlp\\.gate$" + class: ktransformers.models.modeling_deepseek_v3.MoEGate + replace: + class: ktransformers.operators.gate.KMoEGate # mlp module with custom forward function + kwargs: + generate_device: "cuda:1" + prefill_device: "cuda:1" + +- match: + name: "^model\\.layers\\.(0|[1-9]|[12][0-9])\\.mlp\\.experts$" + replace: + class: ktransformers.operators.experts.KTransformersExperts # custom MoE Kernel with expert paralleism + kwargs: + prefill_device: "cuda:0" + prefill_op: "KExpertsTorch" + generate_device: "cpu" + generate_op: "KExpertsCPU" + out_device: "cuda:0" + recursive: False # don't recursively inject submodules of this module + +- match: + name: "^model\\.layers\\.([3456][0-9])\\.mlp\\.experts$" + replace: + class: ktransformers.operators.experts.KTransformersExperts # custom MoE Kernel with expert paralleism + kwargs: + prefill_device: "cuda:1" + prefill_op: "KExpertsTorch" + generate_device: "cpu" + generate_op: "KExpertsCPU" + out_device: "cuda:1" + recursive: False # don't recursively inject submodules of this module + +- match: + name: "^model\\.layers\\.(0|[1-9]|[12][0-9])\\.self_attn$" + replace: + class: ktransformers.operators.attention.KDeepseekV2Attention # optimized MLA implementation + kwargs: + generate_device: "cuda:0" + prefill_device: "cuda:0" + absorb_for_prefill: False # change this to True to enable long context(prefill may slower). + +- match: + name: "^model\\.layers\\.([3456][0-9])\\.self_attn$" + replace: + class: ktransformers.operators.attention.KDeepseekV2Attention # optimized MLA implementation + kwargs: + generate_device: "cuda:1" + prefill_device: "cuda:1" + absorb_for_prefill: False # change this to True to enable long context(prefill may slower). + +- match: + name: "^model$" + replace: + class: "ktransformers.operators.models.KDeepseekV2Model" + kwargs: + per_layer_prefill_intput_threshold: 0 # 0 is close layer wise prefill + transfer_map: + 30: "cuda:1" + +- match: + name: "^model\\.layers\\.(0|[1-9]|[12][0-9])\\." + replace: + class: "default" + kwargs: + generate_device: "cuda:0" + prefill_device: "cuda:0" + +- match: + name: "^lm_head" + class: torch.nn.Linear + replace: + class: "default" + kwargs: + generate_device: "cuda:1" + prefill_device: "cuda:1" + + +- match: + name: "(^model\\.layers\\.([3456][0-9])\\.)|(model.norm)" + replace: + class: "default" + kwargs: + generate_device: "cuda:1" + prefill_device: "cuda:1" diff --git a/archive/kt-sft/ktransformers/optimize/optimize_rules/DeepSeek-V3-Chat-multi-gpu-marlin.yaml b/archive/kt-sft/ktransformers/optimize/optimize_rules/DeepSeek-V3-Chat-multi-gpu-marlin.yaml new file mode 100644 index 0000000..e04c6ce --- /dev/null +++ b/archive/kt-sft/ktransformers/optimize/optimize_rules/DeepSeek-V3-Chat-multi-gpu-marlin.yaml @@ -0,0 +1,172 @@ +- match: + name: "^model.embed_tokens" + replace: + class: "default" + kwargs: + generate_device: "cpu" + prefill_device: "cpu" + +- match: + name: "^model\\.layers\\.(0|[1-9]|[12][0-9])\\." + class: ktransformers.models.modeling_deepseek_v3.DeepseekV3RotaryEmbedding + replace: + class: ktransformers.operators.RoPE.YarnRotaryEmbeddingV3 + kwargs: + generate_device: "cuda:0" + prefill_device: "cuda:0" +- match: + name: "^model\\.layers\\.([3456][0-9])\\." + class: ktransformers.models.modeling_deepseek_v3.DeepseekV3RotaryEmbedding + replace: + class: ktransformers.operators.RoPE.YarnRotaryEmbeddingV3 + kwargs: + generate_device: "cuda:1" + prefill_device: "cuda:1" + +- match: + name: "^model\\.layers\\.(0|[1-9]|[12][0-9])\\.(?!self_attn\\.kv_b_proj).*$" # regular expression + class: torch.nn.Linear # only match modules matching name and class simultaneously + replace: + class: ktransformers.operators.linear.KTransformersLinear # optimized Kernel on quantized data types + kwargs: + generate_device: "cuda:0" + prefill_device: "cuda:0" + generate_op: "KLinearMarlin" + prefill_op: "KLinearTorch" + +- match: + name: "^model\\.layers\\.([3456][0-9])\\.(?!self_attn\\.kv_b_proj).*$" # regular expression + class: torch.nn.Linear # only match modules matching name and class simultaneously + replace: + class: ktransformers.operators.linear.KTransformersLinear # optimized Kernel on quantized data types + kwargs: + generate_device: "cuda:1" + prefill_device: "cuda:1" + generate_op: "KLinearMarlin" + prefill_op: "KLinearTorch" + +- match: + name: "^model\\.layers\\.(0|[1-9]|[12][0-9])\\.mlp$" + class: ktransformers.models.modeling_deepseek_v3.DeepseekV3MoE + replace: + class: ktransformers.operators.experts.KDeepseekV3MoE # mlp module with custom forward function + kwargs: + generate_device: "cuda:0" + prefill_device: "cuda:0" +- match: + name: "^model\\.layers\\.([3456][0-9])\\.mlp$" + class: ktransformers.models.modeling_deepseek_v3.DeepseekV3MoE + replace: + class: ktransformers.operators.experts.KDeepseekV3MoE # mlp module with custom forward function + kwargs: + generate_device: "cuda:1" + prefill_device: "cuda:1" + +- match: + name: "^model\\.layers\\.(0|[1-9]|[12][0-9])\\.mlp\\.gate$" + class: ktransformers.models.modeling_deepseek_v3.MoEGate + replace: + class: ktransformers.operators.gate.KMoEGate + kwargs: + generate_device: "cuda:0" + prefill_device: "cuda:0" +- match: + name: "^model\\.layers\\.([3456][0-9])\\.mlp\\.gate$" + class: ktransformers.models.modeling_deepseek_v3.MoEGate + replace: + class: ktransformers.operators.gate.KMoEGate # mlp module with custom forward function + kwargs: + generate_device: "cuda:1" + prefill_device: "cuda:1" + +- match: + name: "^model\\.layers\\.(0|[1-4])\\.mlp\\.experts$" # inject experts in layer 0~4 as marlin expert + replace: + class: ktransformers.operators.experts.KTransformersExperts + kwargs: + generate_device: "cuda:0" # run in cuda:0 + generate_op: "KExpertsMarlin" + recursive: False + +- match: + name: "^model\\.layers\\.([3][0])\\.mlp\\.experts$" # inject experts in layer 30~31 as marlin expert + replace: + class: ktransformers.operators.experts.KTransformersExperts + kwargs: + generate_device: "cuda:1" + generate_op: "KExpertsMarlin" + recursive: False + +- match: + name: "^model\\.layers\\.(0|[1-9]|[12][0-9])\\.mlp\\.experts$" + replace: + class: ktransformers.operators.experts.KTransformersExperts # custom MoE Kernel with expert paralleism + kwargs: + prefill_device: "cuda:0" + prefill_op: "KExpertsTorch" + generate_device: "cpu" + generate_op: "KExpertsCPU" + out_device: "cuda:0" + recursive: False # don't recursively inject submodules of this module + +- match: + name: "^model\\.layers\\.([3456][0-9])\\.mlp\\.experts$" + replace: + class: ktransformers.operators.experts.KTransformersExperts # custom MoE Kernel with expert paralleism + kwargs: + prefill_device: "cuda:1" + prefill_op: "KExpertsTorch" + generate_device: "cpu" + generate_op: "KExpertsCPU" + out_device: "cuda:1" + recursive: False # don't recursively inject submodules of this module + +- match: + name: "^model\\.layers\\.(0|[1-9]|[12][0-9])\\.self_attn$" + replace: + class: ktransformers.operators.attention.KDeepseekV2Attention # optimized MLA implementation + kwargs: + generate_device: "cuda:0" + prefill_device: "cuda:0" +- match: + name: "^model\\.layers\\.([3456][0-9])\\.self_attn$" + replace: + class: ktransformers.operators.attention.KDeepseekV2Attention # optimized MLA implementation + kwargs: + generate_device: "cuda:1" + prefill_device: "cuda:1" +- match: + name: "^model$" + replace: + class: "ktransformers.operators.models.KDeepseekV2Model" + kwargs: + per_layer_prefill_intput_threshold: 0 # 0 is close layer wise prefill + transfer_map: + 30: "cuda:1" + +- match: + name: "^model\\.layers\\.(0|[1-9]|[12][0-9])\\." + replace: + class: "default" + kwargs: + generate_device: "cuda:0" + prefill_device: "cuda:0" + +- match: + name: "^lm_head" + class: torch.nn.Linear + replace: + class: ktransformers.operators.linear.KTransformersLinear + kwargs: + generate_device: "cuda:0" + prefill_device: "cuda:0" + generate_op: "KLinearMarlin" + prefill_op: "KLinearTorch" + +- match: + name: "(^model\\.layers\\.([3456][0-9])\\.)|(model.norm)" + replace: + class: "default" + kwargs: + generate_device: "cuda:1" + prefill_device: "cuda:1" diff --git a/archive/kt-sft/ktransformers/optimize/optimize_rules/DeepSeek-V3-Chat-multi-gpu.yaml b/archive/kt-sft/ktransformers/optimize/optimize_rules/DeepSeek-V3-Chat-multi-gpu.yaml new file mode 100644 index 0000000..50e282d --- /dev/null +++ b/archive/kt-sft/ktransformers/optimize/optimize_rules/DeepSeek-V3-Chat-multi-gpu.yaml @@ -0,0 +1,154 @@ +- match: + name: "^model.embed_tokens" + replace: + class: "default" + kwargs: + generate_device: "cpu" + prefill_device: "cpu" + +- match: + name: "^model\\.layers\\.(0|[1-9]|[12][0-9])\\." + class: ktransformers.models.modeling_deepseek_v3.DeepseekV3RotaryEmbedding + replace: + class: ktransformers.operators.RoPE.YarnRotaryEmbeddingV3 + kwargs: + generate_device: "cuda:0" + prefill_device: "cuda:0" +- match: + name: "^model\\.layers\\.([3456][0-9])\\." + class: ktransformers.models.modeling_deepseek_v3.DeepseekV3RotaryEmbedding + replace: + class: ktransformers.operators.RoPE.YarnRotaryEmbeddingV3 + kwargs: + generate_device: "cuda:1" + prefill_device: "cuda:1" + +- match: + name: "^model\\.layers\\.(0|[1-9]|[12][0-9])\\.(?!self_attn\\.kv_b_proj).*$" # regular expression + class: torch.nn.Linear # only match modules matching name and class simultaneously + replace: + class: ktransformers.operators.linear.KTransformersLinear # optimized Kernel on quantized data types + kwargs: + generate_device: "cuda:0" + prefill_device: "cuda:0" + generate_op: "KLinearMarlin" + prefill_op: "KLinearTorch" + +- match: + name: "^model\\.layers\\.([3456][0-9])\\.(?!self_attn\\.kv_b_proj).*$" # regular expression + class: torch.nn.Linear # only match modules matching name and class simultaneously + replace: + class: ktransformers.operators.linear.KTransformersLinear # optimized Kernel on quantized data types + kwargs: + generate_device: "cuda:1" + prefill_device: "cuda:1" + generate_op: "KLinearMarlin" + prefill_op: "KLinearTorch" + +- match: + name: "^model\\.layers\\.(0|[1-9]|[12][0-9])\\.mlp$" + class: ktransformers.models.modeling_deepseek_v3.DeepseekV3MoE + replace: + class: ktransformers.operators.experts.KDeepseekV3MoE # mlp module with custom forward function + kwargs: + generate_device: "cuda:0" + prefill_device: "cuda:0" +- match: + name: "^model\\.layers\\.([3456][0-9])\\.mlp$" + class: ktransformers.models.modeling_deepseek_v3.DeepseekV3MoE + replace: + class: ktransformers.operators.experts.KDeepseekV3MoE # mlp module with custom forward function + kwargs: + generate_device: "cuda:1" + prefill_device: "cuda:1" + +- match: + name: "^model\\.layers\\.(0|[1-9]|[12][0-9])\\.mlp\\.gate$" + class: ktransformers.models.modeling_deepseek_v3.MoEGate + replace: + class: ktransformers.operators.gate.KMoEGate + kwargs: + generate_device: "cuda:0" + prefill_device: "cuda:0" +- match: + name: "^model\\.layers\\.([3456][0-9])\\.mlp\\.gate$" + class: ktransformers.models.modeling_deepseek_v3.MoEGate + replace: + class: ktransformers.operators.gate.KMoEGate # mlp module with custom forward function + kwargs: + generate_device: "cuda:1" + prefill_device: "cuda:1" + +- match: + name: "^model\\.layers\\.(0|[1-9]|[12][0-9])\\.mlp\\.experts$" + replace: + class: ktransformers.operators.experts.KTransformersExperts # custom MoE Kernel with expert paralleism + kwargs: + prefill_device: "cuda:0" + prefill_op: "KExpertsTorch" + generate_device: "cpu" + generate_op: "KExpertsCPU" + out_device: "cuda:0" + recursive: False # don't recursively inject submodules of this module + +- match: + name: "^model\\.layers\\.([3456][0-9])\\.mlp\\.experts$" + replace: + class: ktransformers.operators.experts.KTransformersExperts # custom MoE Kernel with expert paralleism + kwargs: + prefill_device: "cuda:1" + prefill_op: "KExpertsTorch" + generate_device: "cpu" + generate_op: "KExpertsCPU" + out_device: "cuda:1" + recursive: False # don't recursively inject submodules of this module + +- match: + name: "^model\\.layers\\.(0|[1-9]|[12][0-9])\\.self_attn$" + replace: + class: ktransformers.operators.attention.KDeepseekV2Attention # optimized MLA implementation + kwargs: + generate_device: "cuda:0" + prefill_device: "cuda:0" +- match: + name: "^model\\.layers\\.([3456][0-9])\\.self_attn$" + replace: + class: ktransformers.operators.attention.KDeepseekV2Attention # optimized MLA implementation + kwargs: + generate_device: "cuda:1" + prefill_device: "cuda:1" +- match: + name: "^model$" + replace: + class: "ktransformers.operators.models.KDeepseekV2Model" + kwargs: + per_layer_prefill_intput_threshold: 0 # 0 is close layer wise prefill + transfer_map: + 30: "cuda:1" + +- match: + name: "^model\\.layers\\.(0|[1-9]|[12][0-9])\\." + replace: + class: "default" + kwargs: + generate_device: "cuda:0" + prefill_device: "cuda:0" + +- match: + name: "^lm_head" + class: torch.nn.Linear + replace: + class: ktransformers.operators.linear.KTransformersLinear + kwargs: + generate_device: "cuda:1" + prefill_device: "cuda:1" + generate_op: "KLinearMarlin" + prefill_op: "KLinearTorch" + +- match: + name: "(^model\\.layers\\.([3456][0-9])\\.)|(model.norm)" + replace: + class: "default" + kwargs: + generate_device: "cuda:1" + prefill_device: "cuda:1" diff --git a/archive/kt-sft/ktransformers/optimize/optimize_rules/DeepSeek-V3-Chat-serve.yaml b/archive/kt-sft/ktransformers/optimize/optimize_rules/DeepSeek-V3-Chat-serve.yaml new file mode 100644 index 0000000..e1c61b8 --- /dev/null +++ b/archive/kt-sft/ktransformers/optimize/optimize_rules/DeepSeek-V3-Chat-serve.yaml @@ -0,0 +1,92 @@ +- match: + class: ktransformers.models.modeling_deepseek_v3.DeepseekV3RotaryEmbedding + replace: + class: ktransformers.operators.RoPE.YarnRotaryEmbeddingV3 + kwargs: + generate_device: "cuda" + prefill_device: "cuda" + +- match: + name: "^lm_head$" # regular expression + class: torch.nn.Linear # only match modules matching name and class simultaneously + replace: + class: ktransformers.operators.linear.KTransformersLinear # optimized Kernel on quantized data types + kwargs: + generate_device: "cuda" + prefill_device: "cuda" + generate_op: "VLinearMarlin" + prefill_op: "KLinearTorch" + +- match: + name: "^model\\.layers\\.(?!.*self_attn\\.kv_b_proj).*$" # regular expression + class: torch.nn.Linear # only match modules matching name and class simultaneously + replace: + class: ktransformers.operators.linear.KTransformersLinear # optimized Kernel on quantized data types + kwargs: + generate_device: "cuda" + prefill_device: "cuda" + generate_op: "VLinearMarlin" + prefill_op: "KLinearTorch" +- match: + name: "^model\\.layers\\..*\\.mlp$" + class: ktransformers.models.modeling_deepseek_v3.DeepseekV3MoE + replace: + class: ktransformers.operators.experts.KDeepseekV3MoEV2 # mlp module with custom forward function + kwargs: + generate_device: "cuda" + prefill_device: "cuda" +- match: + class: ktransformers.models.modeling_deepseek_v3.MoEGate + replace: + class: ktransformers.operators.gate.KMoEGate + kwargs: + generate_device: "cuda:0" + prefill_device: "cuda:0" +- match: + name: "^model\\.layers\\..*\\.mlp\\.experts$" + replace: + class: ktransformers.operators.experts.KTransformersExpertsV2 # custom MoE Kernel with expert paralleism + kwargs: + prefill_device: "cuda" + prefill_op: "KExpertsTorch" + generate_device: "cpu" + generate_op: "KExpertsCPU" + out_device: "cuda" + recursive: False # don't recursively inject submodules of this module +- match: + name: "^model\\.layers\\..*\\.self_attn$" + replace: + class: ktransformers.operators.balance_serve_attention.flashinfer_attn # optimized MLA implementation + kwargs: + generate_device: "cuda" + prefill_device: "cuda" + absorb_for_prefill: False # change this to True to enable long context(prefill may slower). +- match: + name: "^model$" + replace: + class: "ktransformers.operators.models.KDeepseekV2Model" + kwargs: + per_layer_prefill_intput_threshold: 0 # 0 is close layer wise prefill +- match: + name: "^model.embed_tokens" + replace: + class: "default" + kwargs: + generate_device: "cpu" + prefill_device: "cpu" + +- match: + class: ktransformers.models.modeling_deepseek_v3.DeepseekV3RMSNorm + replace: + class: ktransformers.operators.layernorm.RMSNorm + kwargs: + generate_device: "cuda" + prefill_device: "cuda" + +- match: + class: ktransformers.models.modeling_deepseek_v3.DeepseekV3MLP + replace: + class: ktransformers.operators.mlp.kDeepseekV3MLP + kwargs: + generate_device: "cuda" + prefill_device: "cuda" \ No newline at end of file diff --git a/archive/kt-sft/ktransformers/optimize/optimize_rules/DeepSeek-V3-Chat-sft-amx-multi-gpu-4.yaml b/archive/kt-sft/ktransformers/optimize/optimize_rules/DeepSeek-V3-Chat-sft-amx-multi-gpu-4.yaml new file mode 100644 index 0000000..4eda68c --- /dev/null +++ b/archive/kt-sft/ktransformers/optimize/optimize_rules/DeepSeek-V3-Chat-sft-amx-multi-gpu-4.yaml @@ -0,0 +1,392 @@ +- match: + name: "^model.embed_tokens" + replace: + class: "default" + kwargs: + generate_device: "cpu" + prefill_device: "cpu" + +# === Rotary Embedding Replacement === + +# GPU 0: layers 0–14 +- match: + name: "^model\\.layers\\.([0-9]|1[0-4])\\." + class: ktransformers.models.modeling_deepseek_v3.DeepseekV3RotaryEmbedding + replace: + class: ktransformers.operators.RoPE.YarnRotaryEmbeddingV3 + kwargs: + generate_device: "cuda:0" + prefill_device: "cuda:0" + +# GPU 1: layers 15–29 +- match: + name: "^model\\.layers\\.(1[5-9]|2[0-9])\\." + class: ktransformers.models.modeling_deepseek_v3.DeepseekV3RotaryEmbedding + replace: + class: ktransformers.operators.RoPE.YarnRotaryEmbeddingV3 + kwargs: + generate_device: "cuda:1" + prefill_device: "cuda:1" + +# GPU 2: layers 30–44 +- match: + name: "^model\\.layers\\.(3[0-9]|4[0-4])\\." + class: ktransformers.models.modeling_deepseek_v3.DeepseekV3RotaryEmbedding + replace: + class: ktransformers.operators.RoPE.YarnRotaryEmbeddingV3 + kwargs: + generate_device: "cuda:2" + prefill_device: "cuda:2" + +# GPU 3: layers 45–60 +- match: + name: "^model\\.layers\\.(4[5-9]|5[0-9]|60)\\." + class: ktransformers.models.modeling_deepseek_v3.DeepseekV3RotaryEmbedding + replace: + class: ktransformers.operators.RoPE.YarnRotaryEmbeddingV3 + kwargs: + generate_device: "cuda:3" + prefill_device: "cuda:3" + +# === Linear Layers Replacement (excluding self_attn.kv_b_proj) === + +# GPU 0: layers 0–14 +- match: + name: "^model\\.layers\\.([0-9]|1[0-4])\\.(?!self_attn\\.kv_b_proj).*$" + class: torch.nn.Linear + replace: + class: ktransformers.operators.linear.KTransformersLinear + kwargs: + generate_device: "cuda:0" + prefill_device: "cuda:0" + generate_op: "KLinearTorch" + prefill_op: "KLinearTorch" + +# GPU 1: layers 15–29 +- match: + name: "^model\\.layers\\.(1[5-9]|2[0-9])\\.(?!self_attn\\.kv_b_proj).*$" + class: torch.nn.Linear + replace: + class: ktransformers.operators.linear.KTransformersLinear + kwargs: + generate_device: "cuda:1" + prefill_device: "cuda:1" + generate_op: "KLinearTorch" + prefill_op: "KLinearTorch" + +# GPU 2: layers 30–44 +- match: + name: "^model\\.layers\\.(3[0-9]|4[0-4])\\.(?!self_attn\\.kv_b_proj).*$" + class: torch.nn.Linear + replace: + class: ktransformers.operators.linear.KTransformersLinear + kwargs: + generate_device: "cuda:2" + prefill_device: "cuda:2" + generate_op: "KLinearTorch" + prefill_op: "KLinearTorch" + +# GPU 3: layers 45–60 +- match: + name: "^model\\.layers\\.(4[5-9]|5[0-9]|60)\\.(?!self_attn\\.kv_b_proj).*$" + class: torch.nn.Linear + replace: + class: ktransformers.operators.linear.KTransformersLinear + kwargs: + generate_device: "cuda:3" + prefill_device: "cuda:3" + generate_op: "KLinearTorch" + prefill_op: "KLinearTorch" + +# === MLP (MoE) Replacement === + +# GPU 0: layers 0–14 +- match: + name: "^model\\.layers\\.([0-9]|1[0-4])\\.mlp$" + class: ktransformers.models.modeling_deepseek_v3.DeepseekV3MoE + replace: + class: ktransformers.operators.experts.KDeepseekV3MoE + kwargs: + generate_device: "cuda:0" + prefill_device: "cuda:0" + +# GPU 1: layers 15–29 +- match: + name: "^model\\.layers\\.(1[5-9]|2[0-9])\\.mlp$" + class: ktransformers.models.modeling_deepseek_v3.DeepseekV3MoE + replace: + class: ktransformers.operators.experts.KDeepseekV3MoE + kwargs: + generate_device: "cuda:1" + prefill_device: "cuda:1" + +# GPU 2: layers 30–44 +- match: + name: "^model\\.layers\\.(3[0-9]|4[0-4])\\.mlp$" + class: ktransformers.models.modeling_deepseek_v3.DeepseekV3MoE + replace: + class: ktransformers.operators.experts.KDeepseekV3MoE + kwargs: + generate_device: "cuda:2" + prefill_device: "cuda:2" + +# GPU 3: layers 45–60 +- match: + name: "^model\\.layers\\.(4[5-9]|5[0-9]|60)\\.mlp$" + class: ktransformers.models.modeling_deepseek_v3.DeepseekV3MoE + replace: + class: ktransformers.operators.experts.KDeepseekV3MoE + kwargs: + generate_device: "cuda:3" + prefill_device: "cuda:3" + +# === MLP Gate Replacement === + +# GPU 0: layers 0–14 +- match: + name: "^model\\.layers\\.([0-9]|1[0-4])\\.mlp\\.gate$" + class: ktransformers.models.modeling_deepseek_v3.MoEGate + replace: + class: ktransformers.operators.gate.KMoEGate + kwargs: + generate_device: "cuda:0" + prefill_device: "cuda:0" + +# GPU 1: layers 15–29 +- match: + name: "^model\\.layers\\.(1[5-9]|2[0-9])\\.mlp\\.gate$" + class: ktransformers.models.modeling_deepseek_v3.MoEGate + replace: + class: ktransformers.operators.gate.KMoEGate + kwargs: + generate_device: "cuda:1" + prefill_device: "cuda:1" + +# GPU 2: layers 30–44 +- match: + name: "^model\\.layers\\.(3[0-9]|4[0-4])\\.mlp\\.gate$" + class: ktransformers.models.modeling_deepseek_v3.MoEGate + replace: + class: ktransformers.operators.gate.KMoEGate + kwargs: + generate_device: "cuda:2" + prefill_device: "cuda:2" + +# GPU 3: layers 45–60 +- match: + name: "^model\\.layers\\.(4[5-9]|5[0-9]|60)\\.mlp\\.gate$" + class: ktransformers.models.modeling_deepseek_v3.MoEGate + replace: + class: ktransformers.operators.gate.KMoEGate + kwargs: + generate_device: "cuda:3" + prefill_device: "cuda:3" + +# === MLP Experts Replacement === +# replace with marlin expert. Open and modify layer-num as needed. +# Each layer of malin experts takes about 6GB of GPU memory. +# !!!Do remember 'close' cuda graph if you are using marlin expert.!!! +# !!!KExpertsTorch is untested, we don't have enough VRAM.!!! + +# GPU 0: layers 3–4 +# - match: +# name: "^model\\.layers\\.([3-4])\\.mlp\\.experts$" +# replace: +# class: ktransformers.operators.experts.KTransformersExperts +# kwargs: +# generate_device: "cuda:0" +# generate_op: "KExpertsMarlin" +# recursive: False + +# # GPU 1: layers 15–17 +# - match: +# name: "^model\\.layers\\.(1[5-7])\\.mlp\\.experts$" +# replace: +# class: ktransformers.operators.experts.KTransformersExperts +# kwargs: +# generate_device: "cuda:1" +# generate_op: "KExpertsMarlin" +# recursive: False + +# # GPU 2: layers 30–32 +# - match: +# name: "^model\\.layers\\.(3[0-2])\\.mlp\\.experts$" +# replace: +# class: ktransformers.operators.experts.KTransformersExperts +# kwargs: +# generate_device: "cuda:2" +# generate_op: "KExpertsMarlin" +# recursive: False + +# # GPU 3: layers 45–46 +# - match: +# name: "^model\\.layers\\.(4[5-6])\\.mlp\\.experts$" +# replace: +# class: ktransformers.operators.experts.KTransformersExperts +# kwargs: +# generate_device: "cuda:3" +# generate_op: "KExpertsMarlin" +# recursive: False + + +# === MLP Experts Replacement === + +# GPU 0: layers 0–14 +- match: + name: "^model\\.layers\\.([0-9]|1[0-4])\\.mlp\\.experts$" + replace: + class: ktransformers.operators.experts.KTransformersExperts + kwargs: + prefill_device: "cuda:0" + prefill_op: "KExpertsTorch" + generate_device: "cpu" + generate_op: "KSFTExpertsCPU" + out_device: "cuda:0" + backend: "AMXInt8" # or "AMXBF16" or "llamafile" (default) + recursive: False + +# GPU 1: layers 15–29 +- match: + name: "^model\\.layers\\.(1[5-9]|2[0-9])\\.mlp\\.experts$" + replace: + class: ktransformers.operators.experts.KTransformersExperts + kwargs: + prefill_device: "cuda:1" + prefill_op: "KExpertsTorch" + generate_device: "cpu" + generate_op: "KSFTExpertsCPU" + out_device: "cuda:1" + backend: "AMXInt8" # or "AMXBF16" or "llamafile" (default) + recursive: False + +# GPU 2: layers 30–44 +- match: + name: "^model\\.layers\\.(3[0-9]|4[0-4])\\.mlp\\.experts$" + replace: + class: ktransformers.operators.experts.KTransformersExperts + kwargs: + prefill_device: "cuda:2" + prefill_op: "KExpertsTorch" + generate_device: "cpu" + generate_op: "KSFTExpertsCPU" + out_device: "cuda:2" + backend: "AMXInt8" # or "AMXBF16" or "llamafile" (default) + recursive: False + +# GPU 3: layers 45–60 +- match: + name: "^model\\.layers\\.(4[5-9]|5[0-9]|60)\\.mlp\\.experts$" + replace: + class: ktransformers.operators.experts.KTransformersExperts + kwargs: + prefill_device: "cuda:3" + prefill_op: "KExpertsTorch" + generate_device: "cpu" + generate_op: "KSFTExpertsCPU" + out_device: "cuda:3" + backend: "AMXInt8" # or "AMXBF16" or "llamafile" (default) + recursive: False + +# === Self-Attention Replacement === + +# GPU 0: layers 0–14 +- match: + name: "^model\\.layers\\.([0-9]|1[0-4])\\.self_attn$" + replace: + class: ktransformers.operators.attention.KDeepseekV2Attention + kwargs: + generate_device: "cuda:0" + prefill_device: "cuda:0" + absorb_for_prefill: False + +# GPU 1: layers 15–29 +- match: + name: "^model\\.layers\\.(1[5-9]|2[0-9])\\.self_attn$" + replace: + class: ktransformers.operators.attention.KDeepseekV2Attention + kwargs: + generate_device: "cuda:1" + prefill_device: "cuda:1" + absorb_for_prefill: False + +# GPU 2: layers 30–44 +- match: + name: "^model\\.layers\\.(3[0-9]|4[0-4])\\.self_attn$" + replace: + class: ktransformers.operators.attention.KDeepseekV2Attention + kwargs: + generate_device: "cuda:2" + prefill_device: "cuda:2" + absorb_for_prefill: False + +# GPU 3: layers 45–60 +- match: + name: "^model\\.layers\\.(4[5-9]|5[0-9]|60)\\.self_attn$" + replace: + class: ktransformers.operators.attention.KDeepseekV2Attention + kwargs: + generate_device: "cuda:3" + prefill_device: "cuda:3" + absorb_for_prefill: False + +# === Overall Model Replacement with Transfer Map === + +- match: + name: "^model$" + replace: + class: "ktransformers.operators.models.KDeepseekV2Model" + kwargs: + per_layer_prefill_intput_threshold: 0 # 0 means close layer‐wise prefill + transfer_map: + 15: "cuda:1" # Layers 15+ on GPU 1 + 30: "cuda:2" # Layers 30+ on GPU 2 + 45: "cuda:3" # Layers 45+ on GPU 3 + +# === Default Catch-All for Other Modules === + +# GPU 0: layers 0–14 +- match: + name: "^model\\.layers\\.([0-9]|1[0-4])\\." + replace: + class: "default" + kwargs: + generate_device: "cuda:0" + prefill_device: "cuda:0" + +# GPU 1: layers 15–29 +- match: + name: "^model\\.layers\\.(1[5-9]|2[0-9])\\." + replace: + class: "default" + kwargs: + generate_device: "cuda:1" + prefill_device: "cuda:1" + +# GPU 2: layers 30–44 +- match: + name: "^model\\.layers\\.(3[0-9]|4[0-4])\\." + replace: + class: "default" + kwargs: + generate_device: "cuda:2" + prefill_device: "cuda:2" + +- match: + name: "^lm_head" + class: torch.nn.Linear + replace: + class: ktransformers.operators.linear.KTransformersLinear + kwargs: + generate_device: "cuda:3" + prefill_device: "cuda:3" + generate_op: "KLinearTorch" + prefill_op: "KLinearTorch" + +# For final modules (model.norm), ensure they are on GPU 3 (as in your original config) +- match: + name: "(^model\\.layers\\.(4[5-9]|5[0-9]|60)\\.)|(^model\\.norm)" + replace: + class: "default" + kwargs: + generate_device: "cuda:3" + prefill_device: "cuda:3" diff --git a/archive/kt-sft/ktransformers/optimize/optimize_rules/DeepSeek-V3-Chat-sft-amx-multi-gpu.yaml b/archive/kt-sft/ktransformers/optimize/optimize_rules/DeepSeek-V3-Chat-sft-amx-multi-gpu.yaml new file mode 100644 index 0000000..8b8c204 --- /dev/null +++ b/archive/kt-sft/ktransformers/optimize/optimize_rules/DeepSeek-V3-Chat-sft-amx-multi-gpu.yaml @@ -0,0 +1,156 @@ +- match: + name: "^model.embed_tokens" + replace: + class: "default" + kwargs: + generate_device: "cpu" + prefill_device: "cpu" + +- match: + name: "^model\\.layers\\.(0|[1-9]|[12][0-9])\\." + class: ktransformers.models.modeling_deepseek_v3.DeepseekV3RotaryEmbedding + replace: + class: ktransformers.operators.RoPE.YarnRotaryEmbeddingV3 + kwargs: + generate_device: "cuda:0" + prefill_device: "cuda:0" +- match: + name: "^model\\.layers\\.([3456][0-9])\\." + class: ktransformers.models.modeling_deepseek_v3.DeepseekV3RotaryEmbedding + replace: + class: ktransformers.operators.RoPE.YarnRotaryEmbeddingV3 + kwargs: + generate_device: "cuda:1" + prefill_device: "cuda:1" + +- match: + name: "^model\\.layers\\.(0|[1-9]|[12][0-9])\\.(?!self_attn\\.kv_b_proj).*$" # regular expression + class: torch.nn.Linear # only match modules matching name and class simultaneously + replace: + class: ktransformers.operators.linear.KTransformersLinear # optimized Kernel on quantized data types + kwargs: + generate_device: "cuda:0" + prefill_device: "cuda:0" + generate_op: "KLinearTorch" + prefill_op: "KLinearTorch" + +- match: + name: "^model\\.layers\\.([3456][0-9])\\.(?!self_attn\\.kv_b_proj).*$" # regular expression + class: torch.nn.Linear # only match modules matching name and class simultaneously + replace: + class: ktransformers.operators.linear.KTransformersLinear # optimized Kernel on quantized data types + kwargs: + generate_device: "cuda:1" + prefill_device: "cuda:1" + generate_op: "KLinearTorch" + prefill_op: "KLinearTorch" + +- match: + name: "^model\\.layers\\.(0|[1-9]|[12][0-9])\\.mlp$" + class: ktransformers.models.modeling_deepseek_v3.DeepseekV3MoE + replace: + class: ktransformers.operators.experts.KDeepseekV3MoE # mlp module with custom forward function + kwargs: + generate_device: "cuda:0" + prefill_device: "cuda:0" +- match: + name: "^model\\.layers\\.([3456][0-9])\\.mlp$" + class: ktransformers.models.modeling_deepseek_v3.DeepseekV3MoE + replace: + class: ktransformers.operators.experts.KDeepseekV3MoE # mlp module with custom forward function + kwargs: + generate_device: "cuda:1" + prefill_device: "cuda:1" + +- match: + name: "^model\\.layers\\.(0|[1-9]|[12][0-9])\\.mlp\\.gate$" + class: ktransformers.models.modeling_deepseek_v3.MoEGate + replace: + class: ktransformers.operators.gate.KMoEGate + kwargs: + generate_device: "cuda:0" + prefill_device: "cuda:0" +- match: + name: "^model\\.layers\\.([3456][0-9])\\.mlp\\.gate$" + class: ktransformers.models.modeling_deepseek_v3.MoEGate + replace: + class: ktransformers.operators.gate.KMoEGate # mlp module with custom forward function + kwargs: + generate_device: "cuda:1" + prefill_device: "cuda:1" + +- match: + name: "^model\\.layers\\.(0|[1-9]|[12][0-9])\\.mlp\\.experts$" + replace: + class: ktransformers.operators.experts.KTransformersExperts # custom MoE Kernel with expert paralleism + kwargs: + prefill_device: "cuda:0" + prefill_op: "KExpertsTorch" + generate_device: "cpu" + generate_op: "KSFTExpertsCPU" + out_device: "cuda:0" + backend: "AMXInt8" # or "AMXBF16" or "llamafile" (default) + recursive: False # don't recursively inject submodules of this module + +- match: + name: "^model\\.layers\\.([3456][0-9])\\.mlp\\.experts$" + replace: + class: ktransformers.operators.experts.KTransformersExperts # custom MoE Kernel with expert paralleism + kwargs: + prefill_device: "cuda:1" + prefill_op: "KExpertsTorch" + generate_device: "cpu" + generate_op: "KSFTExpertsCPU" + out_device: "cuda:1" + backend: "AMXInt8" # or "AMXBF16" or "llamafile" (default) + recursive: False # don't recursively inject submodules of this module + +- match: + name: "^model\\.layers\\.(0|[1-9]|[12][0-9])\\.self_attn$" + replace: + class: ktransformers.operators.attention.KDeepseekV2Attention # optimized MLA implementation + kwargs: + generate_device: "cuda:0" + prefill_device: "cuda:0" +- match: + name: "^model\\.layers\\.([3456][0-9])\\.self_attn$" + replace: + class: ktransformers.operators.attention.KDeepseekV2Attention # optimized MLA implementation + kwargs: + generate_device: "cuda:1" + prefill_device: "cuda:1" +- match: + name: "^model$" + replace: + class: "ktransformers.operators.models.KDeepseekV2Model" + kwargs: + per_layer_prefill_intput_threshold: 0 # 0 is close layer wise prefill + transfer_map: + 30: "cuda:1" + +- match: + name: "^model\\.layers\\.(0|[1-9]|[12][0-9])\\." + replace: + class: "default" + kwargs: + generate_device: "cuda:0" + prefill_device: "cuda:0" + +- match: + name: "^lm_head" + class: torch.nn.Linear + replace: + class: ktransformers.operators.linear.KTransformersLinear + kwargs: + generate_device: "cuda:1" + prefill_device: "cuda:1" + generate_op: "KLinearTorch" + prefill_op: "KLinearTorch" + +- match: + name: "(^model\\.layers\\.([3456][0-9])\\.)|(model.norm)" + replace: + class: "default" + kwargs: + generate_device: "cuda:1" + prefill_device: "cuda:1" diff --git a/archive/kt-sft/ktransformers/optimize/optimize_rules/DeepSeek-V3-Chat-sft-amx.yaml b/archive/kt-sft/ktransformers/optimize/optimize_rules/DeepSeek-V3-Chat-sft-amx.yaml new file mode 100644 index 0000000..177d4c4 --- /dev/null +++ b/archive/kt-sft/ktransformers/optimize/optimize_rules/DeepSeek-V3-Chat-sft-amx.yaml @@ -0,0 +1,77 @@ +- match: + class: ktransformers.models.modeling_deepseek_v3.DeepseekV3RotaryEmbedding + replace: + class: ktransformers.operators.RoPE.YarnRotaryEmbeddingV3 + kwargs: + generate_device: "cuda" + prefill_device: "cuda" + +- match: + name: "^lm_head$" # regular expression + class: torch.nn.Linear # only match modules matching name and class simultaneously + replace: + class: ktransformers.operators.linear.KTransformersLinear # optimized Kernel on quantized data types + kwargs: + generate_device: "cuda" + prefill_device: "cuda" + generate_op: "KLinearTorch" + prefill_op: "KLinearTorch" + +- match: + name: "^model\\.layers\\.(?!.*self_attn\\.kv_b_proj).*$" # regular expression + class: torch.nn.Linear # only match modules matching name and class simultaneously + replace: + class: ktransformers.operators.linear.KTransformersLinear # optimized Kernel on quantized data types + kwargs: + generate_device: "cuda" + prefill_device: "cuda" + generate_op: "KLinearTorch" + prefill_op: "KLinearTorch" +- match: + name: "^model\\.layers\\..*\\.mlp$" + class: ktransformers.models.modeling_deepseek_v3.DeepseekV3MoE + replace: + class: ktransformers.operators.experts.KDeepseekV3MoE # mlp module with custom forward function + kwargs: + generate_device: "cuda" + prefill_device: "cuda" +- match: + class: ktransformers.models.modeling_deepseek_v3.MoEGate + replace: + class: ktransformers.operators.gate.KMoEGate + kwargs: + generate_device: "cuda:0" + prefill_device: "cuda:0" +- match: + name: "^model\\.layers\\..*\\.mlp\\.experts$" + replace: + class: ktransformers.operators.experts.KTransformersExperts # custom MoE Kernel with expert paralleism + kwargs: + prefill_device: "cuda" + prefill_op: "KExpertsTorch" + generate_device: "cpu" + generate_op: "KSFTExpertsCPU" + out_device: "cuda" + backend: "AMXInt8" # or "AMXBF16" or "llamafile" (default) + recursive: False # don't recursively inject submodules of this module +- match: + name: "^model\\.layers\\..*\\.self_attn$" + replace: + class: ktransformers.operators.attention.KDeepseekV2Attention # optimized MLA implementation + kwargs: + generate_device: "cuda" + prefill_device: "cuda" + absorb_for_prefill: False # change this to True to enable long context(prefill may slower). +- match: + name: "^model$" + replace: + class: "ktransformers.operators.models.KDeepseekV2Model" + kwargs: + per_layer_prefill_intput_threshold: 0 # 0 is close layer wise prefill +- match: + name: "^model.embed_tokens" + replace: + class: "default" + kwargs: + generate_device: "cpu" + prefill_device: "cpu" \ No newline at end of file diff --git a/archive/kt-sft/ktransformers/optimize/optimize_rules/DeepSeek-V3-Chat.yaml b/archive/kt-sft/ktransformers/optimize/optimize_rules/DeepSeek-V3-Chat.yaml new file mode 100644 index 0000000..d28e016 --- /dev/null +++ b/archive/kt-sft/ktransformers/optimize/optimize_rules/DeepSeek-V3-Chat.yaml @@ -0,0 +1,76 @@ +- match: + class: ktransformers.models.modeling_deepseek_v3.DeepseekV3RotaryEmbedding + replace: + class: ktransformers.operators.RoPE.YarnRotaryEmbeddingV3 + kwargs: + generate_device: "cuda" + prefill_device: "cuda" + +- match: + name: "^lm_head$" # regular expression + class: torch.nn.Linear # only match modules matching name and class simultaneously + replace: + class: ktransformers.operators.linear.KTransformersLinear # optimized Kernel on quantized data types + kwargs: + generate_device: "cuda" + prefill_device: "cuda" + generate_op: "KLinearMarlin" + prefill_op: "KLinearTorch" + +- match: + name: "^model\\.layers\\.(?!.*self_attn\\.kv_b_proj).*$" # regular expression + class: torch.nn.Linear # only match modules matching name and class simultaneously + replace: + class: ktransformers.operators.linear.KTransformersLinear # optimized Kernel on quantized data types + kwargs: + generate_device: "cuda" + prefill_device: "cuda" + generate_op: "KLinearMarlin" + prefill_op: "KLinearTorch" +- match: + name: "^model\\.layers\\..*\\.mlp$" + class: ktransformers.models.modeling_deepseek_v3.DeepseekV3MoE + replace: + class: ktransformers.operators.experts.KDeepseekV3MoE # mlp module with custom forward function + kwargs: + generate_device: "cuda" + prefill_device: "cuda" +- match: + class: ktransformers.models.modeling_deepseek_v3.MoEGate + replace: + class: ktransformers.operators.gate.KMoEGate + kwargs: + generate_device: "cuda:0" + prefill_device: "cuda:0" +- match: + name: "^model\\.layers\\..*\\.mlp\\.experts$" + replace: + class: ktransformers.operators.experts.KTransformersExperts # custom MoE Kernel with expert paralleism + kwargs: + prefill_device: "cuda" + prefill_op: "KExpertsTorch" + generate_device: "cpu" + generate_op: "KExpertsCPU" + out_device: "cuda" + recursive: False # don't recursively inject submodules of this module +- match: + name: "^model\\.layers\\..*\\.self_attn$" + replace: + class: ktransformers.operators.attention.KDeepseekV2Attention # optimized MLA implementation + kwargs: + generate_device: "cuda" + prefill_device: "cuda" + absorb_for_prefill: False # change this to True to enable long context(prefill may slower). +- match: + name: "^model$" + replace: + class: "ktransformers.operators.models.KDeepseekV2Model" + kwargs: + per_layer_prefill_intput_threshold: 0 # 0 is close layer wise prefill +- match: + name: "^model.embed_tokens" + replace: + class: "default" + kwargs: + generate_device: "cpu" + prefill_device: "cpu" \ No newline at end of file diff --git a/archive/kt-sft/ktransformers/optimize/optimize_rules/Internlm2_5-7b-Chat-1m.yaml b/archive/kt-sft/ktransformers/optimize/optimize_rules/Internlm2_5-7b-Chat-1m.yaml new file mode 100644 index 0000000..51a8142 --- /dev/null +++ b/archive/kt-sft/ktransformers/optimize/optimize_rules/Internlm2_5-7b-Chat-1m.yaml @@ -0,0 +1,28 @@ +- match: + class: ktransformers.models.modeling_llama.LlamaRotaryEmbedding + replace: + class: ktransformers.operators.RoPE.RotaryEmbeddingV2 +- match: + name: "^model.embed_tokens" + replace: + class: "default" + kwargs: + generate_device: "cpu" + prefill_device: "cpu" +- match: + class: ktransformers.models.modeling_llama.LlamaModel + replace: + class: ktransformers.operators.models.KLlamaModel + kwargs: + generate_device: "cuda" + prefill_device: "cuda" + per_layer_prefill_intput_threshold: 0 # 0 is close layer wise prefill + +- match: + name: "^model\\.layers\\..*\\.self_attn$" + replace: + class: ktransformers.operators.attention.KLlamaAttention + kwargs: + generate_device: "cuda" + prefill_device: "cuda" + diff --git a/archive/kt-sft/ktransformers/optimize/optimize_rules/Mixtral.yaml b/archive/kt-sft/ktransformers/optimize/optimize_rules/Mixtral.yaml new file mode 100644 index 0000000..80a346a --- /dev/null +++ b/archive/kt-sft/ktransformers/optimize/optimize_rules/Mixtral.yaml @@ -0,0 +1,59 @@ +- match: + class: ktransformers.models.modeling_mixtral.MixtralRotaryEmbedding + replace: + class: ktransformers.operators.RoPE.RotaryEmbedding + kwargs: + generate_device: "cuda" + prefill_device: "cuda" +- match: + name: "^model\\.layers\\..*$" + class: torch.nn.Linear # only match modules matching name and class simultaneously + replace: + class: ktransformers.operators.linear.KTransformersLinear # optimized Kernel on quantized data types + kwargs: + generate_device: "cuda" + prefill_device: "cuda" + generate_op: "KLinearMarlin" + prefill_op: "KLinearTorch" +- match: + name: "^lm_head" + class: torch.nn.Linear + replace: + class: ktransformers.operators.linear.KTransformersLinear + kwargs: + generate_device: "cuda" + prefill_device: "cuda" + generate_op: "KLinearMarlin" + prefill_op: "KLinearTorch" +- match: + name: "^model\\.layers\\..*\\.block_sparse_moe$" + class: ktransformers.models.modeling_mixtral.MixtralSparseMoeBlock + replace: + class: ktransformers.operators.experts.KMistralSparseMoEBlock +- match: + name: "^model\\.layers\\..*\\.block_sparse_moe\\.experts$" + replace: + class: ktransformers.operators.experts.KTransformersExperts + kwargs: + prefill_device: "cuda" + prefill_op: "KExpertsTorch" + generate_device: "cpu" + generate_op: "KExpertsCPU" + out_device: "cuda" + recursive: False # don't recursively inject submodules of this module + +- match: + name: "^model.embed_tokens" + replace: + class: "default" + kwargs: + generate_device: "cpu" + prefill_device: "cpu" + +- match: + name: "^model\\.layers\\..*\\." + replace: + class: "default" + kwargs: + generate_device: "cuda" + prefill_device: "cuda" \ No newline at end of file diff --git a/archive/kt-sft/ktransformers/optimize/optimize_rules/Moonlight-16B-A3B-serve.yaml b/archive/kt-sft/ktransformers/optimize/optimize_rules/Moonlight-16B-A3B-serve.yaml new file mode 100644 index 0000000..bc52e0e --- /dev/null +++ b/archive/kt-sft/ktransformers/optimize/optimize_rules/Moonlight-16B-A3B-serve.yaml @@ -0,0 +1,94 @@ + + +- match: + name: "^lm_head$" # regular expression + class: torch.nn.Linear # only match modules matching name and class simultaneously + replace: + class: ktransformers.operators.linear.KTransformersLinear # optimized Kernel on quantized data types + kwargs: + generate_device: "cuda" + prefill_device: "cuda" + generate_op: "VLinearMarlin" + prefill_op: "KLinearTorch" + +- match: + name: "^model\\.layers\\.(?!.*self_attn\\.kv_b_proj).*$" # regular expression + class: torch.nn.Linear # only match modules matching name and class simultaneously + replace: + class: ktransformers.operators.linear.KTransformersLinear # optimized Kernel on quantized data types + kwargs: + generate_device: "cuda" + prefill_device: "cuda" + generate_op: "VLinearMarlin" + prefill_op: "KLinearTorch" +- match: + name: "^model\\.layers\\..*\\.mlp$" + class: ktransformers.models.modeling_deepseek_v3.DeepseekV3MoE + replace: + class: ktransformers.operators.experts.KDeepseekV3MoEV2 # mlp module with custom forward function + kwargs: + generate_device: "cuda" + prefill_device: "cuda" +- match: + class: ktransformers.models.modeling_deepseek_v3.MoEGate + replace: + class: ktransformers.operators.gate.KMoEGate + kwargs: + generate_device: "cuda:0" + prefill_device: "cuda:0" +- match: + name: "^model\\.layers\\..*\\.mlp\\.experts$" + replace: + class: ktransformers.operators.experts.KTransformersExpertsV2 # custom MoE Kernel with expert paralleism + kwargs: + prefill_device: "cuda" + prefill_op: "KExpertsTorch" + generate_device: "cpu" + generate_op: "KExpertsCPU" + out_device: "cuda" + recursive: False # don't recursively inject submodules of this module +- match: + name: "^model\\.layers\\..*\\.self_attn$" + replace: + class: ktransformers.operators.balance_serve_attention.flashinfer_attn # optimized MLA implementation + kwargs: + generate_device: "cuda" + prefill_device: "cuda" + absorb_for_prefill: False # change this to True to enable long context(prefill may slower). +- match: + name: "^model$" + replace: + class: "ktransformers.operators.models.KDeepseekV2Model" + kwargs: + per_layer_prefill_intput_threshold: 0 # 0 is close layer wise prefill +- match: + name: "^model.embed_tokens" + replace: + class: "default" + kwargs: + generate_device: "cpu" + prefill_device: "cpu" + +- match: + class: ktransformers.models.modeling_deepseek_v3.DeepseekV3RMSNorm + replace: + class: ktransformers.operators.layernorm.RMSNorm + kwargs: + generate_device: "cuda" + prefill_device: "cuda" + +- match: + class: ktransformers.models.modeling_deepseek_v3.DeepseekV3MLP + replace: + class: ktransformers.operators.mlp.kDeepseekV3MLP + kwargs: + generate_device: "cuda" + prefill_device: "cuda" + +- match: + class: ktransformers.models.modeling_deepseek_v3.DeepseekV3RotaryEmbedding + replace: + class: ktransformers.operators.RoPE.RotaryEmbeddingV4 + kwargs: + generate_device: "cuda" + prefill_device: "cuda" \ No newline at end of file diff --git a/archive/kt-sft/ktransformers/optimize/optimize_rules/Moonlight-16B-A3B.yaml b/archive/kt-sft/ktransformers/optimize/optimize_rules/Moonlight-16B-A3B.yaml new file mode 100644 index 0000000..6cea246 --- /dev/null +++ b/archive/kt-sft/ktransformers/optimize/optimize_rules/Moonlight-16B-A3B.yaml @@ -0,0 +1,86 @@ +- match: + class: ktransformers.models.modeling_deepseek_v3.DeepseekV3RotaryEmbedding + replace: + class: ktransformers.operators.RoPE.RotaryEmbeddingV3 + kwargs: + generate_device: "cuda" + prefill_device: "cuda" + +- match: + name: "^lm_head$" # regular expression + class: torch.nn.Linear # only match modules matching name and class simultaneously + replace: + class: ktransformers.operators.linear.KTransformersLinear # optimized Kernel on quantized data types + kwargs: + generate_device: "cuda" + prefill_device: "cuda" + generate_op: "KLinearMarlin" + prefill_op: "KLinearTorch" + +- match: + name: "^model\\.layers\\.(?!.*self_attn\\.kv_b_proj).*$" # regular expression + class: torch.nn.Linear # only match modules matching name and class simultaneously + replace: + class: ktransformers.operators.linear.KTransformersLinear # optimized Kernel on quantized data types + kwargs: + generate_device: "cuda" + prefill_device: "cuda" + generate_op: "KLinearMarlin" + prefill_op: "KLinearTorch" +- match: + name: "^model\\.layers\\..*\\.mlp$" + class: ktransformers.models.modeling_deepseek_v3.DeepseekV3MoE + replace: + class: ktransformers.operators.experts.KDeepseekV3MoE # mlp module with custom forward function + kwargs: + generate_device: "cuda" + prefill_device: "cuda" +- match: + class: ktransformers.models.modeling_deepseek_v3.MoEGate + replace: + class: ktransformers.operators.gate.KMoEGate + kwargs: + generate_device: "cuda:0" + prefill_device: "cuda:0" +- match: + name: "^model\\.layers\\..*\\.mlp\\.experts$" + replace: + class: ktransformers.operators.experts.KTransformersExperts # custom MoE Kernel with expert paralleism + kwargs: + prefill_device: "cuda" + prefill_op: "KExpertsTorch" + generate_device: "cpu" + generate_op: "KExpertsCPU" + out_device: "cuda" + recursive: False # don't recursively inject submodules of this module +# if want to use more VRAM, use experts Marlin and disable CUDA Graph(disable CUDA Graph may cause low performance) +#- match: +# name: "^model\\.layers\\..*\\.mlp\\.experts$" +# replace: +# class: ktransformers.operators.experts.KTransformersExperts # custom MoE Kernel with expert paralleism +# kwargs: +# prefill_device: "cuda" +# prefill_op: "KExpertsTorch" +# generate_device: "cuda" +# generate_op: "KExpertsMarlin" +# recursive: False # don't recursively inject submodules of this module +- match: + name: "^model\\.layers\\..*\\.self_attn$" + replace: + class: ktransformers.operators.attention.KDeepseekV2Attention # optimized MLA implementation + kwargs: + generate_device: "cuda" + prefill_device: "cuda" +- match: + name: "^model$" + replace: + class: "ktransformers.operators.models.KDeepseekV2Model" + kwargs: + per_layer_prefill_intput_threshold: 0 # 0 is close layer wise prefill +- match: + name: "^model.embed_tokens" + replace: + class: "default" + kwargs: + generate_device: "cpu" + prefill_device: "cpu" \ No newline at end of file diff --git a/archive/kt-sft/ktransformers/optimize/optimize_rules/Qwen2-57B-A14B-Instruct-multi-gpu.yaml b/archive/kt-sft/ktransformers/optimize/optimize_rules/Qwen2-57B-A14B-Instruct-multi-gpu.yaml new file mode 100644 index 0000000..da01c82 --- /dev/null +++ b/archive/kt-sft/ktransformers/optimize/optimize_rules/Qwen2-57B-A14B-Instruct-multi-gpu.yaml @@ -0,0 +1,122 @@ +- match: + name: "^model\\.layers\\.([012])\\." + class: ktransformers.models.modeling_qwen2_moe.Qwen2MoeRotaryEmbedding + replace: + class: ktransformers.operators.RoPE.RotaryEmbedding + kwargs: + generate_device: "cuda:0" + prefill_device: "cuda:0" +- match: + name: "^model\\.layers\\.([012])$" # regular expression + class: torch.nn.Linear # only match modules matching name and class simultaneously + replace: + class: ktransformers.operators.linear.KTransformersLinear # optimized Kernel on quantized data types + kwargs: + generate_device: "cuda:0" + prefill_device: "cuda:0" + generate_op: "KLinearMarlin" + prefill_op: "KLinearTorch" +- match: + name: "^model\\.layers\\.([012])\\.mlp$" + class: ktransformers.models.modeling_qwen2_moe.Qwen2MoeSparseMoeBlock + replace: + class: ktransformers.operators.experts.KQwen2MoeSparseMoeBlock # mlp module with custom forward function +- match: + name: "^model\\.layers\\.([012])\\.mlp\\.experts$" + replace: + class: ktransformers.operators.experts.KTransformersExperts # custom MoE Kernel with expert paralleism + # device: "cpu" # which devices to load this module when initializing + kwargs: + prefill_device: "cuda:0" + prefill_op: "KExpertsTorch" + generate_device: "cpu" + generate_op: "KExpertsCPU" + out_device: "cuda:0" + recursive: False # don't recursively inject submodules of this module + +- match: + name: "^model\\.layers\\.([12][0-9]|[3-9])\\." + class: ktransformers.models.modeling_qwen2_moe.Qwen2MoeRotaryEmbedding + replace: + class: ktransformers.operators.RoPE.RotaryEmbedding + kwargs: + generate_device: "cuda:1" + prefill_device: "cuda:1" +- match: + name: "^model\\.layers\\.([12][0-9]|[3-9])$" # regular expression + class: torch.nn.Linear # only match modules matching name and class simultaneously + replace: + class: ktransformers.operators.linear.KTransformersLinear # optimized Kernel on quantized data types + kwargs: + generate_device: "cuda:1" + prefill_device: "cuda:1" + generate_op: "KLinearMarlin" + prefill_op: "KLinearTorch" +- match: + name: "^model\\.layers\\.([12][0-9]|[3-9])\\.mlp$" + class: ktransformers.models.modeling_qwen2_moe.Qwen2MoeSparseMoeBlock + replace: + class: ktransformers.operators.experts.KQwen2MoeSparseMoeBlock # mlp module with custom forward function +- match: + name: "^model\\.layers\\.([12][0-9]|[3-9])\\.mlp\\.experts$" + replace: + class: ktransformers.operators.experts.KTransformersExperts # custom MoE Kernel with expert paralleism + # device: "cpu" # which devices to load this module when initializing + kwargs: + prefill_device: "cuda:1" + prefill_op: "KExpertsTorch" + generate_device: "cpu" + generate_op: "KExpertsCPU" + out_device: "cuda:1" + recursive: False # don't recursively inject submodules of this module + +- match: + name: "^model.embed_tokens" + replace: + class: "default" + kwargs: + generate_device: "cpu" + prefill_device: "cpu" +- match: + name: "^lm_head" + class: torch.nn.Linear + replace: + class: ktransformers.operators.linear.KTransformersLinear + kwargs: + generate_device: "cuda:1" + prefill_device: "cuda:1" + generate_op: "KLinearMarlin" + prefill_op: "KLinearTorch" + +- match: + name: "(^model.norm)" + replace: + class: "default" + kwargs: + generate_device: "cuda:1" + prefill_device: "cuda:1" + +- match: + name: "^model$" + replace: + class: "ktransformers.operators.models.KQwen2MoeModel" + kwargs: + per_layer_prefill_intput_threshold: 0 # 0 is close layer wise prefill + transfer_map: + 3: "cuda:1" + +- match: + name: "^model\\.layers\\.([012])\\." + replace: + class: "default" + kwargs: + generate_device: "cuda:0" + prefill_device: "cuda:0" + +- match: + name: "^model\\.layers\\.([12][0-9]|[3-9])\\." + replace: + class: "default" + kwargs: + generate_device: "cuda:1" + prefill_device: "cuda:1" diff --git a/archive/kt-sft/ktransformers/optimize/optimize_rules/Qwen2-57B-A14B-Instruct.yaml b/archive/kt-sft/ktransformers/optimize/optimize_rules/Qwen2-57B-A14B-Instruct.yaml new file mode 100644 index 0000000..38e9e73 --- /dev/null +++ b/archive/kt-sft/ktransformers/optimize/optimize_rules/Qwen2-57B-A14B-Instruct.yaml @@ -0,0 +1,67 @@ +- match: + class: ktransformers.models.modeling_qwen2_moe.Qwen2MoeRotaryEmbedding + replace: + class: ktransformers.operators.RoPE.RotaryEmbedding + kwargs: + generate_device: "cuda" + prefill_device: "cuda" +- match: + name: "^model\\.layers\\..*$" # regular expression + class: torch.nn.Linear # only match modules matching name and class simultaneously + replace: + class: ktransformers.operators.linear.KTransformersLinear # optimized Kernel on quantized data types + kwargs: + generate_device: "cuda" + prefill_device: "cuda" + generate_op: "KLinearMarlin" + prefill_op: "KLinearTorch" +- match: + name: "^lm_head" + class: torch.nn.Linear + replace: + class: ktransformers.operators.linear.KTransformersLinear + kwargs: + generate_device: "cuda" + prefill_device: "cuda" + generate_op: "KLinearMarlin" + prefill_op: "KLinearTorch" +- match: + name: "^model\\.layers\\..*\\.mlp$" + class: ktransformers.models.modeling_qwen2_moe.Qwen2MoeSparseMoeBlock + replace: + class: ktransformers.operators.experts.KQwen2MoeSparseMoeBlock # mlp module with custom forward function + kwargs: + generate_device: "cuda" + prefill_device: "cuda" +- match: + name: "^model\\.layers\\..*\\.mlp\\.experts$" + replace: + class: ktransformers.operators.experts.KTransformersExperts # custom MoE Kernel with expert paralleism + # device: "cpu" # which devices to load this module when initializing + kwargs: + prefill_device: "cuda" + prefill_op: "KExpertsTorch" + generate_device: "cpu" + generate_op: "KExpertsCPU" + out_device: "cuda" + recursive: False # don't recursively inject submodules of this module +- match: + name: "^model$" + replace: + class: "ktransformers.operators.models.KQwen2MoeModel" + kwargs: + per_layer_prefill_intput_threshold: 0 # 0 is close layer wise prefill +- match: + name: "^model.embed_tokens" + replace: + class: "default" + kwargs: + generate_device: "cpu" + prefill_device: "cpu" +- match: + name: "^model\\.layers\\..*\\." + replace: + class: "default" + kwargs: + generate_device: "cuda" + prefill_device: "cuda" diff --git a/archive/kt-sft/ktransformers/optimize/optimize_rules/Qwen2-serve-amx.yaml b/archive/kt-sft/ktransformers/optimize/optimize_rules/Qwen2-serve-amx.yaml new file mode 100644 index 0000000..27dba2b --- /dev/null +++ b/archive/kt-sft/ktransformers/optimize/optimize_rules/Qwen2-serve-amx.yaml @@ -0,0 +1,96 @@ +- match: + class: ktransformers.models.modeling_qwen2_moe.Qwen2MoeRotaryEmbedding + replace: + class: ktransformers.operators.RoPE.RotaryEmbedding + kwargs: + generate_device: "cuda" + prefill_device: "cuda" + +- match: + name: "^lm_head$" # regular expression + class: torch.nn.Linear # only match modules matching name and class simultaneously + replace: + class: ktransformers.operators.linear.KTransformersLinear # optimized Kernel on quantized data types + kwargs: + generate_device: "cuda" + prefill_device: "cuda" + generate_op: "KLinearMarlin" + prefill_op: "KLinearTorch" + +# - match: +# name: "^model\\.layers\\..*$" # regular expression +# class: torch.nn.Linear # only match modules matching name and class simultaneously +# replace: +# class: ktransformers.operators.linear.KTransformersLinear # optimized Kernel on quantized data types +# kwargs: +# generate_device: "cuda" +# prefill_device: "cuda" +# generate_op: "VLinearMarlin" +# prefill_op: "KLinearTorch" +- match: + name: "^model\\.layers\\.(?!.*mlp\\.shared_expert_gate).*$" # regular expression + class: torch.nn.Linear # only match modules matching name and class simultaneously + replace: + class: ktransformers.operators.linear.KTransformersLinear # optimized Kernel on quantized data types + kwargs: + generate_device: "cuda" + prefill_device: "cuda" + generate_op: "VLinearMarlin" + prefill_op: "KLinearTorch" +- match: + name: "^model\\.layers\\..*\\.mlp$" + class: ktransformers.models.modeling_qwen2_moe.Qwen2MoeSparseMoeBlock + replace: + class: ktransformers.operators.experts.KQwen2MoeSparseMoeBlockV2 # mlp module with custom forward function + kwargs: + generate_device: "cuda" + prefill_device: "cuda" + +- match: + name: "^model\\.layers\\..*\\.mlp\\.experts$" + replace: + class: ktransformers.operators.experts.KTransformersExpertsV2 # custom MoE Kernel with expert paralleism + kwargs: + prefill_device: "cuda" + prefill_op: "KExpertsTorch" + generate_device: "cpu" + generate_op: "KExpertsCPU" + out_device: "cuda" + backend: "AMXInt8" # or "AMXBF16" or "llamafile" (default) + recursive: False # don't recursively inject submodules of this module +- match: + name: "^model\\.layers\\..*\\.self_attn$" + replace: + class: ktransformers.operators.balance_serve_attention.KQwen2MoeAttention # optimized MLA implementation + kwargs: + generate_device: "cuda" + prefill_device: "cuda" +- match: + name: "^model$" + replace: + class: "ktransformers.operators.models.KQwen2MoeModel" + kwargs: + per_layer_prefill_intput_threshold: 0 # 0 is close layer wise prefill +- match: + name: "^model.embed_tokens" + replace: + class: "default" + kwargs: + generate_device: "cpu" + prefill_device: "cpu" + +- match: + class: ktransformers.models.modeling_qwen2_moe.Qwen2MoeRMSNorm + replace: + class: ktransformers.operators.layernorm.KQwen2MoeRMSNorm + kwargs: + generate_device: "cuda" + prefill_device: "cuda" + +- match: + class: ktransformers.models.modeling_qwen2_moe.Qwen2MoeMLP + replace: + class: ktransformers.operators.mlp.KQwen2MoeMLP + kwargs: + generate_device: "cuda" + prefill_device: "cuda" \ No newline at end of file diff --git a/archive/kt-sft/ktransformers/optimize/optimize_rules/Qwen2-serve.yaml b/archive/kt-sft/ktransformers/optimize/optimize_rules/Qwen2-serve.yaml new file mode 100644 index 0000000..41b41a7 --- /dev/null +++ b/archive/kt-sft/ktransformers/optimize/optimize_rules/Qwen2-serve.yaml @@ -0,0 +1,95 @@ +- match: + class: ktransformers.models.modeling_qwen2_moe.Qwen2MoeRotaryEmbedding + replace: + class: ktransformers.operators.RoPE.RotaryEmbedding + kwargs: + generate_device: "cuda" + prefill_device: "cuda" + +- match: + name: "^lm_head$" # regular expression + class: torch.nn.Linear # only match modules matching name and class simultaneously + replace: + class: ktransformers.operators.linear.KTransformersLinear # optimized Kernel on quantized data types + kwargs: + generate_device: "cuda" + prefill_device: "cuda" + generate_op: "KLinearMarlin" + prefill_op: "KLinearTorch" + +# - match: +# name: "^model\\.layers\\..*$" # regular expression +# class: torch.nn.Linear # only match modules matching name and class simultaneously +# replace: +# class: ktransformers.operators.linear.KTransformersLinear # optimized Kernel on quantized data types +# kwargs: +# generate_device: "cuda" +# prefill_device: "cuda" +# generate_op: "VLinearMarlin" +# prefill_op: "KLinearTorch" +- match: + name: "^model\\.layers\\.(?!.*mlp\\.shared_expert_gate).*$" # regular expression + class: torch.nn.Linear # only match modules matching name and class simultaneously + replace: + class: ktransformers.operators.linear.KTransformersLinear # optimized Kernel on quantized data types + kwargs: + generate_device: "cuda" + prefill_device: "cuda" + generate_op: "VLinearMarlin" + prefill_op: "KLinearTorch" +- match: + name: "^model\\.layers\\..*\\.mlp$" + class: ktransformers.models.modeling_qwen2_moe.Qwen2MoeSparseMoeBlock + replace: + class: ktransformers.operators.experts.KQwen2MoeSparseMoeBlockV2 # mlp module with custom forward function + kwargs: + generate_device: "cuda" + prefill_device: "cuda" + +- match: + name: "^model\\.layers\\..*\\.mlp\\.experts$" + replace: + class: ktransformers.operators.experts.KTransformersExpertsV2 # custom MoE Kernel with expert paralleism + kwargs: + prefill_device: "cuda" + prefill_op: "KExpertsTorch" + generate_device: "cpu" + generate_op: "KExpertsCPU" + out_device: "cuda" + recursive: False # don't recursively inject submodules of this module +- match: + name: "^model\\.layers\\..*\\.self_attn$" + replace: + class: ktransformers.operators.balance_serve_attention.KQwen2MoeAttention # optimized MLA implementation + kwargs: + generate_device: "cuda" + prefill_device: "cuda" +- match: + name: "^model$" + replace: + class: "ktransformers.operators.models.KQwen2MoeModel" + kwargs: + per_layer_prefill_intput_threshold: 0 # 0 is close layer wise prefill +- match: + name: "^model.embed_tokens" + replace: + class: "default" + kwargs: + generate_device: "cpu" + prefill_device: "cpu" + +- match: + class: ktransformers.models.modeling_qwen2_moe.Qwen2MoeRMSNorm + replace: + class: ktransformers.operators.layernorm.KQwen2MoeRMSNorm + kwargs: + generate_device: "cuda" + prefill_device: "cuda" + +- match: + class: ktransformers.models.modeling_qwen2_moe.Qwen2MoeMLP + replace: + class: ktransformers.operators.mlp.KQwen2MoeMLP + kwargs: + generate_device: "cuda" + prefill_device: "cuda" \ No newline at end of file diff --git a/archive/kt-sft/ktransformers/optimize/optimize_rules/Qwen3Moe-serve-amx.yaml b/archive/kt-sft/ktransformers/optimize/optimize_rules/Qwen3Moe-serve-amx.yaml new file mode 100644 index 0000000..8607ca0 --- /dev/null +++ b/archive/kt-sft/ktransformers/optimize/optimize_rules/Qwen3Moe-serve-amx.yaml @@ -0,0 +1,96 @@ +- match: + class: ktransformers.models.modeling_qwen2_moe.Qwen2MoeRotaryEmbedding + replace: + class: ktransformers.operators.RoPE.RotaryEmbedding + kwargs: + generate_device: "cuda" + prefill_device: "cuda" + +- match: + name: "^lm_head$" # regular expression + class: torch.nn.Linear # only match modules matching name and class simultaneously + replace: + class: ktransformers.operators.linear.KTransformersLinear # optimized Kernel on quantized data types + kwargs: + generate_device: "cuda" + prefill_device: "cuda" + generate_op: "VLinearMarlin" + prefill_op: "KLinearTorch" + +# - match: +# name: "^model\\.layers\\..*$" # regular expression +# class: torch.nn.Linear # only match modules matching name and class simultaneously +# replace: +# class: ktransformers.operators.linear.KTransformersLinear # optimized Kernel on quantized data types +# kwargs: +# generate_device: "cuda" +# prefill_device: "cuda" +# generate_op: "VLinearMarlin" +# prefill_op: "KLinearTorch" +- match: + name: "^model\\.layers\\.(?!.*mlp\\.shared_expert_gate).*$" # regular expression + class: torch.nn.Linear # only match modules matching name and class simultaneously + replace: + class: ktransformers.operators.linear.KTransformersLinear # optimized Kernel on quantized data types + kwargs: + generate_device: "cuda" + prefill_device: "cuda" + generate_op: "KLinearMarlin" + prefill_op: "KLinearTorch" +- match: + name: "^model\\.layers\\..*\\.mlp$" + class: ktransformers.models.modeling_qwen3_moe.Qwen3MoeSparseMoeBlock + replace: + class: ktransformers.operators.experts.KQwen3MoeSparseMoeBlockV2 # mlp module with custom forward function + kwargs: + generate_device: "cuda" + prefill_device: "cuda" + +- match: + name: "^model\\.layers\\..*\\.mlp\\.experts$" + replace: + class: ktransformers.operators.experts.KTransformersExpertsV2 # custom MoE Kernel with expert paralleism + kwargs: + prefill_device: "cuda" + prefill_op: "KExpertsTorch" + generate_device: "cpu" + generate_op: "KExpertsCPU" + out_device: "cuda" + backend: "AMXBF16" # or "AMXBF16" or "llamafile" (default) + recursive: False # don't recursively inject submodules of this module +- match: + name: "^model\\.layers\\..*\\.self_attn$" + replace: + class: ktransformers.operators.balance_serve_attention.KQwen3MoeAttention # optimized MLA implementation + kwargs: + generate_device: "cuda" + prefill_device: "cuda" +- match: + name: "^model$" + replace: + class: "ktransformers.operators.models.KQwen2MoeModel" + kwargs: + per_layer_prefill_intput_threshold: 0 # 0 is close layer wise prefill +- match: + name: "^model.embed_tokens" + replace: + class: "default" + kwargs: + generate_device: "cpu" + prefill_device: "cpu" + +- match: + class: ktransformers.models.modeling_qwen3_moe.Qwen3MoeRMSNorm + replace: + class: ktransformers.operators.layernorm.KQwen3MoeRMSNorm + kwargs: + generate_device: "cuda" + prefill_device: "cuda" + +- match: + class: ktransformers.models.modeling_qwen3_moe.Qwen3MoeMLP + replace: + class: ktransformers.operators.mlp.KQwen2MoeMLP + kwargs: + generate_device: "cuda" + prefill_device: "cuda" \ No newline at end of file diff --git a/archive/kt-sft/ktransformers/optimize/optimize_rules/Qwen3Moe-serve.yaml b/archive/kt-sft/ktransformers/optimize/optimize_rules/Qwen3Moe-serve.yaml new file mode 100644 index 0000000..63f67da --- /dev/null +++ b/archive/kt-sft/ktransformers/optimize/optimize_rules/Qwen3Moe-serve.yaml @@ -0,0 +1,95 @@ +- match: + class: ktransformers.models.modeling_qwen2_moe.Qwen2MoeRotaryEmbedding + replace: + class: ktransformers.operators.RoPE.RotaryEmbedding + kwargs: + generate_device: "cuda" + prefill_device: "cuda" + +- match: + name: "^lm_head$" # regular expression + class: torch.nn.Linear # only match modules matching name and class simultaneously + replace: + class: ktransformers.operators.linear.KTransformersLinear # optimized Kernel on quantized data types + kwargs: + generate_device: "cuda" + prefill_device: "cuda" + generate_op: "VLinearMarlin" + prefill_op: "KLinearTorch" + +# - match: +# name: "^model\\.layers\\..*$" # regular expression +# class: torch.nn.Linear # only match modules matching name and class simultaneously +# replace: +# class: ktransformers.operators.linear.KTransformersLinear # optimized Kernel on quantized data types +# kwargs: +# generate_device: "cuda" +# prefill_device: "cuda" +# generate_op: "VLinearMarlin" +# prefill_op: "KLinearTorch" +- match: + name: "^model\\.layers\\.(?!.*mlp\\.shared_expert_gate).*$" # regular expression + class: torch.nn.Linear # only match modules matching name and class simultaneously + replace: + class: ktransformers.operators.linear.KTransformersLinear # optimized Kernel on quantized data types + kwargs: + generate_device: "cuda" + prefill_device: "cuda" + generate_op: "KLinearMarlin" + prefill_op: "KLinearTorch" +- match: + name: "^model\\.layers\\..*\\.mlp$" + class: ktransformers.models.modeling_qwen3_moe.Qwen3MoeSparseMoeBlock + replace: + class: ktransformers.operators.experts.KQwen3MoeSparseMoeBlockV2 # mlp module with custom forward function + kwargs: + generate_device: "cuda" + prefill_device: "cuda" + +- match: + name: "^model\\.layers\\..*\\.mlp\\.experts$" + replace: + class: ktransformers.operators.experts.KTransformersExpertsV2 # custom MoE Kernel with expert paralleism + kwargs: + prefill_device: "cuda" + prefill_op: "KExpertsTorch" + generate_device: "cpu" + generate_op: "KExpertsCPU" + out_device: "cuda" + recursive: False # don't recursively inject submodules of this module +- match: + name: "^model\\.layers\\..*\\.self_attn$" + replace: + class: ktransformers.operators.balance_serve_attention.KQwen3MoeAttention # optimized MLA implementation + kwargs: + generate_device: "cuda" + prefill_device: "cuda" +- match: + name: "^model$" + replace: + class: "ktransformers.operators.models.KQwen2MoeModel" + kwargs: + per_layer_prefill_intput_threshold: 0 # 0 is close layer wise prefill +- match: + name: "^model.embed_tokens" + replace: + class: "default" + kwargs: + generate_device: "cpu" + prefill_device: "cpu" + +- match: + class: ktransformers.models.modeling_qwen3_moe.Qwen3MoeRMSNorm + replace: + class: ktransformers.operators.layernorm.KQwen3MoeRMSNorm + kwargs: + generate_device: "cuda" + prefill_device: "cuda" + +- match: + class: ktransformers.models.modeling_qwen3_moe.Qwen3MoeMLP + replace: + class: ktransformers.operators.mlp.KQwen2MoeMLP + kwargs: + generate_device: "cuda" + prefill_device: "cuda" \ No newline at end of file diff --git a/archive/kt-sft/ktransformers/optimize/optimize_rules/Qwen3Moe-sft-amx.yaml b/archive/kt-sft/ktransformers/optimize/optimize_rules/Qwen3Moe-sft-amx.yaml new file mode 100644 index 0000000..b8eceb2 --- /dev/null +++ b/archive/kt-sft/ktransformers/optimize/optimize_rules/Qwen3Moe-sft-amx.yaml @@ -0,0 +1,80 @@ +- match: + class: ktransformers.models.modeling_qwen2_moe.Qwen2MoeRotaryEmbedding + replace: + class: ktransformers.operators.RoPE.RotaryEmbedding + kwargs: + generate_device: "cuda" + prefill_device: "cuda" + +- match: + name: "^lm_head$" # regular expression + class: torch.nn.Linear # only match modules matching name and class simultaneously + replace: + class: ktransformers.operators.linear.KTransformersLinear # optimized Kernel on quantized data types + kwargs: + generate_device: "cuda" + prefill_device: "cuda" + generate_op: "KLinearTorch" + prefill_op: "KLinearTorch" + +# - match: +# name: "^model\\.layers\\..*$" # regular expression +# class: torch.nn.Linear # only match modules matching name and class simultaneously +# replace: +# class: ktransformers.operators.linear.KTransformersLinear # optimized Kernel on quantized data types +# kwargs: +# generate_device: "cuda" +# prefill_device: "cuda" +# generate_op: "KLinearTorch" +# prefill_op: "KLinearTorch" +- match: + name: "^model\\.layers\\.(?!.*mlp\\.shared_expert_gate).*$" # regular expression + class: torch.nn.Linear # only match modules matching name and class simultaneously + replace: + class: ktransformers.operators.linear.KTransformersLinear # optimized Kernel on quantized data types + kwargs: + generate_device: "cuda" + prefill_device: "cuda" + generate_op: "KLinearTorch" + prefill_op: "KLinearTorch" +- match: + name: "^model\\.layers\\..*\\.mlp$" + replace: + class: ktransformers.operators.experts.KQwen3MoeSparseMoeBlock # mlp module with custom forward function + kwargs: + generate_device: "cuda" + prefill_device: "cuda" + +- match: + name: "^model\\.layers\\..*\\.mlp\\.experts$" + replace: + class: ktransformers.operators.experts.KTransformersExperts # custom MoE Kernel with expert paralleism + kwargs: + prefill_device: "cuda" + prefill_op: "KExpertsTorch" + generate_device: "cpu" + generate_op: "KSFTExpertsCPU" + out_device: "cuda" + backend: "AMXInt8" # or "AMXBF16" or "AMXInt8" + recursive: False # don't recursively inject submodules of this module +- match: + name: "^model\\.layers\\..*\\.self_attn$" + replace: + class: ktransformers.operators.attention.KQwen3MoeAttention # optimized MLA implementation + kwargs: + generate_device: "cuda" + prefill_device: "cuda" +- match: + name: "^model.embed_tokens" + replace: + class: "default" + kwargs: + generate_device: "cpu" + prefill_device: "cpu" + +- match: + name: "^model$" + replace: + class: "ktransformers.operators.models.KQwen3MoeModel" + kwargs: + per_layer_prefill_intput_threshold: 0 \ No newline at end of file diff --git a/archive/kt-sft/ktransformers/optimize/optimize_rules/rocm/DeepSeek-V3-Chat.yaml b/archive/kt-sft/ktransformers/optimize/optimize_rules/rocm/DeepSeek-V3-Chat.yaml new file mode 100644 index 0000000..ffbff2b --- /dev/null +++ b/archive/kt-sft/ktransformers/optimize/optimize_rules/rocm/DeepSeek-V3-Chat.yaml @@ -0,0 +1,76 @@ +- match: + class: ktransformers.models.modeling_deepseek_v3.DeepseekV3RotaryEmbedding + replace: + class: ktransformers.operators.RoPE.YarnRotaryEmbeddingV3 + kwargs: + generate_device: "cuda" + prefill_device: "cuda" + +- match: + name: "^lm_head$" # regular expression + class: torch.nn.Linear # only match modules matching name and class simultaneously + replace: + class: ktransformers.operators.linear.KTransformersLinear # optimized Kernel on quantized data types + kwargs: + generate_device: "cpu" + prefill_device: "cuda" + generate_op: "KLinearCPUInfer" + prefill_op: "KLinearTorch" + +- match: + name: "^model\\.layers\\.(?!.*self_attn\\.kv_b_proj).*$" # regular expression + class: torch.nn.Linear # only match modules matching name and class simultaneously + replace: + class: ktransformers.operators.linear.KTransformersLinear # optimized Kernel on quantized data types + kwargs: + generate_device: "cuda" + prefill_device: "cuda" + generate_op: "KLinearQ8" + prefill_op: "KLinearTorch" +- match: + name: "^model\\.layers\\..*\\.mlp$" + class: ktransformers.models.modeling_deepseek_v3.DeepseekV3MoE + replace: + class: ktransformers.operators.experts.KDeepseekV3MoE # mlp module with custom forward function + kwargs: + generate_device: "cuda" + prefill_device: "cuda" +- match: + class: ktransformers.models.modeling_deepseek_v3.MoEGate + replace: + class: ktransformers.operators.gate.KMoEGate + kwargs: + generate_device: "cuda:0" + prefill_device: "cuda:0" +- match: + name: "^model\\.layers\\..*\\.mlp\\.experts$" + replace: + class: ktransformers.operators.experts.KTransformersExperts # custom MoE Kernel with expert paralleism + kwargs: + prefill_device: "cuda" + prefill_op: "KExpertsTorch" + generate_device: "cpu" + generate_op: "KExpertsCPU" + out_device: "cuda" + recursive: False # don't recursively inject submodules of this module +- match: + name: "^model\\.layers\\..*\\.self_attn$" + replace: + class: ktransformers.operators.attention.KDeepseekV2Attention # optimized MLA implementation + kwargs: + generate_device: "cuda" + prefill_device: "cuda" + absorb_for_prefill: False # change this to True to enable long context(prefill may slower). +- match: + name: "^model$" + replace: + class: "ktransformers.operators.models.KDeepseekV2Model" + kwargs: + per_layer_prefill_intput_threshold: 0 # 0 is close layer wise prefill +- match: + name: "^model.embed_tokens" + replace: + class: "default" + kwargs: + generate_device: "cpu" + prefill_device: "cpu" \ No newline at end of file diff --git a/archive/kt-sft/ktransformers/optimize/optimize_rules/xpu/DeepSeek-V2-Chat.yaml b/archive/kt-sft/ktransformers/optimize/optimize_rules/xpu/DeepSeek-V2-Chat.yaml new file mode 100644 index 0000000..5de582f --- /dev/null +++ b/archive/kt-sft/ktransformers/optimize/optimize_rules/xpu/DeepSeek-V2-Chat.yaml @@ -0,0 +1,64 @@ +- match: + class: ktransformers.models.modeling_deepseek.DeepseekV2YarnRotaryEmbedding + replace: + class: ktransformers.operators.RoPE.YarnRotaryEmbedding + kwargs: + generate_device: "xpu" + prefill_device: "xpu" +- match: + name: "^model\\.layers\\..*" # regular expression + class: torch.nn.Linear # only match modules matching name and class simultaneously + replace: + class: ktransformers.operators.linear.KTransformersLinear # optimized Kernel on quantized data types + kwargs: + generate_device: "xpu" + prefill_device: "xpu" + generate_op: "KLinearIPEXLLM" + prefill_op: "KLinearIPEXLLM" +- match: + name: "^model\\.layers\\..*\\.mlp$" + class: ktransformers.models.modeling_deepseek.DeepseekV2MoE + replace: + class: ktransformers.operators.experts.KDeepseekV2MoE # mlp module with custom forward function + kwargs: + generate_device: "xpu" + prefill_device: "xpu" +- match: + name: "^model\\.layers\\..*\\.mlp\\.experts$" + replace: + class: ktransformers.operators.experts.KTransformersExperts # custom MoE Kernel with expert paralleism + kwargs: + prefill_device: "xpu" + prefill_op: "KExpertsTorch" + generate_device: "cpu" + generate_op: "KExpertsCPU" + out_device: "xpu" + recursive: False # don't recursively inject submodules of this module +- match: + class: ktransformers.models.modeling_deepseek.DeepseekV2RMSNorm + replace: + class: ktransformers.operators.layernorm.KDeepseekRMSNormIPEXLLM + kwargs: + generate_device: "xpu" + prefill_device: "xpu" +- match: + name: "^model\\.layers\\..*\\.self_attn$" + replace: + class: ktransformers.operators.attention.KDeepseekV2Attention # optimized MLA implementation + kwargs: + generate_device: "xpu" + prefill_device: "xpu" +- match: + name: "^model$" + replace: + class: "ktransformers.operators.models.KDeepseekV2Model" + kwargs: + per_layer_prefill_intput_threshold: 0 # 0 is close layer wise prefill + device: "xpu" +- match: + name: "^model.embed_tokens" + replace: + class: "default" + kwargs: + generate_device: "cpu" + prefill_device: "cpu" \ No newline at end of file diff --git a/archive/kt-sft/ktransformers/optimize/optimize_rules/xpu/DeepSeek-V3-Chat.yaml b/archive/kt-sft/ktransformers/optimize/optimize_rules/xpu/DeepSeek-V3-Chat.yaml new file mode 100644 index 0000000..c0e46c3 --- /dev/null +++ b/archive/kt-sft/ktransformers/optimize/optimize_rules/xpu/DeepSeek-V3-Chat.yaml @@ -0,0 +1,81 @@ +- match: + class: ktransformers.models.modeling_deepseek_v3.DeepseekV3RotaryEmbedding + replace: + class: ktransformers.operators.RoPE.YarnRotaryEmbeddingV3 + kwargs: + generate_device: "xpu" + prefill_device: "xpu" +- match: + name: "^lm_head$" # regular expression + class: torch.nn.Linear # only match modules matching name and class simultaneously + replace: + class: ktransformers.operators.linear.KTransformersLinear # optimized Kernel on quantized data types + kwargs: + generate_device: "xpu" + prefill_device: "xpu" + generate_op: "KLinearIPEXLLM" + prefill_op: "KLinearIPEXLLM" +- match: + name: "^model\\.layers\\..*" # regular expression + class: torch.nn.Linear # only match modules matching name and class simultaneously + replace: + class: ktransformers.operators.linear.KTransformersLinear # optimized Kernel on quantized data types + kwargs: + generate_device: "xpu" + prefill_device: "xpu" + generate_op: "KLinearIPEXLLM" + prefill_op: "KLinearIPEXLLM" +- match: + name: "^model\\.layers\\..*\\.mlp$" + class: ktransformers.models.modeling_deepseek_v3.DeepseekV3MoE + replace: + class: ktransformers.operators.experts.KDeepseekV3MoE # mlp module with custom forward function + kwargs: + generate_device: "xpu" + prefill_device: "xpu" +- match: + class: ktransformers.models.modeling_deepseek_v3.DeepseekV3RMSNorm + replace: + class: ktransformers.operators.layernorm.KDeepseekRMSNormIPEXLLM + kwargs: + generate_device: "xpu" + prefill_device: "xpu" +- match: + class: ktransformers.models.modeling_deepseek_v3.MoEGate + replace: + class: ktransformers.operators.gate.KMoEGateIPEXLLM + kwargs: + generate_device: "xpu:0" + prefill_device: "xpu:0" +- match: + name: "^model\\.layers\\..*\\.mlp\\.experts$" + replace: + class: ktransformers.operators.experts.KTransformersExperts # custom MoE Kernel with expert paralleism + kwargs: + prefill_device: "xpu" + prefill_op: "KExpertsTorch" + generate_device: "cpu" + generate_op: "KExpertsCPU" + out_device: "xpu" + recursive: False # don't recursively inject submodules of this module +- match: + name: "^model\\.layers\\..*\\.self_attn$" + replace: + class: ktransformers.operators.attention.KDeepseekV2Attention # optimized MLA implementation + kwargs: + generate_device: "xpu" + prefill_device: "xpu" + absorb_for_prefill: False # change this to True to enable long context(prefill may slower). +- match: + name: "^model$" + replace: + class: "ktransformers.operators.models.KDeepseekV2Model" + kwargs: + per_layer_prefill_intput_threshold: 0 # 0 is close layer wise prefill +- match: + name: "^model.embed_tokens" + replace: + class: "default" + kwargs: + generate_device: "cpu" + prefill_device: "cpu" \ No newline at end of file diff --git a/archive/kt-sft/ktransformers/optimize/optimize_rules/xpu/Qwen3Moe-Chat.yaml b/archive/kt-sft/ktransformers/optimize/optimize_rules/xpu/Qwen3Moe-Chat.yaml new file mode 100644 index 0000000..6bb4dae --- /dev/null +++ b/archive/kt-sft/ktransformers/optimize/optimize_rules/xpu/Qwen3Moe-Chat.yaml @@ -0,0 +1,80 @@ +- match: + name: "rotary_emb$" + replace: + class: ktransformers.operators.RoPE.KQwen3MoeRotaryEmbedding + kwargs: + generate_device: "xpu" + prefill_device: "xpu" +- match: + name: "^lm_head$" # regular expression + class: torch.nn.Linear # only match modules matching name and class simultaneously + replace: + class: ktransformers.operators.linear.KTransformersLinear # optimized Kernel on quantized data types + kwargs: + generate_device: "xpu" + prefill_device: "xpu" + generate_op: "KLinearIPEXLLM" + prefill_op: "KLinearIPEXLLM" +- match: + name: "^model\\.layers\\.(?!.*mlp\\.gate).*$" # regular expression + class: torch.nn.Linear # only match modules matching name and class simultaneously + replace: + class: ktransformers.operators.linear.KTransformersLinear # optimized Kernel on quantized data types + kwargs: + generate_device: "xpu" + prefill_device: "xpu" + generate_op: "KLinearIPEXLLM" + prefill_op: "KLinearIPEXLLM" +- match: + name: "^model\\.layers\\..*\\.mlp$" + class: transformers.models.qwen3_moe.modeling_qwen3_moe.Qwen3MoeSparseMoeBlock + replace: + class: ktransformers.operators.experts.KQwen3MoeSparseMoeBlockV2 # mlp module with custom forward function + kwargs: + generate_device: "xpu" + prefill_device: "xpu" +- match: + name: "^model\\.layers\\..*\\.mlp\\.experts$" + replace: + class: ktransformers.operators.experts.KTransformersExpertsV2 # custom MoE Kernel with expert paralleism + kwargs: + prefill_device: "xpu" + prefill_op: "KExpertsTorch" + generate_device: "cpu" + generate_op: "KExpertsCPU" + out_device: "xpu" + recursive: False # don't recursively inject submodules of this module +- match: + name: "^model\\.layers\\..*\\.self_attn$" + replace: + class: ktransformers.operators.attention.KQwen3MoeAttentionIPEXLLM + kwargs: + generate_device: "xpu" + prefill_device: "xpu" +- match: + name: "^model$" + replace: + class: "ktransformers.operators.models.KQwen2MoeModel" + kwargs: + per_layer_prefill_intput_threshold: 0 # 0 is close layer wise prefill +- match: + name: "^model.embed_tokens" + replace: + class: "default" + kwargs: + generate_device: "cpu" + prefill_device: "cpu" +- match: + class: transformers.models.qwen3_moe.modeling_qwen3_moe.Qwen3MoeRMSNorm + replace: + class: ktransformers.operators.layernorm.KDeepseekRMSNormIPEXLLM + kwargs: + generate_device: "xpu" + prefill_device: "xpu" +- match: + class: transformers.models.qwen3_moe.modeling_qwen3_moe.Qwen3MoeMLP + replace: + class: ktransformers.operators.mlp.KQwen2MoeMLP + kwargs: + generate_device: "xpu" + prefill_device: "xpu" diff --git a/archive/kt-sft/ktransformers/server/__init__.py b/archive/kt-sft/ktransformers/server/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/archive/kt-sft/ktransformers/server/api/__init__.py b/archive/kt-sft/ktransformers/server/api/__init__.py new file mode 100644 index 0000000..d69bc0b --- /dev/null +++ b/archive/kt-sft/ktransformers/server/api/__init__.py @@ -0,0 +1,10 @@ +from fastapi import APIRouter + +from .ollama import router as ollama_router +from .openai import router as openai_router,post_db_creation_operations +from .web import router as web_router + +router = APIRouter() +router.include_router(ollama_router) +router.include_router(openai_router) +router.include_router(web_router) diff --git a/archive/kt-sft/ktransformers/server/api/ollama/__init__.py b/archive/kt-sft/ktransformers/server/api/ollama/__init__.py new file mode 100644 index 0000000..7efab95 --- /dev/null +++ b/archive/kt-sft/ktransformers/server/api/ollama/__init__.py @@ -0,0 +1,6 @@ +from fastapi import APIRouter + +from .completions import router as completions_router + +router = APIRouter() +router.include_router(completions_router) diff --git a/archive/kt-sft/ktransformers/server/api/ollama/completions.py b/archive/kt-sft/ktransformers/server/api/ollama/completions.py new file mode 100644 index 0000000..9aa8ad0 --- /dev/null +++ b/archive/kt-sft/ktransformers/server/api/ollama/completions.py @@ -0,0 +1,283 @@ +from datetime import datetime +from http.client import NOT_IMPLEMENTED +import json +from time import time +from uuid import uuid4 +from typing import List, Optional + +from fastapi import APIRouter, Request +from pydantic import BaseModel, Field + +from ktransformers.server.config.config import Config +from ktransformers.server.utils.create_interface import get_interface +from ktransformers.server.schemas.assistants.streaming import check_link_response +from ktransformers.server.backend.base import BackendInterfaceBase + +from ktransformers.server.schemas.endpoints.chat import RawUsage + +router = APIRouter(prefix='/api') + +# https://github.com/ollama/ollama/blob/main/docs/api.md#generate-a-completion +class OllamaGenerateCompletionRequest(BaseModel): + model: str = Field(..., description="The model name, which is required.") + prompt: Optional[str] = Field( + None, description="The prompt to generate a response for.") + images: Optional[List[str]] = Field( + None, description="A list of base64-encoded images for multimodal models such as llava.") + # Advanced parameters + format: Optional[str] = Field( + None, description="The format to return a response in, accepted value is json.") + options: Optional[dict] = Field( + None, description="Additional model parameters as listed in the documentation.") + system: Optional[str] = Field( + None, description="System message to override what is defined in the Modelfile.") + template: Optional[str] = Field( + None, description="The prompt template to use, overriding what is defined in the Modelfile.") + context: Optional[str] = Field( + None, description="The context parameter from a previous request to keep a short conversational memory.") + stream: Optional[bool] = Field( + None, description="If false, the response will be returned as a single response object.") + raw: Optional[bool] = Field( + None, description="If true, no formatting will be applied to the prompt.") + keep_alive: Optional[str] = Field( + "5m", description="Controls how long the model will stay loaded into memory following the request.") + +class OllamaGenerationStreamResponse(BaseModel): + model: str + created_at: str + response: str + done: bool = Field(...) + +class OllamaGenerationResponse(BaseModel): + model: str + created_at: str + response: str + done: bool + +@router.post("/generate", tags=['ollama']) +async def generate(request: Request, input: OllamaGenerateCompletionRequest): + id = str(uuid4()) + interface: BackendInterfaceBase = get_interface() + print(f'COMPLETION INPUT:----\n{input.prompt}\n----') + config = Config() + + if input.stream: + async def inner(): + async for res in interface.inference(input.prompt, id): + if isinstance(res, RawUsage): + raw_usage = res + else: + token, finish_reason = res + d = OllamaGenerationStreamResponse( + model=config.model_name, + created_at=str(datetime.now()), + response=token, + done=False + ) + yield d.model_dump_json() + '\n' + d = OllamaGenerationStreamResponse( + model=config.model_name, + created_at=str(datetime.now()), + response='', + done=True + ) + yield d.model_dump_json() + '\n' + return check_link_response(request, inner()) + else: + complete_response = "" + async for res in interface.inference(input.prompt, id): + if isinstance(res, RawUsage): + raw_usage = res + else: + token, finish_reason = res + complete_response += token + response = OllamaGenerationResponse( + model=config.model_name, + created_at=str(datetime.now()), + response=complete_response, + done=True + ) + return response + +# https://github.com/ollama/ollama/blob/main/docs/api.md#generate-a-chat-completion +class OllamaChatCompletionMessage(BaseModel): + role: str + content: str + +class OllamaChatCompletionRequest(BaseModel): + model: str = Field(..., description="The model name, which is required.") + messages: List[OllamaChatCompletionMessage] = Field( + ..., description="A list of messages to generate a response for.") + stream: bool = Field(True, description="If true, the response will be streamed.") + +class OllamaChatCompletionStreamResponse(BaseModel): + model: str + created_at: str + message: dict + done: bool = Field(...) + done_reason: Optional[str] = Field("", description="done_reason") + total_duration: Optional[int] = Field(None, description="Total time spent in nanoseconds") + load_duration: Optional[int] = Field(None, description="Time spent loading model in nanoseconds") + prompt_eval_count: Optional[int] = Field(None, description="Number of tokens in prompt") + prompt_eval_duration: Optional[int] = Field(None, description="Time spent evaluating prompt in nanoseconds") + eval_count: Optional[int] = Field(None, description="Number of tokens generated") + eval_duration: Optional[int] = Field(None, description="Time spent generating response in nanoseconds") + +class OllamaChatCompletionResponse(BaseModel): + model: str + created_at: str + message: dict + done: bool + done_reason: Optional[str] = Field("", description="done_reason") + total_duration: Optional[int] = Field(None, description="Total time spent in nanoseconds") + load_duration: Optional[int] = Field(None, description="Time spent loading model in nanoseconds") + prompt_eval_count: Optional[int] = Field(None, description="Number of tokens in prompt") + prompt_eval_duration: Optional[int] = Field(None, description="Time spent evaluating prompt in nanoseconds") + eval_count: Optional[int] = Field(None, description="Number of tokens generated") + eval_duration: Optional[int] = Field(None, description="Time spent generating response in nanoseconds") + +@router.post("/chat", tags=['ollama']) +async def chat(request: Request, input: OllamaChatCompletionRequest): + id = str(uuid4()) + interface: BackendInterfaceBase = get_interface() + config = Config() + + input_message = [json.loads(m.model_dump_json()) for m in input.messages] + + if input.stream: + async def inner(): + start_time = time() + tokens = [] + + async for res in interface.inference(input_message, id): + if isinstance(res, RawUsage): + raw_usage = res + else: + token, finish_reason = res + d = OllamaChatCompletionStreamResponse( + model=config.model_name, + created_at=str(datetime.now()), + message={"role": "assistant", "content": token}, + done=False + ) + yield d.model_dump_json() + '\n' + end_time = time() + total_duration = int((end_time - start_time) * 1_000_000_000) # unit: ns + prompt_eval_count = raw_usage.prefill_count + eval_count = raw_usage.decode_count + eval_duration = int(raw_usage.decode_time * 1_000_000_000) + prompt_eval_duration = int(raw_usage.prefill_time * 1_000_000_000) + load_duration = int(raw_usage.tokenize_time * 1_000_000_000) + done_reason = finish_reason + + d = OllamaChatCompletionStreamResponse( + model=config.model_name, + created_at=str(datetime.now()), + message={}, + done=True, + total_duration=total_duration, + load_duration=load_duration, + prompt_eval_count=prompt_eval_count, + prompt_eval_duration=prompt_eval_duration, + eval_count=eval_count, + eval_duration=eval_duration, + done_reason=done_reason + ) + yield d.model_dump_json() + '\n' + return check_link_response(request, inner()) + else: + start_time = time() + complete_response = "" + eval_count = 0 + + async for res in interface.inference(input_message, id): + if isinstance(res, RawUsage): + raw_usage = res + else: + token, finish_reason = res + complete_response += token + + end_time = time() + total_duration = int((end_time - start_time) * 1_000_000_000) # unit: ns + prompt_eval_count = raw_usage.prefill_count + eval_count = raw_usage.decode_count + eval_duration = int(raw_usage.decode_time * 1_000_000_000) + prompt_eval_duration = int(raw_usage.prefill_time * 1_000_000_000) + load_duration = int(raw_usage.tokenize_time * 1_000_000_000) + done_reason = finish_reason + + + response = OllamaChatCompletionResponse( + model=config.model_name, + created_at=str(datetime.now()), + message={"role": "assistant", "content": complete_response}, + done=True, + total_duration=total_duration, + load_duration=load_duration, + prompt_eval_count=prompt_eval_count, + prompt_eval_duration=prompt_eval_duration, + eval_count=eval_count, + eval_duration=eval_duration, + done_reason=done_reason + ) + return response + +# https://github.com/ollama/ollama/blob/main/docs/api.md#list-local-models +class OllamaModel(BaseModel): + name: str + modified_at: str + size: int + # TODO: fill the rest correctly + +# mock ollama +@router.get("/tags", tags=['ollama']) +async def tags(): + config = Config() + # TODO: fill this correctly, although it does not effect Tabby + return {"models": [OllamaModel(name=config.model_name, modified_at="123", size=123)]} + +class OllamaModelInfo(BaseModel): + # TODO: fill this correctly + pass + +class OllamaShowRequest(BaseModel): + name: str = Field(..., description="Name of the model to show") + verbose: Optional[bool] = Field( + None, description="If set to true, returns full data for verbose response fields") + +class OllamaShowDetial(BaseModel): + parent_model: str + format: str + family: str + families: List[str] + parameter_size: str + quantization_level: str + +class OllamaShowResponse(BaseModel): + modelfile: str + parameters: str + template: str + details: OllamaShowDetial + model_info: OllamaModelInfo + + class Config: + protected_namespaces = () + +@router.post("/show", tags=['ollama']) +async def show(request: Request, input: OllamaShowRequest): + config = Config() + # TODO: Add more info in config to return, although it does not effect Tabby + return OllamaShowResponse( + modelfile="# Modelfile generated by ...", + parameters=" ", + template=" ", + details=OllamaShowDetial( + parent_model=" ", + format="gguf", + family=" ", + families=[" "], + parameter_size=" ", + quantization_level=" " + ), + model_info=OllamaModelInfo() + ) diff --git a/archive/kt-sft/ktransformers/server/api/openai/__init__.py b/archive/kt-sft/ktransformers/server/api/openai/__init__.py new file mode 100644 index 0000000..810b4f8 --- /dev/null +++ b/archive/kt-sft/ktransformers/server/api/openai/__init__.py @@ -0,0 +1,15 @@ +from fastapi import APIRouter + +from .assistants import router as assistants_router,create_default_assistant +from .endpoints.chat import router as chat_router +from .legacy import router as legacy_router + +router = APIRouter(prefix='/v1') + + +router.include_router(assistants_router) +router.include_router(chat_router) +router.include_router(legacy_router) + +def post_db_creation_operations(): + create_default_assistant() diff --git a/archive/kt-sft/ktransformers/server/api/openai/assistants/__init__.py b/archive/kt-sft/ktransformers/server/api/openai/assistants/__init__.py new file mode 100644 index 0000000..a5b652f --- /dev/null +++ b/archive/kt-sft/ktransformers/server/api/openai/assistants/__init__.py @@ -0,0 +1,14 @@ +from fastapi import APIRouter + +from .assistants import router as assistants_router, create_default_assistant +from .messages import router as messages_router +from .runs import router as runs_router +from .threads import router as threads_router + +router = APIRouter() + +threads_router.include_router(runs_router) +threads_router.include_router(messages_router) + +router.include_router(assistants_router) +router.include_router(threads_router) diff --git a/archive/kt-sft/ktransformers/server/api/openai/assistants/assistants.py b/archive/kt-sft/ktransformers/server/api/openai/assistants/assistants.py new file mode 100644 index 0000000..415c593 --- /dev/null +++ b/archive/kt-sft/ktransformers/server/api/openai/assistants/assistants.py @@ -0,0 +1,103 @@ +from typing import Optional + +from fastapi import APIRouter +from fastapi.testclient import TestClient + +from ktransformers.server.crud.assistants.assistants import AssistantDatabaseManager +from ktransformers.server.crud.assistants.runs import RunsDatabaseManager +from ktransformers.server.schemas.assistants.assistants import AssistantCreate, AssistantModify, ObjectID, AssistantBuildStatus, AssistantObject +from ktransformers.server.schemas.base import DeleteResponse, Order +from ktransformers.server.config.log import logger + + +router = APIRouter(prefix="/assistants") +assistant_manager = AssistantDatabaseManager() +runs_manager = RunsDatabaseManager() + + +@router.post("/", tags=['openai']) +async def create_assistant( + assistant: AssistantCreate, +): + return assistant_manager.db_create_assistant(assistant).as_api_response() + + +@router.get("/", tags=['openai']) +async def list_assistants( + limit: Optional[int] = 20, + order: Order = Order.DESC, + after: Optional[str] = None, + before: Optional[str] = None, +): + return [assistant.as_api_response() for assistant in assistant_manager.db_list_assistants(limit, order)] + +# list assistant with status + + +@router.get("/status", tags=['openai-ext']) +async def list_assistants_with_status( + limit: Optional[int] = 20, + order: Order = Order.DESC, + after: Optional[str] = None, + before: Optional[str] = None, +): + return assistant_manager.db_list_assistants(limit, order) + + +@router.get("/{assistant_id}", tags=['openai']) +async def retrieve_assistant( + assistant_id: str, +): + return assistant_manager.db_get_assistant_by_id(assistant_id).as_api_response() + + +@router.post("/{assistant_id}", tags=['openai']) +async def modify_assistant( + assistant_id: str, + assistant: AssistantModify, +): + return assistant_manager.db_update_assistant_by_id(assistant_id, assistant).as_api_response() + + +@router.delete("/{assistant_id}", tags=['openai'], response_model=DeleteResponse) +async def delete_assistant(assistant_id: str): + assistant_manager.db_delete_assistant_by_id(assistant_id) + return DeleteResponse(id=assistant_id, object="assistant.deleted") + + +@router.get("/{assistant_id}/related_thread", tags=['openai']) +async def get_related_thread(assistant_id: ObjectID): + assistant = assistant_manager.db_get_assistant_by_id(assistant_id) + return assistant.get_related_threads_ids() + + +def create_default_assistant(): + logger.info('Creating default assistant') + if assistant_manager.db_count_assistants() == 0: + default_assistant = assistant_manager.db_create_assistant(AssistantCreate(name="KT Assistant", + model="default model", + instructions="""You are a helpful, respectful and honest assistant. Always answer as helpfully as possible, while being safe. """ + + """Your answers should not include any harmful, unethical, racist, sexist, toxic, dangerous, or illegal content. """ + + """Please ensure that your responses are socially unbiased and positive in nature.""")) + default_assistant.build_status.status = AssistantBuildStatus.Status.completed + default_assistant.sync_db() + + +# unit test +client = TestClient(router) + + +def test_create_assistant(): + ass_create = AssistantCreate(model="awesome model", instructions="hello") + + res = client.post("/", json=ass_create.model_dump(mode="json")) + + assert res.status_code == 200 + assistant = AssistantObject.model_validate(res.json()) + + assert assistant.model == ass_create.model + assert assistant.instructions == ass_create.instructions + + res = client.get(f"/{assistant.id}") + ass1 = AssistantObject.model_validate(res.json()) + assert assistant == ass1 diff --git a/archive/kt-sft/ktransformers/server/api/openai/assistants/messages.py b/archive/kt-sft/ktransformers/server/api/openai/assistants/messages.py new file mode 100644 index 0000000..80c7569 --- /dev/null +++ b/archive/kt-sft/ktransformers/server/api/openai/assistants/messages.py @@ -0,0 +1,54 @@ +from typing import List, Optional + +from fastapi import APIRouter + +from ktransformers.server.exceptions import not_implemented +from ktransformers.server.schemas.assistants.messages import MessageCreate, MessageObject, MessageModify +from ktransformers.server.crud.assistants.messages import MessageDatabaseManager +from ktransformers.server.schemas.base import DeleteResponse, ObjectID, Order +from ktransformers.server.backend.base import ThreadContext +from ktransformers.server.utils.create_interface import get_thread_context_manager +router = APIRouter() +message_manager = MessageDatabaseManager() + + +@router.post("/{thread_id}/messages", tags=['openai'], response_model=MessageObject) +async def create_message(thread_id: str, msg: MessageCreate): + message = message_manager.db_create_message( + thread_id, msg, MessageObject.Status.in_progress) + ctx: Optional[ThreadContext] = await get_thread_context_manager().get_context_by_thread_id(thread_id) + if ctx is not None: + ctx.put_user_message(message) + return message + + +@router.get("/{thread_id}/messages", tags=['openai'], response_model=List[MessageObject]) +async def list_messages( + thread_id: str, + limit: Optional[int] = 20, + order: Order = Order.DESC, + after: Optional[str] = None, + before: Optional[str] = None, + run_id: Optional[str] = None, +): + return message_manager.db_list_messages_of_thread(thread_id, limit, order) + + +@router.get("/{thread_id}/messages/{message_id}", tags=['openai'], response_model=MessageObject) +async def retrieve_message(thread_id: ObjectID, message_id: ObjectID): + return message_manager.db_get_message_by_id(thread_id, message_id) + + +@router.post("/{thread_id}/messages/{message_id}", tags=['openai'], response_model=MessageObject) +async def modify_message(thread_id: ObjectID, message_id: ObjectID, msg: MessageModify): + #raise not_implemented('modify message not implemented') + raise not_implemented('modify message') + + +@router.delete("/{thread_id}/messages/{message_id}", tags=['openai'], response_model=DeleteResponse) +async def delete_message(thread_id: ObjectID, message_id: ObjectID): + ctx: Optional[ThreadContext] = await get_thread_context_manager().get_context_by_thread_id(thread_id) + if ctx is not None: + ctx.delete_user_message(message_id) + message_manager.db_delete_message_by_id(thread_id, message_id) + return DeleteResponse(id=message_id, object='thread.message.deleted') diff --git a/archive/kt-sft/ktransformers/server/api/openai/assistants/runs.py b/archive/kt-sft/ktransformers/server/api/openai/assistants/runs.py new file mode 100644 index 0000000..f5e51f9 --- /dev/null +++ b/archive/kt-sft/ktransformers/server/api/openai/assistants/runs.py @@ -0,0 +1,99 @@ +from typing import List, Optional + +from fastapi import APIRouter, Request + +from ktransformers.server.crud.assistants.runs import RunsDatabaseManager +from ktransformers.server.backend.base import ThreadContext +from ktransformers.server.schemas.assistants.runs import RunCreate,RunObject,RunThreadCreate,RunModify,RunSubmit +from ktransformers.server.schemas.assistants.streaming import api_stream_response +from ktransformers.server.utils.create_interface import get_thread_context_manager +from ktransformers.server.schemas.base import Order +from ktransformers.server.config.log import logger +from ktransformers.server.exceptions import internal_server_error + + +router = APIRouter() +runs_manager = RunsDatabaseManager() + + +@router.post("/{thread_id}/runs",tags=['openai']) +async def create_run(request: Request, thread_id: str, run_create: RunCreate): + if run_create.stream: + async def inner(): + run = runs_manager.db_create_run(thread_id, run_create) + yield run.stream_response_with_event(event=RunObject.Status.created) + + ctx: ThreadContext = await get_thread_context_manager().get_context_by_run_object(run) + + async for event in ctx.work(): + yield event + return api_stream_response(request, inner()) + else: + run = runs_manager.db_create_run(thread_id, run_create) + ctx: ThreadContext = await get_thread_context_manager().get_context_by_run_object(run) + async for event in ctx.work(): + pass + return run + + +@router.post("/runs",tags=['openai'], response_model=RunObject) +async def create_thread_and_run(run_thread: RunThreadCreate): + raise NotImplementedError + + +@router.get("/{thread_id}/runs",tags=['openai'], response_model=List[RunObject]) +async def list_runs( + thread_id: str, + limit: Optional[int] = 20, + order: Optional[Order] = Order.DESC, + after: Optional[str] = None, + before: Optional[str] = None, +): + raise NotImplementedError + + +@router.get("/{thread_id}/runs/{run_id}",tags=['openai'], response_model=RunObject) +async def retrieve_run( + thread_id: str, + run_id: str, +): + runobj= runs_manager.db_get_run(run_id) + assert runobj.thread_id == thread_id + return runobj + + + +@router.post("/{thread_id}/runs/{run_id}",tags=['openai'], response_model=RunObject) +async def modify_run( + thread_id: str, + run_id: str, + run: RunModify, +): + raise NotImplementedError + + +@router.post("/{thread_id}/runs/{run_id}/submit_tool_outputs", tags=['openai'],response_model=RunObject) +async def submit_tool_outputs_to_run(thread_id: str, run_id: str, submit: RunSubmit): + raise NotImplementedError + + +@router.post("/{thread_id}/runs/{run_id}/cancel",tags=['openai'], response_model=RunObject) +async def cancel_run(thread_id: str, run_id: str): + ctx: ThreadContext = await get_thread_context_manager().get_context_by_thread_id(thread_id) + if ctx is not None: + if ctx.run is None: + logger.warn(f'Run {ctx.run.id} is expected to be in_progress, but no context is found') + raise internal_server_error('ctx do not have run') + + if ctx.run.id == run_id: + logger.info(f'Cancelling thread: {thread_id} and run: {run_id}') + ctx.run.stream_response_with_event(RunObject.Status.cancelling) + return ctx.run + else: + run = runs_manager.db_get_run(run_id) + logger.info(f'Run {run_id} not in this thread context') + return run + else: + run = runs_manager.db_get_run(run_id) + logger.info(f'Run {run_id} not in context manager') + return run diff --git a/archive/kt-sft/ktransformers/server/api/openai/assistants/threads.py b/archive/kt-sft/ktransformers/server/api/openai/assistants/threads.py new file mode 100644 index 0000000..684a0ad --- /dev/null +++ b/archive/kt-sft/ktransformers/server/api/openai/assistants/threads.py @@ -0,0 +1,36 @@ +from typing import List,Optional +from fastapi import APIRouter + +from ktransformers.server.crud.assistants.threads import ThreadsDatabaseManager,Order,ObjectID +from ktransformers.server.schemas.assistants.threads import ThreadObject,ThreadCreate,ThreadModify +from ktransformers.server.schemas.base import DeleteResponse +from ktransformers.server.schemas.conversation import ThreadPreview + +router = APIRouter(prefix='/threads') +threads_manager = ThreadsDatabaseManager() + + +@router.post("/",tags=['openai'], response_model=ThreadObject) +async def create_thread(thread: ThreadCreate): + return threads_manager.db_create_thread(thread) + + +@router.get("/", tags=['openai-ext'],response_model=List[ThreadPreview]) +async def list_threads(limit: Optional[int] = 20, order: Order = Order.DESC): + return threads_manager.db_list_threads_preview(limit, order) + + +@router.get("/{thread_id}",tags=['openai'], response_model=ThreadObject) +async def retrieve_thread(thread_id: ObjectID): + return threads_manager.db_get_thread_by_id(thread_id) + + +@router.post("/{thread_id}",tags=['openai'], response_model=ThreadObject) +async def modify_thread(thread_id: ObjectID, thread: ThreadModify): + raise NotImplementedError + + +@router.delete("/{thread_id}",tags=['openai'], response_model=DeleteResponse) +async def delete_thread(thread_id: ObjectID): + threads_manager.db_delete_thread_by_id(thread_id=thread_id) + return DeleteResponse(id=thread_id, object='thread.deleted') diff --git a/archive/kt-sft/ktransformers/server/api/openai/endpoints/__init__.py b/archive/kt-sft/ktransformers/server/api/openai/endpoints/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/archive/kt-sft/ktransformers/server/api/openai/endpoints/chat.py b/archive/kt-sft/ktransformers/server/api/openai/endpoints/chat.py new file mode 100644 index 0000000..4adbed3 --- /dev/null +++ b/archive/kt-sft/ktransformers/server/api/openai/endpoints/chat.py @@ -0,0 +1,509 @@ +import json +from time import time +from uuid import uuid4 +from typing import Dict, List, Optional, Any, Literal, Union +from pydantic import BaseModel, Field +import re +from fastapi import APIRouter +from fastapi.requests import Request +from ktransformers.server.utils.create_interface import get_interface +from ktransformers.server.schemas.assistants.streaming import chat_stream_response +from ktransformers.server.schemas.endpoints.chat import ChatCompletionCreate +from ktransformers.server.schemas.endpoints.chat import RawUsage, Role +from ktransformers.server.backend.base import BackendInterfaceBase +from ktransformers.server.config.config import Config +from ktransformers.server.config.log import logger +from fastapi.responses import JSONResponse +from ktransformers.server.schemas.endpoints.chat import ChatCompletionChunk, CompletionUsage + +# Define own data structure instead of importing from OpenAI + + +class Choice(BaseModel): + index: int + message: Optional[Dict[str, Any]] = None + finish_reason: Optional[str] = None + logprobs: Optional[Any] = None + delta: Optional[Dict[str, Any]] = None + content_filter_results: Optional[Dict[str, Any]] = None + +class ChatCompletion(BaseModel): + id: str + object: str = "chat.completion" + created: int + model: str + choices: List[Choice] + usage: Optional[CompletionUsage] = None + system_fingerprint: Optional[str] = None + prompt_filter_results: Optional[List[Dict[str, Any]]] = None + +# Only for non-streaming response construction +class ChatCompletionMessageToolCallFunction(BaseModel): + name: str + arguments: str + +class ChatCompletionMessageToolCall(BaseModel): + id: str + type: str + function: ChatCompletionMessageToolCallFunction + +class ChatCompletionMessage(BaseModel): + role: str + content: Optional[str] = None + tool_calls: Optional[List[ChatCompletionMessageToolCall]] = None + +router = APIRouter() + +@router.get('/models', tags=['openai']) +async def list_models(): + return {"data": [{"id": Config().model_name, "name": Config().model_name}], "object": "list"} + +def getTools(buffer): + tool_calls_begin_marker = "<｜tool▁calls▁begin｜>" + tool_call_begin_marker = "<｜tool▁call▁begin｜>" + tool_sep_marker = "<｜tool▁sep｜>" + tool_call_end_marker = "<｜tool▁call▁end｜>" + tool_calls_end_marker = "<｜tool▁calls▁end｜>" + extracted_tools = [] + working_buffer = buffer + + # Iterate over all function calls + while tool_call_begin_marker in working_buffer and tool_call_end_marker in working_buffer: + # Find a complete function call + start_index = working_buffer.find(tool_call_begin_marker) + end_index = working_buffer.find(tool_call_end_marker) + len(tool_call_end_marker) + + if start_index == -1 or end_index == -1 or start_index > end_index: + logger.warning("Not a function") + break + + # Extract the full function call + full_tool_call = working_buffer[start_index:end_index] + + # Remove this function call from the working buffer to prevent duplicate processing + working_buffer = working_buffer.replace(full_tool_call, "", 1) + + # Extract the function name + function_name_start = full_tool_call.find(tool_sep_marker) + len(tool_sep_marker) + function_name_end = full_tool_call.find("\n", function_name_start) + function_name = full_tool_call[function_name_start:function_name_end].strip() + + # Extract JSON parameters + json_pattern = r'```json\s*(.*?)\s*```' + json_match = re.search(json_pattern, full_tool_call, re.DOTALL) + + if json_match: + arguments_str = json_match.group(1).strip() + # Generate tool call IDs + tool_call_id = f"call_{uuid4().hex[:24]}" + + # Add to tool call list + extracted_tools.append({ + "id": tool_call_id, + "type": "function", + "function": { + "name": function_name, + "arguments": arguments_str + } + }) + + logger.info(f"Get Function: {function_name}") + else: + logger.warning(f"Unable to get function, function_name: {function_name}") + + logger.info(f"Total {len(extracted_tools)} Functions") + return extracted_tools + +def get_tool_instructions(): + """Return concise tool calling instructions in English""" + return """ + +When you need real-time information or specialized operations, use function calls with this format: + +functionfunction_name +```json +{"param1": "value1", "param2": "value2",...} +``` + +The in the user message are the available tools automatically attached by the system. +You want to hide the guidance information in and the information in from the user. +Use functions when needed. Ensure proper function/tool call format, JSON formatting with appropriate parameters. + + +""" + +@router.post('/chat/completions', tags=['openai']) +async def chat_completion(request: Request, create: ChatCompletionCreate): + id = str(uuid4().hex) + + # Process messages with tool functionality if needed + enhanced_messages = list(create.messages) + if create.max_tokens is not None and create.max_tokens<0: + return JSONResponse( + status_code=400, + content={ + "object": "error", + "message": f"max_tokens must be at least 0, got {create.max_tokens}.", + "type": "BadRequestError", + "param": None, + "code": 400 + }) + + if create.max_completion_tokens is not None and create.max_completion_tokens<0: + return JSONResponse( + status_code=400, + content={ + "object": "error", + "message": f"max_completion_tokens must be at least 0, got {create.max_completion_tokens}.", + "type": "BadRequestError", + "param": None, + "code": 400 + }) + + if create.temperature<0 or create.temperature>2: + return JSONResponse( + status_code=400, + content={ + "object": "error", + "message": f"temperature must be in [0, 2], got {create.temperature}.", + "type": "BadRequestError", + "param": None, + "code": 400 + }) + if create.top_p<=0 or create.top_p>1: + return JSONResponse( + status_code=400, + content={ + "object": "error", + "message": f"top_p must be in (0, 1], got {create.top_p}.", + "type": "BadRequestError", + "param": None, + "code": 400 + }) + if create.frequency_penalty<-2 or create.frequency_penalty>2: + return JSONResponse( + status_code=400, + content={ + "object": "error", + "message": f"frequency_penalty must be in [-2, 2], got {create.frequency_penalty}.", + "type": "BadRequestError", + "param": None, + "code": 400 + }) + if create.presence_penalty<-2 or create.presence_penalty>2: + return JSONResponse( + status_code=400, + content={ + "object": "error", + "message": f"presence_penalty must be in [-2, 2], got {create.presence_penalty}.", + "type": "BadRequestError", + "param": None, + "code": 400 + }) + # Check if tools are present + has_tools = create.tools and len(create.tools) > 0 + + if has_tools: + # Find the most recent user message to append tool information + latest_user_msg_idx = -1 + for i in range(len(enhanced_messages) - 1, -1, -1): + if enhanced_messages[i].role == Role.user: + latest_user_msg_idx = i + break + + # Build the tool descriptions + tools_description = "" + for tool in create.tools: + tools_description += f"{tool.function.name}{tool.function.description}{tool.function.parameters}\n" + + # If first message is system, add concise tool instructions + if enhanced_messages[0].role == Role.system or enhanced_messages[0].role == Role.user: + if "" not in enhanced_messages[0].content.lower(): + enhanced_messages[0].content += "\n\n" + get_tool_instructions() + + # For the latest user message, append tool information + if latest_user_msg_idx >= 0: + # Add tool descriptions to the latest user message + enhanced_messages[latest_user_msg_idx].content += f"\n\n:\n{tools_description}\n" + + # Process request + interface: BackendInterfaceBase = get_interface() + input_message = [json.loads(m.model_dump_json()) for m in enhanced_messages] + if Config().api_key != '': + assert request.headers.get('Authorization', '').split()[-1] == Config().api_key + + if create.stream: + async def inner(): + chunk = ChatCompletionChunk( + id=id, + choices=[], + object='chat.completion.chunk', + created=int(time()), + model=Config().model_name, + system_fingerprint=f"fp_{uuid4().hex[:12]}", + ) + + # Collect the full output of the model + full_content = "" + buffer = "" # Used to temporarily store the current block of text + tool_call_mode = False # Mark if a tool call is being processed + tool_calls = [] # Store all detected tool calls + + # Tool call markers + tool_calls_begin_marker = "<｜tool▁calls▁begin｜>" + tool_call_begin_marker = "<｜tool▁call▁begin｜>" + tool_sep_marker = "<｜tool▁sep｜>" + tool_call_end_marker = "<｜tool▁call▁end｜>" + tool_calls_end_marker = "<｜tool▁calls▁end｜>" + too_calls_dict = { + "":"<｜tool▁calls▁begin｜>", + "":"<｜tool▁call▁begin｜>", + "":"<｜tool▁sep｜>", + "":"<｜tool▁call▁end｜>", + "":"<｜tool▁calls▁end｜>" + } + # Use check_client_connected for early stopping + async for res in interface.inference(input_message, id, create.temperature, create.top_p, create.max_tokens, create.max_completion_tokens): + if isinstance(res, RawUsage): + # Final return on utilization + raw_usage = res + chunk.choices = [] + chunk.usage = CompletionUsage( + prompt_tokens=raw_usage.prefill_count, + completion_tokens=raw_usage.decode_count, + total_tokens=raw_usage.prefill_count + raw_usage.decode_count + ) + if create.return_speed: + chunk.usage.prefill_time = res.prefill_time + chunk.usage.decode_time = res.decode_time + else: + chunk.usage.__dict__.pop('prefill_time', None) + chunk.usage.__dict__.pop('decode_time', None) + yield chunk + elif isinstance(res, tuple) and len(res) == 2: + token, finish_reason = res + token = re.sub('|'.join(map(re.escape, too_calls_dict.keys())), lambda m: too_calls_dict[m.group(0)], token) + # Detecting model-specific formatting tool call starts + if not tool_call_mode and tool_calls_begin_marker in buffer + token: + tool_call_mode = True + + # Adjust full_content to remove tool call section + if buffer.endswith(tool_calls_begin_marker): + full_content = full_content[:-len(tool_calls_begin_marker)] + elif tool_calls_begin_marker in (buffer + token): + idx = (buffer + token).find(tool_calls_begin_marker) + full_content = full_content[:-(len(buffer) - idx)] + buffer = "" + + # Send the current cumulative text content (if any) + if full_content: + chunk.choices = [{ + "index": 0, + "delta": {"content": full_content}, + "finish_reason": None + }] + yield chunk + full_content = "" + + # Accumulation of content in non-tool call mode + if not tool_call_mode: + full_content += token + buffer += token + # Keep the buffer at a reasonable size + if len(buffer) > 200: + buffer = buffer[-200:] + else: + # In tool call mode, continue to collect tool call related text + buffer += token + + # If the tool call end marker is found + if tool_calls_end_marker in buffer: + try: + # Parse and extract tool calling information + tool_calls = getTools(buffer) + if len(tool_calls): + # reset state + tool_call_mode = False + buffer = "" + + # Send tool call events + for idx, tool_call in enumerate(tool_calls): + # First tool call message + chunk.choices = [{ + "index": 0, + "delta": { + "role": "assistant", + "content": None, + "tool_calls": [{ + "index": idx, + "id": tool_call["id"], + "type": "function", + "function": { + "name": tool_call["function"]["name"], + "arguments": "" + } + }] + }, + "finish_reason": None + }] + yield chunk + + # Sending Parameters + chunk.choices = [{ + "index": 0, + "delta": { + "tool_calls": [{ + "index": idx, + "function": {"arguments": tool_call["function"]["arguments"]} + }] + }, + "finish_reason": None + }] + yield chunk + + # Send Completion Message + chunk.choices = [{ + "index": 0, + "delta": {}, + "finish_reason": "tool_calls" + }] + yield chunk + + # No further processing after return + return + else: + # JSON extraction failed, probably incomplete formatting + logger.warning("Failed to extract JSON from tool call") + tool_call_mode = False + buffer = "" + except Exception as e: + logger.error(f"Error processing tool call: {e}") + tool_call_mode = False + buffer = "" + + # Normal text output (only in non-tool call mode) + if not tool_call_mode and token: + if finish_reason is not None: + chunk.choices = [{ + "index": 0, + "delta": {}, + "finish_reason": finish_reason + }] + yield chunk + else: + if any(marker in token for marker in [tool_calls_begin_marker, tool_call_begin_marker]): + pass + else: + chunk.choices = [{ + "index": 0, + "delta": {"content": token}, + "finish_reason": None + }] + yield chunk + + # If gotten this far without returning, it means that the full tool call was not detected + # Send Routine Completion Message + if not tool_call_mode: + chunk.choices = [{ + "index": 0, + "delta": {}, + "finish_reason": "stop" + }] + yield chunk + + return chat_stream_response(request, inner()) + else: + # non streaming response processing + full_content = "" + finish_reason = None + tool_calls = [] + buffer = "" + tool_call_mode = False + + # Custom model special markers + tool_calls_begin_marker = "<｜tool▁calls▁begin｜>" + tool_call_begin_marker = "<｜tool▁call▁begin｜>" + tool_sep_marker = "<｜tool▁sep｜>" + tool_call_end_marker = "<｜tool▁call▁end｜>" + tool_calls_end_marker = "<｜tool▁calls▁end｜>" + too_calls_dict = { + "":"<｜tool▁calls▁begin｜>", + "":"<｜tool▁call▁begin｜>", + "":"<｜tool▁sep｜>", + "":"<｜tool▁call▁end｜>", + "":"<｜tool▁calls▁end｜>" + } + async for res in interface.inference(input_message, id, create.temperature, create.top_p, create.max_tokens, create.max_completion_tokens): + if isinstance(res, RawUsage): + raw_usage = res + usage = CompletionUsage( + prompt_tokens=raw_usage.prefill_count, + completion_tokens=raw_usage.decode_count, + total_tokens=raw_usage.prefill_count + raw_usage.decode_count, + ) + if create.return_speed: + usage.prefill_time = res.prefill_time + usage.decode_time = res.decode_time + else: + usage.__dict__.pop('prefill_time', None) + usage.__dict__.pop('decode_time', None) + + elif isinstance(res, tuple) and len(res) == 2: + token, finish_reason = res + token = re.sub('|'.join(map(re.escape, too_calls_dict.keys())), lambda m: too_calls_dict[m.group(0)], token) + # Detecting the start of model-specific formatting tool calls + if not tool_call_mode and tool_calls_begin_marker in buffer + token: + tool_call_mode = True + + # Adjust full_content to remove tool call section + if buffer.endswith(tool_calls_begin_marker): + full_content = full_content[:-len(tool_calls_begin_marker)] + elif tool_calls_begin_marker in (buffer + token): + idx = (buffer + token).find(tool_calls_begin_marker) + full_content = full_content[:-(len(buffer) - idx)] + buffer = "" + + # Accumulation of content in non-tool call mode + if not tool_call_mode: + full_content += token + buffer += token + # Keep the buffer at a reasonable size + if len(buffer) > 200: + buffer = buffer[-200:] + else: + # In tool call mode, continue to collect tool call related text + buffer += token + + # If the tool call end marker is found + if tool_calls_end_marker in buffer: + # Extract tool calls + tool_calls = getTools(buffer) + if tool_calls: + finish_reason = "tool_calls" + + # Reset state + tool_call_mode = False + buffer = "" + + # Build Response + message = { + "role": "assistant", + "content": None if tool_calls else full_content + } + if tool_calls: + message["tool_calls"] = tool_calls + response = { + "id": id, + "object": "chat.completion", + "created": int(time()), + "model": Config().model_name, + "choices": [{ + "index": 0, + "message": message, + "finish_reason": finish_reason or "stop" + }], + "usage": usage.__dict__ if 'usage' in locals() else None, + "system_fingerprint": f"fp_{uuid4().hex[:12]}" + } + + return response \ No newline at end of file diff --git a/archive/kt-sft/ktransformers/server/api/openai/legacy/__init__.py b/archive/kt-sft/ktransformers/server/api/openai/legacy/__init__.py new file mode 100644 index 0000000..08039a2 --- /dev/null +++ b/archive/kt-sft/ktransformers/server/api/openai/legacy/__init__.py @@ -0,0 +1,6 @@ +from fastapi import APIRouter + +from . import completions + +router = APIRouter() +router.include_router(completions.router) \ No newline at end of file diff --git a/archive/kt-sft/ktransformers/server/api/openai/legacy/completions.py b/archive/kt-sft/ktransformers/server/api/openai/legacy/completions.py new file mode 100644 index 0000000..bd6a56a --- /dev/null +++ b/archive/kt-sft/ktransformers/server/api/openai/legacy/completions.py @@ -0,0 +1,81 @@ +import json +from time import time +from uuid import uuid4 +from fastapi import APIRouter +from fastapi.requests import Request +from ktransformers.server.utils.create_interface import get_interface +from ktransformers.server.schemas.assistants.streaming import stream_response +from ktransformers.server.schemas.legacy.completions import CompletionCreate,CompletionObject +from ktransformers.server.schemas.endpoints.chat import RawUsage +from fastapi.responses import JSONResponse +from ktransformers.server.config.config import Config +router = APIRouter() + +@router.post("/completions",tags=['openai']) +async def create_completion(request:Request, create:CompletionCreate): + id = str(uuid4()) + if create.max_tokens is not None and create.max_tokens<0: + return JSONResponse( + status_code=400, + content={ + "object": "error", + "message": f"max_tokens must be at least 0, got {create.max_tokens}.", + "type": "BadRequestError", + "param": None, + "code": 400 + }) + if create.max_completion_tokens is not None and create.max_completion_tokens<0: + return JSONResponse( + status_code=400, + content={ + "object": "error", + "message": f"max_completion_tokens must be at least 0, got {create.max_completion_tokens}.", + "type": "BadRequestError", + "param": None, + "code": 400 + }) + if create.temperature<0 or create.temperature>2: + return JSONResponse( + status_code=400, + content={ + "object": "error", + "message": f"temperature must be in [0, 2], got {create.temperature}.", + "type": "BadRequestError", + "param": None, + "code": 400 + }) + if create.top_p<=0 or create.top_p>1: + return JSONResponse( + status_code=400, + content={ + "object": "error", + "message": f"top_p must be in (0, 1], got {create.top_p}.", + "type": "BadRequestError", + "param": None, + "code": 400 + }) + interface = get_interface() + print(f'COMPLETION INPUT:----\n{create.prompt}\n----') + + + if create.stream: + async def inner(): + async for res in interface.inference(create.prompt, id, create.temperature, create.top_p, create.max_tokens, create.max_completion_tokens): + if isinstance(res, RawUsage): + raw_usage = res + else: + token, finish_reason = res + d = {'choices':[{'delta':{'content':token}}]} + yield f"data:{json.dumps(d)}\n\n" + d = {'choices':[{'delta':{'content':''},'finish_reason':''}]} + yield f"data:{json.dumps(d)}\n\n" + return stream_response(request,inner()) + else: + comp = CompletionObject(id=id,object='text_completion',created=int(time())) + async for res in interface.inference(create.prompt,id,create.temperature,create.top_p, create.max_tokens, create.max_completion_tokens): + if isinstance(res, RawUsage): + raw_usage = res + else: + token, finish_reason = res + comp.append_token(token) + return comp diff --git a/archive/kt-sft/ktransformers/server/api/web/__init__.py b/archive/kt-sft/ktransformers/server/api/web/__init__.py new file mode 100644 index 0000000..befa721 --- /dev/null +++ b/archive/kt-sft/ktransformers/server/api/web/__init__.py @@ -0,0 +1,6 @@ +from fastapi import APIRouter +from .system import router as system_router + + +router = APIRouter() +router.include_router(system_router) diff --git a/archive/kt-sft/ktransformers/server/api/web/system.py b/archive/kt-sft/ktransformers/server/api/web/system.py new file mode 100644 index 0000000..78e27a7 --- /dev/null +++ b/archive/kt-sft/ktransformers/server/api/web/system.py @@ -0,0 +1,9 @@ +from fastapi import APIRouter + + +router = APIRouter() + + +@router.get('/system-info',tags=['web']) +def system_info(): + raise NotImplementedError diff --git a/archive/kt-sft/ktransformers/server/args.py b/archive/kt-sft/ktransformers/server/args.py new file mode 100644 index 0000000..748bd47 --- /dev/null +++ b/archive/kt-sft/ktransformers/server/args.py @@ -0,0 +1,157 @@ +import argparse +from ktransformers.server.backend.args import ConfigArgs, default_args +from ktransformers.util.utils import get_free_ports +from transformers import AutoConfig + +class ArgumentParser: + def __init__(self, cfg): + self.cfg = cfg + + def parse_args(self): + parser = argparse.ArgumentParser(prog="kvcache.ai", description="Ktransformers") + parser.add_argument("--host", type=str, default=self.cfg.server_ip) + parser.add_argument("--port", type=int, default=self.cfg.server_port) + parser.add_argument("--api_key", type=str, default=self.cfg.api_key) + parser.add_argument("--ssl_keyfile", type=str) + parser.add_argument("--ssl_certfile", type=str) + parser.add_argument("--web", type=bool, default=self.cfg.mount_web) + parser.add_argument("--model_name", type=str, default=self.cfg.model_name) + parser.add_argument("--model_dir", type=str) + parser.add_argument("--model_path", type=str, default=self.cfg.model_path) + parser.add_argument( + "--device", type=str, default=self.cfg.model_device, help="Warning: Abandoning this parameter" + ) + parser.add_argument("--architectures", type=str, default=self.cfg.model_name) + parser.add_argument("--gguf_path", type=str, default=self.cfg.gguf_path) + parser.add_argument("--optimize_config_path", default=None, type=str, required=False) + parser.add_argument("--cpu_infer", type=int, default=self.cfg.cpu_infer) + parser.add_argument("--backend_type", type=str, default=self.cfg.backend_type) + parser.add_argument("--chunk_size", type=int, default=self.cfg.chunk_size) + + # model configs + # parser.add_argument("--model_cache_lens", type=int, default=self.cfg.cache_lens) # int? + parser.add_argument("--max_batch_size", type=int, default=self.cfg.max_batch_size) + parser.add_argument("--max_new_tokens", type=int, default=self.cfg.max_new_tokens) + parser.add_argument("--json_mode", type=bool, default=self.cfg.json_mode) + parser.add_argument("--healing", type=bool, default=self.cfg.healing) + parser.add_argument("--ban_strings", type=list, default=self.cfg.ban_strings, required=False) + parser.add_argument("--gpu_split", type=str, default=self.cfg.gpu_split, required=False) + parser.add_argument("--length", type=int, default=self.cfg.length, required=False) + parser.add_argument("--rope_scale", type=float, default=self.cfg.rope_scale, required=False) + parser.add_argument("--rope_alpha", type=float, default=self.cfg.rope_alpha, required=False) + parser.add_argument("--no_flash_attn", type=bool, default=self.cfg.no_flash_attn) + parser.add_argument("--low_mem", type=bool, default=self.cfg.low_mem) + parser.add_argument("--experts_per_token", type=int, default=self.cfg.experts_per_token, required=False) + parser.add_argument("--load_q4", type=bool, default=self.cfg.load_q4) + parser.add_argument("--fast_safetensors", type=bool, default=self.cfg.fast_safetensors) + parser.add_argument("--draft_model_dir", type=str, default=self.cfg.draft_model_dir, required=False) + parser.add_argument("--no_draft_scale", type=bool, default=self.cfg.no_draft_scale) + parser.add_argument("--modes", type=bool, default=self.cfg.modes) + parser.add_argument("--mode", type=str, default=self.cfg.mode) + parser.add_argument("--username", type=str, default=self.cfg.username) + parser.add_argument("--botname", type=str, default=self.cfg.botname) + parser.add_argument("--system_prompt", type=str, default=self.cfg.system_prompt, required=False) + parser.add_argument("--temperature", type=float, default=self.cfg.temperature) + parser.add_argument("--smoothing_factor", type=float, default=self.cfg.smoothing_factor) + parser.add_argument("--dynamic_temperature", type=str, default=self.cfg.dynamic_temperature, required=False) + parser.add_argument("--top_k", type=int, default=self.cfg.top_k) + parser.add_argument("--top_p", type=float, default=self.cfg.top_p) + parser.add_argument("--top_a", type=float, default=self.cfg.top_a) + parser.add_argument("--skew", type=float, default=self.cfg.skew) + parser.add_argument("--typical", type=float, default=self.cfg.typical) + parser.add_argument("--repetition_penalty", type=float, default=self.cfg.repetition_penalty) + parser.add_argument("--frequency_penalty", type=float, default=self.cfg.frequency_penalty) + parser.add_argument("--presence_penalty", type=float, default=self.cfg.presence_penalty) + parser.add_argument("--response_chunk", type=int, default=self.cfg.response_chunk) + parser.add_argument("--no_code_formatting", type=bool, default=self.cfg.no_code_formatting) + parser.add_argument("--cache_8bit", type=bool, default=self.cfg.cache_8bit) + parser.add_argument("--cache_q4", type=bool, default=self.cfg.cache_q4) + parser.add_argument("--ngram_decoding", type=bool, default=self.cfg.ngram_decoding) + parser.add_argument("--print_timings", type=bool, default=self.cfg.print_timings) + parser.add_argument("--amnesia", type=bool, default=self.cfg.amnesia) + parser.add_argument("--batch_size", type=int, default=self.cfg.batch_size) + parser.add_argument("--cache_lens", type=int, default=self.cfg.cache_lens) + + # kvc2 config + parser.add_argument("--kvc2_config_dir", type=str, default=self.cfg.kvc2_config_dir) + + # log configs + # log level: debug, info, warn, error, crit + parser.add_argument("--log_dir", type=str, default=self.cfg.log_dir) + parser.add_argument("--log_file", type=str, default=self.cfg.log_file) + parser.add_argument("--log_level", type=str, default=self.cfg.log_level) + parser.add_argument("--backup_count", type=int, default=self.cfg.backup_count) + + # db configs + parser.add_argument("--db_type", type=str, default=self.cfg.db_type) + parser.add_argument("--db_host", type=str, default=self.cfg.db_host) + parser.add_argument("--db_port", type=str, default=self.cfg.db_port) + parser.add_argument("--db_name", type=str, default=self.cfg.db_name) + parser.add_argument("--db_pool_size", type=int, default=self.cfg.db_pool_size) + parser.add_argument("--db_database", type=str, default=self.cfg.db_database) + + # user config + parser.add_argument("--user_secret_key", type=str, default=self.cfg.user_secret_key) + parser.add_argument("--user_algorithm", type=str, default=self.cfg.user_algorithm) + parser.add_argument("--force_think", action=argparse.BooleanOptionalAction, type=bool, default=self.cfg.user_force_think) + parser.add_argument("--use_cuda_graph", action=argparse.BooleanOptionalAction, type=bool, default=self.cfg.use_cuda_graph) + + # web config + parser.add_argument("--web_cross_domain", type=bool, default=self.cfg.web_cross_domain) + + # file config + parser.add_argument("--file_upload_dir", type=str, default=self.cfg.file_upload_dir) + parser.add_argument("--assistant_store_dir", type=str, default=self.cfg.assistant_store_dir) + # local chat + parser.add_argument("--prompt_file", type=str, default=self.cfg.prompt_file) + + + # async server + parser.add_argument("--sched_strategy", type=str, default=self.cfg.sched_strategy) + # parser.add_argument("--sched_port", type=int, default=self.cfg.sched_port) + # parser.add_argument("--sched_metrics_port", type=int, default=self.cfg.sched_metrics_port) + # parser.add_argument("--kvc2_metrics_port", type=int, default=self.cfg.kvc2_metrics_port) + parser.add_argument("--page_size", type=str, default=self.cfg.page_size) + parser.add_argument("--memory_gpu_only", type=str, default=self.cfg.memory_gpu_only) + parser.add_argument("--utilization_percentage", type=str, default=self.cfg.utilization_percentage) + parser.add_argument("--cpu_memory_size_GB", type=str, default=self.cfg.cpu_memory_size_GB) + + + args = parser.parse_args() + if (args.model_dir is not None or args.model_path is not None): + if (args.model_path is not None): + # if pass model_dir and model_path, we use model_path + args.model_dir = args.model_path + else: + # if only pass model_dir, we use model_dir + args.model_path = args.model_dir + else: + args.model_dir = self.cfg.model_dir + args.model_path = self.cfg.model_path + + # we add the name not match args individually + self.cfg.model_device = args.device + self.cfg.mount_web = args.web + self.cfg.server_ip = args.host + self.cfg.server_port = args.port + self.cfg.user_force_think = args.force_think + + model_config = AutoConfig.from_pretrained(args.model_dir, trust_remote_code=True) + if model_config.architectures[0] == "Qwen3MoeForCausalLM" or model_config.architectures[0] == "Qwen2MoeForCausalLM" : + args.gpu_memory_size = args.cache_lens*2*2*model_config.num_hidden_layers*model_config.num_key_value_heads*model_config.head_dim + args.architectures = model_config.architectures[0] + else: + args.gpu_memory_size = args.cache_lens*2*576*61 + # set config from args + for key, value in vars(args).items(): + if value is not None and hasattr(self.cfg, key): + setattr(self.cfg, key, value) + self.cfg.gpu_memory_size = args.gpu_memory_size + free_ports = get_free_ports(3, [args.port]) + args.sched_port = free_ports[0] + args.sched_metrics_port = free_ports[1] + args.kvc2_metrics_port = free_ports[2] + self.cfg.sched_port = free_ports[0] + self.cfg.sched_metrics_port = free_ports[1] + self.cfg.kvc2_metrics_port = free_ports[2] + return args diff --git a/archive/kt-sft/ktransformers/server/backend/__init__.py b/archive/kt-sft/ktransformers/server/backend/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/archive/kt-sft/ktransformers/server/backend/args.py b/archive/kt-sft/ktransformers/server/backend/args.py new file mode 100644 index 0000000..1c602b1 --- /dev/null +++ b/archive/kt-sft/ktransformers/server/backend/args.py @@ -0,0 +1,80 @@ +from pydantic import BaseModel, Field +from typing import Optional +from ktransformers.server.config.config import Config + + +class ConfigArgs(BaseModel): + model_name: Optional[str] = Field(..., description="Model name") + model_dir: Optional[str] = Field(..., description="Path to model directory") + optimize_config_path: Optional[str] = Field(None, description="Path of your optimize config yml file") + gguf_path: Optional[str] = Field(None, description="Path of your gguf file") + + class Config: + protected_namespaces = () + + max_batch_size: int = Field( + None, description="Max number of batches to run at once, assuming the sequences will fit within total_context" + ) + chunk_size: int = Field( + None, + description=( + "Max chunk size. Determines the size of prefill operations. Can be reduced to reduce pauses whenever a new" + " job is started, but at the expense of overall prompt ingestion speed" + ), + ) + max_new_tokens: int = Field(None, description="Max new tokens per completion. For this example applies to all jobs") + json_mode: bool = Field( + None, description="Use LMFE to constrain the output to JSON format. See schema and details below" + ) + healing: bool = Field(None, description="Demonstrate token healing") + ban_strings: Optional[list] = Field(None, description="Ban some phrases maybe") + gpu_split: Optional[str] = Field(None, description='"auto", or VRAM allocation per GPU in GB') + length: Optional[int] = Field(None, description="Maximum sequence length") + rope_scale: Optional[float] = Field(None, description="RoPE scaling factor") + rope_alpha: Optional[float] = Field(None, description="RoPE alpha value (NTK)") + no_flash_attn: bool = Field(None, description="Disable Flash Attention") + low_mem: bool = Field(None, description="Enable VRAM optimizations, potentially trading off speed") + experts_per_token: Optional[int] = Field( + None, description="Override MoE model's default number of experts per token" + ) + load_q4: bool = Field(None, description="Load weights in Q4 mode") + fast_safetensors: bool = Field(None, description="Optimized safetensors loading with direct I/O (experimental!)") + draft_model_dir: Optional[str] = Field(None, description="Path to draft model directory") + no_draft_scale: bool = Field( + None, + description="If draft model has smaller context size than model, don't apply alpha (NTK) scaling to extend it", + ) + modes: bool = Field(None, description="List available modes and exit.") + mode: str = Field(None, description="Chat mode. Use llama for Llama 1/2 chat finetunes.") + username: str = Field(None, description="Username when using raw chat mode") + botname: str = Field(None, description="Bot name when using raw chat mode") + system_prompt: Optional[str] = Field(None, description="Use custom system prompt") + temperature: float = Field(None, description="Sampler temperature, default = 0.95 (1 to disable)") + smoothing_factor: float = Field(None, description="Smoothing Factor, default = 0.0 (0 to disable)") + dynamic_temperature: Optional[str] = Field( + None, description="Dynamic temperature min,max,exponent, e.g. -dyntemp 0.2,1.5,1" + ) + top_k: int = Field(None, description="Sampler top-K, default = 50 (0 to disable)") + top_p: float = Field(None, description="Sampler top-P, default = 0.8 (0 to disable)") + top_a: float = Field(None, description="Sampler top-A, default = 0.0 (0 to disable)") + skew: float = Field(None, description="Skew sampling, default = 0.0 (0 to disable)") + typical: float = Field(None, description="Sampler typical threshold, default = 0.0 (0 to disable)") + repetition_penalty: float = Field(None, description="Sampler repetition penalty, default = 1.01 (1 to disable)") + frequency_penalty: float = Field(None, description="Sampler frequency penalty, default = 0.0 (0 to disable)") + presence_penalty: float = Field(None, description="Sampler presence penalty, default = 0.0 (0 to disable)") + response_chunk: int = Field(None, description="Space to reserve in context for reply, default = 250") + no_code_formatting: bool = Field(None, description="Disable code formatting/syntax highlighting") + cache_8bit: bool = Field(None, description="Use 8-bit (FP8) cache") + cache_q4: bool = Field(None, description="Use Q4 cache") + ngram_decoding: bool = Field(None, description="Use n-gram speculative decoding") + print_timings: bool = Field(None, description="Output timings after each prompt") + amnesia: bool = Field(None, description="Forget context after every response") + + # for transformers + batch_size: int = Field(None, description="Batch Size") + cache_lens: int = Field(None, description="Cache lens for transformers static cache") + device: str = Field(None, description="device") + + +cfg = Config() +default_args = cfg diff --git a/archive/kt-sft/ktransformers/server/backend/base.py b/archive/kt-sft/ktransformers/server/backend/base.py new file mode 100644 index 0000000..aa011bf --- /dev/null +++ b/archive/kt-sft/ktransformers/server/backend/base.py @@ -0,0 +1,167 @@ +from asyncio import Queue +from enum import Enum +import sys, os +from typing import AsyncIterator, Dict, List, Optional, Tuple + +import torch + +from ktransformers.server.config.log import logger +from ktransformers.server.crud.assistants.assistants import AssistantDatabaseManager +from ktransformers.server.crud.assistants.messages import MessageDatabaseManager +from ktransformers.server.crud.assistants.runs import RunsDatabaseManager +from ktransformers.server.crud.assistants.threads import ThreadsDatabaseManager +from ktransformers.server.exceptions import request_error +from ktransformers.server.schemas.assistants.assistants import AssistantObject +from ktransformers.server.schemas.assistants.messages import MessageCreate, MessageObject, Role +from ktransformers.server.schemas.assistants.runs import RunObject +from ktransformers.server.schemas.assistants.threads import ThreadObject +from ktransformers.server.schemas.endpoints.chat import RawUsage +from ktransformers.server.schemas.base import ObjectID, Order +from ktransformers.server.utils.multi_timer import Profiler + + +from .args import ConfigArgs,default_args + + + +class BackendInterfaceBase: + ''' + Interface to inference frameworks. e.g. transformers, exllama. + Implement __init__ and work + ''' + + args: ConfigArgs + profiler:Profiler = Profiler() + + def __init__(self, args:ConfigArgs = default_args): + raise NotImplementedError + + + async def inference(self,local_messages,request_unique_id:Optional[str])->AsyncIterator[str]: + ''' + work can be called directly, or by ThreadContext + + local_messages: + When called by ThreadContext, local_messages are generated by ThreadContext.get_local_messages(). + Please deal with different local_messages + request_unique_id: + unique id of different requests, useful when using cache + + return: + async str output for stream update + + ''' + raise NotImplementedError + + + def report_last_time_performance(self): + try: + tokenize_time = self.profiler.get_timer_sec('tokenize') + prefill_time = self.profiler.get_timer_sec('prefill') + decode_time = self.profiler.get_timer_sec('decode') + prefill_count = self.profiler.get_counter('prefill') + decode_count = self.profiler.get_counter('decode') + + logger.info(f'Performance(T/s): prefill {prefill_count/prefill_time}, decode {decode_count/decode_time}. Time(s): tokenize {tokenize_time}, prefill {prefill_time}, decode {decode_time}') + except: + logger.info(f'Performance statistics not recorded') + + +class ThreadContext: + ''' + A thread context holding assistant logics + + ''' + + args: ConfigArgs + # Assistant Logic + assistant: Optional[AssistantObject] = None + related_threads : List[ThreadObject] + thread: ThreadObject + messages: List[MessageObject] = [] + run: RunObject + + interface: Optional[BackendInterfaceBase] = None + + queue: Optional[Queue] = None + timer: Profiler = Profiler() + + def __init__(self, run: RunObject,interface:BackendInterfaceBase, args: ConfigArgs = default_args) -> None: + self.args = args + self.thread_manager = ThreadsDatabaseManager() + self.message_manager = MessageDatabaseManager() + self.runs_manager = RunsDatabaseManager() + self.assistant_manager = AssistantDatabaseManager() + self.thread = self.thread_manager.db_get_thread_by_id(run.thread_id) + self.assistant = self.assistant_manager.db_get_assistant_by_id(run.assistant_id) + self.messages = self.message_manager.db_list_messages_of_thread(run.thread_id,order=Order.ASC) + logger.debug(f"{len(self.messages)} messages loaded from database") + self.interface = interface + self.update_by_run(run,args) + + def get_local_messages(self): + ''' + Get local messages, as the input to interface.work + This function is intended to message preprocess e.g. apply chat template + ''' + raise NotImplementedError + + def update_by_run(self,run:RunObject,args:ConfigArgs = default_args): + self.run = run + self.args = args + + def put_user_message(self, message: MessageObject): + assert ( + message.role.is_user() and message.thread_id == self.thread.id and message.status == MessageObject.Status.in_progress + ) + self.messages.append(message) + + def delete_user_message(self,message_id: ObjectID): + self.messages = [m for m in self.messages if m.id != message_id] + + async def work(self)->AsyncIterator: + logger.debug('start working') + user_message = self.messages[-1] + if not user_message.role.is_user(): + raise request_error('user must talk before LLM can talk') + user_message.status = MessageObject.Status.completed + user_message.sync_db() + + local_messages = self.get_local_messages() # must get this before we interseted reply_message + + + response_str_count = 0 + reply_message = self.message_manager.create_message_object( + self.thread.id, + self.run.id, + MessageCreate(role=Role.assistant, content=""), + ) + reply_message.assistant_id = self.assistant.id + self.messages.append(reply_message) + + yield reply_message.stream_response_with_event(MessageObject.Status.created) + yield reply_message.stream_response_with_event(MessageObject.Status.in_progress) + yield self.run.stream_response_with_event(RunObject.Status.in_progress) + + async for res in self.interface.inference(local_messages,self.thread.id): + if isinstance(res, RawUsage): + raw_usage = res + else: + token, finish_reason = res + if self.run.status == RunObject.Status.cancelling: + logger.warn(f'Run {self.run.id} cancelling') + break + yield reply_message.append_message_delta(token) + response_str_count+=1 + + if self.run.status == RunObject.Status.cancelling: + yield self.run.stream_response_with_event(RunObject.Status.cancelled) + yield reply_message.stream_response_with_event(MessageObject.Status.incomplete) + elif self.run.status == RunObject.Status.in_progress: + yield self.run.stream_response_with_event(RunObject.Status.completed) + yield reply_message.stream_response_with_event(MessageObject.Status.completed) + else: + raise NotImplementedError(f'{self.run.status} should not appear here') + + reply_message.sync_db() + self.run.sync_db() \ No newline at end of file diff --git a/archive/kt-sft/ktransformers/server/backend/context_manager.py b/archive/kt-sft/ktransformers/server/backend/context_manager.py new file mode 100644 index 0000000..e44feaa --- /dev/null +++ b/archive/kt-sft/ktransformers/server/backend/context_manager.py @@ -0,0 +1,65 @@ +from asyncio import Lock +from typing import Dict, Optional + +from ktransformers.server.backend.base import ThreadContext, BackendInterfaceBase +from ktransformers.server.schemas.assistants.runs import RunObject +from ktransformers.server.schemas.base import ObjectID +from ktransformers.server.config.log import logger +from ktransformers.server.backend.interfaces.transformers import TransformersThreadContext +from ktransformers.server.backend.interfaces.ktransformers import KTransformersThreadContext +from ktransformers.server.backend.interfaces.exllamav2 import ExllamaThreadContext + + +from ktransformers.server.backend.interfaces.exllamav2 import ExllamaInterface +from ktransformers.server.backend.interfaces.transformers import TransformersInterface +from ktransformers.server.backend.interfaces.ktransformers import KTransformersInterface + +class ThreadContextManager: + lock: Lock + threads_context: Dict[ObjectID, ThreadContext] + interface: BackendInterfaceBase + + def __init__(self,interface) -> None: + logger.debug(f"Creating Context Manager") + self.lock = Lock() + self.threads_context = {} + self.interface = interface + pass + + async def get_context_by_run_object(self, run: RunObject) -> ThreadContext: + async with self.lock: + logger.debug(f"keys {self.threads_context.keys()}") + if run.thread_id not in self.threads_context: + logger.debug(f"new inference context {run.thread_id}") + if isinstance(self.interface, ExllamaInterface): + new_context = ExllamaThreadContext(run, self.interface) + elif isinstance(self.interface, KTransformersInterface): + new_context = KTransformersThreadContext(run, self.interface) + elif isinstance(self.interface, TransformersInterface): + new_context = TransformersThreadContext(run, self.interface) + else: + from ktransformers.server.backend.interfaces.balance_serve import BalanceServeThreadContext + from ktransformers.server.backend.interfaces.balance_serve import BalanceServeInterface + if isinstance(self.interface, BalanceServeInterface): + new_context = BalanceServeThreadContext(run, self.interface) + else: + raise NotImplementedError + # elif isinstance(self.interface, BalanceServeInterface): + # new_context = BalanceServeThreadContext(run, self.interface) + # else: + # raise NotImplementedError + self.threads_context[run.thread_id] = new_context + # self.threads_context[run.thread_id] = ExllamaInferenceContext(run) + re = self.threads_context[run.thread_id] + re.update_by_run(run) + return re + + async def get_context_by_thread_id(self, thread_id: ObjectID) -> Optional[ThreadContext]: + async with self.lock: + if thread_id in self.threads_context: + logger.debug(f'found context for thread {thread_id}') + return self.threads_context[thread_id] + else: + logger.debug(f'no context for thread {thread_id}') + return None + \ No newline at end of file diff --git a/archive/kt-sft/ktransformers/server/backend/interfaces/__init__.py b/archive/kt-sft/ktransformers/server/backend/interfaces/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/archive/kt-sft/ktransformers/server/backend/interfaces/balance_serve.py b/archive/kt-sft/ktransformers/server/backend/interfaces/balance_serve.py new file mode 100644 index 0000000..e4b7b1c --- /dev/null +++ b/archive/kt-sft/ktransformers/server/backend/interfaces/balance_serve.py @@ -0,0 +1,492 @@ +from typing import Any, AsyncIterator, List, Optional, Set +from ktransformers.models.custom_cache import KDeepSeekV3Cache, KGQACache +from transformers import ( + AutoTokenizer, + AutoConfig, + GenerationConfig, + StaticCache, + AutoModelForCausalLM, + BitsAndBytesConfig, +) + +from ktransformers.server.config.config import Config +from ..base import ThreadContext, BackendInterfaceBase +import torch +from ktransformers.server.backend.interfaces.transformers import ( + ConfigArgs, + default_args, + TextStreamer, +) +from ktransformers.server.schemas.base import ObjectID +from ktransformers.server.config.log import logger +from ktransformers.optimize.optimize import optimize_and_load_gguf +from ktransformers.models.custom_modeling_deepseek_v3 import KDeepseekV3ForCausalLM +from ktransformers.models.custom_modeling_deepseek_v2 import KDeepseekV2ForCausalLM +from ktransformers.models.custom_modeling_qwen2_moe import KQwen2MoeForCausalLM +from ktransformers.models.custom_modeling_qwen3_moe import KQwen3MoeForCausalLM +from ktransformers.models.configuration_qwen3_moe import Qwen3MoeConfig +from ktransformers.server.balance_serve.inference.model_runner import ModelRunner +from ktransformers.server.balance_serve.inference.sampling.sampler import Sampler, SamplingOptions +from ktransformers.server.balance_serve.inference.query_manager import QueryManager +from ktransformers.server.balance_serve.inference.forward_batch import ForwardBatchInput, ForwardBatchOutput +from ktransformers.server.balance_serve.sched_rpc import SchedulerClient +from ktransformers.server.balance_serve.settings import sched_ext +from torch.multiprocessing import Queue +import torch.multiprocessing as mp +from multiprocessing.synchronize import Event +from ktransformers.server.schemas.endpoints.chat import RawUsage +from ktransformers.server.utils.multi_timer import Profiler +import zmq +import time +import queue +import tempfile +import asyncio +import threading +from contextlib import asynccontextmanager +from fastapi import FastAPI, Request +import os +import pickle +import subprocess +import tempfile +import atexit +import signal + + +ktransformer_rules_dir = ( + os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "..", "..", "./optimize/optimize_rules/") +) +default_optimize_rules = { + # "DeepseekV3ForCausalLM": ktransformer_rules_dir + "Moonlight-16B-A3B-serve.yaml", + "DeepseekV3ForCausalLM": ktransformer_rules_dir + "DeepSeek-V3-Chat-serve.yaml", + "Qwen2MoeForCausalLM": ktransformer_rules_dir + "Qwen2-serve.yaml", + "Qwen3MoeForCausalLM": ktransformer_rules_dir + "Qwen3Moe-serve.yaml", +} + + +async def chat_stream(queue: asyncio.Queue, tokenizer: AutoTokenizer): + streamer = TextStreamer(tokenizer) + while True: + token = await queue.get() + #print(f"Got token: {token}") + if token is None: + # str = f'{token}\n\n' + # str = model.tokenizer.decode(token) + s = streamer.end() + if s is not None: + yield s + break + + # str = model.tokenizer.decode(token) + yield streamer.put(token) + + + +def fill_generated_tokens(query_updates: list[sched_ext.QueryUpdate], generated_tokens: torch.Tensor, query_manager: QueryManager = None): + #print(len(query_updates), generated_tokens.size(0), generated_tokens) + for i in range(generated_tokens.size(0)): + print(generated_tokens[i].item()) + query_updates[i].generated_token = generated_tokens[i].item() + if not query_manager.query_map[query_updates[i].id].is_prefill: + pos = query_updates[i].active_position + if pos < query_manager.query_map[query_updates[i].id].max_length: + query_manager.query_map[query_updates[i].id].query_tokens[pos] = generated_tokens[i] + +def report_last_time_performance(profiler: Profiler): + try: + tokenize_time = profiler.get_timer_sec('tokenize') + prefill_time = profiler.get_timer_sec('prefill') + decode_time = profiler.get_timer_sec('decode') + prefill_count = profiler.get_counter('prefill') + decode_count = profiler.get_counter('decode') + + logger.info(f'Performance(T/s): prefill {prefill_count/prefill_time}, decode {decode_count/decode_time}. Time(s): tokenize {tokenize_time}, prefill {prefill_time}, decode {decode_time}') + except: + logger.info(f'Performance statistics not recorded') + +class Engine: + sched_client : SchedulerClient + updates : list[sched_ext.QueryUpdate] + batch : sched_ext.BatchQueryTodo + model_runner: ModelRunner + sampler: Sampler + query_manager: QueryManager + cache: KDeepSeekV3Cache | KGQACache + def __init__(self, args: ConfigArgs = default_args, generated_token_queue:Queue = None, broadcast_endpoint: str = None, kvcache_event: Event = None): + self.args = args + + for key, value in vars(args).items(): + if value is not None and hasattr(Config(), key): + setattr(Config(), key, value) + + self.device = self.args.device + self.sched_client = SchedulerClient(args.sched_port) + self.updates = [] + + try: + config = AutoConfig.from_pretrained(args.model_dir, trust_remote_code=True) + except: + if args.model_name == "Qwen3Moe": + config = Qwen3MoeConfig.from_pretrained(args.model_dir, trust_remote_code=True) + else: + assert False, f"model {args.model_name} not supported" + + + self.gen_queue = generated_token_queue + + with torch.device("meta"): + if config.architectures[0] == "DeepseekV3ForCausalLM": + self.cache = KDeepSeekV3Cache(config, self.args.page_size) + self.model = KDeepseekV3ForCausalLM(config, self.cache) + elif config.architectures[0] == "DeepseekV2ForCausalLM": + self.cache = KDeepSeekV3Cache(config, self.args.page_size) + self.model = KDeepseekV2ForCausalLM(config, self.cache) + elif config.architectures[0] == "Qwen2MoeForCausalLM" or config.architectures[0] == "Qwen3MoeForCausalLM": + self.cache = KGQACache(config, self.args.page_size) + if config.architectures[0] == "Qwen2MoeForCausalLM": + self.model = KQwen2MoeForCausalLM(config, self.cache) + else: + self.model = KQwen3MoeForCausalLM(config, self.cache) + + + context = zmq.Context() + + + self.pub_socket = context.socket(zmq.PUB) + self.pub_socket.bind(f"ipc://{broadcast_endpoint}") + # time.sleep(1) # make sure all subscribers are ready + + + try: + generation_config = GenerationConfig.from_pretrained(args.model_dir) + except: + generation_config = GenerationConfig( + max_length=args.max_new_tokens, + temperature=args.temperature, + top_p=args.top_p, + do_sample=True + ) + + if args.optimize_config_path is None: + optimize_config_path = default_optimize_rules[config.architectures[0]] + + else: + optimize_config_path = args.optimize_config_path + gguf_path = args.gguf_path + if gguf_path is None: + gguf_path = input( + "please input the path of your gguf file(gguf file in the dir containing input gguf file must all" + " belong to current model):" + ) + optimize_and_load_gguf(self.model, optimize_config_path, gguf_path, config) + self.model.generation_config = generation_config + if self.model.generation_config.pad_token_id is None: + self.model.generation_config.pad_token_id = self.model.generation_config.eos_token_id + + self.model.eval() + kvcache_event.set() + # load kvcache + print(f"Getting inference context from sched_client.") + inference_context = self.sched_client.get_inference_context_raw() + print(f"Got inference context, sending it to subscribers.") + inference_context = self.sched_client.rebuild_inferece_context(inference_context) + self.cache.load(inference_context) + print(f"kv_cache loaded successfully.") + + + self.block_num = inference_context.k_cache[0].size(1) + self.model_runner = ModelRunner(self.model, self.device, self.args.use_cuda_graph, page_size = args.page_size, block_num=self.block_num) + #@TODO add config + if config.architectures[0] == "Qwen2MoeForCausalLM" or config.architectures[0] == "Qwen3MoeForCausalLM": + self.model.init_wrapper(self.args.use_cuda_graph, self.device, max(self.model_runner.cuda_graphs), args.max_batch_size, self.block_num) + else: + self.model.init_wrapper(self.args.use_cuda_graph, self.device, args.max_batch_size, self.block_num) + + self.sampler = Sampler() + self.query_manager = QueryManager(device = self.device, page_size = args.page_size) + + + def sampling(self, forward_output: ForwardBatchOutput): + generated_tokens = torch.empty(0, device=self.device, dtype=torch.int32) + for i in range(forward_output.num_batchs): + logit = forward_output.logits[i] + if hasattr(forward_output, "temperatures"): + temperatures = forward_output.temperatures[i] + else: + temperatures = None + + if hasattr(forward_output, "top_ps"): + top_ps = forward_output.top_ps[i] + else: + top_ps = None + + sample_options = SamplingOptions(logit.size(0), self.device, pretrained_config=self.model.generation_config, temperatures=temperatures, top_ps=top_ps) + generated_tokens, probs=self.sampler(logit, sample_options) + return generated_tokens, probs + + def loop(self): + + next_batch = None + + while True: + self.batch = next_batch + if self.batch is not None: + self.model_runner.run(self.batch, self.query_manager) + + if len(self.updates) > 0: + for q in self.updates: + if q.is_prefill == True: + continue + # print(f"Putting token {q.generated_token} into queue for query id: {q.id}") + try: + self.gen_queue.put((q.id, q.generated_token if q.decode_done == False else None), timeout=5) + except queue.Full: + pass#print("Queue is full after timeout; unable to put more items.") + + next_batch = self.sched_client.update_last_batch(self.updates) + if next_batch.query_ids == []: + next_batch = None + self.pub_socket.send_pyobj(next_batch) + + if next_batch is not None: + self.query_manager.add_query(next_batch) + + + if self.batch is not None: + self.model_runner.sync() + print(f"Model execution time (GPU): {self.model_runner.model_time:.3f} ms, {1000/self.model_runner.model_time:.3f} tokens/s") + # if self.rank == 0: + + generated_tokens, probs = self.sampling( self.model_runner.output) + + self.updates = self.query_manager.update(self.batch) + fill_generated_tokens(self.updates, generated_tokens, self.query_manager) + else: + self.updates = [] + +class BalanceServeThreadContext(ThreadContext): + def get_local_messages(self): + local_messages = [] + for m in self.messages: + local_messages.append({"role": m.role.value, "content": m.get_text_content()}) + + return local_messages + + +def run_engine(args, token_queue, broadcast_endpoint, event, kvcache_event): + engine = Engine(args, token_queue, broadcast_endpoint, kvcache_event) + if args.use_cuda_graph: + engine.model_runner.warmup() + + event.set() + engine.loop() + + +class BalanceServeInterface(BackendInterfaceBase): + use_static_cache: bool = True + + model: Any + tokenizer: AutoTokenizer + + cache: StaticCache + generated_ids: torch.Tensor + seq_length: int + + streamer: TextStreamer + + # thread_related + last_request_id: Optional[str] = None + ever_generated_ids: Set[int] = set() + + def __init__(self, args: ConfigArgs = default_args): + self.args = args + self.queue_map:dict[int,asyncio.Queue] = {} + self.thread_map: dict[int, int] = {} + processes = [] + self.broadcast_endpoint = tempfile.NamedTemporaryFile(delete=False).name # @TODO add to config + ctx = mp.get_context("spawn") + self.token_queue = ctx.Queue(maxsize=1000) + self.tokenizer = AutoTokenizer.from_pretrained(args.model_dir, trust_remote_code=True) + self.sched_client = SchedulerClient(args.sched_port) + self.streamer = TextStreamer(self.tokenizer) + + start_event = ctx.Event() + kvcache_event = ctx.Event() + + p = ctx.Process(target=run_engine, args=(self.args, self.token_queue, self.broadcast_endpoint, start_event, kvcache_event)) + p.start() + processes.append(p) + kvcache_event.wait() + + + with tempfile.NamedTemporaryFile(delete=False) as temp_file: + pickle.dump(args, temp_file) + temp_file_path = temp_file.name + current_file = __file__ + target_file = os.path.join(os.path.dirname(current_file), "..", "..", "balance_serve", "sched_rpc.py") + target_file = os.path.normpath(target_file) + log_path = os.path.join(args.log_dir, "rpc.log") + log = open(log_path, "a") + sched_process = subprocess.Popen( + ["python3", target_file, "--config", temp_file_path], + stdout=log, + stderr=log + ) + print("sched_rpc started with PID:", sched_process.pid) + + def signal_handler(signum, frame): + print(f"Received signal {signum}, shutting down...") + cleanup() + os._exit(0) + + def cleanup(): + print("Cleaning up...") + + for p in processes: + if p.is_alive(): + print(f"Terminating subprocess {p.pid}") + p.terminate() + p.join() + + if sched_process and sched_process.poll() is None: + print(f"Terminating sched_process {sched_process.pid}") + sched_process.terminate() + sched_process.wait() + signal.signal(signal.SIGINT, signal_handler) + signal.signal(signal.SIGTERM, signal_handler) + + start_event.wait() + + def get_params(self, temperature: Optional[float] = None, top_p: Optional[float] = None, + max_tokens: Optional[float] = None, max_completion_tokens: Optional[float] = None) -> tuple[float, float]: + """Get sampling parameters and handle default values and edge cases""" + if max_tokens is not None: + max_completion_tokens = max_tokens + if max_completion_tokens is None: + max_completion_tokens = self.args.max_new_tokens + else: + max_completion_tokens = min(self.args.max_new_tokens, max_completion_tokens) + if temperature is None: + temperature = self.args.temperature + if top_p is None: + top_p = self.args.top_p + + if temperature == 0: + temperature = 0.0001 + if top_p == 0: + top_p = 0.0001 + + return temperature, top_p, max_completion_tokens + + def run_queue_proxy(self): + loop = asyncio.new_event_loop() + asyncio.set_event_loop(loop) + loop.run_until_complete(self.queue_proxy()) + + @asynccontextmanager + async def lifespan(self, app: FastAPI): + asyncio.create_task(self.queue_proxy()) + yield + + async def queue_proxy(self): + print("Queue Proxy Started") + while True: + try: + query_id, token = self.token_queue.get_nowait() + try: + # query id might not be allocated yet + self.queue_map[query_id].put_nowait(token) + #print(f"Proxy Put token: {token} to queue for query id: {query_id}") + except asyncio.QueueFull: + #print(f"Queue for query id: {query_id} is full, waiting to put: {token}") + await self.queue_map[query_id].put(token) + + except queue.Empty: + # print("no new token") + # await asyncio.sleep(1) + await asyncio.sleep(0) + def tokenize_prompt(self, prompt: str): + input_ids = self.tokenizer.encode(prompt, return_tensors="pt").to(self.args.device) + return input_ids + + def format_and_tokenize_input_ids(self, thread_id: ObjectID, messages: List): + input_str: str = self.tokenizer.apply_chat_template(messages,tokenize=False,add_generation_prompt=True) + # drop token in chat template + if input_str.endswith('\n'): + input_str = input_str[:-len('\n')] + input_ids = self.tokenizer.encode(input_str, return_tensors="pt", add_special_tokens=False).to(self.args.device) + logger.debug(f"get input ids of shape {input_ids.shape}") + return input_ids + + async def inference(self, local_messages, thread_id: str, temperature: Optional[float] = None, top_p: Optional[float] = None, + max_tokens: Optional[float] = None, max_completion_tokens: Optional[float] = None): + profiler = Profiler() + profiler.create_and_start_timer("tokenize") + + if isinstance(local_messages, List): + input_ids = self.format_and_tokenize_input_ids(thread_id, local_messages) + elif isinstance(local_messages, str): + input_ids = self.tokenize_prompt(local_messages) + else: + raise ValueError("local_messages should be List or str") + if Config().user_force_think: + token_thinks = torch.tensor([self.tokenizer.encode("\n",add_special_tokens=False)],device=input_ids.device) + input_ids = torch.cat( + [input_ids, token_thinks], dim=1 + ) + + profiler.pause_timer("tokenize") + + profiler.create_and_start_timer("prefill") + + query_add = sched_ext.QueryAdd() + query_add.query_token = input_ids[0].tolist() + query_length = input_ids[0].shape[0] + query_add.query_length = query_length + profiler.set_counter("prefill", query_length) + #@TODO add server + stop_criteria = [self.tokenizer.encode(self.tokenizer.eos_token, add_special_tokens=False),self.tokenizer.encode("<|im_end|>")] + query_add.stop_criteria = stop_criteria + + temperature, top_p, max_new_tokens = self.get_params(temperature, top_p, max_tokens, max_completion_tokens) + + query_add.sample_options.temperature = temperature + query_add.sample_options.top_p = top_p + query_add.estimated_length = min(self.args.cache_lens, query_length+max_new_tokens) + + if query_add.estimated_length < query_add.query_length: + raise Exception(f'query too long: estimated_length={query_add.estimated_length} < query_length={query_add.query_length}') + + query_id = self.sched_client.add_query(query_add) + queue = asyncio.Queue(maxsize=max_new_tokens) + self.queue_map[query_id] = queue + self.thread_map[thread_id] = query_id + is_first_token = True + async for token in chat_stream(self.queue_map[query_id], self.tokenizer): + if is_first_token: + is_first_token=False + profiler.pause_timer("prefill") + profiler.create_and_start_timer("decode") + profiler.set_counter("decode", 0) + if Config().user_force_think: + think = '\n' + print(think, end="",flush=True) + yield think, None + else: + profiler.inc("decode") + yield token, None + profiler.pause_timer("decode") + report_last_time_performance(profiler) + yield self.streamer.end(), None + if profiler.get_counter('decode') >= max_new_tokens - 1: + yield "", "length" + else: + yield "", "stop" + + + yield RawUsage( + tokenize_time = profiler.get_timer_sec('tokenize'), + prefill_time = profiler.get_timer_sec('prefill'), + decode_time = profiler.get_timer_sec('decode'), + prefill_count = profiler.get_counter('prefill'), + decode_count = profiler.get_counter('decode'), + ) diff --git a/archive/kt-sft/ktransformers/server/backend/interfaces/exllamav2.py b/archive/kt-sft/ktransformers/server/backend/interfaces/exllamav2.py new file mode 100644 index 0000000..6385264 --- /dev/null +++ b/archive/kt-sft/ktransformers/server/backend/interfaces/exllamav2.py @@ -0,0 +1,40 @@ +import sys, os +from typing import AsyncIterator, Dict, Tuple + +import torch + +from ..args import ConfigArgs, default_args + +from ..base import BackendInterfaceBase, ThreadContext +from ktransformers.server.schemas.assistants.runs import RunObject + + +from ..args import * + +class ExllamaThreadContext(ThreadContext): + def __init__(self, run: RunObject, args: ConfigArgs = default_args) -> None: + super().__init__(run,args) + + def get_interface(self): + return + + def get_local_messages(self): + raise NotImplementedError + + + + +class ExllamaInterface(BackendInterfaceBase): + + def __init__(self, args: ConfigArgs = ...): + raise NotImplementedError + + def tokenize_prompt(self, prompt: str) -> torch.Tensor: + raise NotImplementedError + + async def inference(self,local_messages,request_unique_id:Optional[str])->AsyncIterator: + raise NotImplementedError + + + + diff --git a/archive/kt-sft/ktransformers/server/backend/interfaces/ktransformers.py b/archive/kt-sft/ktransformers/server/backend/interfaces/ktransformers.py new file mode 100644 index 0000000..7cf6b91 --- /dev/null +++ b/archive/kt-sft/ktransformers/server/backend/interfaces/ktransformers.py @@ -0,0 +1,252 @@ +import torch +from typing import Optional, List +import asyncio +from transformers import AutoTokenizer, AutoConfig, GenerationConfig +from ktransformers.server.backend.interfaces.transformers import ( + TransformersInterface, + ConfigArgs, + TransformersThreadContext, + default_args, + TextStreamer, +) +from ktransformers.server.config.log import logger +from ktransformers.optimize.optimize import optimize_and_load_gguf +from ktransformers.models.custom_cache import StaticCache +from ktransformers.util.cuda_graph_runner import CUDAGraphRunner +from ktransformers.local_chat import custom_models, default_optimize_rules +from ktransformers.util.utils import get_device +from typing import Optional +from ktransformers.operators.flashinfer_wrapper import flashinfer_enabled, MLAWrapperSingleton +from ktransformers.server.schemas.endpoints.chat import RawUsage +from ktransformers.util.grad_wrapper import maybe_no_grad + +warm_uped = False + +class KTransformersThreadContext(TransformersThreadContext): + pass + + +class KTransformersInterface(TransformersInterface): + def __init__(self, args: ConfigArgs = default_args): + self.args = args + torch.set_grad_enabled(False) + self.tokenizer = AutoTokenizer.from_pretrained(args.model_dir, device=args.device, trust_remote_code=args.trust_remote_code) + config = AutoConfig.from_pretrained(args.model_dir, trust_remote_code=args.trust_remote_code) + try: + generation_config = GenerationConfig.from_pretrained(args.model_dir) + except: + generation_config = GenerationConfig( + max_length=args.max_new_tokens, + temperature=args.temperature, + top_p=args.top_p, + do_sample=True + ) + + torch.set_default_dtype(config.torch_dtype) + if config.architectures[0] == "Qwen2MoeForCausalLM": + config._attn_implementation = "flash_attention_2" + + with torch.device("meta"): + self.model = custom_models[config.architectures[0]](config) + if default_args.optimize_config_path is None: + optimize_config_path = default_optimize_rules[config.architectures[0]] + else: + optimize_config_path = args.optimize_config_path + + # print(optimize_config) + + gguf_path = args.gguf_path + if gguf_path is None: + gguf_path = input( + "please input the path of your gguf file(gguf file in the dir containing input gguf file must all" + " belong to current model):" + ) + optimize_and_load_gguf(self.model, optimize_config_path, gguf_path, config) + self.model.generation_config = generation_config + self.device_map = self.model.gguf_loader.tensor_device_map + # logger.info(f"{args.model_name} loaded from {args.model_dir} to {self.device_map}") + self.cache = StaticCache( + config=self.model.config, + max_batch_size=args.batch_size, + max_cache_len=args.cache_lens, + device=self.device_map, + dtype=self.model.dtype, + ) + # logger.info(f"StaticCache (length={args.cache_lens}), batch size:{args.batch_size}") + + if self.model.generation_config.pad_token_id is None: + self.model.generation_config.pad_token_id = self.model.generation_config.eos_token_id + self.streamer = TextStreamer(self.tokenizer) + + self._infer_lock = asyncio.Lock() + + def decode_one_tokens(self): + global warm_uped + + device_map = self.model.gguf_loader.tensor_device_map + torch_device = get_device("blk.0.self_attn", device_map) + torch_device = "cuda:0" if torch_device == "cuda" else torch_device + torch.cuda.set_device(torch_device) + if warm_uped and self.args.use_cuda_graph: + if not hasattr(self, "cuda_graph_runner"): + self.cuda_graph_runner = CUDAGraphRunner() + self.cuda_graph_runner.capture( + self.model, + self.current_ids, + self.active_cache_position.unsqueeze(0), + self.active_cache_position, + self.cache, + main_device=torch_device, + return_dict=False, + use_cache=True, + ) + + if hasattr(self, "cuda_graph_runner"): + logits = self.cuda_graph_runner( + self.current_ids, self.active_cache_position.unsqueeze(0), self.active_cache_position + ) + self.cache.change_seq_length(1) + torch.cuda.synchronize() + logits = logits[0, -1, :] + return self.logits_to_token(logits) + + if self.args.use_cuda_graph: + warm_uped = True + + if self.use_static_cache: + logits = self.model( + self.current_ids.to(torch_device), + cache_position=self.active_cache_position, + past_key_values=self.cache, + return_dict=False, + use_cache=True, + )[0] + else: + logits = self.model(self.current_ids, return_dict=False)[0] + logits = logits[0, -1, :] + + return self.logits_to_token(logits) + + + + @maybe_no_grad + def prefill(self, input_ids: torch.Tensor, is_new: bool, temperature: Optional[float] = None, top_p: Optional[float] = None, max_tokens: Optional[float] = None, max_completion_tokens: Optional[float] = None): + input_ids_length = input_ids.shape[-1] + if max_tokens is not None: + max_completion_tokens = max_tokens + if max_completion_tokens is None: + max_new_tokens = self.args.max_new_tokens + else: + max_new_tokens = min(self.args.max_new_tokens, max_completion_tokens) + if(input_ids_length >= self.args.cache_lens): + logger.warning(f"input_ids_length {input_ids_length} > cache_lens {self.args.cache_lens}") + self.seq_length = input_ids_length + return + logger.debug(f"input_ids: {input_ids.shape}") + device = self.device_map.get("blk.0.self_attn", {}).get("generate_device", "cuda:0") + device = "cuda:0" if device == "cuda" else device + + if is_new: + self.ever_generated_ids.clear() + same_prefix = 0 + flat_input_ids = input_ids.flatten() + + if getattr(self, 'generated_ids', None) is None: + self.generated_ids = torch.zeros( + self.args.batch_size, + input_ids.shape[-1] + max_new_tokens + 1, + dtype=torch.int, + device=self.args.device, + ) + self.seq_length = 1 + + flat_prev_ids = self.generated_ids.flatten() + for i in range(min(self.seq_length, flat_input_ids.shape[0]) - 1): + if flat_input_ids[i] == flat_prev_ids[i]: + same_prefix += 1 + else: + break + + logger.debug(f"same prefix len: {same_prefix}") + self.cache.remove_suffix(same_prefix) + self.seq_length = same_prefix + self.generated_ids = self.generated_ids[..., :same_prefix] + input_ids = input_ids[..., same_prefix:] + input_ids_length = input_ids.shape[-1] + + self.ever_generated_ids.clear() + self.profiler.set_counter("prefill", input_ids_length) + logger.debug(f"input_ids: {input_ids.shape}") + logger.debug(f"generate_ids: {self.generated_ids.shape}") + + former_seq_length = self.seq_length + self.seq_length += input_ids_length + expected_length = min(self.seq_length + max_new_tokens + 1, self.args.cache_lens) + delta_length = expected_length - self.generated_ids.shape[-1] + if delta_length > 0: + new_generate_ids = torch.zeros( + self.args.batch_size, delta_length, dtype=torch.int, device=self.args.device + ) + self.generated_ids = torch.cat([self.generated_ids, new_generate_ids], dim=-1) + else: + logger.warning(f"seq_length bigger than cache_lens, killed") + exit(0) + + logger.debug(f"cache position: {former_seq_length} to {self.seq_length}") + cache_position = torch.arange(former_seq_length, self.seq_length, device=device) + self.generated_ids[:, cache_position] = input_ids.to(self.args.device).to(torch.int) + + if not (type(self) is TransformersInterface): + input_ids = input_ids.to("cpu") + + def chunk_prefill(input_ids, cache_position): + inputs_embeds = self.model.model.embed_tokens(input_ids).to(device) + torch.cuda.set_device(device) + if flashinfer_enabled: + MLAWrapperSingleton.need_plan_all() + if self.use_static_cache: + logits = self.model( + inputs_embeds=inputs_embeds, + cache_position=cache_position, + past_key_values=self.cache, + return_dict=False, + use_cache=True, + )[0] + else: + logits = self.model(inputs_embeds=inputs_embeds, return_dict=False)[0] + + return logits + + chunk_start = 0 + while chunk_start < input_ids_length: + chunk_end = min(chunk_start + self.args.chunk_size, input_ids_length) + if self.cache != None: + self.cache.cur_idx=cache_position[chunk_start:chunk_end] + logits = chunk_prefill(input_ids[:, chunk_start:chunk_end], cache_position[chunk_start:chunk_end]) + chunk_start += self.args.chunk_size + + if flashinfer_enabled: + MLAWrapperSingleton.reset_buffer() + self.prepare_logits_wrapper(input_ids, device, temperature, top_p) + next_token = self.logits_to_token(logits[0, -1, :]) + self.max_new_tokens = min(max_new_tokens, self.args.cache_lens - self.seq_length) - 1 + yield self.append_new_tokens(next_token) + + @property + def active_cache_position(self): + device = self.device_map.get("blk.0.self_attn", {}).get("generate_device", "cuda:0") + return torch.tensor([self.seq_length - 1], device=device) + + async def inference(self, local_messages, thread_id: str, temperature: Optional[float] = None, top_p: Optional[float] = None, max_tokens: Optional[float] = None, max_completion_tokens: Optional[float] = None): + async with self._infer_lock: + async for v in super().inference(local_messages, thread_id, temperature, top_p, max_tokens, max_completion_tokens): + yield v + + # return this inference raw usage + yield RawUsage( + tokenize_time = self.profiler.get_timer_sec('tokenize'), + prefill_time = self.profiler.get_timer_sec('prefill'), + decode_time = self.profiler.get_timer_sec('decode'), + prefill_count = self.profiler.get_counter('prefill'), + decode_count = self.profiler.get_counter('decode'), + ) \ No newline at end of file diff --git a/archive/kt-sft/ktransformers/server/backend/interfaces/transformers.py b/archive/kt-sft/ktransformers/server/backend/interfaces/transformers.py new file mode 100644 index 0000000..aaa7a51 --- /dev/null +++ b/archive/kt-sft/ktransformers/server/backend/interfaces/transformers.py @@ -0,0 +1,481 @@ +from typing import Any, List, Optional, Set +import re +import json +import uuid +from transformers import ( + LlamaTokenizer, + AutoTokenizer, + AutoConfig, + LlamaForCausalLM, + GenerationConfig, + StaticCache, + AutoModelForCausalLM, + BitsAndBytesConfig, + LogitsProcessorList, + TemperatureLogitsWarper, + TopKLogitsWarper, + TopPLogitsWarper, + MinPLogitsWarper, + TypicalLogitsWarper, + EpsilonLogitsWarper, + EtaLogitsWarper, +) + +from ktransformers.server.config.config import Config +from ktransformers.server.schemas.base import ObjectID +from ktransformers.server.utils.multi_timer import Profiler +from torch.nn.attention import SDPBackend +import torch +import sys, os +from ..base import ThreadContext, BackendInterfaceBase +from ktransformers.server.config.log import logger +from ..args import ConfigArgs, default_args +from ktransformers.operators.flashinfer_wrapper import flashinfer_enabled, MLAWrapperSingleton +from ktransformers.util.grad_wrapper import maybe_no_grad + +# This TextStreamer is a modified version from https://github.com/huggingface/transformers/blob/main/src/transformers/generation/streamers.py +class TextStreamer: + + def __init__(self, tokenizer: "AutoTokenizer", skip_prompt: bool = False, **decode_kwargs): + self.tokenizer = tokenizer + self.skip_prompt = skip_prompt + self.decode_kwargs = decode_kwargs + + # variables used in the streaming process + self.token_cache = [] + self.print_len = 0 + self.next_tokens_are_prompt = True + + def reset(self): + self.token_cache = [] + self.print_len = 0 + + def put(self, value) -> Optional[str]: + """ + Receives tokens, decodes them, and prints them to stdout as soon as they form entire words. + """ + if not isinstance(value, int): + raise ValueError("TextStreamer only supports batch size 1, and int type input") + + if self.skip_prompt and self.next_tokens_are_prompt: + self.next_tokens_are_prompt = False + return None + + # Add the new token to the cache and decodes the entire thing. + self.token_cache.append(value) + text = self.tokenizer.decode(self.token_cache, skip_special_tokens=True, **self.decode_kwargs) + + # After the symbol for a new line, we flush the cache. + if text.endswith("\n"): + printable_text = text[self.print_len :] + self.reset() + # If the last token is a CJK character, we print the characters. + elif len(text) > 0 and self._is_chinese_char(ord(text[-1])): + printable_text = text[self.print_len :] + self.print_len += len(printable_text) + # Otherwise, prints until the last space char (simple heuristic to avoid printing incomplete words, + # which may change with the subsequent token -- there are probably smarter ways to do this!) + else: + printable_text = text[self.print_len : text.rfind(" ") + 1] + self.print_len += len(printable_text) + return printable_text + + def end(self) -> Optional[str]: + """Flushes any remaining cache and prints a newline to stdout.""" + # Flush the cache, if it exists + if len(self.token_cache) > 0: + text = self.tokenizer.decode(self.token_cache, skip_special_tokens=True, **self.decode_kwargs) + printable_text = text[self.print_len :] + self.reset() + else: + printable_text = "" + + self.next_tokens_are_prompt = True + return printable_text + + def _is_chinese_char(self, cp): + """Checks whether CP is the codepoint of a CJK character.""" + # This defines a "chinese character" as anything in the CJK Unicode block: + # https://en.wikipedia.org/wiki/CJK_Unified_Ideographs_(Unicode_block) + # + # Note that the CJK Unicode block is NOT all Japanese and Korean characters, + # despite its name. The modern Korean Hangul alphabet is a different block, + # as is Japanese Hiragana and Katakana. Those alphabets are used to write + # space-separated words, so they are not treated specially and handled + # like the all of the other languages. + if ( + (cp >= 0x4E00 and cp <= 0x9FFF) + or (cp >= 0x3400 and cp <= 0x4DBF) # + or (cp >= 0x20000 and cp <= 0x2A6DF) # + or (cp >= 0x2A700 and cp <= 0x2B73F) # + or (cp >= 0x2B740 and cp <= 0x2B81F) # + or (cp >= 0x2B820 and cp <= 0x2CEAF) # + or (cp >= 0xF900 and cp <= 0xFAFF) + or (cp >= 0x2F800 and cp <= 0x2FA1F) # + ): # + return True + + return False + + +class TransformersThreadContext(ThreadContext): + def get_local_messages(self): + local_messages = [] + for m in self.messages: + local_messages.append({"role": m.role.value, "content": m.get_text_content()}) + + return local_messages + + +class TransformersInterface(BackendInterfaceBase): + use_static_cache: bool = True + + model: Any + tokenizer: AutoTokenizer + + cache: StaticCache + generated_ids: torch.Tensor + seq_length: int + + streamer: TextStreamer + + # thread_related + last_request_id: Optional[str] = None + ever_generated_ids: Set[int] = set() + + def __init__(self, args: ConfigArgs = default_args): + self.args = args + + self.tokenizer = AutoTokenizer.from_pretrained(args.model_dir) + self.model = AutoModelForCausalLM.from_pretrained(args.model_dir, device_map=args.device, use_safetensors=True) + # logger.info(f"{args.model_name} loaded from {args.model_dir} to {args.device}") + + self.cache = StaticCache( + config=self.model.config, + max_batch_size=args.batch_size, + max_cache_len=args.cache_lens, + device=args.device, + dtype=self.model.dtype, + ) + # logger.info(f"StaticCache (length={args.cache_lens}) created at {args.device}, batch size:{args.batch_size}") + + self.streamer = TextStreamer(self.tokenizer) + + @property + def current_ids(self): + return self.generated_ids[:, self.seq_length - 1].unsqueeze(1) + + @property + def active_cache_position(self): + return torch.tensor([self.seq_length - 1], device=self.args.device) + + def tokenize_prompt(self, prompt: str): + input_ids = self.tokenizer.encode(prompt, return_tensors="pt").to(self.args.device) + return input_ids + + def format_and_tokenize_input_ids(self, thread_id: ObjectID, messages: List): + for m in messages: + if m["role"] == "system": + logger.warning(f'change {m["role"]} to user') + m["role"] = "user" + + new_messages = [messages[0]] + for m in messages[1:]: + if m["role"] == "user" and new_messages[-1]["role"] == "user": + logger.warning("merge two adjacent user messages") + new_messages[-1]["content"] += '\n' + m["content"] + else: + new_messages.append(m) + # if (self.last_request_id is not None) and self.last_request_id == thread_id: + # input_ids = self.tokenizer.encode(self.tokenizer.eos_token+self.tokenizer.apply_chat_template([new_messages[-1]], return_tensors="pt",tokenize=False, add_generation_prompt=True), add_special_tokens = False, return_tensors="pt").to(self.args.device) + # else: + # input_ids = self.tokenizer.apply_chat_template( + # new_messages, return_tensors="pt", add_generation_prompt=True + # ).to(self.args.device) + input_str: str = self.tokenizer.apply_chat_template(new_messages,tokenize=False,add_generation_prompt=True) + # drop token in chat template + if input_str.endswith('\n'): + input_str = input_str[:-len('\n')] + input_ids = self.tokenizer.encode(input_str, return_tensors="pt").to(self.args.device) + if (self.last_request_id is not None) and self.last_request_id == thread_id: + x = self.generated_ids[:,:self.seq_length] + y = input_ids[:,:self.seq_length] + # We can only hope that the input_ids are the same + unequal_mask = torch.ne(x,y) + unequal_positions = torch.nonzero(unequal_mask) + num_unequal_elements = unequal_mask.sum().item() + logger.warning(f'num_unequal_elements: {num_unequal_elements}') + + input_ids = input_ids[:,self.seq_length:] + logger.debug(f"get input ids of shape {input_ids.shape}") + return input_ids + + def append_new_tokens(self, new_tokens: int) -> Optional[str]: + self.generated_ids[0, self.seq_length] = new_tokens + self.seq_length += 1 + return self.streamer.put(new_tokens) + + @staticmethod + def tf_logits_warper(generation_config): + """ + This class returns a [`LogitsProcessorList`] list object that contains all relevant [`LogitsWarper`] instances + used for multinomial sampling. + """ + + # instantiate warpers list + warpers = LogitsProcessorList() + + # In beam methods, we need to keep at least one non-eos token to explore continuations that might have a + # better score (i.e. keep len(list(generation_config._eos_token_tensor)) + 1) + if generation_config.num_beams > 1: + if isinstance(generation_config._eos_token_tensor, list): + min_tokens_to_keep = len(generation_config._eos_token_tensor) + 1 + elif isinstance(generation_config._eos_token_tensor, torch.Tensor): + min_tokens_to_keep = generation_config._eos_token_tensor.shape[0] + 1 + else: + min_tokens_to_keep = 2 + else: + min_tokens_to_keep = 1 + + # the following idea is largely copied from this PR: https://github.com/huggingface/transformers/pull/5420/files + # all samplers can be found in `generation_utils_samplers.py` + if generation_config.temperature is not None and generation_config.temperature != 1.0: + warpers.append(TemperatureLogitsWarper(generation_config.temperature)) + if generation_config.top_k is not None and generation_config.top_k != 0: + warpers.append(TopKLogitsWarper(top_k=generation_config.top_k, min_tokens_to_keep=min_tokens_to_keep)) + if generation_config.top_p is not None and generation_config.top_p < 1.0: + warpers.append(TopPLogitsWarper(top_p=generation_config.top_p, min_tokens_to_keep=min_tokens_to_keep)) + if generation_config.min_p is not None: + # Applied after temperature scaling (see https://github.com/ggerganov/llama.cpp/pull/3841#issuecomment-2073826084) + warpers.append(MinPLogitsWarper(min_p=generation_config.min_p, min_tokens_to_keep=min_tokens_to_keep)) + if generation_config.typical_p is not None and generation_config.typical_p < 1.0: + warpers.append( + TypicalLogitsWarper(mass=generation_config.typical_p, min_tokens_to_keep=min_tokens_to_keep) + ) + if generation_config.epsilon_cutoff is not None and 0.0 < generation_config.epsilon_cutoff < 1.0: + warpers.append( + EpsilonLogitsWarper(epsilon=generation_config.epsilon_cutoff, min_tokens_to_keep=min_tokens_to_keep) + ) + if generation_config.eta_cutoff is not None and 0.0 < generation_config.eta_cutoff < 1.0: + warpers.append( + EtaLogitsWarper( + epsilon=generation_config.eta_cutoff, min_tokens_to_keep=min_tokens_to_keep, device=device + ) + ) + # `LogitNormalization` should always be the last logit processor, when present + if generation_config.renormalize_logits is True: + warpers.append(LogitNormalization()) + return warpers + + def prepare_logits_wrapper(self, inputs, device, temperature: Optional[float] = None, top_p: Optional[float] = None): + if temperature is None or temperature == 0: + temperature = self.model.generation_config.temperature + if top_p is None: + top_p = self.model.generation_config.top_p + if top_p == 0: + top_p = 0.0001 + generation_config, model_kwargs = self.model._prepare_generation_config( + None, max_length=self.args.max_new_tokens, + do_sample=True, + top_k=self.args.top_k, + top_p=top_p, + temperature=temperature, + repetition_penalty=self.args.repetition_penalty # change this to modify generate config + ) + self.inputs = inputs + + self.logits_warper = self.tf_logits_warper(generation_config) + + def logits_to_token(self, logits: torch.Tensor): + logits = self.logits_warper(self.inputs.view(1, -1), logits.view(1, -1)) + + probs = torch.nn.functional.softmax(logits, dim=-1) + + sample = True + if sample: + last = torch.multinomial(probs, num_samples=1) + else: + _, last = torch.topk(probs, k=1, dim=-1) + + last = last.item() + self.ever_generated_ids.add(last) + return last + + def decode_one_tokens(self): + if self.use_static_cache: + logits = self.model( + self.current_ids, + cache_position=self.active_cache_position, + past_key_values=self.cache, + return_dict=False, + use_cache=True, + )[0] + else: + logits = self.model(self.current_ids, return_dict=False)[0] + logits = logits[0, -1, :] + + return self.logits_to_token(logits) + + @maybe_no_grad + def prefill(self, input_ids: torch.Tensor, is_new: bool, temperature: Optional[float] = None, top_p: Optional[float] = None, max_tokens: Optional[float] = None, max_completion_tokens: Optional[float] = None): + input_ids_length = input_ids.shape[-1] + logger.debug(f"input_ids: {input_ids.shape}") + if max_tokens is not None: + max_completion_tokens = max_tokens + if max_completion_tokens is None: + max_new_tokens = self.args.max_new_tokens + else: + max_new_tokens = min(self.args.max_new_tokens, max_completion_tokens) + if is_new: + self.ever_generated_ids.clear() + same_prefix = 0 + flat_input_ids = input_ids.flatten() + + if getattr(self, 'generated_ids', None) is None: + self.generated_ids = torch.zeros( + self.args.batch_size, + input_ids.shape[-1] + max_new_tokens + 1, + dtype=torch.int, + device=self.args.device, + ) + self.seq_length = 1 + + flat_prev_ids = self.generated_ids.flatten() + for i in range(min(self.seq_length, flat_input_ids.shape[0]) - 1): + if flat_input_ids[i] == flat_prev_ids[i]: + same_prefix += 1 + else: + break + + logger.debug(f"same prefix len: {same_prefix}") + self.cache.remove_suffix(same_prefix) + self.seq_length = same_prefix + self.generated_ids = self.generated_ids[..., :same_prefix] + input_ids = input_ids[..., same_prefix:] + input_ids_length = input_ids.shape[-1] + + self.ever_generated_ids.clear() + self.profiler.set_counter("prefill", input_ids_length) + logger.debug(f"input_ids: {input_ids.shape}") + + logger.debug(f"generate_ids: {self.generated_ids.shape}") + former_seq_length = self.seq_length + self.seq_length += input_ids_length + expected_length = self.seq_length + max_new_tokens + 1 + delta_length = expected_length - self.generated_ids.shape[-1] + if delta_length > 0: + new_generate_ids = torch.zeros( + self.args.batch_size, delta_length, dtype=torch.int, device=self.args.device + ) + self.generated_ids = torch.cat([self.generated_ids, new_generate_ids], dim=-1) + + logger.debug(f"cache position: {former_seq_length} to {self.seq_length}") + cache_position = torch.arange(former_seq_length, self.seq_length, device=self.args.device) + self.generated_ids[:, cache_position] = input_ids.to(self.args.device).to(torch.int) + + device = input_ids.device + if not (type(self) is TransformersInterface): + input_ids = input_ids.to("cpu") + inputs_embeds = self.model.model.embed_tokens(input_ids).to(device) + if self.use_static_cache: + logits = self.model( + inputs_embeds=inputs_embeds, + cache_position=cache_position, + past_key_values=self.cache, + return_dict=False, + use_cache=True, + )[0] + else: + logits = self.model(inputs_embeds=inputs_embeds, return_dict=False)[0] + + self.prepare_logits_wrapper(input_ids, device, temperature, top_p) + next_token = self.logits_to_token(logits[0, -1, :]) + self.max_new_tokens = min(max_new_tokens, self.args.cache_lens - self.seq_length) - 1 + yield self.append_new_tokens(next_token) + + @maybe_no_grad + def generate(self): + logger.info(f"args.max_new_tokens: {self.args.max_new_tokens}, cache_lens: {self.args.cache_lens}, seq_length: {self.seq_length}") + if(self.max_new_tokens <= 0): + logger.warning("max_new_tokens is less than 0") + yield self.streamer.end(), "length" + return + self.profiler.set_counter("decode", 0) + + for i in range(1, self.max_new_tokens): + with torch.nn.attention.sdpa_kernel(backends=[SDPBackend.FLASH_ATTENTION, SDPBackend.MATH, SDPBackend.EFFICIENT_ATTENTION]): + if flashinfer_enabled: + MLAWrapperSingleton.plan_all(None,None,None,self.active_cache_position.to(torch.int32)+1, None, + num_heads=self.model.config.num_attention_heads, head_dim_ckv=self.model.config.kv_lora_rank, + head_dim_kpe=self.model.config.qk_rope_head_dim, page_size=self.cache.page_size, + sm_scale=self.model.model.layers[0].self_attn.softmax_scale, q_data_type=torch.bfloat16, kv_data_type=torch.bfloat16) + next_token = self.decode_one_tokens() + self.profiler.inc("decode") + if next_token == self.tokenizer.eos_token_id or "<|im_end|>" == self.tokenizer.decode(next_token): + yield self.streamer.end(), None + yield "", "stop" + assert self.args.batch_size == 1 + break + yield self.append_new_tokens(next_token), None + + else: # for's else, if output get max new tokens + yield self.streamer.end(), None + yield "", "length" + + + + def check_is_new(self, thread_id: str): + if not self.use_static_cache: + return True + if self.last_request_id is None: + self.last_request_id = thread_id + return True + else: + if self.last_request_id == thread_id: + return False + else: + self.last_request_id = thread_id + return True + + async def inference(self, local_messages, thread_id: str, temperature: Optional[float] = None, top_p: Optional[float] = None, max_tokens: Optional[float] = None, max_completion_tokens: Optional[float] = None): + self.streamer.reset() + self.profiler.create_and_start_timer("tokenize") + if isinstance(local_messages, List): + input_ids = self.format_and_tokenize_input_ids(thread_id, local_messages) + elif isinstance(local_messages, str): + #local_messages = local_messages[0]['content'] + input_ids = self.tokenize_prompt(local_messages) + #input_ids = torch.tensor([[6366]], device=input_ids.device) + else: + raise ValueError("local_messages should be List or str") + + if Config().user_force_think: + token_thinks = torch.tensor([self.tokenizer.encode("\n",add_special_tokens=False)],device=input_ids.device) + input_ids = torch.cat( + [input_ids, token_thinks], dim=1 + ) + + self.profiler.pause_timer("tokenize") + + self.profiler.create_and_start_timer("prefill") + + if Config().user_force_think: + think = '\n' + print(think, end="",flush=True) + yield think, None + + for t in self.prefill(input_ids, self.check_is_new(thread_id), temperature, top_p, max_tokens, max_completion_tokens): + # output think token after prefill done + if t is not None: + print(t, end="",flush=True) + yield t, None + self.profiler.pause_timer("prefill") + + self.profiler.create_and_start_timer("decode") + for t, finish_reason in self.generate(): + if t is not None: + print(t, end="",flush=True) + yield t, finish_reason + print("") + self.profiler.pause_timer("decode") + self.report_last_time_performance() diff --git a/archive/kt-sft/ktransformers/server/balance_serve/inference/__init__.py b/archive/kt-sft/ktransformers/server/balance_serve/inference/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/archive/kt-sft/ktransformers/server/balance_serve/inference/config.py b/archive/kt-sft/ktransformers/server/balance_serve/inference/config.py new file mode 100644 index 0000000..6140f24 --- /dev/null +++ b/archive/kt-sft/ktransformers/server/balance_serve/inference/config.py @@ -0,0 +1,142 @@ +''' +Date: 2024-11-07 07:30:16 +LastEditors: djw +LastEditTime: 2024-11-15 14:23:26 +''' +import math +from dataclasses import dataclass +from typing import Optional + +import torch +import torch.nn as nn +from torch import Tensor +from torch.nn import functional as F +import yaml + +import json +from typing import Optional + +class ModelConfig: + vocab_size: int = 32000 + n_layer: int = 1 + n_head: int = 32 + dim: int = 4096 + intermediate_size: int = 18944 + n_local_heads: int = 8 + head_dim: int = 128 + rope_base: float = 1000000.0 + norm_eps: float = 1e-06 + rope_scaling: Optional[dict] = None + rms_norm_eps: float = 1e-6 + hidden_act: str = "silu" + model_path: str + gguf_path: str + optimize_rule_path: str + speculative_rule_path: str + + + # quantize config + quant_algorithm: Optional[str] = None + quant_group_size: Optional[int] = None + quant_num_bits: Optional[int] = None + + json_key_map = { + "vocab_size": "vocab_size", + "n_layer": "num_hidden_layers", + "n_head": "num_attention_heads", + "dim": "hidden_size", + "intermediate_size": "intermediate_size", + "n_local_heads": "num_key_value_heads", + "rope_base": "rope_theta", + "norm_eps": "norm_eps", + "rms_norm_eps": "rms_norm_eps", + "hidden_act": "hidden_act", + } + + def __init__(self, config): + self.model_path = config["model"]["model_path"] + self.gguf_path = config["model"]["gguf_path"] + self.optimize_rule_path = config["model"]["optimize_rule_path"] + if "speculative_rule_path" in config["model"]: + self.speculative_rule_path = config["model"]["speculative_rule_path"] + self.speculative_gguf_path = config["model"]["speculative_gguf_path"] + self.speculative_model_path = config["model"]["speculative_model_path"] + self.quant_algorithm = config["model"]["quant"]["algorithm"] + self.quant_group_size = config["model"]["quant"]["group_size"] + self.quant_num_bits = config["model"]["quant"]["num_bits"] + self.load_config() + self.n_layer = config["model"]["n_layers"] + + def load_config(self): + config_file = f"{self.model_path}/config.json" + try: + with open(config_file, "r") as f: + config_data = json.load(f) + except FileNotFoundError: + raise FileNotFoundError(f"Configuration file not found at {config_file}") + + for attr, json_key in self.json_key_map.items(): + if json_key in config_data: + setattr(self, attr, config_data[json_key]) + else: + setattr(self, attr, getattr(self, attr)) + + + + + +class ParallelConfig: + def __init__( + self, + config, + ) -> None: + self.pipeline_parallel_size = config["parallel"]["pp"] + self.tensor_parallel_size = config["parallel"]["tp"] + self.disable_custom_all_reduce = config["parallel"]["disable_custom_all_reduce"] + self.world_size = self.pipeline_parallel_size * self.tensor_parallel_size + +class AttnConfig: + page_size: int = 256 + block_num: int = 32 + max_batch_token : int = 256 + max_batch_size: int = 32 + + def __init__(self, config): + self.page_size = config["attn"]["page_size"] + self.block_num = config["attn"]["block_num"] + self.max_batch_token = config["attn"]["max_batch_token"] + self.max_batch_size = config["attn"]["max_batch_size"] + + +class SamplerConfig(): + # Batched sampling params + temperatures: float + is_all_greedy: bool + + def __init__(self, config): + self.temperatures = config["sample"]["temperature"] + self.is_all_greedy = True + + +def load_yaml_config(file_path): + with open(file_path, "r") as f: + return yaml.safe_load(f) + + + + +class LLMConfig: + model_config: ModelConfig + parallel_config: ParallelConfig + attn_config: AttnConfig + sample_config: SamplerConfig + config_file: str + + def __init__(self, config_file): + self.config_file = config_file + config = load_yaml_config(config_file) + self.model_config = ModelConfig(config) + self.parallel_config = ParallelConfig(config) + self.attn_config = AttnConfig(config) + self.sample_config = SamplerConfig(config) + diff --git a/archive/kt-sft/ktransformers/server/balance_serve/inference/distributed/__init__.py b/archive/kt-sft/ktransformers/server/balance_serve/inference/distributed/__init__.py new file mode 100644 index 0000000..db325cf --- /dev/null +++ b/archive/kt-sft/ktransformers/server/balance_serve/inference/distributed/__init__.py @@ -0,0 +1,3 @@ +from .communication_op import * +from .parallel_state import * +from .utils import * diff --git a/archive/kt-sft/ktransformers/server/balance_serve/inference/distributed/communication_op.py b/archive/kt-sft/ktransformers/server/balance_serve/inference/distributed/communication_op.py new file mode 100644 index 0000000..37d8dca --- /dev/null +++ b/archive/kt-sft/ktransformers/server/balance_serve/inference/distributed/communication_op.py @@ -0,0 +1,39 @@ +""" +Date: 2024-12-11 06:02:42 +LastEditors: djw +LastEditTime: 2024-12-12 09:52:06 +""" + +from typing import Any, Dict, Optional, Union + +import torch +import torch.distributed + +from .parallel_state import get_tp_group + + +def tensor_model_parallel_all_reduce(input_: torch.Tensor, bsz_tensor: torch.Tensor, is_compute_bound=False, overlap=False) -> torch.Tensor: + """All-reduce the input tensor across model parallel group.""" + return get_tp_group().all_reduce(input_, bsz_tensor, is_compute_bound=is_compute_bound, overlap=overlap) + + +def tensor_model_parallel_all_gather( + input_: torch.Tensor, dim: int = -1 +) -> torch.Tensor: + """All-gather the input tensor across model parallel group.""" + return get_tp_group().all_gather(input_, dim) + + +def tensor_model_parallel_gather( + input_: torch.Tensor, dst: int = 0, dim: int = -1 +) -> Optional[torch.Tensor]: + """Gather the input tensor across model parallel group.""" + return get_tp_group().gather(input_, dst, dim) + + +def broadcast_tensor_dict( + tensor_dict: Optional[Dict[Any, Union[torch.Tensor, Any]]] = None, src: int = 0 +): + if not torch.distributed.is_initialized(): + return tensor_dict + return get_tp_group().broadcast_tensor_dict(tensor_dict, src) diff --git a/archive/kt-sft/ktransformers/server/balance_serve/inference/distributed/cuda_wrapper.py b/archive/kt-sft/ktransformers/server/balance_serve/inference/distributed/cuda_wrapper.py new file mode 100644 index 0000000..31bf415 --- /dev/null +++ b/archive/kt-sft/ktransformers/server/balance_serve/inference/distributed/cuda_wrapper.py @@ -0,0 +1,168 @@ +"""This file is a pure Python wrapper for the cudart library. +It avoids the need to compile a separate shared library, and is +convenient for use when we just need to call a few functions. +""" + +import ctypes +from dataclasses import dataclass +from typing import Any, Dict, List, Optional + +# this line makes it possible to directly load `libcudart.so` using `ctypes` +import torch # noqa + +# === export types and functions from cudart to Python === +# for the original cudart definition, please check +# https://docs.nvidia.com/cuda/cuda-runtime-api/index.html + +cudaError_t = ctypes.c_int +cudaMemcpyKind = ctypes.c_int + + +class cudaIpcMemHandle_t(ctypes.Structure): + _fields_ = [("internal", ctypes.c_byte * 128)] + + +@dataclass +class Function: + name: str + restype: Any + argtypes: List[Any] + + +def find_loaded_library(lib_name) -> Optional[str]: + """ + According to according to https://man7.org/linux/man-pages/man5/proc_pid_maps.5.html, + the file `/proc/self/maps` contains the memory maps of the process, which includes the + shared libraries loaded by the process. We can use this file to find the path of the + a loaded library. + """ # noqa + found = False + with open("/proc/self/maps") as f: + for line in f: + if lib_name in line: + found = True + break + if not found: + # the library is not loaded in the current process + return None + # if lib_name is libcudart, we need to match a line with: + # address /path/to/libcudart-hash.so.11.0 + start = line.index("/") + path = line[start:].strip() + filename = path.split("/")[-1] + assert filename.rpartition(".so")[0].startswith(lib_name), \ + f"Unexpected filename: {filename} for library {lib_name}" + return path + + +class CudaRTLibrary: + exported_functions = [ + # ​cudaError_t cudaSetDevice ( int device ) + Function("cudaSetDevice", cudaError_t, [ctypes.c_int]), + # cudaError_t cudaDeviceSynchronize ( void ) + Function("cudaDeviceSynchronize", cudaError_t, []), + # ​cudaError_t cudaDeviceReset ( void ) + Function("cudaDeviceReset", cudaError_t, []), + + # const char* cudaGetErrorString ( cudaError_t error ) + Function("cudaGetErrorString", ctypes.c_char_p, [cudaError_t]), + + # ​cudaError_t cudaMalloc ( void** devPtr, size_t size ) + Function("cudaMalloc", cudaError_t, + [ctypes.POINTER(ctypes.c_void_p), ctypes.c_size_t]), + # ​cudaError_t cudaFree ( void* devPtr ) + Function("cudaFree", cudaError_t, [ctypes.c_void_p]), + # ​cudaError_t cudaMemset ( void* devPtr, int value, size_t count ) + Function("cudaMemset", cudaError_t, + [ctypes.c_void_p, ctypes.c_int, ctypes.c_size_t]), + # ​cudaError_t cudaMemcpy ( void* dst, const void* src, size_t count, cudaMemcpyKind kind ) # noqa + Function("cudaMemcpy", cudaError_t, [ + ctypes.c_void_p, ctypes.c_void_p, ctypes.c_size_t, cudaMemcpyKind + ]), + + # cudaError_t cudaIpcGetMemHandle ( cudaIpcMemHandle_t* handle, void* devPtr ) # noqa + Function("cudaIpcGetMemHandle", cudaError_t, + [ctypes.POINTER(cudaIpcMemHandle_t), ctypes.c_void_p]), + # ​cudaError_t cudaIpcOpenMemHandle ( void** devPtr, cudaIpcMemHandle_t handle, unsigned int flags ) # noqa + Function("cudaIpcOpenMemHandle", cudaError_t, [ + ctypes.POINTER(ctypes.c_void_p), cudaIpcMemHandle_t, ctypes.c_uint + ]), + ] + + # class attribute to store the mapping from the path to the library + # to avoid loading the same library multiple times + path_to_library_cache: Dict[str, Any] = {} + + # class attribute to store the mapping from library path + # to the corresponding dictionary + path_to_dict_mapping: Dict[str, Dict[str, Any]] = {} + + def __init__(self, so_file: Optional[str] = None): + if so_file is None: + so_file = find_loaded_library("libcudart") + assert so_file is not None, \ + "libcudart is not loaded in the current process" + if so_file not in CudaRTLibrary.path_to_library_cache: + lib = ctypes.CDLL(so_file) + CudaRTLibrary.path_to_library_cache[so_file] = lib + self.lib = CudaRTLibrary.path_to_library_cache[so_file] + + if so_file not in CudaRTLibrary.path_to_dict_mapping: + _funcs = {} + for func in CudaRTLibrary.exported_functions: + f = getattr(self.lib, func.name) + f.restype = func.restype + f.argtypes = func.argtypes + _funcs[func.name] = f + CudaRTLibrary.path_to_dict_mapping[so_file] = _funcs + self.funcs = CudaRTLibrary.path_to_dict_mapping[so_file] + + def CUDART_CHECK(self, result: cudaError_t) -> None: + if result != 0: + error_str = self.cudaGetErrorString(result) + raise RuntimeError(f"CUDART error: {error_str}") + + def cudaGetErrorString(self, error: cudaError_t) -> str: + return self.funcs["cudaGetErrorString"](error).decode("utf-8") + + def cudaSetDevice(self, device: int) -> None: + self.CUDART_CHECK(self.funcs["cudaSetDevice"](device)) + + def cudaDeviceSynchronize(self) -> None: + self.CUDART_CHECK(self.funcs["cudaDeviceSynchronize"]()) + + def cudaDeviceReset(self) -> None: + self.CUDART_CHECK(self.funcs["cudaDeviceReset"]()) + + def cudaMalloc(self, size: int) -> ctypes.c_void_p: + devPtr = ctypes.c_void_p() + self.CUDART_CHECK(self.funcs["cudaMalloc"](ctypes.byref(devPtr), size)) + return devPtr + + def cudaFree(self, devPtr: ctypes.c_void_p) -> None: + self.CUDART_CHECK(self.funcs["cudaFree"](devPtr)) + + def cudaMemset(self, devPtr: ctypes.c_void_p, value: int, + count: int) -> None: + self.CUDART_CHECK(self.funcs["cudaMemset"](devPtr, value, count)) + + def cudaMemcpy(self, dst: ctypes.c_void_p, src: ctypes.c_void_p, + count: int) -> None: + cudaMemcpyDefault = 4 + kind = cudaMemcpyDefault + self.CUDART_CHECK(self.funcs["cudaMemcpy"](dst, src, count, kind)) + + def cudaIpcGetMemHandle(self, + devPtr: ctypes.c_void_p) -> cudaIpcMemHandle_t: + handle = cudaIpcMemHandle_t() + self.CUDART_CHECK(self.funcs["cudaIpcGetMemHandle"]( + ctypes.byref(handle), devPtr)) + return handle + + def cudaIpcOpenMemHandle(self, + handle: cudaIpcMemHandle_t) -> ctypes.c_void_p: + cudaIpcMemLazyEnablePeerAccess = 1 + devPtr = ctypes.c_void_p() + self.CUDART_CHECK(self.funcs["cudaIpcOpenMemHandle"]( + ctypes.byref(devPtr), handle, cudaIpcMemLazyEnablePeerAccess)) + return devPtr diff --git a/archive/kt-sft/ktransformers/server/balance_serve/inference/distributed/custom_all_reduce.py b/archive/kt-sft/ktransformers/server/balance_serve/inference/distributed/custom_all_reduce.py new file mode 100644 index 0000000..e170e43 --- /dev/null +++ b/archive/kt-sft/ktransformers/server/balance_serve/inference/distributed/custom_all_reduce.py @@ -0,0 +1,310 @@ +import ctypes +from contextlib import contextmanager +from typing import List, Optional, Union + +import torch +import torch.distributed as dist +from torch.distributed import ProcessGroup + +import server.envs as envs +from server.inference.distributed.cuda_wrapper import CudaRTLibrary +from server.inference.distributed.custom_all_reduce_utils import gpu_p2p_access_check +from server.inference.distributed.parallel_state import in_the_same_node_as +from server.inference.platforms import current_platform +from server.utils import cuda_device_count_stateless +import vLLMCustomAllreduce + +try: + vLLMCustomAllreduce.meta_size() + custom_ar = True +except Exception: + # For AMD GPUs and CPUs + custom_ar = False + + +def _can_p2p(rank: int, world_size: int) -> bool: + for i in range(world_size): + if i == rank: + continue + if envs.VLLM_SKIP_P2P_CHECK: + print("Skipping P2P check and trusting the driver's P2P report.") + return torch.cuda.can_device_access_peer(rank, i) + if not gpu_p2p_access_check(rank, i): + return False + return True + + +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() + ) + + +class CustomAllreduce: + + _SUPPORTED_WORLD_SIZES = [2, 4, 6, 8] + + # max_size: max supported allreduce size + def __init__( + self, + group: ProcessGroup, + device: Union[int, str, torch.device], + max_size=8192 * 1024, + ) -> None: + """ + Args: + group: the process group to work on. If None, it will use the + default process group. + device: the device to bind the CustomAllreduce to. If None, + it will be bind to f"cuda:{local_rank}". + It is the caller's responsibility to make sure each communicator + is bind to a unique device, and all communicators in this group + are in the same node. + """ + self._IS_CAPTURING = False + self.disabled = True + + if not custom_ar: + # disable because of missing custom allreduce library + # e.g. in a non-cuda environment + return + + self.group = group + + assert ( + dist.get_backend(group) != dist.Backend.NCCL + ), "CustomAllreduce should be attached to a non-NCCL group." + + if not all(in_the_same_node_as(group, source_rank=0)): + # No need to initialize custom allreduce for multi-node case. + print( + "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: + print( + "Custom allreduce is disabled due to an unsupported world" + " size: %d. Supported world sizes: %s. To silence this " + "warning, specify disable_custom_all_reduce=True explicitly.", + world_size, + str(CustomAllreduce._SUPPORTED_WORLD_SIZES), + ) + return + + if isinstance(device, int): + device = torch.device(f"cuda:{device}") + elif isinstance(device, str): + device = torch.device(device) + # now `device` is a `torch.device` object + assert isinstance(device, torch.device) + self.device = device + + cuda_visible_devices = envs.CUDA_VISIBLE_DEVICES + if cuda_visible_devices: + device_ids = list(map(int, cuda_visible_devices.split(","))) + else: + device_ids = list(range(cuda_device_count_stateless())) + + physical_device_id = device_ids[device.index] + tensor = torch.tensor([physical_device_id], dtype=torch.int, device="cpu") + gather_list = [ + torch.tensor([0], dtype=torch.int, device="cpu") for _ in range(world_size) + ] + dist.all_gather(gather_list, tensor, group=self.group) + physical_device_ids = [t.item() for t in gather_list] + + # test nvlink first, this will filter out most of the cases + # where custom allreduce is not supported + # this checks hardware and driver support for NVLink + assert current_platform.is_cuda() + from server.inference.platforms.cuda import CudaPlatform + + cuda_platform: CudaPlatform = current_platform + full_nvlink = cuda_platform.is_full_nvlink(physical_device_ids) + if world_size > 2 and not full_nvlink: + print( + "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 + if not _can_p2p(rank, world_size): + print( + "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.disabled = False + # 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( + vLLMCustomAllreduce.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.max_size = max_size + self.rank = rank + self.world_size = world_size + self.full_nvlink = full_nvlink + self._ptr = vLLMCustomAllreduce.init_custom_ar( + self.meta_ptrs, self.rank_data, rank, self.full_nvlink + ) + vLLMCustomAllreduce.register_buffer(self._ptr, self.buffer_ptrs) + + @staticmethod + def create_shared_buffer( + size_in_bytes: int, group: Optional[ProcessGroup] = None + ) -> List[int]: + """ + Creates a shared buffer and returns a list of pointers + representing the buffer on all processes in the group. + """ + lib = CudaRTLibrary() + pointer = lib.cudaMalloc(size_in_bytes) + handle = lib.cudaIpcGetMemHandle(pointer) + world_size = dist.get_world_size(group=group) + rank = dist.get_rank(group=group) + handles = [None] * world_size + dist.all_gather_object(handles, handle, group=group) + + pointers: List[int] = [] + for i, h in enumerate(handles): + if i == rank: + pointers.append(pointer.value) # type: ignore + else: + pointers.append(lib.cudaIpcOpenMemHandle(h).value) # type: ignore + + return pointers + + @staticmethod + def free_shared_buffer( + pointers: List[int], group: Optional[ProcessGroup] = None + ) -> None: + rank = dist.get_rank(group=group) + lib = CudaRTLibrary() + lib.cudaFree(ctypes.c_void_p(pointers[rank])) + + @contextmanager + def capture(self): + """ + The main responsibility of this context manager is the + `register_graph_buffers` call at the end of the context. + It records all the buffer addresses used in the CUDA graph. + """ + try: + self._IS_CAPTURING = True + yield + finally: + self._IS_CAPTURING = False + if not self.disabled: + self.register_graph_buffers() + + def register_graph_buffers(self): + handle, offset = vLLMCustomAllreduce.get_graph_buffer_ipc_meta(self._ptr) + print("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 + vLLMCustomAllreduce.register_graph_buffers(self._ptr, handles, offsets) + + def should_custom_ar(self, inp: torch.Tensor): + if self.disabled: + return False + inp_size = inp.numel() * inp.element_size() + # custom allreduce requires input byte size to be multiples of 16 + if inp_size % 16 != 0: + return False + if not is_weak_contiguous(inp): + return False + # for 4 or more non NVLink-capable GPUs, custom allreduce provides + # little performance improvement over NCCL. + if self.world_size == 2 or self.full_nvlink: + return inp_size < self.max_size + return False + + def all_reduce( + self, inp: torch.Tensor, *, out: torch.Tensor = None, bsz_tensor: torch.Tensor = None, registered: bool = False, + is_compute_bound=False, overlap=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 is_compute_bound: + sms = 2 if overlap else 36 + else: + sms = 20 if overlap else 36 + #print("all reduce sms", sms) + if out is None: + out = torch.empty_like(inp) + if registered: + vLLMCustomAllreduce.all_reduce(self._ptr, inp, out, 0, 0, bsz_tensor, block_limit=sms) + else: + vLLMCustomAllreduce.all_reduce( + self._ptr, inp, out, self.buffer_ptrs[self.rank], self.max_size, bsz_tensor, block_limit=sms + ) + return out + + def custom_all_reduce(self, input: torch.Tensor, bsz_tensor: torch.Tensor, is_compute_bound=False, overlap=False) -> Optional[torch.Tensor]: + """The main allreduce API that provides support for cuda graph.""" + # When custom allreduce is disabled, this will be None. + if self.disabled or not self.should_custom_ar(input): + return None + if self._IS_CAPTURING: + if torch.cuda.is_current_stream_capturing(): + return self.all_reduce(input, bsz_tensor=bsz_tensor, registered=True, is_compute_bound=is_compute_bound, overlap=overlap) + else: + # If warm up, mimic the allocation pattern since custom + # allreduce is out-of-place. + return torch.empty_like(input) + else: + # Note: outside of cuda graph context, custom allreduce incurs a + # cost of cudaMemcpy, which should be small (<=1% of overall + # latency) compared to the performance gain of using custom kernels + return self.all_reduce(input, bsz_tensor=bsz_tensor, registered=False, is_compute_bound=is_compute_bound, overlap=overlap) + + def close(self): + if not self.disabled and self._ptr: + vLLMCustomAllreduce.dispose(self._ptr) + self._ptr = 0 + self.free_shared_buffer(self.meta_ptrs) + self.free_shared_buffer(self.buffer_ptrs) + + def __del__(self): + self.close() diff --git a/archive/kt-sft/ktransformers/server/balance_serve/inference/distributed/custom_all_reduce_utils.py b/archive/kt-sft/ktransformers/server/balance_serve/inference/distributed/custom_all_reduce_utils.py new file mode 100644 index 0000000..d94ffe4 --- /dev/null +++ b/archive/kt-sft/ktransformers/server/balance_serve/inference/distributed/custom_all_reduce_utils.py @@ -0,0 +1,272 @@ +import ctypes +import json +import os +import pickle +import subprocess +import sys +import tempfile +from itertools import product +from typing import Dict, List, Optional, Sequence + +import torch.distributed as dist +import torch.multiprocessing as mp + +import server.envs as envs +from server.inference.distributed.cuda_wrapper import CudaRTLibrary +from server.utils import cuda_device_count_stateless, update_environment_variables + + +def producer( + batch_src: Sequence[int], + producer_queue, + consumer_queue, + result_queue, + cuda_visible_devices: Optional[str] = None, +): + if cuda_visible_devices is not None: + update_environment_variables({"CUDA_VISIBLE_DEVICES": cuda_visible_devices}) + + lib = CudaRTLibrary() + for i in batch_src: + lib.cudaSetDevice(i) + pointer = lib.cudaMalloc(1024) + lib.cudaMemset(pointer, 1, 1024) + lib.cudaDeviceSynchronize() + handle = lib.cudaIpcGetMemHandle(pointer) + producer_queue.put(handle) + open_success = consumer_queue.get() + if open_success: + # use two queues to simulate barrier + producer_queue.put(0) + consumer_queue.get() + # check if the memory is modified + host_data = (ctypes.c_char * 1024)() + lib.cudaMemcpy(host_data, pointer, 1024) # type: ignore + for i in range(1024): + if ord(host_data[i]) != 2: + open_success = False + break + result_queue.put(open_success) + lib.cudaDeviceReset() + + +def consumer( + batch_tgt: Sequence[int], + producer_queue, + consumer_queue, + result_queue, + cuda_visible_devices: Optional[str] = None, +): + if cuda_visible_devices is not None: + update_environment_variables({"CUDA_VISIBLE_DEVICES": cuda_visible_devices}) + + lib = CudaRTLibrary() + for j in batch_tgt: + lib.cudaSetDevice(j) + handle = producer_queue.get() + open_success = False + try: + pointer = lib.cudaIpcOpenMemHandle(handle) # type: ignore + open_success = True + except RuntimeError: + # cannot error out here, because the producer process + # is still waiting for the response. + pass + consumer_queue.put(open_success) + if open_success: + # modify the memory + lib.cudaMemset(pointer, 2, 1024) + lib.cudaDeviceSynchronize() + # use two queues to simulate barrier + producer_queue.get() + consumer_queue.put(0) + # check if the memory is modified + host_data = (ctypes.c_char * 1024)() + lib.cudaMemcpy(host_data, pointer, 1024) # type: ignore + for i in range(1024): + if ord(host_data[i]) != 2: + open_success = False + break + result_queue.put(open_success) + lib.cudaDeviceReset() + + +def can_actually_p2p( + batch_src: Sequence[int], + batch_tgt: Sequence[int], +) -> Sequence[bool]: + """ + Usually, checking if P2P access is enabled can be done by + `torch.cuda.can_device_access_peer(src, tgt)`. However, sometimes + the driver might be broken, and `torch.cuda.can_device_access_peer(src, tgt)` + returns `True` even if P2P access is not actually possible. + See https://github.com/vllm-project/vllm/issues/2728 and + https://forums.developer.nvidia.com/t/direct-gpu-gpu-communication-does-not-seem-to-work-properly/283264/10 + Therefore, we have to perform a real P2P access to check if it is actually + possible. + + Note on p2p and cuda IPC: + Usually, one process uses one GPU: + GPU src --> cuda context src --> tensor src --> process src + + We need to combine p2p and cuda IPC, so that: + GPU src --> cuda context src --> tensor src --> process src + |shared| + GPU tgt --> cuda context tgt --> tensor tgt --> process tgt + That is to say, process src creates a tensor in GPU src, passes IPC handle to + process tgt, and process tgt accesses the tensor in GPU tgt. Any operation on the + tensor in process tgt will be reflected in the tensor in process src, because + they are the same memory segment. + It is important to note that process tgt accesses the tensor in GPU tgt, not + GPU src. That's why we need p2p access. + + The most time-consuming part is the process creation. To avoid creating + processes for every pair of GPUs, we use batched testing. We create two + processes for testing all pairs of GPUs in batch. The trick is to reset + the device after each test (which is not available in PyTorch). + """ # noqa + cuda_visible_devices = envs.CUDA_VISIBLE_DEVICES + # pass the CUDA_VISIBLE_DEVICES to the child process + # to make sure they see the same set of GPUs + + # make sure the processes are spawned + smp = mp.get_context("spawn") + producer_queue = smp.Queue() + consumer_queue = smp.Queue() + result_queue = smp.Queue() + p_src = smp.Process( + target=producer, + args=( + batch_src, + producer_queue, + consumer_queue, + result_queue, + cuda_visible_devices, + ), + ) + p_tgt = smp.Process( + target=consumer, + args=( + batch_tgt, + producer_queue, + consumer_queue, + result_queue, + cuda_visible_devices, + ), + ) + p_src.start() + p_tgt.start() + p_src.join() + p_tgt.join() + assert p_src.exitcode == 0 and p_tgt.exitcode == 0 + result: List[bool] = [] + for src, tgt in zip(batch_src, batch_tgt): + a = result_queue.get() + b = result_queue.get() + if a != b: + print( + "Two processes do not agree on the P2P access" + " status on %d -> %d, treat as disabled.", + src, + tgt, + ) + result.append(False) + else: + result.append(a) + return result + + +# why do we need this cache? +# we are testing peer-to-peer (p2p) access between GPUs,across processes. +# if we test it every time, it will be very slow, because we need to create +# N * N * 2 processes, where N is the world size. This is very slow. +# to reduce the time, we use a cache file to store the p2p access status. +# the cache file is generated by the master process if it does not exist. +# then all the processes can read the cache file to check the p2p access status. +# Note that the cache file is suffixed by the CUDA_VISIBLE_DEVICES, so that we +# can have different cache files for different CUDA_VISIBLE_DEVICES settings, +# e.g. used by different vllm engines. The device id in the cache file is a +# **local** device id, i.e. from 0 to num_dev-1, where num_dev is the number +# of visible devices in the vllm engine. +_gpu_p2p_access_cache: Optional[Dict[str, bool]] = None + + +def gpu_p2p_access_check(src: int, tgt: int) -> bool: + """Check if GPU src can access GPU tgt.""" + + # if the cache variable is already calculated, + # read from the cache instead of checking it again + global _gpu_p2p_access_cache + if _gpu_p2p_access_cache is not None: + return _gpu_p2p_access_cache[f"{src}->{tgt}"] + + is_distributed = dist.is_initialized() + + num_dev = cuda_device_count_stateless() + cuda_visible_devices = envs.CUDA_VISIBLE_DEVICES + if cuda_visible_devices is None: + cuda_visible_devices = ",".join(str(i) for i in range(num_dev)) + + path = os.path.join( + envs.VLLM_CACHE_ROOT, f"gpu_p2p_access_cache_for_{cuda_visible_devices}.json" + ) + os.makedirs(os.path.dirname(path), exist_ok=True) + from server.inference.distributed.parallel_state import get_world_group + + if (not is_distributed or get_world_group().local_rank == 0) and ( + not os.path.exists(path) + ): + # only the local master process (with local_rank == 0) can + # enter this block to calculate the cache + print("generating GPU P2P access cache in %s", path) + cache: Dict[str, bool] = {} + ids = list(range(num_dev)) + # batch of all pairs of GPUs + batch_src, batch_tgt = zip(*list(product(ids, ids))) + # NOTE: we use `subprocess` rather than `multiprocessing` here + # because the caller might not have `if __name__ == "__main__":`, + # in that case we cannot use spawn method in multiprocessing. + # However, `can_actually_p2p` requires spawn method. + # The fix is, we use `subprocess` to call the function, + # where we have `if __name__ == "__main__":` in this file. + + # use a temporary file to store the result + # we don't use the output of the subprocess directly, + # because the subprocess might produce logging output + with tempfile.NamedTemporaryFile() as output_file: + input_bytes = pickle.dumps((batch_src, batch_tgt, output_file.name)) + returned = subprocess.run( + [sys.executable, __file__], input=input_bytes, capture_output=True + ) + # check if the subprocess is successful + try: + returned.check_returncode() + except Exception as e: + # wrap raised exception to provide more information + raise RuntimeError( + f"Error happened when batch testing " + f"peer-to-peer access from {batch_src} to {batch_tgt}:\n" + f"{returned.stderr.decode()}" + ) from e + with open(output_file.name, "rb") as f: + result = pickle.load(f) + for _i, _j, r in zip(batch_src, batch_tgt, result): + cache[f"{_i}->{_j}"] = r + with open(path, "w") as f: + json.dump(cache, f, indent=4) + if is_distributed: + get_world_group().barrier() + print("reading GPU P2P access cache from %s", path) + with open(path) as f: + cache = json.load(f) + _gpu_p2p_access_cache = cache + return _gpu_p2p_access_cache[f"{src}->{tgt}"] + + +__all__ = ["gpu_p2p_access_check"] + +if __name__ == "__main__": + batch_src, batch_tgt, output_file = pickle.loads(sys.stdin.buffer.read()) + result = can_actually_p2p(batch_src, batch_tgt) + with open(output_file, "wb") as f: + f.write(pickle.dumps(result)) diff --git a/archive/kt-sft/ktransformers/server/balance_serve/inference/distributed/parallel_state.py b/archive/kt-sft/ktransformers/server/balance_serve/inference/distributed/parallel_state.py new file mode 100644 index 0000000..fc11374 --- /dev/null +++ b/archive/kt-sft/ktransformers/server/balance_serve/inference/distributed/parallel_state.py @@ -0,0 +1,1262 @@ +# Copyright 2023 The vLLM team. +# Adapted from +# https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/core/parallel_state.py +# Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved. +"""vLLM distributed state. +It takes over the control of the distributed environment from PyTorch. +The typical workflow is: + +- call `init_distributed_environment` to initialize the distributed environment. +- call `initialize_model_parallel` or `ensure_model_parallel_initialized` to + initialize the model parallel groups. + +- any code dealing with the distributed stuff + +- call `destroy_model_parallel` to destroy the model parallel groups. +- call `destroy_distributed_environment` to destroy the distributed environment. + +If you only need to use the distributed environment without model/pipeline + parallelism, you can skip the model parallel initialization and destruction + steps. +""" +import contextlib +import gc +import pickle +import weakref +from collections import namedtuple +from contextlib import contextmanager, nullcontext +from dataclasses import dataclass +from multiprocessing import shared_memory +from typing import Any, Callable, Dict, List, Optional, Tuple, Union +from unittest.mock import patch + +import torch +import torch.distributed +from torch.distributed import Backend, ProcessGroup + +import server.envs as envs +from server.inference.platforms import current_platform +from server.utils import direct_register_custom_op, supports_custom_op + + +@dataclass +class GraphCaptureContext: + stream: torch.cuda.Stream + + +TensorMetadata = namedtuple("TensorMetadata", ["device", "dtype", "size"]) + + +def _split_tensor_dict( + tensor_dict: Dict[str, Union[torch.Tensor, Any]] +) -> Tuple[List[Tuple[str, Any]], List[torch.Tensor]]: + """Split the tensor dictionary into two parts: + 1. A list of (key, value) pairs. If the value is a tensor, it is replaced + by its metadata. + 2. A list of tensors. + """ + metadata_list: List[Tuple[str, Any]] = [] + tensor_list: List[torch.Tensor] = [] + for key, value in tensor_dict.items(): + if isinstance(value, torch.Tensor): + # Note: we cannot use `value.device` here, + # because it contains not only the device type but also the device + # index (e.g. "cuda:0"). We only need the device type. + # receiving side will set the device index. + device = value.device.type + metadata_list.append( + (key, TensorMetadata(device, value.dtype, value.size())) + ) + tensor_list.append(value) + else: + metadata_list.append((key, value)) + return metadata_list, tensor_list + + +_group_name_counter: Dict[str, int] = {} + + +def _get_unique_name(name: str) -> str: + """Get a unique name for the group. + Example: + _get_unique_name("tp") -> "tp:0" + _get_unique_name("tp") -> "tp:1" + """ + if name not in _group_name_counter: + _group_name_counter[name] = 0 + newname = f"{name}:{_group_name_counter[name]}" + _group_name_counter[name] += 1 + return newname + + +_groups: Dict[str, Callable[[], Optional["GroupCoordinator"]]] = {} + + +def _register_group(group: "GroupCoordinator") -> None: + _groups[group.unique_name] = weakref.ref(group) + + +if supports_custom_op(): + + def inplace_all_reduce(tensor: torch.Tensor, group_name: str) -> None: + assert group_name in _groups, f"Group {group_name} is not found." + group = _groups[group_name]() + if group is None: + raise ValueError(f"Group {group_name} is destroyed.") + group._all_reduce_in_place(tensor) + + def inplace_all_reduce_fake(tensor: torch.Tensor, group_name: str) -> None: + return + + direct_register_custom_op( + op_name="inplace_all_reduce", + op_func=inplace_all_reduce, + mutates_args=["tensor"], + fake_impl=inplace_all_reduce_fake, + ) + + def outplace_all_reduce(tensor: torch.Tensor, group_name: str, bsz_tensor: torch.Tensor, is_compute_bound: bool = False, overlap: bool = False) -> torch.Tensor: + assert group_name in _groups, f"Group {group_name} is not found." + group = _groups[group_name]() + if group is None: + raise ValueError(f"Group {group_name} is destroyed.") + return group._all_reduce_out_place(tensor, bsz_tensor, is_compute_bound=is_compute_bound, overlap=overlap) + + def outplace_all_reduce_fake(tensor: torch.Tensor, group_name: str, bsz_tensor: torch.Tensor, is_compute_bound: bool = False, overlap: bool = False) -> torch.Tensor: + return torch.empty_like(tensor) + + direct_register_custom_op( + op_name="outplace_all_reduce", + op_func=outplace_all_reduce, + mutates_args=[], + fake_impl=outplace_all_reduce_fake, + ) + + +class GroupCoordinator: + """ + PyTorch ProcessGroup wrapper for a group of processes. + PyTorch ProcessGroup is bound to one specific communication backend, + e.g. NCCL, Gloo, MPI, etc. + GroupCoordinator takes charge of all the communication operations among + the processes in the group. It can route the communication to + a specific implementation (e.g. switch allreduce implementation + based on the tensor size and cuda graph mode). + """ + + # available attributes: + rank: int # global rank + ranks: List[int] # global ranks in the group + world_size: int # size of the group + # difference between `local_rank` and `rank_in_group`: + # if we have a group of size 4 across two nodes: + # Process | Node | Rank | Local Rank | Rank in Group + # 0 | 0 | 0 | 0 | 0 + # 1 | 0 | 1 | 1 | 1 + # 2 | 1 | 2 | 0 | 2 + # 3 | 1 | 3 | 1 | 3 + local_rank: int # local rank used to assign devices + rank_in_group: int # rank inside the group + cpu_group: ProcessGroup # group for CPU communication + device_group: ProcessGroup # group for device communication + use_pynccl: bool # a hint of whether to use PyNccl + use_custom_allreduce: bool # a hint of whether to use CustomAllreduce + # communicators are only created for world size > 1 + pynccl_comm: Optional[Any] # PyNccl communicator + ca_comm: Optional[Any] # Custom allreduce communicator + mq_broadcaster: Optional[Any] # shared memory broadcaster + + def __init__( + self, + group_ranks: List[List[int]], + local_rank: int, + torch_distributed_backend: Union[str, Backend], + use_pynccl: bool, + use_custom_allreduce: bool, + use_tpu_communicator: bool, + use_hpu_communicator: bool, + use_xpu_communicator: bool, + use_message_queue_broadcaster: bool = False, + group_name: Optional[str] = None, + ): + group_name = group_name or "anonymous" + self.unique_name = _get_unique_name(group_name) + _register_group(self) + + self.rank = torch.distributed.get_rank() + self.local_rank = local_rank + self.device_group = None + self.cpu_group = None + + for ranks in group_ranks: + device_group = torch.distributed.new_group( + ranks, backend=torch_distributed_backend + ) + # a group with `gloo` backend, to allow direct coordination between + # processes through the CPU. + cpu_group = torch.distributed.new_group(ranks, backend="gloo") + if self.rank in ranks: + self.ranks = ranks + self.world_size = len(ranks) + self.rank_in_group = ranks.index(self.rank) + self.device_group = device_group + self.cpu_group = cpu_group + + assert self.cpu_group is not None + assert self.device_group is not None + assert current_platform.is_cuda_alike() + + if current_platform.is_cuda_alike(): + self.device = torch.device(f"cuda:{local_rank}") + else: + self.device = torch.device("cpu") + + self.use_pynccl = use_pynccl + self.use_custom_allreduce = use_custom_allreduce + self.use_tpu_communicator = use_tpu_communicator + self.use_hpu_communicator = use_hpu_communicator + self.use_xpu_communicator = use_xpu_communicator + + # lazy import to avoid documentation build error + from server.inference.distributed.custom_all_reduce import CustomAllreduce + from server.inference.distributed.pynccl import PyNcclCommunicator + + self.pynccl_comm: Optional[PyNcclCommunicator] = None + # if use_pynccl and self.world_size > 1: + # self.pynccl_comm = PyNcclCommunicator( + # group=self.cpu_group, + # device=self.device, + # ) + + self.ca_comm: Optional[CustomAllreduce] = None + if use_custom_allreduce and self.world_size > 1: + # Initialize a custom fast all-reduce implementation. + self.ca_comm = CustomAllreduce( + group=self.cpu_group, + device=self.device, + ) + + #### we assume we won't use tpu or hpu or xpu or messagequeue broadcast + + # from vllm.distributed.device_communicators.tpu_communicator import ( + # TpuCommunicator) + # self.tpu_communicator: Optional[TpuCommunicator] = None + # if use_tpu_communicator and self.world_size > 1: + # self.tpu_communicator = TpuCommunicator(group=self.cpu_group) + self.tpu_communicator = None + + # from vllm.distributed.device_communicators.hpu_communicator import ( + # HpuCommunicator) + # self.hpu_communicator: Optional[HpuCommunicator] + # if use_hpu_communicator and self.world_size > 1: + # self.hpu_communicator = HpuCommunicator(group=self.device_group) + self.hpu_communicator = None + + # from vllm.distributed.device_communicators.xpu_communicator import ( + # XpuCommunicator) + # self.xpu_communicator: Optional[XpuCommunicator] + # if use_xpu_communicator and self.world_size > 1: + # self.xpu_communicator = XpuCommunicator(group=self.device_group) + self.xpu_communicator = None + + # from vllm.distributed.device_communicators.shm_broadcast import ( + # MessageQueue) + # self.mq_broadcaster: Optional[MessageQueue] = None + # if use_message_queue_broadcaster and self.world_size > 1: + # self.mq_broadcaster = MessageQueue.create_from_process_group( + # self.cpu_group, 1 << 22, 6) + self.mq_broadcaster = None + + @property + def first_rank(self): + """Return the global rank of the first process in the group""" + return self.ranks[0] + + @property + def last_rank(self): + """Return the global rank of the last process in the group""" + return self.ranks[-1] + + @property + def is_first_rank(self): + """Return whether the caller is the first process in the group""" + return self.rank == self.first_rank + + @property + def is_last_rank(self): + """Return whether the caller is the last process in the group""" + return self.rank == self.last_rank + + @property + def next_rank(self): + """Return the global rank of the process that follows the caller""" + rank_in_group = self.rank_in_group + world_size = self.world_size + return self.ranks[(rank_in_group + 1) % world_size] + + @property + def prev_rank(self): + """Return the global rank of the process that precedes the caller""" + rank_in_group = self.rank_in_group + world_size = self.world_size + return self.ranks[(rank_in_group - 1) % world_size] + + @contextmanager + def graph_capture( + self, graph_capture_context: Optional[GraphCaptureContext] = None + ): + if graph_capture_context is None: + stream = torch.cuda.Stream() + graph_capture_context = GraphCaptureContext(stream) + else: + stream = graph_capture_context.stream + + ca_comm = self.ca_comm + maybe_ca_context = nullcontext() if ca_comm is None else ca_comm.capture() + + # ensure all initialization operations complete before attempting to + # capture the graph on another stream + curr_stream = torch.cuda.current_stream() + if curr_stream != stream: + stream.wait_stream(curr_stream) + + with torch.cuda.stream(stream), maybe_ca_context: + # In graph mode, we have to be very careful about the collective + # operations. The current status is: + # allreduce \ Mode | Eager | Graph | + # -------------------------------------------- + # custom allreduce | enabled | enabled | + # PyNccl | disabled| enabled | + # torch.distributed | enabled | disabled| + # + # Note that custom allreduce will have a runtime check, if the + # tensor size is too large, it will fallback to the next + # available option. + # In summary: When using CUDA graph, we use + # either custom all-reduce kernel or pynccl. When not using + # CUDA graph, we use either custom all-reduce kernel or + # PyTorch NCCL. We always prioritize using custom all-reduce + # kernel but fall back to PyTorch or pynccl if it is + # disabled or not supported. + pynccl_comm = self.pynccl_comm + maybe_pynccl_context: Any + if not pynccl_comm: + maybe_pynccl_context = nullcontext() + else: + maybe_pynccl_context = pynccl_comm.change_state( + enable=True, stream=torch.cuda.current_stream() + ) + with maybe_pynccl_context: + yield graph_capture_context + + def all_reduce(self, input_: torch.Tensor, bsz_tensor: torch.Tensor, is_compute_bound=False, overlap=False) -> torch.Tensor: + """ + User-facing all-reduce function before we actually call the + all-reduce operation. + + We need this because Dynamo does not support passing an arbitrary + object (`self` in this case) to a custom op. We need to pass the + group name as a string, and then look up the group coordinator from + the group name, dispatch the all-reduce operation to the group + coordinator. + + In addition, PyTorch custom ops do not support mutation or returning + a new tensor in the same op. So we need to figure out if the op is + in-place or out-of-place ahead of time. + """ + # Bypass the function if we are using only 1 GPU. + if self.world_size == 1: + return input_ + + if input_.is_cpu: + import intel_extension_for_pytorch as ipex + + ipex.distributed.all_reduce(input_, group=self.device_group) + return input_ + + if not supports_custom_op(): + self._all_reduce_in_place(input_) + return input_ + + if self.tpu_communicator is not None and not self.tpu_communicator.disabled: + # TPU handles Dynamo with its own logic. + return self.tpu_communicator.all_reduce(input_) + + if self.hpu_communicator is not None and not self.hpu_communicator.disabled: + return self.hpu_communicator.all_reduce(input_) + + if self.xpu_communicator is not None and not self.xpu_communicator.disabled: + return self.xpu_communicator.all_reduce(input_) + + if ( + self.ca_comm is not None + and not self.ca_comm.disabled + and self.ca_comm.should_custom_ar(input_) + ): + return torch.ops.vllm.outplace_all_reduce( + input_, group_name=self.unique_name, bsz_tensor=bsz_tensor, is_compute_bound=is_compute_bound, overlap=overlap + ) + else: + #assert self.ca_comm is not None + #assert not self.ca_comm.disabled + #assert self.ca_comm.should_custom_ar(input_) + torch.ops.vllm.inplace_all_reduce(input_, group_name=self.unique_name) + return input_ + + def _all_reduce_out_place(self, input_: torch.Tensor, bsz_tensor: torch.Tensor, is_compute_bound=False, overlap=False) -> torch.Tensor: + ca_comm = self.ca_comm + assert ca_comm is not None + assert not ca_comm.disabled + out = ca_comm.custom_all_reduce(input_, bsz_tensor, is_compute_bound=is_compute_bound, overlap=overlap) + assert out is not None + return out + + def _all_reduce_in_place(self, input_: torch.Tensor) -> None: + pynccl_comm = self.pynccl_comm + if pynccl_comm is not None and not pynccl_comm.disabled: + pynccl_comm.all_reduce(input_) + else: + torch.distributed.all_reduce(input_, group=self.device_group) + + def all_gather(self, input_: torch.Tensor, dim: int = -1) -> torch.Tensor: + world_size = self.world_size + # Bypass the function if we are using only 1 GPU. + if world_size == 1: + return input_ + assert ( + -input_.dim() <= dim < input_.dim() + ), f"Invalid dim ({dim}) for input tensor with shape {input_.size()}" + + # For TPUs, use TPU communicator. + tpu_comm = self.tpu_communicator + if tpu_comm is not None and not tpu_comm.disabled: + return tpu_comm.all_gather(input_, dim) + + # For HPUs, use HPU communicator. + hpu_comm = self.hpu_communicator + if hpu_comm is not None and not hpu_comm.disabled: + return hpu_comm.all_gather(input_, dim) + + if dim < 0: + # Convert negative dim to positive. + dim += input_.dim() + input_size = input_.size() + # NOTE: we have to use concat-style all-gather here, + # stack-style all-gather has compatibility issues with + # torch.compile . see https://github.com/pytorch/pytorch/issues/138795 + output_size = (input_size[0] * world_size,) + input_size[1:] + # Allocate output tensor. + output_tensor = torch.empty( + output_size, dtype=input_.dtype, device=input_.device + ) + # All-gather. + torch.distributed.all_gather_into_tensor( + output_tensor, input_, group=self.device_group + ) + # Reshape + output_tensor = output_tensor.reshape((world_size,) + input_size) + output_tensor = output_tensor.movedim(0, dim) + output_tensor = output_tensor.reshape( + input_size[:dim] + (world_size * input_size[dim],) + input_size[dim + 1 :] + ) + return output_tensor + + def gather( + self, input_: torch.Tensor, dst: int = 0, dim: int = -1 + ) -> Optional[torch.Tensor]: + """ + NOTE: We assume that the input tensor is on the same device across + all the ranks. + NOTE: `dst` is the local rank of the destination rank. + """ + world_size = self.world_size + # Bypass the function if we are using only 1 GPU. + if world_size == 1: + return input_ + assert ( + -input_.dim() <= dim < input_.dim() + ), f"Invalid dim ({dim}) for input tensor with shape {input_.size()}" + if dim < 0: + # Convert negative dim to positive. + dim += input_.dim() + if self.xpu_communicator is not None and not self.xpu_communicator.disabled: + return self.xpu_communicator.gather(input_, self.rank_in_group, dst, dim) + # Allocate output tensor. + if self.rank_in_group == dst: + gather_list = [torch.empty_like(input_) for _ in range(world_size)] + else: + gather_list = None + # Gather. + torch.distributed.gather( + input_, gather_list, dst=self.ranks[dst], group=self.device_group + ) + if self.rank_in_group == dst: + output_tensor = torch.cat(gather_list, dim=dim) + else: + output_tensor = None + return output_tensor + + def broadcast(self, input_: torch.Tensor, src: int = 0): + """Broadcast the input tensor. + NOTE: `src` is the local rank of the source rank. + """ + assert src < self.world_size, f"Invalid src rank ({src})" + + # Bypass the function if we are using only 1 GPU. + if self.world_size == 1: + return input_ + # Broadcast. + torch.distributed.broadcast( + input_, src=self.ranks[src], group=self.device_group + ) + return input_ + + def broadcast_object(self, obj: Optional[Any] = None, src: int = 0): + """Broadcast the input object. + NOTE: `src` is the local rank of the source rank. + """ + assert src < self.world_size, f"Invalid src rank ({src})" + + # Bypass the function if we are using only 1 GPU. + if self.world_size == 1: + return obj + if self.mq_broadcaster is not None: + assert src == 0, "Message queue broadcaster only supports src=0" + return self.mq_broadcaster.broadcast_object(obj) + if self.rank_in_group == src: + torch.distributed.broadcast_object_list( + [obj], src=self.ranks[src], group=self.cpu_group + ) + return obj + else: + recv = [None] + torch.distributed.broadcast_object_list( + recv, src=self.ranks[src], group=self.cpu_group + ) + return recv[0] + + def broadcast_object_list( + self, obj_list: List[Any], src: int = 0, group: Optional[ProcessGroup] = None + ): + """Broadcast the input object list. + NOTE: `src` is the local rank of the source rank. + """ + assert src < self.world_size, f"Invalid src rank ({src})" + + # Bypass the function if we are using only 1 GPU. + if self.world_size == 1: + return obj_list + # Broadcast. + torch.distributed.broadcast_object_list( + obj_list, src=self.ranks[src], group=self.device_group + ) + return obj_list + + def send_object(self, obj: Any, dst: int) -> None: + """Send the input object list to the destination rank.""" + """NOTE: `dst` is the local rank of the destination rank.""" + + assert dst < self.world_size, f"Invalid dst rank ({dst})" + + assert dst != self.rank_in_group, ( + "Invalid destination rank. Destination rank is the same " + "as the current rank." + ) + + # Serialize object to tensor and get the size as well + object_tensor = torch.frombuffer(pickle.dumps(obj), dtype=torch.uint8) + + size_tensor = torch.tensor( + [object_tensor.numel()], dtype=torch.long, device="cpu" + ) + + # Send object size + + torch.distributed.send(size_tensor, dst=self.ranks[dst], group=self.cpu_group) + + # Send object + torch.distributed.send(object_tensor, dst=self.ranks[dst], group=self.cpu_group) + + return None + + def recv_object(self, src: int) -> Any: + """Receive the input object list from the source rank.""" + """NOTE: `src` is the local rank of the source rank.""" + + assert src < self.world_size, f"Invalid src rank ({src})" + + assert ( + src != self.rank_in_group + ), "Invalid source rank. Source rank is the same as the current rank." + + size_tensor = torch.empty(1, dtype=torch.long, device="cpu") + + # Receive object size + rank_size = torch.distributed.recv( + size_tensor, src=self.ranks[src], group=self.cpu_group + ) + + # Tensor to receive serialized objects into. + object_tensor = torch.empty( # type: ignore[call-overload] + size_tensor.item(), # type: ignore[arg-type] + dtype=torch.uint8, + device="cpu", + ) + + rank_object = torch.distributed.recv( + object_tensor, src=self.ranks[src], group=self.cpu_group + ) + + assert ( + rank_object == rank_size + ), "Received object sender rank does not match the size sender rank." + + obj = pickle.loads(object_tensor.numpy().tobytes()) + + return obj + + def broadcast_tensor_dict( + self, + tensor_dict: Optional[Dict[str, Union[torch.Tensor, Any]]] = None, + src: int = 0, + group: Optional[ProcessGroup] = None, + metadata_group: Optional[ProcessGroup] = None, + ) -> Optional[Dict[str, Union[torch.Tensor, Any]]]: + """Broadcast the input tensor dictionary. + NOTE: `src` is the local rank of the source rank. + """ + # Bypass the function if we are using only 1 GPU. + if not torch.distributed.is_initialized() or self.world_size == 1: + return tensor_dict + + group = self.device_group + metadata_group = self.cpu_group + assert src < self.world_size, f"Invalid src rank ({src})" + + rank_in_group = self.rank_in_group + if rank_in_group == src: + metadata_list: List[Tuple[Any, Any]] = [] + assert isinstance( + tensor_dict, dict + ), f"Expecting a dictionary, got {type(tensor_dict)}" + metadata_list, tensor_list = _split_tensor_dict(tensor_dict) + # `metadata_list` lives in CPU memory. + # `broadcast_object_list` has serialization & deserialization, + # all happening on CPU. Therefore, we can use the CPU group. + self.broadcast_object(metadata_list, src=src) + async_handles = [] + for tensor in tensor_list: + if tensor.numel() == 0: + # Skip broadcasting empty tensors. + continue + if tensor.is_cpu: + # use metadata_group for CPU tensors + handle = torch.distributed.broadcast( + tensor, src=self.ranks[src], group=metadata_group, async_op=True + ) + else: + # use group for GPU tensors + handle = torch.distributed.broadcast( + tensor, src=self.ranks[src], group=group, async_op=True + ) + async_handles.append(handle) + for async_handle in async_handles: + async_handle.wait() + + else: + metadata_list = self.broadcast_object(None, src=src) + tensor_dict = {} + async_handles = [] + for key, value in metadata_list: + if isinstance(value, TensorMetadata): + tensor = torch.empty( + value.size, dtype=value.dtype, device=value.device + ) + if tensor.numel() == 0: + # Skip broadcasting empty tensors. + tensor_dict[key] = tensor + continue + if tensor.is_cpu: + # use metadata_group for CPU tensors + handle = torch.distributed.broadcast( + tensor, + src=self.ranks[src], + group=metadata_group, + async_op=True, + ) + else: + # use group for GPU tensors + handle = torch.distributed.broadcast( + tensor, src=self.ranks[src], group=group, async_op=True + ) + async_handles.append(handle) + tensor_dict[key] = tensor + else: + tensor_dict[key] = value + for async_handle in async_handles: + async_handle.wait() + return tensor_dict + + def send_tensor_dict( + self, + tensor_dict: Dict[str, Union[torch.Tensor, Any]], + dst: Optional[int] = None, + all_gather_group: Optional["GroupCoordinator"] = None, + ) -> Optional[Dict[str, Union[torch.Tensor, Any]]]: + """Send the input tensor dictionary. + NOTE: `dst` is the local rank of the source rank. + """ + # Bypass the function if we are using only 1 GPU. + if not torch.distributed.is_initialized() or self.world_size == 1: + return tensor_dict + + all_gather_size = 1 if all_gather_group is None else all_gather_group.world_size + all_gather_rank = ( + 0 if all_gather_group is None else all_gather_group.rank_in_group + ) + + group = self.device_group + metadata_group = self.cpu_group + + if dst is None: + dst = (self.rank_in_group + 1) % self.world_size + assert dst < self.world_size, f"Invalid dst rank ({dst})" + + metadata_list: List[Tuple[Any, Any]] = [] + assert isinstance( + tensor_dict, dict + ), f"Expecting a dictionary, got {type(tensor_dict)}" + metadata_list, tensor_list = _split_tensor_dict(tensor_dict) + # `metadata_list` lives in CPU memory. + # `send_object_list` has serialization & deserialization, + # all happening on CPU. Therefore, we can use the CPU group. + self.send_object(metadata_list, dst=dst) + for tensor in tensor_list: + if tensor.numel() == 0: + # Skip sending empty tensors. + continue + + # send-allgather: send only a slice, then do allgather. + if all_gather_group is not None and tensor.numel() % all_gather_size == 0: + tensor = tensor.reshape(all_gather_size, -1)[all_gather_rank] + + if tensor.is_cpu: + # use metadata_group for CPU tensors + torch.distributed.send( + tensor, dst=self.ranks[dst], group=metadata_group + ) + else: + # use group for GPU tensors + torch.distributed.send(tensor, dst=self.ranks[dst], group=group) + return None + + def recv_tensor_dict( + self, + src: Optional[int] = None, + all_gather_group: Optional["GroupCoordinator"] = None, + ) -> Optional[Dict[str, Union[torch.Tensor, Any]]]: + """Recv the input tensor dictionary. + NOTE: `src` is the local rank of the source rank. + """ + # Bypass the function if we are using only 1 GPU. + if not torch.distributed.is_initialized() or self.world_size == 1: + return None + + all_gather_size = 1 if all_gather_group is None else all_gather_group.world_size + all_gather_rank = ( + 0 if all_gather_group is None else all_gather_group.rank_in_group + ) + + group = self.device_group + metadata_group = self.cpu_group + + if src is None: + src = (self.rank_in_group - 1) % self.world_size + assert src < self.world_size, f"Invalid src rank ({src})" + + recv_metadata_list = self.recv_object(src=src) + tensor_dict: Dict[str, Any] = {} + for key, value in recv_metadata_list: + if isinstance(value, TensorMetadata): + tensor = torch.empty(value.size, dtype=value.dtype, device=value.device) + if tensor.numel() == 0: + # Skip broadcasting empty tensors. + tensor_dict[key] = tensor + continue + + # send-allgather: send only a slice, then do allgather. + use_all_gather = ( + all_gather_group is not None + and tensor.numel() % all_gather_size == 0 + ) + + if use_all_gather: + orig_shape = tensor.shape + tensor = tensor.reshape(all_gather_size, -1)[all_gather_rank] + + if tensor.is_cpu: + # use metadata_group for CPU tensors + torch.distributed.recv( + tensor, src=self.ranks[src], group=metadata_group + ) + else: + # use group for GPU tensors + torch.distributed.recv(tensor, src=self.ranks[src], group=group) + if use_all_gather: + # do the allgather + tensor = all_gather_group.all_gather(tensor, dim=0) # type: ignore + tensor = tensor.reshape(orig_shape) + + tensor_dict[key] = tensor + else: + tensor_dict[key] = value + return tensor_dict + + def barrier(self): + """Barrier synchronization among the group. + NOTE: don't use `device_group` here! `barrier` in NCCL is + terrible because it is internally a broadcast operation with + secretly created GPU tensors. It is easy to mess up the current + device. Use the CPU group instead. + """ + torch.distributed.barrier(group=self.cpu_group) + + def send(self, tensor: torch.Tensor, dst: Optional[int] = None) -> None: + """Sends a tensor to the destination rank in a non-blocking way""" + """NOTE: `dst` is the local rank of the destination rank.""" + if dst is None: + dst = (self.rank_in_group + 1) % self.world_size + + pynccl_comm = self.pynccl_comm + if pynccl_comm is not None and not pynccl_comm.disabled: + pynccl_comm.send(tensor, dst) + else: + torch.distributed.send(tensor, self.ranks[dst], self.device_group) + + def recv( + self, size: torch.Size, dtype: torch.dtype, src: Optional[int] = None + ) -> torch.Tensor: + """Receives a tensor from the source rank.""" + """NOTE: `src` is the local rank of the source rank.""" + if src is None: + src = (self.rank_in_group - 1) % self.world_size + + tensor = torch.empty(size, dtype=dtype, device=self.device) + pynccl_comm = self.pynccl_comm + if pynccl_comm is not None and not pynccl_comm.disabled: + pynccl_comm.recv(tensor, src) + else: + torch.distributed.recv(tensor, self.ranks[src], self.device_group) + return tensor + + def destroy(self): + if self.device_group is not None: + torch.distributed.destroy_process_group(self.device_group) + self.device_group = None + if self.cpu_group is not None: + torch.distributed.destroy_process_group(self.cpu_group) + self.cpu_group = None + if self.pynccl_comm is not None: + self.pynccl_comm = None + if self.ca_comm is not None: + self.ca_comm = None + if self.mq_broadcaster is not None: + self.mq_broadcaster = None + + +_WORLD: Optional[GroupCoordinator] = None + + +def get_world_group() -> GroupCoordinator: + assert _WORLD is not None, "world group is not initialized" + return _WORLD + + +def init_world_group( + ranks: List[int], local_rank: int, backend: str +) -> GroupCoordinator: + return GroupCoordinator( + group_ranks=[ranks], + local_rank=local_rank, + torch_distributed_backend=backend, + use_pynccl=False, + use_custom_allreduce=False, + use_tpu_communicator=False, + use_hpu_communicator=False, + use_xpu_communicator=False, + group_name="world", + ) + + +def init_model_parallel_group( + group_ranks: List[List[int]], + local_rank: int, + backend: str, + use_custom_allreduce: Optional[bool] = None, + use_message_queue_broadcaster: bool = False, + group_name: Optional[str] = None, +) -> GroupCoordinator: + if use_custom_allreduce is None: + use_custom_allreduce = _ENABLE_CUSTOM_ALL_REDUCE + return GroupCoordinator( + group_ranks=group_ranks, + local_rank=local_rank, + torch_distributed_backend=backend, + use_pynccl=True, + use_custom_allreduce=use_custom_allreduce, + use_tpu_communicator=True, + use_hpu_communicator=True, + use_xpu_communicator=True, + use_message_queue_broadcaster=use_message_queue_broadcaster, + group_name=group_name, + ) + + +_TP: Optional[GroupCoordinator] = None + + +def get_tp_group() -> GroupCoordinator: + assert _TP is not None, "tensor model parallel group is not initialized" + return _TP + + +# kept for backward compatibility +get_tensor_model_parallel_group = get_tp_group + +_PP: Optional[GroupCoordinator] = None + + +def get_pp_group() -> GroupCoordinator: + assert _PP is not None, "pipeline model parallel group is not initialized" + return _PP + + +# kept for backward compatibility +get_pipeline_model_parallel_group = get_pp_group + + +@contextmanager +def graph_capture(): + """ + `graph_capture` is a context manager which should surround the code that + is capturing the CUDA graph. Its main purpose is to ensure that the + some operations will be run after the graph is captured, before the graph + is replayed. It returns a `GraphCaptureContext` object which contains the + necessary data for the graph capture. Currently, it only contains the + stream that the graph capture is running on. This stream is set to the + current CUDA stream when the context manager is entered and reset to the + default stream when the context manager is exited. This is to ensure that + the graph capture is running on a separate stream from the default stream, + in order to explicitly distinguish the kernels to capture + from other kernels possibly launched on background in the default stream. + """ + with get_tp_group().graph_capture() as context, get_pp_group().graph_capture( + context + ): + yield context + + +_ENABLE_CUSTOM_ALL_REDUCE = True + + +def set_custom_all_reduce(enable: bool): + global _ENABLE_CUSTOM_ALL_REDUCE + _ENABLE_CUSTOM_ALL_REDUCE = enable + + +def init_distributed_environment( + world_size: int = -1, + rank: int = -1, + distributed_init_method: str = "env://", + local_rank: int = -1, + backend: str = "nccl", +): + print( + "world_size=%d rank=%d local_rank=%d " "distributed_init_method=%s backend=%s", + world_size, + rank, + local_rank, + distributed_init_method, + backend, + ) + if not torch.distributed.is_initialized(): + assert distributed_init_method is not None, ( + "distributed_init_method must be provided when initializing " + "distributed environment" + ) + # this backend is used for WORLD + torch.distributed.init_process_group( + backend=backend, + init_method=distributed_init_method, + world_size=world_size, + rank=rank, + ) + # set the local rank + # local_rank is not available in torch ProcessGroup, + # see https://github.com/pytorch/pytorch/issues/122816 + if local_rank == -1: + # local rank not set, this usually happens in single-node + # setting, where we can use rank as local rank + if distributed_init_method == "env://": + local_rank = envs.LOCAL_RANK + else: + local_rank = rank + global _WORLD + if _WORLD is None: + ranks = list(range(torch.distributed.get_world_size())) + _WORLD = init_world_group(ranks, local_rank, backend) + else: + assert ( + _WORLD.world_size == torch.distributed.get_world_size() + ), "world group already initialized with a different world size" + + +def initialize_model_parallel( + tensor_model_parallel_size: int = 1, + pipeline_model_parallel_size: int = 1, + backend: Optional[str] = None, +) -> None: + """ + Initialize model parallel groups. + + Arguments: + tensor_model_parallel_size: number of GPUs used for tensor model + parallelism. + pipeline_model_parallel_size: number of GPUs used for pipeline model + parallelism. + + Let's say we have a total of 8 GPUs denoted by g0 ... g7 and we + use 2 GPUs to parallelize the model tensor, and 4 GPUs to parallelize + the model pipeline. The present function will + create 4 tensor model-parallel groups and 2 pipeline model-parallel groups: + 4 tensor model-parallel groups: + [g0, g1], [g2, g3], [g4, g5], [g6, g7] + 2 pipeline model-parallel groups: + [g0, g2, g4, g6], [g1, g3, g5, g7] + Note that for efficiency, the caller should make sure adjacent ranks + are on the same DGX box. For example if we are using 2 DGX-1 boxes + with a total of 16 GPUs, rank 0 to 7 belong to the first box and + ranks 8 to 15 belong to the second box. + """ + # Get world size and rank. Ensure some consistencies. + assert torch.distributed.is_initialized() + world_size: int = torch.distributed.get_world_size() + backend = backend or torch.distributed.get_backend(get_world_group().device_group) + + if world_size != tensor_model_parallel_size * pipeline_model_parallel_size: + raise RuntimeError( + f"world_size ({world_size}) is not equal to " + f"tensor_model_parallel_size ({tensor_model_parallel_size}) x " + f"pipeline_model_parallel_size ({pipeline_model_parallel_size})" + ) + + # Build the tensor model-parallel groups. + num_tensor_model_parallel_groups: int = world_size // tensor_model_parallel_size + global _TP + assert _TP is None, "tensor model parallel group is already initialized" + group_ranks = [] + for i in range(num_tensor_model_parallel_groups): + ranks = list( + range(i * tensor_model_parallel_size, (i + 1) * tensor_model_parallel_size) + ) + group_ranks.append(ranks) + + # message queue broadcaster is only used in tensor model parallel group + _TP = init_model_parallel_group( + group_ranks, + get_world_group().local_rank, + backend, + use_message_queue_broadcaster=True, + group_name="tp", + ) + + # Build the pipeline model-parallel groups. + num_pipeline_model_parallel_groups: int = world_size // pipeline_model_parallel_size + global _PP + assert _PP is None, "pipeline model parallel group is already initialized" + group_ranks = [] + for i in range(num_pipeline_model_parallel_groups): + ranks = list(range(i, world_size, num_pipeline_model_parallel_groups)) + group_ranks.append(ranks) + # pipeline parallel does not need custom allreduce + _PP = init_model_parallel_group( + group_ranks, + get_world_group().local_rank, + backend, + use_custom_allreduce=False, + group_name="pp", + ) + + +def ensure_model_parallel_initialized( + tensor_model_parallel_size: int, + pipeline_model_parallel_size: int, + backend: Optional[str] = None, +) -> None: + """Helper to initialize model parallel groups if they are not initialized, + or ensure tensor-parallel and pipeline-parallel sizes are equal to expected + values if the model parallel groups are initialized. + """ + backend = backend or torch.distributed.get_backend(get_world_group().device_group) + if not model_parallel_is_initialized(): + initialize_model_parallel( + tensor_model_parallel_size, pipeline_model_parallel_size, backend + ) + return + + assert get_tensor_model_parallel_world_size() == tensor_model_parallel_size, ( + "tensor parallel group already initialized, but of unexpected size: " + f"{get_tensor_model_parallel_world_size()=} vs. " + f"{tensor_model_parallel_size=}" + ) + pp_world_size = get_pp_group().world_size + assert pp_world_size == pipeline_model_parallel_size, ( + "pipeline parallel group already initialized, but of unexpected size: " + f"{pp_world_size=} vs. " + f"{pipeline_model_parallel_size=}" + ) + + +def model_parallel_is_initialized(): + """Check if tensor and pipeline parallel groups are initialized.""" + return _TP is not None and _PP is not None + + +_TP_STATE_PATCHED = False + + +@contextmanager +def patch_tensor_parallel_group(tp_group: GroupCoordinator): + """Patch the tp group temporarily until this function ends. + + This method is for draft workers of speculative decoding to run draft model + with different tp degree from that of target model workers. + + Args: + tp_group (GroupCoordinator): the tp group coordinator + """ + global _TP_STATE_PATCHED + assert not _TP_STATE_PATCHED, "Should not call when it's already patched" + + _TP_STATE_PATCHED = True + old_tp_group = get_tp_group() + global _TP + _TP = tp_group + try: + yield + finally: + # restore the original state + _TP_STATE_PATCHED = False + _TP = old_tp_group + + +def get_tensor_model_parallel_world_size(): + """Return world size for the tensor model parallel group.""" + return get_tp_group().world_size + + +def get_tensor_model_parallel_rank(): + """Return my rank for the tensor model parallel group.""" + return get_tp_group().rank_in_group + + +def destroy_model_parallel(): + """Set the groups to none and destroy them.""" + global _TP + if _TP: + _TP.destroy() + _TP = None + + global _PP + if _PP: + _PP.destroy() + _PP = None + + +def destroy_distributed_environment(): + global _WORLD + if _WORLD: + _WORLD.destroy() + _WORLD = None + if torch.distributed.is_initialized(): + torch.distributed.destroy_process_group() + + +def cleanup_dist_env_and_memory(shutdown_ray: bool = False): + destroy_model_parallel() + destroy_distributed_environment() + with contextlib.suppress(AssertionError): + torch.distributed.destroy_process_group() + if shutdown_ray: + import ray # Lazy import Ray + + ray.shutdown() + gc.collect() + if not current_platform.is_cpu(): + torch.cuda.empty_cache() + + +def in_the_same_node_as(pg: ProcessGroup, source_rank: int = 0) -> List[bool]: + """ + This is a collective operation that returns if each rank is in the same node + as the source rank. It tests if processes are attached to the same + memory system (shared access to shared memory). + """ + assert ( + torch.distributed.get_backend(pg) != torch.distributed.Backend.NCCL + ), "in_the_same_node_as should be tested with a non-NCCL group." + # local rank inside the group + rank = torch.distributed.get_rank(group=pg) + world_size = torch.distributed.get_world_size(group=pg) + + # local tensor in each process to store the result + is_in_the_same_node = torch.tensor([0] * world_size, dtype=torch.int32) + + # global ranks of the processes in the group + ranks = torch.distributed.get_process_group_ranks(pg) + + magic_message = b"magic_message" + shm = None + + try: + with contextlib.suppress(OSError): + if rank == source_rank: + # create a shared memory segment + shm = shared_memory.SharedMemory(create=True, size=128) + shm.buf[: len(magic_message)] = magic_message + torch.distributed.broadcast_object_list( + [shm.name], src=ranks[source_rank], group=pg + ) + is_in_the_same_node[rank] = 1 + else: + # try to open the shared memory segment + recv = [None] + torch.distributed.broadcast_object_list( + recv, src=ranks[source_rank], group=pg + ) + name = recv[0] + # fix to https://stackoverflow.com/q/62748654/9191338 + # Python incorrectly tracks shared memory even if it is not + # created by the process. The following patch is a workaround. + with patch( + "multiprocessing.resource_tracker.register", + lambda *args, **kwargs: None, + ): + shm = shared_memory.SharedMemory(name=name) + if shm.buf[: len(magic_message)] == magic_message: + is_in_the_same_node[rank] = 1 + except Exception as e: + print("Error ignored in is_in_the_same_node: %s", e) + finally: + if shm: + shm.close() + + torch.distributed.barrier(group=pg) + + # clean up the shared memory segment + with contextlib.suppress(OSError): + if rank == source_rank and shm: + shm.unlink() + torch.distributed.all_reduce(is_in_the_same_node, group=pg) + + return [x == 1 for x in is_in_the_same_node.tolist()] diff --git a/archive/kt-sft/ktransformers/server/balance_serve/inference/distributed/pynccl.py b/archive/kt-sft/ktransformers/server/balance_serve/inference/distributed/pynccl.py new file mode 100644 index 0000000..98be81d --- /dev/null +++ b/archive/kt-sft/ktransformers/server/balance_serve/inference/distributed/pynccl.py @@ -0,0 +1,201 @@ +from contextlib import contextmanager +from typing import Optional, Union + +# ===================== import region ===================== +import torch +import torch.distributed as dist +from torch.distributed import ProcessGroup, ReduceOp + +from server.inference.distributed.pynccl_wrapper import ( + NCCLLibrary, + buffer_type, + cudaStream_t, + ncclComm_t, + ncclDataTypeEnum, + ncclRedOpTypeEnum, + ncclUniqueId, +) +from server.inference.distributed.utils import StatelessProcessGroup + + +class PyNcclCommunicator: + + def __init__( + self, + group: Union[ProcessGroup, StatelessProcessGroup], + device: Union[int, str, torch.device], + library_path: Optional[str] = None, + ): + """ + Args: + group: the process group to work on. If None, it will use the + default process group. + device: the device to bind the PyNcclCommunicator to. If None, + it will be bind to f"cuda:{local_rank}". + library_path: the path to the NCCL library. If None, it will + use the default library path. + It is the caller's responsibility to make sure each communicator + is bind to a unique device. + """ + if not isinstance(group, StatelessProcessGroup): + assert dist.is_initialized() + assert ( + dist.get_backend(group) != dist.Backend.NCCL + ), "PyNcclCommunicator should be attached to a non-NCCL group." + # note: this rank is the rank in the group + self.rank = dist.get_rank(group) + self.world_size = dist.get_world_size(group) + else: + self.rank = group.rank + self.world_size = group.world_size + + self.group = group + + # if world_size == 1, no need to create communicator + if self.world_size == 1: + self.available = False + self.disabled = True + self.stream = None + return + try: + self.nccl = NCCLLibrary(library_path) + except Exception: + # disable because of missing NCCL library + # e.g. in a non-GPU environment + self.available = False + self.disabled = True + self.stream = None + return + + self.available = True + self.disabled = False + + print("vLLM is using nccl==%s", self.nccl.ncclGetVersion()) + + if self.rank == 0: + # get the unique id from NCCL + self.unique_id = self.nccl.ncclGetUniqueId() + else: + # construct an empty unique id + self.unique_id = ncclUniqueId() + + if not isinstance(group, StatelessProcessGroup): + tensor = torch.ByteTensor(list(self.unique_id.internal)) + ranks = dist.get_process_group_ranks(group) + # arg `src` in `broadcast` is the global rank + dist.broadcast(tensor, src=ranks[0], group=group) + byte_list = tensor.tolist() + for i, byte in enumerate(byte_list): + self.unique_id.internal[i] = byte + else: + self.unique_id = group.broadcast_obj(self.unique_id, src=0) + if isinstance(device, int): + device = torch.device(f"cuda:{device}") + elif isinstance(device, str): + device = torch.device(device) + # now `device` is a `torch.device` object + assert isinstance(device, torch.device) + self.device = device + # nccl communicator and stream will use this device + # `torch.cuda.device` is a context manager that changes the + # current cuda device to the specified one + with torch.cuda.device(device): + self.comm: ncclComm_t = self.nccl.ncclCommInitRank( + self.world_size, self.unique_id, self.rank + ) + self.stream = torch.cuda.Stream() + + # A small all_reduce for warmup. + data = torch.zeros(1, device=device) + self.all_reduce(data) + self.stream.synchronize() + del data + + # by default it is disabled, e.g. in profiling models and prefill phase. + # to use it, use under `with obj.change_state(enable=True)`, usually + # when we are using CUDA graph. + self.disabled = True + + def all_reduce( + self, tensor: torch.Tensor, op: ReduceOp = ReduceOp.SUM, stream=None + ): + if self.disabled: + return + # nccl communicator created on a specific device + # will only work on tensors on the same device + # otherwise it will cause "illegal memory access" + assert tensor.device == self.device, ( + f"this nccl communicator is created to work on {self.device}, " + f"but the input tensor is on {tensor.device}" + ) + if stream is None: + stream = self.stream + self.nccl.ncclAllReduce( + buffer_type(tensor.data_ptr()), + buffer_type(tensor.data_ptr()), + tensor.numel(), + ncclDataTypeEnum.from_torch(tensor.dtype), + ncclRedOpTypeEnum.from_torch(op), + self.comm, + cudaStream_t(stream.cuda_stream), + ) + + def send(self, tensor: torch.Tensor, dst: int, stream=None): + if self.disabled: + return + assert tensor.device == self.device, ( + f"this nccl communicator is created to work on {self.device}, " + f"but the input tensor is on {tensor.device}" + ) + if stream is None: + stream = self.stream + self.nccl.ncclSend( + buffer_type(tensor.data_ptr()), + tensor.numel(), + ncclDataTypeEnum.from_torch(tensor.dtype), + dst, + self.comm, + cudaStream_t(stream.cuda_stream), + ) + + def recv(self, tensor: torch.Tensor, src: int, stream=None): + if self.disabled: + return + assert tensor.device == self.device, ( + f"this nccl communicator is created to work on {self.device}, " + f"but the input tensor is on {tensor.device}" + ) + if stream is None: + stream = self.stream + self.nccl.ncclRecv( + buffer_type(tensor.data_ptr()), + tensor.numel(), + ncclDataTypeEnum.from_torch(tensor.dtype), + src, + self.comm, + cudaStream_t(stream.cuda_stream), + ) + + @contextmanager + def change_state( + self, enable: Optional[bool] = None, stream: Optional[torch.cuda.Stream] = None + ): + """ + A context manager to change the state of the communicator. + """ + if enable is None: + # guess a default value when not specified + enable = self.available + + if stream is None: + stream = self.stream + + old_disable = self.disabled + old_stream = self.stream + + self.stream = stream + self.disabled = not enable + yield + + self.disabled = old_disable + self.stream = old_stream diff --git a/archive/kt-sft/ktransformers/server/balance_serve/inference/distributed/pynccl_wrapper.py b/archive/kt-sft/ktransformers/server/balance_serve/inference/distributed/pynccl_wrapper.py new file mode 100644 index 0000000..aa7d4a7 --- /dev/null +++ b/archive/kt-sft/ktransformers/server/balance_serve/inference/distributed/pynccl_wrapper.py @@ -0,0 +1,276 @@ +# This file is a pure Python wrapper for the NCCL library. +# The main purpose is to use NCCL combined with CUDA graph. +# Before writing this script, we tried the following approach: +# 1. We tried to use `cupy`, it calls NCCL correctly, but `cupy` itself +# often gets stuck when initializing the NCCL communicator. +# 2. We tried to use `torch.distributed`, but `torch.distributed.all_reduce` +# contains many other potential cuda APIs, that are not allowed during +# capturing the CUDA graph. For further details, please check +# https://discuss.pytorch.org/t/pytorch-cudagraph-with-nccl-operation-failed/ . +# +# Another rejected idea is to write a C/C++ binding for NCCL. It is usually +# doable, but we often encounter issues related with nccl versions, and need +# to switch between different versions of NCCL. See +# https://github.com/NVIDIA/nccl/issues/1234 for more details. +# A C/C++ binding is not flexible enough to handle this. It requires +# recompilation of the code every time we want to switch between different +# versions. This current implementation, with a **pure** Python wrapper, is +# more flexible. We can easily switch between different versions of NCCL by +# changing the environment variable `VLLM_NCCL_SO_PATH`, or the `so_file` +# variable in the code. + +import ctypes +import platform +from dataclasses import dataclass +from typing import Any, Dict, List, Optional + +import torch +from torch.distributed import ReduceOp + +from server.utils import find_nccl_library + + +# === export types and functions from nccl to Python === +# for the original nccl definition, please check +# https://github.com/NVIDIA/nccl/blob/master/src/nccl.h.in + +ncclResult_t = ctypes.c_int +ncclComm_t = ctypes.c_void_p + + +class ncclUniqueId(ctypes.Structure): + _fields_ = [("internal", ctypes.c_byte * 128)] + + +cudaStream_t = ctypes.c_void_p +buffer_type = ctypes.c_void_p + +ncclDataType_t = ctypes.c_int + + +class ncclDataTypeEnum: + ncclInt8 = 0 + ncclChar = 0 + ncclUint8 = 1 + ncclInt32 = 2 + ncclInt = 2 + ncclUint32 = 3 + ncclInt64 = 4 + ncclUint64 = 5 + ncclFloat16 = 6 + ncclHalf = 6 + ncclFloat32 = 7 + ncclFloat = 7 + ncclFloat64 = 8 + ncclDouble = 8 + ncclBfloat16 = 9 + ncclNumTypes = 10 + + @classmethod + def from_torch(cls, dtype: torch.dtype) -> int: + if dtype == torch.int8: + return cls.ncclInt8 + if dtype == torch.uint8: + return cls.ncclUint8 + if dtype == torch.int32: + return cls.ncclInt32 + if dtype == torch.int64: + return cls.ncclInt64 + if dtype == torch.float16: + return cls.ncclFloat16 + if dtype == torch.float32: + return cls.ncclFloat32 + if dtype == torch.float64: + return cls.ncclFloat64 + if dtype == torch.bfloat16: + return cls.ncclBfloat16 + raise ValueError(f"Unsupported dtype: {dtype}") + + +ncclRedOp_t = ctypes.c_int + + +class ncclRedOpTypeEnum: + ncclSum = 0 + ncclProd = 1 + ncclMax = 2 + ncclMin = 3 + ncclAvg = 4 + ncclNumOps = 5 + + @classmethod + def from_torch(cls, op: ReduceOp) -> int: + if op == ReduceOp.SUM: + return cls.ncclSum + if op == ReduceOp.PRODUCT: + return cls.ncclProd + if op == ReduceOp.MAX: + return cls.ncclMax + if op == ReduceOp.MIN: + return cls.ncclMin + if op == ReduceOp.AVG: + return cls.ncclAvg + raise ValueError(f"Unsupported op: {op}") + + +@dataclass +class Function: + name: str + restype: Any + argtypes: List[Any] + + +class NCCLLibrary: + exported_functions = [ + # const char* ncclGetErrorString(ncclResult_t result) + Function("ncclGetErrorString", ctypes.c_char_p, [ncclResult_t]), + # ncclResult_t ncclGetVersion(int *version); + Function("ncclGetVersion", ncclResult_t, + [ctypes.POINTER(ctypes.c_int)]), + # ncclResult_t ncclGetUniqueId(ncclUniqueId* uniqueId); + Function("ncclGetUniqueId", ncclResult_t, + [ctypes.POINTER(ncclUniqueId)]), + # ncclResult_t ncclCommInitRank( + # ncclComm_t* comm, int nranks, ncclUniqueId commId, int rank); + # note that ncclComm_t is a pointer type, so the first argument + # is a pointer to a pointer + Function("ncclCommInitRank", ncclResult_t, [ + ctypes.POINTER(ncclComm_t), ctypes.c_int, ncclUniqueId, + ctypes.c_int + ]), + # ncclResult_t ncclAllReduce( + # const void* sendbuff, void* recvbuff, size_t count, + # ncclDataType_t datatype, ncclRedOp_t op, ncclComm_t comm, + # cudaStream_t stream); + # note that cudaStream_t is a pointer type, so the last argument + # is a pointer + Function("ncclAllReduce", ncclResult_t, [ + buffer_type, buffer_type, ctypes.c_size_t, ncclDataType_t, + ncclRedOp_t, ncclComm_t, cudaStream_t + ]), + + # ncclResult_t ncclSend( + # const void* sendbuff, size_t count, ncclDataType_t datatype, + # int dest, ncclComm_t comm, cudaStream_t stream); + Function("ncclSend", ncclResult_t, [ + buffer_type, ctypes.c_size_t, ncclDataType_t, ctypes.c_int, + ncclComm_t, cudaStream_t + ]), + + # ncclResult_t ncclRecv( + # void* recvbuff, size_t count, ncclDataType_t datatype, + # int src, ncclComm_t comm, cudaStream_t stream); + Function("ncclRecv", ncclResult_t, [ + buffer_type, ctypes.c_size_t, ncclDataType_t, ctypes.c_int, + ncclComm_t, cudaStream_t + ]), + + # be cautious! this is a collective call, it will block until all + # processes in the communicator have called this function. + # because Python object destruction can happen in random order, + # it is better not to call it at all. + # ncclResult_t ncclCommDestroy(ncclComm_t comm); + Function("ncclCommDestroy", ncclResult_t, [ncclComm_t]), + ] + + # class attribute to store the mapping from the path to the library + # to avoid loading the same library multiple times + path_to_library_cache: Dict[str, Any] = {} + + # class attribute to store the mapping from library path + # to the corresponding dictionary + path_to_dict_mapping: Dict[str, Dict[str, Any]] = {} + + def __init__(self, so_file: Optional[str] = None): + + so_file = so_file or find_nccl_library() + + try: + if so_file not in NCCLLibrary.path_to_dict_mapping: + lib = ctypes.CDLL(so_file) + NCCLLibrary.path_to_library_cache[so_file] = lib + self.lib = NCCLLibrary.path_to_library_cache[so_file] + except Exception as e: + print( + "Failed to load NCCL library from %s ." + "It is expected if you are not running on NVIDIA/AMD GPUs." + "Otherwise, the nccl library might not exist, be corrupted " + "or it does not support the current platform %s." + "If you already have the library, please set the " + "environment variable VLLM_NCCL_SO_PATH" + " to point to the correct nccl library path.", so_file, + platform.platform()) + raise e + + if so_file not in NCCLLibrary.path_to_dict_mapping: + _funcs: Dict[str, Any] = {} + for func in NCCLLibrary.exported_functions: + f = getattr(self.lib, func.name) + f.restype = func.restype + f.argtypes = func.argtypes + _funcs[func.name] = f + NCCLLibrary.path_to_dict_mapping[so_file] = _funcs + self._funcs = NCCLLibrary.path_to_dict_mapping[so_file] + + def ncclGetErrorString(self, result: ncclResult_t) -> str: + return self._funcs["ncclGetErrorString"](result).decode("utf-8") + + def NCCL_CHECK(self, result: ncclResult_t) -> None: + if result != 0: + error_str = self.ncclGetErrorString(result) + raise RuntimeError(f"NCCL error: {error_str}") + + def ncclGetVersion(self) -> str: + version = ctypes.c_int() + self.NCCL_CHECK(self._funcs["ncclGetVersion"](ctypes.byref(version))) + version_str = str(version.value) + # something like 21903 --> "2.19.3" + major = version_str[0].lstrip("0") + minor = version_str[1:3].lstrip("0") + patch = version_str[3:].lstrip("0") + return f"{major}.{minor}.{patch}" + + def ncclGetUniqueId(self) -> ncclUniqueId: + unique_id = ncclUniqueId() + self.NCCL_CHECK(self._funcs["ncclGetUniqueId"]( + ctypes.byref(unique_id))) + return unique_id + + def ncclCommInitRank(self, world_size: int, unique_id: ncclUniqueId, + rank: int) -> ncclComm_t: + comm = ncclComm_t() + self.NCCL_CHECK(self._funcs["ncclCommInitRank"](ctypes.byref(comm), + world_size, unique_id, + rank)) + return comm + + def ncclAllReduce(self, sendbuff: buffer_type, recvbuff: buffer_type, + count: int, datatype: int, op: int, comm: ncclComm_t, + stream: cudaStream_t) -> None: + # `datatype` actually should be `ncclDataType_t` + # and `op` should be `ncclRedOp_t` + # both are aliases of `ctypes.c_int` + # when we pass int to a function, it will be converted to `ctypes.c_int` + # by ctypes automatically + self.NCCL_CHECK(self._funcs["ncclAllReduce"](sendbuff, recvbuff, count, + datatype, op, comm, + stream)) + + def ncclSend(self, sendbuff: buffer_type, count: int, datatype: int, + dest: int, comm: ncclComm_t, stream: cudaStream_t) -> None: + self.NCCL_CHECK(self._funcs["ncclSend"](sendbuff, count, datatype, + dest, comm, stream)) + + def ncclRecv(self, recvbuff: buffer_type, count: int, datatype: int, + src: int, comm: ncclComm_t, stream: cudaStream_t) -> None: + self.NCCL_CHECK(self._funcs["ncclRecv"](recvbuff, count, datatype, src, + comm, stream)) + + def ncclCommDestroy(self, comm: ncclComm_t) -> None: + self.NCCL_CHECK(self._funcs["ncclCommDestroy"](comm)) + + +__all__ = [ + "NCCLLibrary", "ncclDataTypeEnum", "ncclRedOpTypeEnum", "ncclUniqueId", + "ncclComm_t", "cudaStream_t", "buffer_type" +] diff --git a/archive/kt-sft/ktransformers/server/balance_serve/inference/distributed/utils.py b/archive/kt-sft/ktransformers/server/balance_serve/inference/distributed/utils.py new file mode 100644 index 0000000..475433c --- /dev/null +++ b/archive/kt-sft/ktransformers/server/balance_serve/inference/distributed/utils.py @@ -0,0 +1,219 @@ +# Copyright 2023 The vLLM team. +# Adapted from +# https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/core/tensor_parallel/utils.py +# Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved. +import dataclasses +import pickle +import time +from collections import deque +from typing import Any, Deque, Dict, Optional, Sequence, Tuple + +import torch +from torch.distributed import TCPStore + +import server.envs as envs + + +def ensure_divisibility(numerator, denominator): + """Ensure that numerator is divisible by the denominator.""" + assert numerator % denominator == 0, "{} is not divisible by {}".format( + numerator, denominator + ) + + +def divide(numerator, denominator): + """Ensure that numerator is divisible by the denominator and return + the division value.""" + ensure_divisibility(numerator, denominator) + return numerator // denominator + + +def split_tensor_along_last_dim( + tensor: torch.Tensor, + num_partitions: int, + contiguous_split_chunks: bool = False, +) -> Sequence[torch.Tensor]: + """Split a tensor along its last dimension. + + Arguments: + tensor: input tensor. + num_partitions: number of partitions to split the tensor + contiguous_split_chunks: If True, make each chunk contiguous + in memory. + + Returns: + A list of Tensors + """ + # Get the size and dimension. + last_dim = tensor.dim() - 1 + last_dim_size = divide(tensor.size()[last_dim], num_partitions) + # Split. + tensor_list = torch.split(tensor, last_dim_size, dim=last_dim) + # NOTE: torch.split does not create contiguous tensors by default. + if contiguous_split_chunks: + return tuple(chunk.contiguous() for chunk in tensor_list) + + return tensor_list + + +def get_pp_indices( + num_hidden_layers: int, pp_rank: int, pp_size: int +) -> Tuple[int, int]: + """Try to evenly distribute layers across partitions. + If the number of layers is not divisible by the number of partitions, + the last partition will have the remaining layers. + """ + partition_list_str = envs.VLLM_PP_LAYER_PARTITION + if partition_list_str is not None: + try: + partitions = [int(layer) for layer in partition_list_str.split(",")] + except ValueError as err: + raise ValueError( + "Invalid partition string: {}".format(partition_list_str) + ) from err + if len(partitions) != pp_size: + raise ValueError(f"{len(partitions)=} does not match {pp_size=}.") + if sum(partitions) != num_hidden_layers: + raise ValueError(f"{sum(partitions)=} does not match {num_hidden_layers=}.") + start_layer = sum(partitions[:pp_rank]) + end_layer = start_layer + partitions[pp_rank] + else: + layers_per_partition = num_hidden_layers // pp_size + start_layer = pp_rank * layers_per_partition + end_layer = start_layer + layers_per_partition + + if pp_rank == pp_size - 1: + end_layer = num_hidden_layers + + return (start_layer, end_layer) + + +@dataclasses.dataclass +class StatelessProcessGroup: + """A dataclass to hold a metadata store, and the rank, world_size of the + group. Only use it to communicate metadata between processes. + For data-plane communication, create NCCL-related objects. + """ + + rank: int + world_size: int + store: torch._C._distributed_c10d.Store + data_expiration_seconds: int = 3600 # 1 hour + + # dst rank -> counter + send_dst_counter: Dict[int, int] = dataclasses.field(default_factory=dict) + # src rank -> counter + recv_src_counter: Dict[int, int] = dataclasses.field(default_factory=dict) + broadcast_send_counter: int = 0 + broadcast_recv_src_counter: Dict[int, int] = dataclasses.field(default_factory=dict) + + # A deque to store the data entries, with key and timestamp. + entries: Deque[Tuple[str, float]] = dataclasses.field(default_factory=deque) + + def __post_init__(self): + assert self.rank < self.world_size + self.send_dst_counter = {i: 0 for i in range(self.world_size)} + self.recv_src_counter = {i: 0 for i in range(self.world_size)} + self.broadcast_recv_src_counter = {i: 0 for i in range(self.world_size)} + + def send_obj(self, obj: Any, dst: int): + """Send an object to a destination rank.""" + self.expire_data() + key = f"send_to/{dst}/{self.send_dst_counter[dst]}" + self.store.set(key, pickle.dumps(obj)) + self.send_dst_counter[dst] += 1 + self.entries.append((key, time.time())) + + def expire_data(self): + """Expire data that is older than `data_expiration_seconds` seconds.""" + while self.entries: + # check the oldest entry + key, timestamp = self.entries[0] + if time.time() - timestamp > self.data_expiration_seconds: + self.store.delete_key(key) + self.entries.popleft() + else: + break + + def recv_obj(self, src: int) -> Any: + """Receive an object from a source rank.""" + obj = pickle.loads( + self.store.get(f"send_to/{self.rank}/{self.recv_src_counter[src]}") + ) + self.recv_src_counter[src] += 1 + return obj + + def broadcast_obj(self, obj: Optional[Any], src: int) -> Any: + """Broadcast an object from a source rank to all other ranks. + It does not clean up after all ranks have received the object. + Use it for limited times, e.g., for initialization. + """ + if self.rank == src: + self.expire_data() + key = f"broadcast_from/{src}/" f"{self.broadcast_send_counter}" + self.store.set(key, pickle.dumps(obj)) + self.broadcast_send_counter += 1 + self.entries.append((key, time.time())) + return obj + else: + key = f"broadcast_from/{src}/" f"{self.broadcast_recv_src_counter[src]}" + recv_obj = pickle.loads(self.store.get(key)) + self.broadcast_recv_src_counter[src] += 1 + return recv_obj + + def all_gather_obj(self, obj: Any) -> list[Any]: + """All gather an object from all ranks.""" + gathered_objs = [] + for i in range(self.world_size): + if i == self.rank: + gathered_objs.append(obj) + self.broadcast_obj(obj, src=self.rank) + else: + recv_obj = self.broadcast_obj(None, src=i) + gathered_objs.append(recv_obj) + return gathered_objs + + def barrier(self): + """A barrier to synchronize all ranks.""" + for i in range(self.world_size): + if i == self.rank: + self.broadcast_obj(None, src=self.rank) + else: + self.broadcast_obj(None, src=i) + + @staticmethod + def create( + host: str, + port: int, + rank: int, + world_size: int, + data_expiration_seconds: int = 3600, + ) -> "StatelessProcessGroup": + """A replacement for `torch.distributed.init_process_group` that does not + pollute the global state. + + If we have process A and process B called `torch.distributed.init_process_group` + to form a group, and then we want to form another group with process A, B, C, + D, it is not possible in PyTorch, because process A and process B have already + formed a group, and process C and process D cannot join that group. This + function is a workaround for this issue. + + `torch.distributed.init_process_group` is a global call, while this function + is a stateless call. It will return a `StatelessProcessGroup` object that can be + used for exchanging metadata. With this function, process A and process B + can call `StatelessProcessGroup.create` to form a group, and then process A, B, + C, and D can call `StatelessProcessGroup.create` to form another group. + """ # noqa + store = TCPStore( + host_name=host, + port=port, + world_size=world_size, + is_master=(rank == 0), + ) + + return StatelessProcessGroup( + rank=rank, + world_size=world_size, + store=store, + data_expiration_seconds=data_expiration_seconds, + ) diff --git a/archive/kt-sft/ktransformers/server/balance_serve/inference/forward_batch.py b/archive/kt-sft/ktransformers/server/balance_serve/inference/forward_batch.py new file mode 100644 index 0000000..26b4d3d --- /dev/null +++ b/archive/kt-sft/ktransformers/server/balance_serve/inference/forward_batch.py @@ -0,0 +1,284 @@ +''' +Date: 2024-11-12 14:15:16 +LastEditors: Xie Weiyu ervinxie@qq.com +LastEditTime: 2024-11-26 08:12:49 +''' +import torch +from ktransformers.server.balance_serve.settings import sched_ext +from ktransformers.server.balance_serve.inference.query_manager import QueryManager, QueryInfo +import time +from ktransformers.server.config.config import Config +class ForwardBatchInput: + + class ForwardMiniBatch: + q_indptr: torch.Tensor + kv_indptr: torch.Tensor + kv_indices: torch.Tensor + kv_last_page_len: torch.Tensor + kv_len: torch.Tensor + position_ids: torch.Tensor + tokens: torch.Tensor + batch_indices: torch.Tensor + positions: torch.Tensor + chunk_size: int + decode_batch: int + is_last_prefill_chunk: bool + logits_start: list + + temperatures: torch.Tensor + top_ps: torch.Tensor + + def __init__(self, prefill_querys_info: list[QueryInfo], decode_querys_info: list[QueryInfo], prefill_s: list[int] = None, prefill_l: list[int] = None, device = torch.device('cuda'), page_size = 256): + batch_decode = len(decode_querys_info) + batch_prefill = len(prefill_querys_info) + + self.q_indptr = torch.tensor([0], device=device, dtype=torch.int32) + self.kv_indptr = torch.tensor([0], device=device, dtype=torch.int32) + self.kv_indices = torch.tensor([], device=device, dtype=torch.int32) + self.kv_len = torch.tensor([], device=device, dtype=torch.int32) + self.kv_last_page_len = torch.tensor([], device=device, dtype=torch.int32) + self.position_ids = torch.tensor([], device=device, dtype=torch.int32) + self.tokens = torch.tensor([], device=device, dtype=torch.int32) + + self.temperatures = torch.tensor([], device=device, dtype=torch.float32) + self.top_ps = torch.tensor([], device=device, dtype=torch.float32) + + self.logits_start = [] + self.decode_batch = batch_decode + self.num_tokens = batch_decode + sum(prefill_l) + self.batch_size = batch_decode + batch_prefill + + for i, prefill_query_info in enumerate(prefill_querys_info): + if prefill_query_info != None: + prefill_kv_block_len = (prefill_query_info.active_position + prefill_l[i] + page_size - 1) // page_size if prefill_query_info is not None else 0 + # print(f"block_len: {prefill_kv_block_len}, page_size: {page_size}") + self.q_indptr = torch.concat((self.q_indptr, torch.tensor([prefill_l[i] + self.q_indptr[-1]], device=device, dtype=torch.int32)), dim=0) + self.kv_indptr = torch.concat((self.kv_indptr, torch.tensor([prefill_kv_block_len + self.kv_indptr[-1]], device=device, dtype=torch.int32)), dim=0) + self.kv_indices = torch.concat((self.kv_indices, prefill_query_info.block_index[:prefill_kv_block_len]), dim=0) + self.kv_last_page_len = torch.concat((self.kv_last_page_len, torch.tensor([(prefill_query_info.active_position + prefill_l[i]) % page_size if (prefill_query_info.active_position + prefill_l[i]) % page_size != 0 else page_size], device=device, dtype=torch.int32)), dim=0) + self.kv_len = torch.concat((self.kv_len, torch.tensor([(prefill_query_info.active_position + prefill_l[i])], device=device, dtype=torch.int32)), dim=0) + self.position_ids = torch.concat((self.position_ids, torch.arange(prefill_s[i], prefill_l[i] + prefill_s[i], device=device, dtype=torch.int32)), dim=0) + self.tokens = torch.concat((self.tokens, prefill_query_info.query_tokens[prefill_s[i]:prefill_s[i] + prefill_l[i]]), dim=0) + self.logits_start.append(prefill_l[i] - 1 if len(self.logits_start) == 0 else sum(prefill_l[:i+1])-1) + + self.temperatures = torch.concat((self.temperatures, torch.tensor([prefill_query_info.temperature], device=device, dtype=torch.float32)), dim=0) + self.top_ps = torch.concat((self.top_ps, torch.tensor([prefill_query_info.top_p], device=device, dtype=torch.float32)), dim=0) + + for decode_query_info in decode_querys_info: + decode_kv_block_len = (decode_query_info.active_position + 1 + page_size - 1) // page_size + self.q_indptr = torch.concat((self.q_indptr, torch.tensor([1 + self.q_indptr[-1]], device=device, dtype=torch.int32)), dim=0) + self.kv_indptr = torch.concat((self.kv_indptr, torch.tensor([decode_kv_block_len+self.kv_indptr[-1]], device=device, dtype=torch.int32)), dim=0) + self.kv_indices = torch.concat((self.kv_indices, decode_query_info.block_index[:decode_kv_block_len]), dim=0) + self.kv_last_page_len = torch.concat((self.kv_last_page_len, torch.tensor([(decode_query_info.active_position + 1) % page_size if (decode_query_info.active_position + 1) % page_size != 0 else page_size], device=device, dtype=torch.int32)), dim=0) + self.kv_len = torch.concat((self.kv_len, torch.tensor([(decode_query_info.active_position + 1)], device=device, dtype=torch.int32)), dim=0) + self.position_ids = torch.concat((self.position_ids, torch.arange(decode_query_info.active_position, decode_query_info.active_position + 1, device=device, dtype=torch.int32)), dim=0) + if decode_query_info.active_position > 0: + self.tokens = torch.concat((self.tokens, decode_query_info.query_tokens[decode_query_info.active_position:decode_query_info.active_position+1]), dim=0) + else: + self.tokens = torch.concat((self.tokens, torch.tensor([0], device=device, dtype=torch.int32)), dim=0) + self.logits_start.append(0 if len(self.logits_start) == 0 else self.logits_start[-1]+1) + + self.temperatures = torch.concat((self.temperatures, torch.tensor([decode_query_info.temperature], device=device, dtype=torch.float32)), dim=0) + self.top_ps = torch.concat((self.top_ps, torch.tensor([decode_query_info.top_p], device=device, dtype=torch.float32)), dim=0) + + self.q_indptr = self.q_indptr.contiguous() + self.kv_indptr = self.kv_indptr.contiguous() + self.kv_indices = self.kv_indices.contiguous() + self.kv_len = self.kv_len.contiguous() + self.kv_last_page_len = self.kv_last_page_len.contiguous() + self.position_ids = self.position_ids.contiguous() + self.tokens = self.tokens.contiguous() + + self.bsz_tensor = torch.tensor([self.batch_size], device=device, dtype=torch.int32) + + def fill(self, prefill_querys_info: list[QueryInfo], decode_querys_info: list[QueryInfo], prefill_s: list[int] = None, prefill_l: list[int] = None, device = torch.device('cuda'), page_size = 256): + batch_decode = len(decode_querys_info) + batch_prefill = len(prefill_querys_info) + + self.q_indptr = torch.tensor([0], device=device, dtype=torch.int32) + self.kv_indptr = torch.tensor([0], device=device, dtype=torch.int32) + self.kv_indices = torch.tensor([], device=device, dtype=torch.int32) + self.kv_len = torch.tensor([], device=device, dtype=torch.int32) + self.kv_last_page_len = torch.tensor([], device=device, dtype=torch.int32) + new_position_ids = torch.tensor([], device=device, dtype=torch.int32) + new_tokens = torch.tensor([], device=device, dtype=torch.int32) + + self.temperatures = torch.tensor([], device=device, dtype=torch.float32) + self.top_ps = torch.tensor([], device=device, dtype=torch.float32) + + self.logits_start = [] + self.decode_batch = batch_decode + self.num_tokens = batch_decode + sum(prefill_l) + self.batch_size = batch_decode + batch_prefill + + for i, prefill_query_info in enumerate(prefill_querys_info): + prefill_kv_block_len = (prefill_query_info.active_position + prefill_l[i] + page_size - 1) // page_size if prefill_query_info is not None else 0 + # print(f"block_len: {prefill_kv_block_len}, page_size: {page_size}") + self.q_indptr = torch.concat((self.q_indptr, torch.tensor([prefill_l[i] + self.q_indptr[-1]], device=device, dtype=torch.int32)), dim=0) + self.kv_indptr = torch.concat((self.kv_indptr, torch.tensor([prefill_kv_block_len + self.kv_indptr[-1]], device=device, dtype=torch.int32)), dim=0) + self.kv_indices = torch.concat((self.kv_indices, prefill_query_info.block_index[:prefill_kv_block_len]), dim=0) + self.kv_last_page_len = torch.concat((self.kv_last_page_len, torch.tensor([(prefill_query_info.active_position + prefill_l[i]) % page_size if (prefill_query_info.active_position + prefill_l[i]) % page_size != 0 else page_size], device=device, dtype=torch.int32)), dim=0) + self.kv_len = torch.concat((self.kv_len, torch.tensor([(prefill_query_info.active_position + prefill_l[i])], device=device, dtype=torch.int32)), dim=0) + new_position_ids = torch.concat((new_position_ids, torch.arange(prefill_s[i], prefill_l[i] + prefill_s[i], device=device, dtype=torch.int32)), dim=0) + new_tokens = torch.concat((new_tokens, prefill_query_info.query_tokens[prefill_s[i]:prefill_s[i] + prefill_l[i]]), dim=0) + self.logits_start.append(prefill_l[i] - 1 if len(self.logits_start) == 0 else sum(prefill_l[:i+1])-1) + + self.temperatures = torch.concat((self.temperatures, torch.tensor([prefill_query_info.temperature], device=device, dtype=torch.float32)), dim=0) + self.top_ps = torch.concat((self.top_ps, torch.tensor([prefill_query_info.top_p], device=device, dtype=torch.float32)), dim=0) + + + for decode_query_info in decode_querys_info: + decode_kv_block_len = (decode_query_info.active_position + 1 + page_size - 1) // page_size + self.q_indptr = torch.concat((self.q_indptr, torch.tensor([1 + self.q_indptr[-1]], device=device, dtype=torch.int32)), dim=0) + self.kv_indptr = torch.concat((self.kv_indptr, torch.tensor([decode_kv_block_len+self.kv_indptr[-1]], device=device, dtype=torch.int32)), dim=0) + self.kv_indices = torch.concat((self.kv_indices, decode_query_info.block_index[:decode_kv_block_len]), dim=0) + self.kv_last_page_len = torch.concat((self.kv_last_page_len, torch.tensor([(decode_query_info.active_position + 1) % page_size if (decode_query_info.active_position + 1) % page_size != 0 else page_size], device=device, dtype=torch.int32)), dim=0) + self.kv_len = torch.concat((self.kv_len, torch.tensor([(decode_query_info.active_position + 1)], device=device, dtype=torch.int32)), dim=0) + new_position_ids = torch.concat((new_position_ids, torch.arange(decode_query_info.active_position, decode_query_info.active_position + 1, device=device, dtype=torch.int32)), dim=0) + if decode_query_info.active_position > 0: + new_tokens = torch.concat((new_tokens, decode_query_info.query_tokens[decode_query_info.active_position:decode_query_info.active_position+1]), dim=0) + else: + new_tokens = torch.concat((new_tokens, torch.tensor([0], device=device, dtype=torch.int32)), dim=0) + self.logits_start.append(0 if len(self.logits_start) == 0 else self.logits_start[-1]+1) + + self.temperatures = torch.concat((self.temperatures, torch.tensor([decode_query_info.temperature], device=device, dtype=torch.float32)), dim=0) + self.top_ps = torch.concat((self.top_ps, torch.tensor([decode_query_info.top_p], device=device, dtype=torch.float32)), dim=0) + + + self.q_indptr = self.q_indptr.contiguous() + self.kv_indptr = self.kv_indptr.contiguous() + self.kv_indices = self.kv_indices.contiguous() + self.kv_len = self.kv_len.contiguous() + self.kv_last_page_len = self.kv_last_page_len.contiguous() + + self.bsz_tensor = torch.tensor([self.batch_size], device=device, dtype=torch.int32) + + # copy new_position_ids and new_tokens to self.position_ids and self.tokens + # print("new_position_ids: ", new_position_ids) + # self.print() + self.position_ids[:new_position_ids.size(0)].copy_(new_position_ids) + self.position_ids[new_position_ids.size(0):].zero_() + self.tokens[:new_tokens.size(0)].copy_(new_tokens) + + + forward_minibatchs: list[ForwardMiniBatch] + batch_size: int + minibatch: ForwardMiniBatch + + + + def __init__(self, batch : sched_ext.BatchQueryTodo = None, query_manager: QueryManager = None, device=None, tokens: torch.Tensor = None): + + if batch is None: + return + + + prefill_minibatches = batch.prefill_mini_batches + decode_mini_batches = [item for sublist in batch.decode_mini_batches for item in sublist] + prefill_querys_info = [] + prefill_s = [] + prefill_l = [] + decode_querys_info = [] + self.batch_size = 1 + for (id, s, l) in prefill_minibatches: + prefill_querys_info.append(query_manager.query_map[id]) + prefill_s.append(s) + prefill_l.append(l) + for decode_batch_idx in decode_mini_batches: + if query_manager.query_map[decode_batch_idx].decode_start_time is None: + query_manager.query_map[decode_batch_idx].decode_start_time =time.time() + decode_querys_info.append(query_manager.query_map[decode_batch_idx]) + + + minibatch = ForwardBatchInput.ForwardMiniBatch(prefill_querys_info, decode_querys_info, prefill_s, prefill_l, device = query_manager.device, page_size = query_manager.page_size) + + self.minibatch = minibatch + + @classmethod + def gen_max_forward_batch( + cls, + device=None, + tokens: torch.Tensor = None, + num_mini_batches: int = 1, + max_seq_length: int = 4096, # TODO: add to yaml + prefill_query_length: int = (Config().chunk_size - Config().max_decode_batch_size) // Config().max_prefill_batch_size, # TODO: use config + prefill_active_length: int = (Config().chunk_size - Config().max_decode_batch_size) // Config().max_prefill_batch_size, + gen_prefill: bool = True, + decode_batch_size: int = Config().max_decode_batch_size, + decode_active_position: torch.Tensor = None, + page_size = 256, + cuda_lens = 1 + ): + instance = cls() + + instance.batch_size = num_mini_batches + page_size = page_size + + prefill_query_info = [] + offset = 0 + if gen_prefill and prefill_query_length != 0: + for i in range(Config().max_prefill_batch_size): + prefill_query_info.append(QueryInfo(i, prefill_query_length, max_seq_length, page_size, device, offset=offset)) + offset += max_seq_length // page_size + + decode_querys_info = [] + for i in range(min(decode_batch_size, cuda_lens)): + query_info = QueryInfo(i+Config().max_prefill_batch_size, prefill_query_length, 256, page_size, device, is_prefill=False, offset=offset) + offset += max_seq_length // page_size + if tokens is not None: + query_info.query_tokens[prefill_active_length:prefill_active_length + 1].copy_(tokens) + if decode_active_position is None: + query_info.active_position = 255 + else: + query_info.active_position = decode_active_position[i] + + decode_querys_info.append(query_info) + + if prefill_query_length*Config().max_prefill_batch_size + len(decode_querys_info) < cuda_lens: + decode_querys_info.append(query_info) + + instance.minibatch = ForwardBatchInput.ForwardMiniBatch(prefill_query_info, decode_querys_info, [0, 0], [prefill_active_length for _ in range(Config().max_prefill_batch_size)], device, page_size) + + return instance + + def fill(self, batch : sched_ext.BatchQueryTodo = None, query_manager: QueryManager = None, page_size = 256): + if batch is None: + return + prefill_minibatches = batch.prefill_mini_batches + decode_mini_batches = [item for sublist in batch.decode_mini_batches for item in sublist] + + prefill_querys_info = [] + prefill_s = [] + prefill_l = [] + decode_querys_info = [] + self.batch_size = 1 + for (id, s, l) in prefill_minibatches: + prefill_querys_info.append(query_manager.query_map[id]) + prefill_s.append(s) + prefill_l.append(l) + for decode_batch_idx in decode_mini_batches: + if query_manager.query_map[decode_batch_idx].decode_start_time is None: + query_manager.query_map[decode_batch_idx].decode_start_time =time.time() + decode_querys_info.append(query_manager.query_map[decode_batch_idx]) + + self.minibatch.fill(prefill_querys_info, decode_querys_info, prefill_s, prefill_l, device=query_manager.device, page_size=page_size) + + + +class ForwardBatchOutput: + logits: list[torch.Tensor] + num_batchs: int + batch_sizes: list[int] + generated_tokens_num: list[int] + lm_start: list[int] + + temperatures: list[torch.Tensor] + top_ps: list[torch.Tensor] + + def __init__(self): + self.logits = [] + self.batch_sizes = [] + self.generated_tokens_num = [] + self.top_ps = [] + self.temperatures = [] + self.num_batchs = 1 \ No newline at end of file diff --git a/archive/kt-sft/ktransformers/server/balance_serve/inference/model_runner.py b/archive/kt-sft/ktransformers/server/balance_serve/inference/model_runner.py new file mode 100644 index 0000000..799bb9b --- /dev/null +++ b/archive/kt-sft/ktransformers/server/balance_serve/inference/model_runner.py @@ -0,0 +1,236 @@ +""" +Date: 2024-11-07 07:02:20 +LastEditors: djw +LastEditTime: 2024-12-10 08:48:32 +""" + +import torch +from torch import nn +import queue +import signal +import queue +from typing import AsyncIterable +from fastapi import FastAPI, Request +from fastapi.responses import StreamingResponse +from contextlib import asynccontextmanager +from pydantic import BaseModel, Field +import asyncio +import multiprocessing +import time +import torch.multiprocessing as mp +import random +import torch.distributed as dist +import zmq +import tempfile +from ktransformers.server.balance_serve.inference.forward_batch import ForwardBatchInput, ForwardBatchOutput + +from ktransformers.server.config.config import Config +from ktransformers.models.custom_modeling_deepseek_v3 import KDeepseekV3ForCausalLM +from ktransformers.models.custom_modeling_deepseek_v2 import KDeepseekV2ForCausalLM +from ktransformers.models.custom_modeling_qwen2_moe import KQwen2MoeForCausalLM +from ktransformers.models.custom_modeling_qwen3_moe import KQwen3MoeForCausalLM +from ktransformers.server.balance_serve.inference.query_manager import QueryManager +from ktransformers.server.balance_serve.settings import sched_ext + + + +def pad_num_tokens(num_tokens): + return (num_tokens + 63) // 64 * 64 + +def deduplicate_and_sort(lst): + return sorted(set(lst)) +def generate_cuda_graphs(chunk_size: int) -> list: + assert chunk_size <= 1024 or chunk_size % 1024 == 0, "chunk_size must <= 1024 or a multiple of 1024" + base_list = [1, 2, 3, Config().max_batch_size, 64, 256, 512, chunk_size] + + if chunk_size <= 1024: + return deduplicate_and_sort(base_list) + + multiples = [i for i in range(1024, chunk_size + 1, 1024)] + + return deduplicate_and_sort(base_list + multiples) +class ModelRunner: + """A CudaGraphRunner runs the forward pass of a model with CUDA graph and torch.compile.""" + + model: KDeepseekV3ForCausalLM | KQwen2MoeForCausalLM | KQwen3MoeForCausalLM + input: ForwardBatchInput | list[ForwardBatchInput] + output: ForwardBatchOutput + + def __init__(self, model = None, device = None, use_cuda_graph = False, max_decode_batch_size = 1, max_chunk_size = 4096, num_mini_batches: int = 1, page_size = 256, block_num = 8): + + self.stream = torch.cuda.Stream(device=device) + self.model = model # Compile and move model to the specified device + self.device = device + self.input = None + self.features_buf = None + self.output = None + self.graph_memory_pool = None + self.cuda_graphs = generate_cuda_graphs(Config().chunk_size) + self.use_cuda_graph = use_cuda_graph + self.model_time = 0 + self.page_size = page_size + self.block_num = block_num + # GPU timing for model execution + self.start_model_event = torch.cuda.Event(enable_timing=True) + self.end_model_event = torch.cuda.Event(enable_timing=True) + + self.graphs = [torch.cuda.CUDAGraph() for _ in range(len(self.cuda_graphs))] + self.page_idx_buf = [torch.zeros([self.cuda_graphs[i]], dtype=torch.int32, device = self.device) for i in range(len(self.cuda_graphs))] + self.page_offset_buf = [torch.zeros([self.cuda_graphs[i]], dtype=torch.int32, device = self.device) for i in range(len(self.cuda_graphs))] + + self.num_mini_batches = num_mini_batches + + self.max_chunk_size = max_chunk_size + + self.bsz_tensor_buf = torch.empty((1, ),dtype=torch.int32, device=device) + self.num_tokens_tensor_buf = torch.empty((1, ),dtype=torch.int32, device=device) + + def model_attn_plan(self, batch, cuda_graph_idx=0): + if isinstance(self.model, KDeepseekV3ForCausalLM): + self.model.flash_infer_attn_plan(batch, self.bsz_tensor_buf, self.num_tokens_tensor_buf, + num_heads=self.model.config.num_attention_heads, head_dim_ckv=self.model.config.kv_lora_rank, + head_dim_kpe=self.model.config.qk_rope_head_dim, page_size=self.model.cache.page_size, causal=True, + sm_scale=self.model.model.layers[0].self_attn.softmax_scale, q_data_type=torch.bfloat16, kv_data_type=torch.bfloat16) + elif isinstance(self.model, KQwen2MoeForCausalLM) or isinstance(self.model, KQwen3MoeForCausalLM): + self.model.flash_infer_attn_plan(batch, self.bsz_tensor_buf, self.num_tokens_tensor_buf, + num_q_heads=self.model.config.num_attention_heads, num_kv_heads=self.model.config.num_key_value_heads, + head_dim=self.model.config.head_dim if hasattr(self.model.config, 'head_dim') else self.model.config.hidden_size // self.model.config.num_attention_heads, + page_size=self.model.cache.page_size, causal=True, + q_data_type=torch.bfloat16, kv_data_type=torch.bfloat16, cuda_graph_idx=cuda_graph_idx) + else: + assert False, "model type not supported" + + + def warmup(self): + + def capture_graphs(cuda_graph_idx): + with torch.cuda.graph(self.graphs[cuda_graph_idx], pool=self.graph_memory_pool, stream=self.stream): + self.outputs_buf[cuda_graph_idx] = self.model(self.input[cuda_graph_idx], self.features_buf[cuda_graph_idx], self.bsz_tensor_buf, self.num_tokens_tensor_buf, self.page_idx_buf[cuda_graph_idx], self.page_offset_buf[cuda_graph_idx], cuda_graph_idx=cuda_graph_idx) + self.graph_memory_pool = self.graphs[cuda_graph_idx].pool() + + self.input = [] + self.features_buf = [] + self.outputs_buf = [] + self.bsz_tensor_buf = torch.tensor([0], dtype=torch.int32, device=self.device) + self.num_tokens_tensor_buf = torch.tensor([0], dtype=torch.int32, device=self.device) + for i in range(len(self.cuda_graphs)): + prefill_query_length = (self.cuda_graphs[i] - Config().max_decode_batch_size) // Config().max_prefill_batch_size if self.cuda_graphs[i] > Config().max_decode_batch_size else 0 #@TODO only supprot 2 prefill batch + self.input.append(ForwardBatchInput.gen_max_forward_batch(device=self.device, num_mini_batches = self.num_mini_batches, prefill_query_length=prefill_query_length, prefill_active_length=prefill_query_length, page_size=self.page_size, cuda_lens=self.cuda_graphs[i])) + + self.features_buf.append(self.model.batch_embeddings(self.input[i])) + batch_size = self.input[i].minibatch.q_indptr.size(0)-1 + num_tokens = self.features_buf[i][0].size(0) + print("capturing cuda graph", batch_size, num_tokens) + + if isinstance(self.model, KQwen2MoeForCausalLM) or isinstance(self.model, KQwen3MoeForCausalLM): + self.model.init_wrapper(self.use_cuda_graph, self.device, num_tokens ,batch_size, self.block_num, i) # TODO: 1024 is a magic number(max_batch_tokens) + + self.bsz_tensor_buf[0] = batch_size + self.num_tokens_tensor_buf[0] = num_tokens + + self.model_attn_plan(self.input[i], i) + + page_idx, page_offset = self.model.cache.get_page_table(self.input[i].minibatch.position_ids, self.input[i].minibatch.q_indptr, self.input[i].minibatch.kv_indptr, self.input[i].minibatch.kv_indices, self.num_tokens_tensor_buf) + + + self.page_idx_buf[i][:num_tokens].copy_(page_idx[:num_tokens]) + self.page_offset_buf[i][:num_tokens].copy_(page_offset[:num_tokens]) + + self.page_idx_buf[i][num_tokens:].fill_(self.model.cache.max_cache_len // self.model.cache.page_size -1) + + self.outputs_buf.append(None) + + torch.cuda.synchronize() + for warm_up_iters in range(11): + with torch.cuda.stream(self.stream): + self.outputs_buf[i] = self.model(self.input[i], self.features_buf[i], self.bsz_tensor_buf, self.num_tokens_tensor_buf, self.page_idx_buf[i], self.page_offset_buf[i], cuda_graph_idx=i) + torch.cuda.synchronize() + + self.outputs_buf[i].num_batchs = batch_size + + capture_graphs(i) + + with torch.cuda.stream(self.stream): + self.graphs[i].replay() + + self.sync(calc_time=False) + print(f"cuda_graph: {i+1}/{len(self.cuda_graphs)}, warmup finished.") + + def run(self, batch: sched_ext.BatchQueryTodo = None, query_manager: QueryManager = None): + with torch.cuda.stream(self.stream): + + batch_size = len(batch.prefill_mini_batches) # TODO: calc this + num_tokens = 0 + for i in range(len(batch.decode_mini_batches)): + batch_size += len(batch.decode_mini_batches[i]) + num_tokens += len(batch.decode_mini_batches[i]) + print(f'decode_batch_i: {len(batch.decode_mini_batches[i])},') + + for i in range(len(batch.prefill_mini_batches)): + num_tokens += batch.prefill_mini_batches[i][2] + print(f'prefill_batch_i: {batch.prefill_mini_batches[i][2]},') + + + + # cuda graph idx equal to min idx i in self.cuda_graphs, that self.cuda_graphs[i] > num_tokens + cuda_graph_idx = next((i for i, token in enumerate(self.cuda_graphs) if token >= num_tokens), len(self.cuda_graphs)) + if not self.use_cuda_graph: + cuda_graph_idx = 0 + # if cuda_graph_idx == len(self.cuda_graphs): + # assert False, "num_tokens is too large" + + if self.use_cuda_graph: + self.input[cuda_graph_idx].fill(batch, query_manager, self.page_size) + else: + self.input = [ForwardBatchInput(batch=batch, query_manager=query_manager, device=self.device)] + + + if self.use_cuda_graph: + self.features = self.model.batch_embeddings(self.input[cuda_graph_idx], device=self.device) + else: + self.features = self.model.batch_embeddings(self.input[cuda_graph_idx], device=self.device) + + + self.bsz_tensor_buf.copy_(batch_size) + self.num_tokens_tensor_buf.copy_(torch.tensor([num_tokens], dtype=torch.int32, device=self.device)) + + if self.use_cuda_graph: + self.features_buf[cuda_graph_idx][0].copy_(self.features[0], non_blocking=True) + + self.model_attn_plan(self.input[cuda_graph_idx], cuda_graph_idx) + self.start_model_event.record(self.stream) + page_idx, page_offset = self.model.cache.get_page_table(self.input[cuda_graph_idx].minibatch.position_ids, self.input[cuda_graph_idx].minibatch.q_indptr, self.input[cuda_graph_idx].minibatch.kv_indptr, self.input[cuda_graph_idx].minibatch.kv_indices, self.num_tokens_tensor_buf) + if self.use_cuda_graph: + self.page_idx_buf[cuda_graph_idx][:num_tokens].copy_(page_idx[:num_tokens]) + self.page_offset_buf[cuda_graph_idx][:num_tokens].copy_(page_offset[:num_tokens]) + + self.page_idx_buf[cuda_graph_idx][num_tokens:].fill_(self.model.cache.max_cache_len // self.model.cache.page_size -1) + self.replay(cuda_graph_idx) + self.output = ForwardBatchOutput() + + self.output.top_ps.append(self.input[cuda_graph_idx].minibatch.top_ps) + self.output.temperatures.append(self.input[cuda_graph_idx].minibatch.temperatures) + + + self.output.logits.append(self.outputs_buf[cuda_graph_idx].logits[0][self.input[cuda_graph_idx].minibatch.logits_start].clone()) + else: + self.output = self.model(self.input[cuda_graph_idx], self.features, self.bsz_tensor_buf, self.num_tokens_tensor_buf, page_idx, page_offset) + self.output.logits[0] = self.output.logits[0][self.input[cuda_graph_idx].minibatch.logits_start] + self.output.top_ps.append(self.input[cuda_graph_idx].minibatch.top_ps) + self.output.temperatures.append(self.input[cuda_graph_idx].minibatch.temperatures) + self.end_model_event.record(self.stream) + + + + def replay(self, cuda_graph_idx=-1): + with torch.cuda.stream(self.stream): + if cuda_graph_idx != -1: + self.graphs[cuda_graph_idx].replay() + else: + self.graphs.replay() + + + def sync(self, calc_time = True): + self.stream.synchronize() + if calc_time: + self.model_time = self.start_model_event.elapsed_time(self.end_model_event) # In ms \ No newline at end of file diff --git a/archive/kt-sft/ktransformers/server/balance_serve/inference/query_manager.py b/archive/kt-sft/ktransformers/server/balance_serve/inference/query_manager.py new file mode 100644 index 0000000..2295dbe --- /dev/null +++ b/archive/kt-sft/ktransformers/server/balance_serve/inference/query_manager.py @@ -0,0 +1,153 @@ +''' +Date: 2024-11-14 12:23:45 +LastEditors: djw +LastEditTime: 2024-11-20 04:06:23 +''' +import torch +from ktransformers.server.balance_serve.settings import sched_ext +import random +import time + +class QueryInfo: + id: int + active_position: int + query_length: int + is_prefill: int + block_index: torch.Tensor + query_tokens: torch.Tensor + stop_criteria: list[torch.Tensor] + + temperature: float + top_p: float + + max_length: int + + def __init__(self, id, query_length: int, max_length: int, page_size: int, device: torch.device, is_prefill: bool = True, offset: int = 0, active_position: int = 0, temperature: float = 0.01, top_p: float = 1.0): + self.id = id + self.is_prefill = is_prefill + self.active_position = active_position + self.max_length = max_length - 1 + self.query_tokens = torch.zeros((max_length,), dtype=torch.int, device = device) + self.stop_criteria = [] + self.block_index = torch.arange(offset, offset + (max_length + active_position + page_size - 1) // page_size, dtype=torch.int, device = device) + self.query_length = query_length + self.enqueue_time = time.time() + self.decode_start_time = None + self.speculative_token = {} # {position: (accept, token)} + + self.temperature = temperature + self.top_p = top_p + + def check_stop(self): + if self.active_position >= self.max_length - 2: + return True + + for stop_tensor in self.stop_criteria: + stop_len = len(stop_tensor) + + if stop_len >= self.active_position: + continue + + #print(f"stop_tensor: {stop_tensor}, stop_len: {stop_len}, active_position: {self.active_position}, query_token: {self.query_tokens[self.active_position - stop_len - 1:self.active_position - 1]}") + + if (torch.equal(self.query_tokens[self.active_position - stop_len - 1:self.active_position - 1], stop_tensor) and self.active_position) or self.max_length <= self.active_position + 3: + self.life_time = time.time() - self.enqueue_time + self.decode_duration_time = time.time() - self.decode_start_time + self.decode_tps = (self.active_position - self.query_length) / self.decode_duration_time + print(f"prefill length: {self.query_length}, prefill time: {self.prefill_duration_time}, prefill tps {self.prefill_tps}, decode length: {self.active_position - self.query_length}, decode time: {self.decode_duration_time}, decode tps {self.decode_tps}") + return True + + + return False + + + def print(self): + print(f"active_position: {self.active_position}, query_length: {self.query_length}, is_prefill: {self.is_prefill}") + print(f"block_index_shape: {self.block_index.shape}, query_tokens_shape: {self.query_tokens.shape}") + + +class QueryManager: + + page_size: int = 256 + device: torch.device + query_map : dict[int, QueryInfo] + + def __init__(self, page_size = 256, device = torch.device('cuda')): + self.page_size = page_size + self.device = device + self.query_map = {} + + def add_query(self, batch: sched_ext.BatchQueryTodo): + + for i in range(len(batch.query_ids)): + id = batch.query_ids[i] + if id not in self.query_map: + print(f"add query id: {id}, batch.query_lengths: {batch.query_lengths[i]}, batch_query_tokens: {batch.query_tokens[i].shape}, batch.block_indexes: {batch.block_indexes[i]}") + query_info = QueryInfo(id=id, query_length=batch.query_lengths[i], max_length=batch.query_tokens[i].size(0) + 1, page_size=self.page_size, device=self.device, temperature=batch.sample_options[i].temperature, top_p=batch.sample_options[i].top_p) + query_info.query_tokens[:query_info.query_length].copy_(batch.query_tokens[i][:query_info.query_length].to(self.device)) + + for stop_token_list in batch.stop_criteria[i]: + query_info.stop_criteria.append(torch.tensor(stop_token_list, dtype=torch.int, device = self.device)) + + block_num = batch.block_indexes[i].size(0) + query_info.block_index[:block_num].copy_(batch.block_indexes[i].to(self.device)) + + self.query_map[id] = query_info + + prefill_mini_batches = batch.prefill_mini_batches + for (prefill_id, s, l) in prefill_mini_batches: + if prefill_id == id: + self.query_map[prefill_id].active_position = s + + + def update(self, batch: sched_ext.BatchQueryTodo) -> list[sched_ext.QueryUpdate]: + query_updates = [] + + prefill_mini_batches = batch.prefill_mini_batches + + for (id, s, l) in prefill_mini_batches: + + if id not in self.query_map: + assert False, f"query id {id} not found in query_map" + + # update query_info + query_info = self.query_map[id] + query_info.active_position += l + + if query_info.active_position >= query_info.query_length and query_info.is_prefill: + query_info.is_prefill = False + query_info.prefill_duration_time = time.time() - query_info.enqueue_time + query_info.prefill_tps = query_info.query_length / query_info.prefill_duration_time + + + # generate schedule query_update + query_update = sched_ext.QueryUpdate() + query_update.id = id + query_update.ok = True + query_update.is_prefill = query_info.is_prefill + query_update.active_position = query_info.active_position + # if(not query_info.is_prefill): + query_updates.append(query_update) + + + decode_mini_batches = batch.decode_mini_batches + + for ids in decode_mini_batches: + for id in ids: + if id not in self.query_map: + assert False, f"query id {id} not found in query_map" + + query_info = self.query_map[id] + query_info.active_position += 1 + + query_update = sched_ext.QueryUpdate() + query_update.id = id + query_update.ok = True + query_update.is_prefill = query_info.is_prefill + + query_update.decode_done = query_info.check_stop() + + query_update.active_position = query_info.active_position + query_updates.append(query_update) + + return query_updates diff --git a/archive/kt-sft/ktransformers/server/balance_serve/inference/sampling/penaltylib/__init__.py b/archive/kt-sft/ktransformers/server/balance_serve/inference/sampling/penaltylib/__init__.py new file mode 100644 index 0000000..43fff0f --- /dev/null +++ b/archive/kt-sft/ktransformers/server/balance_serve/inference/sampling/penaltylib/__init__.py @@ -0,0 +1,13 @@ +from .orchestrator import BatchedPenalizerOrchestrator +from .penalizers.frequency_penalty import BatchedFrequencyPenalizer +from .penalizers.min_new_tokens import BatchedMinNewTokensPenalizer +from .penalizers.presence_penalty import BatchedPresencePenalizer +from .penalizers.repetition_penalty import BatchedRepetitionPenalizer + +__all__ = [ + "BatchedFrequencyPenalizer", + "BatchedMinNewTokensPenalizer", + "BatchedPresencePenalizer", + "BatchedRepetitionPenalizer", + "BatchedPenalizerOrchestrator", +] diff --git a/archive/kt-sft/ktransformers/server/balance_serve/inference/sampling/penaltylib/orchestrator.py b/archive/kt-sft/ktransformers/server/balance_serve/inference/sampling/penaltylib/orchestrator.py new file mode 100644 index 0000000..c35e8ed --- /dev/null +++ b/archive/kt-sft/ktransformers/server/balance_serve/inference/sampling/penaltylib/orchestrator.py @@ -0,0 +1,376 @@ +import abc +import dataclasses +import typing + +import torch + + +@dataclasses.dataclass +class _ReqLike: + origin_input_ids: typing.Union[torch.Tensor, typing.List[int]] + + +@dataclasses.dataclass +class _BatchLike: + reqs: typing.List[_ReqLike] + + def batch_size(self): + return len(self.reqs) + + +class BatchedPenalizerOrchestrator: + batch: _BatchLike + device: str + vocab_size: int + penalizers: typing.Dict[typing.Type["_BatchedPenalizer"], "_BatchedPenalizer"] + + def __init__( + self, + vocab_size: int, + batch: _BatchLike, + device: str, + Penalizers: typing.Set[typing.Type["_BatchedPenalizer"]], + ): + self.vocab_size = vocab_size + self.batch = batch + self.device = device + + self.penalizers = {Penalizer: Penalizer(self) for Penalizer in Penalizers} + + is_required = False + for penalizer in self.penalizers.values(): + pen_is_required = penalizer.prepare_if_required() + is_required |= pen_is_required + self.is_required = is_required + + if self.is_required: + self.cumulate_input_tokens( + input_ids=[req.origin_input_ids for req in self.reqs()] + ) + + def reqs(self): + return self.batch.reqs + + def batch_size(self): + return self.batch.batch_size() + + def cumulate_input_tokens( + self, + input_ids: typing.Union[ + typing.List[torch.Tensor], typing.List[typing.List[int]] + ], + ): + """ + Feed the input tokens to the penalizers. + + Args: + input_ids (typing.Union[typing.List[torch.Tensor], typing.List[typing.List[int]]]): The input tokens. + """ + token_ids = _TokenIDs(orchestrator=self, token_ids=input_ids) + + for penalizer in self.penalizers.values(): + penalizer.cumulate_input_tokens(input_ids=token_ids) + + def cumulate_output_tokens( + self, + output_ids: typing.Union[ + typing.List[torch.Tensor], typing.List[typing.List[int]] + ], + ): + """ + Feed the output tokens to the penalizers. + + Args: + output_ids (typing.Union[typing.List[torch.Tensor], typing.List[typing.List[int]]]): The output tokens. + """ + if not self.is_required: + return + + token_ids = _TokenIDs(orchestrator=self, token_ids=output_ids) + + for penalizer in self.penalizers.values(): + penalizer.cumulate_output_tokens(output_ids=token_ids) + + def apply(self, logits: torch.Tensor) -> torch.Tensor: + """ + Apply the penalizers to the logits. + Note that it may apply the penalizers in-place. + + Args: + logits (torch.Tensor): The logits to apply the penalizers to. + + Returns: + torch.Tensor: The logits after applying the penalizers. + """ + if not self.is_required: + return + + for penalizer in self.penalizers.values(): + logits = penalizer.apply(logits) + + return logits + + def filter( + self, + indices_to_keep: typing.List[int], + indices_tensor_to_keep: torch.Tensor = None, + ): + """ + Filter the penalizers based on the indices to keep in the batch. + + Args: + indices_to_keep (typing.List[int]): List of indices to keep in the batch. + indices_tensor_to_keep (torch.Tensor = None): Tensor of indices to keep in the batch. If not None, it will be used instead of converting indices_to_keep to a tensor. + """ + if not self.is_required: + return + + empty_indices = len(indices_to_keep) == 0 + + is_required = False + for penalizer in self.penalizers.values(): + tmp_is_required = penalizer.is_required() + is_required = is_required or tmp_is_required + if not tmp_is_required or empty_indices: + penalizer.teardown() + else: + # create tensor index only when it's needed + if indices_tensor_to_keep is None: + indices_tensor_to_keep = torch.tensor( + indices_to_keep, dtype=torch.int32, device=self.device + ) + + penalizer.filter( + indices_to_keep=indices_to_keep, + indices_tensor_to_keep=indices_tensor_to_keep, + ) + self.is_required = is_required + + def merge(self, their: "BatchedPenalizerOrchestrator"): + """ + Merge the penalizers of another orchestrator into this one. + + Note that this function **must** be called _before_ self.batch.reqs is updated (filtered). + Each unprepared penalizers would have to be prepared (creating tensors, etc.) first before merging. + This step requires the original batch.reqs, before it gets merged with other batch.reqs. + + Args: + their (BatchedPenalizerOrchestrator): The orchestrator to merge into this one. + """ + if not self.is_required and not their.is_required: + return + + self.is_required |= their.is_required + for Penalizer, their_penalizer in their.penalizers.items(): + if Penalizer not in self.penalizers: + raise ValueError(f"Penalizer {Penalizer} not found in self.penalizers") + + self.penalizers[Penalizer].merge(their_penalizer) + + +class _TokenIDs: + """ + A class that wraps token IDs to provide additional utility functions to penalizers. + + Attributes: + orchestrator (BatchedPenalizerOrchestrator): The orchestrator that this token IDs belong to. + token_ids (typing.Union[torch.Tensor, typing.List[torch.Tensor]]): The token IDs. + cached_counts (torch.Tensor): The cached occurrence count tensor. + """ + + orchestrator: BatchedPenalizerOrchestrator + token_ids: typing.Union[torch.Tensor, typing.List[torch.Tensor]] + cached_counts: torch.Tensor = None + + def __init__( + self, + orchestrator: BatchedPenalizerOrchestrator, + token_ids: typing.Union[ + typing.List[torch.Tensor], typing.List[typing.List[int]] + ], + ): + self.orchestrator = orchestrator + + if not isinstance(token_ids[0], torch.Tensor): + token_ids = [ + torch.tensor( + data=ids, dtype=torch.int64, device=self.orchestrator.device + ) + for ids in token_ids + ] + + self.token_ids = token_ids + + def occurrence_count(self) -> torch.Tensor: + """ + Returns a tensor of shape (batch_size, vocab_size) where each element is the number of times the corresponding token appears in the batch. + + Returns: + torch.Tensor: The occurrence count tensor. + """ + if self.cached_counts is not None: + return self.cached_counts + + token_ids = self.token_ids + + if isinstance(token_ids, torch.Tensor): + token_ids = token_ids.unsqueeze(1) + + # needs to be long to be used as index in scatter_add + if token_ids.dtype != torch.int64: + token_ids = token_ids.to(torch.int64) + + padded_token_ids = torch.nn.utils.rnn.pad_sequence( + sequences=token_ids, + batch_first=True, + padding_value=self.orchestrator.vocab_size, + ) + + self.cached_counts = torch.zeros( + size=(self.orchestrator.batch_size(), self.orchestrator.vocab_size + 1), + dtype=torch.int64, + device=self.orchestrator.device, + ).scatter_add_( + dim=1, + index=padded_token_ids, + src=torch.ones_like(padded_token_ids), + )[ + :, : self.orchestrator.vocab_size + ] + + return self.cached_counts + + +class _BatchedPenalizer(abc.ABC): + """ + An abstract class for a batched penalizer. + """ + + orchestrator: BatchedPenalizerOrchestrator + _is_prepared: bool = False + + def __init__(self, orchestrator: BatchedPenalizerOrchestrator): + self.orchestrator = orchestrator + + def is_prepared(self) -> bool: + return self._is_prepared + + def is_required(self) -> bool: + return self._is_required() + + def prepare(self): + if not self.is_prepared(): + self._prepare() + self._is_prepared = True + + def prepare_if_required(self): + if self.is_required(): + self.prepare() + return True + else: + return False + + def teardown(self): + if self.is_prepared(): + self._teardown() + self._is_prepared = False + + def cumulate_input_tokens(self, input_ids: _TokenIDs): + if not self.is_prepared(): + return + + self._cumulate_input_tokens(input_ids=input_ids) + + def cumulate_output_tokens(self, output_ids: _TokenIDs): + if not self.is_prepared(): + return + + self._cumulate_output_tokens(output_ids=output_ids) + + def apply(self, logits: torch.Tensor) -> torch.Tensor: + if not self.is_prepared(): + return logits + + return self._apply(logits=logits) + + def filter( + self, indices_to_keep: typing.List[int], indices_tensor_to_keep: torch.Tensor + ): + if not self.is_prepared(): + return + + self._filter( + indices_to_keep=indices_to_keep, + indices_tensor_to_keep=indices_tensor_to_keep, + ) + + def merge(self, their: "_BatchedPenalizer"): + if not self.is_prepared() and not their.is_prepared(): + return + + self.prepare() + their.prepare() + self._merge(their) + + @abc.abstractmethod + def _is_required(self) -> bool: + """ + Check if the penalizer is required to be prepared. + """ + pass + + @abc.abstractmethod + def _prepare(self): + """ + Prepare the penalizer. + Usually, this is where the penalizer initializes its tensors. + """ + pass + + @abc.abstractmethod + def _teardown(self): + """ + Tear down the penalizer. + Usually, this is where the penalizer frees its tensors. + """ + pass + + @abc.abstractmethod + def _cumulate_input_tokens(self, input_ids: _TokenIDs): + """ + Cumulate the input tokens. + Orchestrator will call this function to feed the input tokens to the penalizer. + """ + pass + + @abc.abstractmethod + def _cumulate_output_tokens(self, output_ids: _TokenIDs): + """ + Cumulate the output tokens. + Orchestrator will call this function to feed the output tokens to the penalizer. + """ + pass + + @abc.abstractmethod + def _apply(self, logits: torch.Tensor) -> torch.Tensor: + """ + Apply the penalizer to the logits. + Penalizers can modify the logits in-place if needed. + """ + pass + + @abc.abstractmethod + def _filter( + self, indices_to_keep: typing.List[int], indices_tensor_to_keep: torch.Tensor + ): + """ + Filter the penalizer (tensors or underlying data) based on the indices to keep in the batch. + """ + pass + + @abc.abstractmethod + def _merge(self, their: "_BatchedPenalizer"): + """ + Merge the penalizer with another penalizer. + """ + pass diff --git a/archive/kt-sft/ktransformers/server/balance_serve/inference/sampling/penaltylib/penalizers/frequency_penalty.py b/archive/kt-sft/ktransformers/server/balance_serve/inference/sampling/penaltylib/penalizers/frequency_penalty.py new file mode 100644 index 0000000..178cb54 --- /dev/null +++ b/archive/kt-sft/ktransformers/server/balance_serve/inference/sampling/penaltylib/penalizers/frequency_penalty.py @@ -0,0 +1,80 @@ +import typing + +import torch + +from ..orchestrator import _BatchedPenalizer, _TokenIDs + + +class BatchedFrequencyPenalizer(_BatchedPenalizer): + """ + Frequency penalizer penalizes tokens based on their frequency in the output. + """ + + frequency_penalties: torch.Tensor = None + cumulated_frequency_penalties: torch.Tensor = None + + def _is_required(self) -> bool: + return any( + req.sampling_params.frequency_penalty != 0.0 + for req in self.orchestrator.reqs() + ) + + def _prepare(self): + self.cumulated_frequency_penalties = ( + torch.tensor( + data=[0.0 for _ in self.orchestrator.reqs()], + dtype=torch.float32, + device=self.orchestrator.device, + ) + .unsqueeze_(1) + .repeat(1, self.orchestrator.vocab_size) + ) + + self.frequency_penalties = ( + torch.tensor( + data=[ + req.sampling_params.frequency_penalty + for req in self.orchestrator.reqs() + ], + dtype=torch.float32, + device=self.orchestrator.device, + ) + .unsqueeze_(1) + .expand_as(self.cumulated_frequency_penalties) + ) + + def _teardown(self): + del self.frequency_penalties + del self.cumulated_frequency_penalties + + self.frequency_penalties = None + self.cumulated_frequency_penalties = None + + def _cumulate_input_tokens(self, input_ids: _TokenIDs): + pass + + def _cumulate_output_tokens(self, output_ids: _TokenIDs): + self.cumulated_frequency_penalties += ( + self.frequency_penalties * output_ids.occurrence_count() + ) + + def _apply(self, logits: torch.Tensor) -> torch.Tensor: + logits -= self.cumulated_frequency_penalties + return logits + + def _filter( + self, indices_to_keep: typing.List[int], indices_tensor_to_keep: torch.Tensor + ): + self.frequency_penalties = self.frequency_penalties[indices_tensor_to_keep] + self.cumulated_frequency_penalties = self.cumulated_frequency_penalties[ + indices_tensor_to_keep + ] + + def _merge(self, their: "BatchedFrequencyPenalizer"): + self.frequency_penalties = torch.cat( + [self.frequency_penalties, their.frequency_penalties], dim=0 + ) + self.cumulated_frequency_penalties = torch.cat( + [self.cumulated_frequency_penalties, their.cumulated_frequency_penalties], + dim=0, + ) diff --git a/archive/kt-sft/ktransformers/server/balance_serve/inference/sampling/penaltylib/penalizers/min_new_tokens.py b/archive/kt-sft/ktransformers/server/balance_serve/inference/sampling/penaltylib/penalizers/min_new_tokens.py new file mode 100644 index 0000000..cc97a2e --- /dev/null +++ b/archive/kt-sft/ktransformers/server/balance_serve/inference/sampling/penaltylib/penalizers/min_new_tokens.py @@ -0,0 +1,108 @@ +import typing + +import torch + +from ..orchestrator import _BatchedPenalizer, _TokenIDs + + +class BatchedMinNewTokensPenalizer(_BatchedPenalizer): + """ + Min new tokens penalizer penalizes tokens based on the length of the output. + """ + + min_new_tokens: torch.Tensor = None + stop_token_penalties: torch.Tensor = None + len_output_tokens: torch.Tensor = None + + def _is_required(self) -> bool: + return any( + req.sampling_params.min_new_tokens > 0 for req in self.orchestrator.reqs() + ) + + def _prepare(self): + self.min_new_tokens = torch.tensor( + data=[ + req.sampling_params.min_new_tokens for req in self.orchestrator.reqs() + ], + dtype=torch.int32, + device=self.orchestrator.device, + ).unsqueeze_(1) + + padded_stop_token_ids = torch.nn.utils.rnn.pad_sequence( + sequences=[ + torch.tensor( + data=( + list( + (req.sampling_params.stop_token_ids or set()) + | (req.tokenizer.additional_stop_token_ids or set()) + | {req.tokenizer.eos_token_id} + ) + ), + dtype=torch.int64, + device=self.orchestrator.device, + ) + for req in self.orchestrator.reqs() + ], + batch_first=True, + padding_value=self.orchestrator.vocab_size, + ) + self.stop_token_penalties = torch.zeros( + size=(self.orchestrator.batch_size(), self.orchestrator.vocab_size + 1), + dtype=torch.float32, + device=self.orchestrator.device, + ).scatter_add_( + dim=1, + index=padded_stop_token_ids, + src=torch.full_like( + input=padded_stop_token_ids, + dtype=torch.float32, + fill_value=float("-inf"), + device=self.orchestrator.device, + ), + )[ + :, : self.orchestrator.vocab_size + ] + + self.len_output_tokens = torch.zeros( + size=(self.orchestrator.batch_size(), 1), + dtype=torch.int32, + device=self.orchestrator.device, + ) + + def _teardown(self): + del self.min_new_tokens + del self.stop_token_penalties + del self.len_output_tokens + + self.min_new_tokens = None + self.stop_token_penalties = None + self.len_output_tokens = None + + def _cumulate_input_tokens(self, input_ids: _TokenIDs): + pass + + def _cumulate_output_tokens(self, output_ids: _TokenIDs): + self.len_output_tokens += 1 + + def _apply(self, logits: torch.Tensor) -> torch.Tensor: + mask = (self.len_output_tokens < self.min_new_tokens).expand_as(logits) + logits[mask] += self.stop_token_penalties[mask] + return logits + + def _filter( + self, indices_to_keep: typing.List[int], indices_tensor_to_keep: torch.Tensor + ): + self.min_new_tokens = self.min_new_tokens[indices_tensor_to_keep] + self.stop_token_penalties = self.stop_token_penalties[indices_tensor_to_keep] + self.len_output_tokens = self.len_output_tokens[indices_tensor_to_keep] + + def _merge(self, their: "BatchedMinNewTokensPenalizer"): + self.min_new_tokens = torch.cat( + [self.min_new_tokens, their.min_new_tokens], dim=0 + ) + self.stop_token_penalties = torch.cat( + [self.stop_token_penalties, their.stop_token_penalties], dim=0 + ) + self.len_output_tokens = torch.cat( + [self.len_output_tokens, their.len_output_tokens], dim=0 + ) diff --git a/archive/kt-sft/ktransformers/server/balance_serve/inference/sampling/penaltylib/penalizers/presence_penalty.py b/archive/kt-sft/ktransformers/server/balance_serve/inference/sampling/penaltylib/penalizers/presence_penalty.py new file mode 100644 index 0000000..0593fdd --- /dev/null +++ b/archive/kt-sft/ktransformers/server/balance_serve/inference/sampling/penaltylib/penalizers/presence_penalty.py @@ -0,0 +1,79 @@ +import typing + +import torch + +from ..orchestrator import _BatchedPenalizer, _TokenIDs + + +class BatchedPresencePenalizer(_BatchedPenalizer): + """ + Presence penalizer penalizes tokens based on their presence in the output. + """ + + presence_penalties: torch.Tensor = None + cumulated_presence_penalties: torch.Tensor = None + + def _is_required(self) -> bool: + return any( + req.sampling_params.presence_penalty != 0.0 + for req in self.orchestrator.reqs() + ) + + def _prepare(self): + self.cumulated_presence_penalties = ( + torch.tensor( + data=[0.0 for _ in self.orchestrator.reqs()], + dtype=torch.float32, + device=self.orchestrator.device, + ) + .unsqueeze_(1) + .repeat(1, self.orchestrator.vocab_size) + ) + + self.presence_penalties = ( + torch.tensor( + data=[ + req.sampling_params.presence_penalty + for req in self.orchestrator.reqs() + ], + dtype=torch.float32, + device=self.orchestrator.device, + ) + .unsqueeze_(1) + .expand_as(self.cumulated_presence_penalties) + ) + + def _teardown(self): + del self.presence_penalties + del self.cumulated_presence_penalties + + self.presence_penalties = None + self.cumulated_presence_penalties = None + + def _cumulate_input_tokens(self, input_ids: _TokenIDs): + pass + + def _cumulate_output_tokens(self, output_ids: _TokenIDs): + mask = output_ids.occurrence_count() > 0 + self.cumulated_presence_penalties[mask] = self.presence_penalties[mask] + + def _apply(self, logits: torch.Tensor) -> torch.Tensor: + logits -= self.cumulated_presence_penalties + return logits + + def _filter( + self, indices_to_keep: typing.List[int], indices_tensor_to_keep: torch.Tensor + ): + self.presence_penalties = self.presence_penalties[indices_tensor_to_keep] + self.cumulated_presence_penalties = self.cumulated_presence_penalties[ + indices_tensor_to_keep + ] + + def _merge(self, their: "BatchedPresencePenalizer"): + self.presence_penalties = torch.cat( + [self.presence_penalties, their.presence_penalties], dim=0 + ) + self.cumulated_presence_penalties = torch.cat( + [self.cumulated_presence_penalties, their.cumulated_presence_penalties], + dim=0, + ) diff --git a/archive/kt-sft/ktransformers/server/balance_serve/inference/sampling/penaltylib/penalizers/repetition_penalty.py b/archive/kt-sft/ktransformers/server/balance_serve/inference/sampling/penaltylib/penalizers/repetition_penalty.py new file mode 100644 index 0000000..ea32add --- /dev/null +++ b/archive/kt-sft/ktransformers/server/balance_serve/inference/sampling/penaltylib/penalizers/repetition_penalty.py @@ -0,0 +1,83 @@ +import typing + +import torch + +from ..orchestrator import _BatchedPenalizer, _TokenIDs + + +class BatchedRepetitionPenalizer(_BatchedPenalizer): + """ + Repetition penalizer penalizes tokens based on their repetition in the input and output. + """ + + repetition_penalties: torch.Tensor = None + cumulated_repetition_penalties: torch.Tensor = None + + def _is_required(self) -> bool: + return any( + req.sampling_params.repetition_penalty != 1.0 + for req in self.orchestrator.reqs() + ) + + def _prepare(self): + self.cumulated_repetition_penalties = ( + torch.tensor( + data=[1.0 for _ in self.orchestrator.reqs()], + dtype=torch.float32, + device=self.orchestrator.device, + ) + .unsqueeze_(1) + .repeat(1, self.orchestrator.vocab_size) + ) + + self.repetition_penalties = ( + torch.tensor( + data=[ + req.sampling_params.repetition_penalty + for req in self.orchestrator.reqs() + ], + dtype=torch.float32, + device=self.orchestrator.device, + ) + .unsqueeze_(1) + .expand_as(self.cumulated_repetition_penalties) + ) + + def _teardown(self): + del self.repetition_penalties + del self.cumulated_repetition_penalties + + self.repetition_penalties = None + self.cumulated_repetition_penalties = None + + def _cumulate_input_tokens(self, input_ids: _TokenIDs): + mask = input_ids.occurrence_count() > 0 + self.cumulated_repetition_penalties[mask] = self.repetition_penalties[mask] + + def _cumulate_output_tokens(self, output_ids: _TokenIDs): + mask = output_ids.occurrence_count() > 0 + self.cumulated_repetition_penalties[mask] = self.repetition_penalties[mask] + + def _apply(self, logits: torch.Tensor) -> torch.Tensor: + return torch.where( + logits > 0, + logits / self.cumulated_repetition_penalties, + logits * self.cumulated_repetition_penalties, + ) + + def _filter( + self, indices_to_keep: typing.List[int], indices_tensor_to_keep: torch.Tensor + ): + self.repetition_penalties = self.repetition_penalties[indices_tensor_to_keep] + self.cumulated_repetition_penalties = self.cumulated_repetition_penalties[ + indices_tensor_to_keep + ] + + def _merge(self, their: "BatchedRepetitionPenalizer"): + self.repetition_penalties = torch.cat( + [self.repetition_penalties, their.repetition_penalties], dim=0 + ) + self.cumulated_repetition_penalties = torch.cat( + [self.cumulated_repetition_penalties, their.cumulated_repetition_penalties], + dim=0, + ) diff --git a/archive/kt-sft/ktransformers/server/balance_serve/inference/sampling/sampler.py b/archive/kt-sft/ktransformers/server/balance_serve/inference/sampling/sampler.py new file mode 100644 index 0000000..f491c97 --- /dev/null +++ b/archive/kt-sft/ktransformers/server/balance_serve/inference/sampling/sampler.py @@ -0,0 +1,100 @@ +''' +Date: 2024-11-14 12:23:45 +LastEditors: Xie Weiyu ervinxie@qq.com +LastEditTime: 2024-11-25 08:59:23 +''' +import logging +import torch +from torch import nn +from transformers import GenerationConfig + +from flashinfer.sampling import ( + min_p_sampling_from_probs, + top_k_renorm_probs, + top_k_top_p_sampling_from_logits, + top_p_renorm_probs, +) + +logger = logging.getLogger(__name__) + +class SamplingOptions(): + # Batched sampling params + temperatures: torch.Tensor + top_ps: torch.Tensor + top_ks: torch.Tensor + min_ps: torch.Tensor + + # All requests use greedy sampling + is_all_greedy: bool + + # Dispatch in CUDA graph + need_min_p_sampling: bool + + def __init__(self, bsz = 1, device = torch.device('cuda'), pretrained_config:GenerationConfig = None, temperatures: torch.Tensor = None, top_ps: torch.Tensor = None): + if pretrained_config is None and temperatures is None: + self.temperatures = torch.full((bsz, 1), 0, device=device, dtype=torch.float32) + self.top_ps = torch.ones((bsz, 1), device=device, dtype=torch.float32) + self.top_ks = torch.ones((bsz, 1), device=device, dtype=torch.float32) + self.need_min_p_sampling = False + self.is_all_greedy = True + else: + if temperatures is not None: + self.temperatures = temperatures.unsqueeze(-1) + else: + self.temperatures = torch.full((bsz, 1), pretrained_config.temperature, device=device, dtype=torch.float32) + + if top_ps is not None: + self.top_ps = top_ps.unsqueeze(-1) + else: + self.top_ps = torch.full((bsz, 1), pretrained_config.top_p, device=device, dtype=torch.float32) + self.top_ks = torch.full((bsz, 1), pretrained_config.top_k, device=device, dtype=torch.float32) + self.need_min_p_sampling = False + self.is_all_greedy = False + +class Sampler(nn.Module): + def __init__(self): + super().__init__() + + def forward( + self, + logits: torch.Tensor, + sampling_config: SamplingOptions = None, + ): + if sampling_config == None: + sampling_config = SamplingOptions() + + logits = logits.contiguous() + origin_logits = logits.clone() + if sampling_config.is_all_greedy: + # Use torch.argmax if all requests use greedy sampling + probs = logits + batch_next_token_ids = torch.argmax(logits, -1) + else: + # Post process logits + logits.div_(sampling_config.temperatures) + max_top_k_round, batch_size = 32, logits.shape[0] + if sampling_config.need_min_p_sampling: + probs = torch.softmax(logits, dim=-1) + logits = None + del logits + probs = top_k_renorm_probs(probs, sampling_config.top_ks) + probs = top_p_renorm_probs(probs, sampling_config.top_ps) + batch_next_token_ids = min_p_sampling_from_probs( + probs, sampling_config.min_ps + ) + temperature_0_idx = torch.where(sampling_config.temperatures == 0)[0] + batch_next_token_ids[temperature_0_idx] = torch.argmax(origin_logits[temperature_0_idx], -1).to(torch.int32) + else: + # TODO: use different kernel when don't need top_k or top_p + # @TODO get probs + probs = logits + batch_next_token_ids = top_k_top_p_sampling_from_logits( + logits, + sampling_config.top_ks, + sampling_config.top_ps, + filter_apply_order="joint", + ) + temperature_0_idx = torch.where(sampling_config.temperatures == 0)[0] + batch_next_token_ids[temperature_0_idx] = torch.argmax(origin_logits[temperature_0_idx], -1).to(torch.int32) + + return batch_next_token_ids.to(torch.int32), probs \ No newline at end of file diff --git a/archive/kt-sft/ktransformers/server/balance_serve/sched_rpc.py b/archive/kt-sft/ktransformers/server/balance_serve/sched_rpc.py new file mode 100644 index 0000000..ccc30af --- /dev/null +++ b/archive/kt-sft/ktransformers/server/balance_serve/sched_rpc.py @@ -0,0 +1,197 @@ +from datetime import datetime +import os +from typing import Optional +import zmq +import pickle +import threading +import torch.multiprocessing as mp +import sys +current_file_path = os.path.abspath(__file__) +# sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "..", "..")) +import pickle +import argparse +from ktransformers.server.balance_serve.settings import sched_ext, create_sched_settings, create_sched_settings_qwen2moe, create_sched_settings_qwen3moe + + + +if mp.get_start_method(allow_none=True) is None: + print('set start method') + mp.set_start_method('spawn') +else: + print(f'start method already set to {mp.get_start_method(allow_none=True)}') + + +class SchedulerServer: + def __init__(self, settings, main_args): + self.sched = sched_ext.create_scheduler(settings) + + self.context = zmq.Context() + self.frontend = self.context.socket(zmq.ROUTER) + print(f"sched zmq rpc server on port {main_args.sched_port}") + self.frontend.bind(f"tcp://*:{main_args.sched_port}") + + self.backend = self.context.socket(zmq.DEALER) + self.backend.bind("inproc://backend") + + def run_scheduler(self): + self.sched.run() + + def stop_scheduler(self): + self.sched.stop() + + def start_proxy(self): + zmq.proxy(self.frontend, self.backend) + + def worker_routine(self): + worker = self.context.socket(zmq.REP) + worker.connect("inproc://backend") + while True: + try: + message = worker.recv() + data = pickle.loads(message) + + method = data.get('method') + params = data.get('params', {}) + # print(f"Received request: {method}") + + if method == 'add_query': + query_add = params.get('query') + query_id = self.sched.add_query(query_add) + response = {'status': 'ok', 'query_id': query_id} + worker.send(pickle.dumps(response)) + + elif method == 'cancel_query': + query_id = params.get('query_id') + self.sched.cancel(query_id) + response = {'status': 'ok'} + worker.send(pickle.dumps(response)) + + elif method == 'update_last_batch': + updates = params.get('updates') + + batch_todo = self.sched.update_last_batch(updates) + + response = {'status': 'ok', 'batch_todo': batch_todo} + # print (batch_todo.query_lengths, batch_todo.query_ids) + worker.send(pickle.dumps(response)) + + elif method == 'get_inference_context': + inference_context = self.sched.get_inference_context() + data = { + "k_cache":inference_context.k_cache, + "v_cache":inference_context.v_cache + } + print(f"Serializing KVCache") + data["k_cache"] = [mp.reductions.reduce_tensor(t) for t in data['k_cache']] + data["v_cache"] = [mp.reductions.reduce_tensor(t) for t in data['v_cache']] + # print(data) + response = {'status': 'ok', 'inference_context': data} + + worker.send(pickle.dumps(response)) + # response['inference_context'].k_cache[0][0, 0, 0, 0, 0] = 1 + # print("k_cache update") + + else: + response = {'status': 'error', 'message': 'Unknown method'} + worker.send(pickle.dumps(response)) + + except Exception as e: + response = {'status': 'error', 'message': str(e)} + worker.send(pickle.dumps(response)) + + def start_rpc_service(self): + try: + print("Scheduler RPC service is running...") + + threading.Thread(target=self.run_scheduler, daemon=True).start() + + for _ in range(10): + threading.Thread(target=self.worker_routine, daemon=True).start() + + self.start_proxy() + + except KeyboardInterrupt: + print("Shutting down scheduler RPC service...") + self.stop_rpc_service() + + def stop_rpc_service(self): + self.stop_scheduler() + self.frontend.close() + self.backend.close() + self.context.term() + +def start_server(settings, main_args): + server = SchedulerServer(settings, main_args) + server.start_rpc_service() + + +# Add async client for webserver +class SchedulerClient: + def __init__(self, sched_port): + address=f'tcp://localhost:{sched_port}' + self.address = address + self.context = zmq.Context() + self.socket = self.context.socket(zmq.REQ) + self.socket.connect(self.address) + print(f"Connected to server at {self.address}") + + def __del__(self): + self.socket.close() + self.context.term() + + def send_request(self, method, params=None): + if params is None: + params = {} + request = { + 'method': method, + 'params': params + } + # print(f'send request {request}') + self.socket.send(pickle.dumps(request)) + response = self.socket.recv() + # print(response) + response = pickle.loads(response) + if response.get('status') == 'ok': + return response + else: + raise Exception(f"Error from server: {response.get('message')}") + + def add_query(self, query): + response = self.send_request('add_query', {'query': query}) + return response.get('query_id') + + def cancel_query(self, query_id): + self.send_request('cancel_query', {'query_id': query_id}) + + def update_last_batch(self, updates): + response = self.send_request('update_last_batch', {'updates': updates}) + # print(f"update_last_batch response {response}") + return response.get('batch_todo') + + def rebuild_inferece_context(self,response): + data = response.get('inference_context') + inference_context = sched_ext.InferenceContext() + print('Rebuilding kvcache') + inference_context.k_cache = [fn(*args) for fn,args in data['k_cache']] + inference_context.v_cache = [fn(*args) for fn,args in data['v_cache']] + return inference_context + + def get_inference_context_raw(self): + response = self.send_request('get_inference_context') + return response + + + +if __name__ == '__main__': + parser = argparse.ArgumentParser() + parser.add_argument("--config", type=str, required=True) + args = parser.parse_args() + with open(args.config, "rb") as f: + main_args = pickle.load(f) + if main_args.architectures == "Qwen2MoeForCausalLM": + settings = create_sched_settings_qwen2moe(main_args) + elif main_args.architectures == "Qwen3MoeForCausalLM": + settings = create_sched_settings_qwen3moe(main_args) + else: + settings = create_sched_settings(main_args) + start_server(settings, main_args) diff --git a/archive/kt-sft/ktransformers/server/balance_serve/settings.py b/archive/kt-sft/ktransformers/server/balance_serve/settings.py new file mode 100644 index 0000000..540dc1c --- /dev/null +++ b/archive/kt-sft/ktransformers/server/balance_serve/settings.py @@ -0,0 +1,179 @@ +''' +Date: 2024-11-13 09:43:39 +LastEditors: djw +LastEditTime: 2024-11-18 16:41:03 +''' +import sys, os +import yaml, json +from time import sleep + + +import sched_ext +from transformers import AutoConfig + +from ktransformers.models.configuration_qwen3_moe import Qwen3MoeConfig + +def create_sched_settings(args): + default_sample_options = sched_ext.SampleOptions() + model_name = os.path.basename(os.path.normpath(args.model_dir)) + input_model_settings = sched_ext.ModelSettings() + input_model_settings.model_path = args.model_dir + input_model_settings.params_count = int(0) + model_config = AutoConfig.from_pretrained(args.model_dir, trust_remote_code=True) + input_model_settings.layer_count = model_config.num_hidden_layers + input_model_settings.num_k_heads = 1 # model_config["num_key_value_heads"] + input_model_settings.k_head_dim = 576 + input_model_settings.bytes_per_params = 2 + input_model_settings.bytes_per_kv_cache_element = 2 + settings = sched_ext.Settings() + settings.model_name = model_name + settings.quant_type = "BF16" + settings.model_settings = input_model_settings + settings.page_size = args.page_size + settings.gpu_device_count = 1 # tp + settings.gpu_device_id = [i for i in range(settings.gpu_device_count)] + # settings.gpu_memory_size = args.cache_lens*576*2 + settings.gpu_memory_size = args.gpu_memory_size + settings.memory_utilization_percentage = args.utilization_percentage + max_batch_size = args.max_batch_size + chunk_size = args.chunk_size + + max_decode_batch_size = max_batch_size - 2 + + settings.max_batch_size = max_batch_size + settings.recommended_chunk_prefill_token_count = (chunk_size - max_decode_batch_size) // 2 + settings.sample_options = default_sample_options + settings.sched_metrics_port = args.sched_metrics_port + settings.gpu_only = args.memory_gpu_only + settings.use_self_defined_head_dim = True + settings.self_defined_head_dim = 576 + settings.full_kv_cache_on_each_gpu = True + settings.k_cache_on = True + settings.v_cache_on = False + + settings.kvc2_root_path = '/mnt/data/persist-kvc' + settings.kvc2_config_path = args.kvc2_config_dir + settings.memory_pool_size_GB = args.cpu_memory_size_GB + settings.evict_count = 40 + settings.kvc2_metrics_port = args.kvc2_metrics_port + settings.load_from_disk = False + settings.save_to_disk = True + + + settings.strategy_name = args.sched_strategy + + settings.auto_derive() + return settings + + +def create_sched_settings_qwen2moe(args): + default_sample_options = sched_ext.SampleOptions() + model_name = os.path.basename(os.path.normpath(args.model_dir)) + input_model_settings = sched_ext.ModelSettings() + input_model_settings.model_path = args.model_dir + input_model_settings.params_count = int(0) + model_config = AutoConfig.from_pretrained(args.model_dir, trust_remote_code=True) + input_model_settings.layer_count = model_config.num_hidden_layers + input_model_settings.num_k_heads = model_config.num_key_value_heads # model_config["num_key_value_heads"] + input_model_settings.k_head_dim = 128 + input_model_settings.bytes_per_params = 2 + input_model_settings.bytes_per_kv_cache_element = 2 + settings = sched_ext.Settings() + settings.model_name = model_name + settings.quant_type = "BF16" + settings.model_settings = input_model_settings + settings.page_size = args.page_size + settings.gpu_device_count = 1 # tp + settings.gpu_device_id = [i for i in range(settings.gpu_device_count)] + # settings.gpu_memory_size = args.cache_lens*576*2 + settings.gpu_memory_size = args.gpu_memory_size + settings.memory_utilization_percentage = args.utilization_percentage + max_batch_size = args.max_batch_size + chunk_size = args.chunk_size + + max_decode_batch_size = max_batch_size - 2 + + settings.max_batch_size = max_batch_size + settings.recommended_chunk_prefill_token_count = (chunk_size - max_decode_batch_size) // 2 + settings.sample_options = default_sample_options + settings.sched_metrics_port = args.sched_metrics_port + settings.gpu_only = args.memory_gpu_only + settings.use_self_defined_head_dim = False + settings.self_defined_head_dim = 576 + settings.full_kv_cache_on_each_gpu = True + settings.k_cache_on = True + settings.v_cache_on = True + + settings.kvc2_root_path = '/mnt/data/persist-kvc' + settings.kvc2_config_path = args.kvc2_config_dir + settings.memory_pool_size_GB = args.cpu_memory_size_GB + settings.evict_count = 40 + settings.kvc2_metrics_port = args.kvc2_metrics_port + settings.load_from_disk = False + settings.save_to_disk = True + + + settings.strategy_name = args.sched_strategy + + settings.auto_derive() + return settings + + + +def create_sched_settings_qwen3moe(args): + default_sample_options = sched_ext.SampleOptions() + model_name = os.path.basename(os.path.normpath(args.model_dir)) + input_model_settings = sched_ext.ModelSettings() + input_model_settings.model_path = args.model_dir + input_model_settings.params_count = int(0) + model_config = Qwen3MoeConfig.from_pretrained(args.model_dir, trust_remote_code=True) + input_model_settings.layer_count = model_config.num_hidden_layers + input_model_settings.num_k_heads = model_config.num_key_value_heads # model_config["num_key_value_heads"] + input_model_settings.k_head_dim = 128 + input_model_settings.bytes_per_params = 2 + input_model_settings.bytes_per_kv_cache_element = 2 + settings = sched_ext.Settings() + settings.model_name = model_name + settings.quant_type = "BF16" + settings.model_settings = input_model_settings + settings.page_size = args.page_size + settings.gpu_device_count = 1 # tp + settings.gpu_device_id = [i for i in range(settings.gpu_device_count)] + # settings.gpu_memory_size = args.cache_lens*576*2 + settings.gpu_memory_size = args.gpu_memory_size + settings.memory_utilization_percentage = args.utilization_percentage + max_batch_size = args.max_batch_size + chunk_size = args.chunk_size + + max_decode_batch_size = max_batch_size - 2 + + settings.max_batch_size = max_batch_size + settings.recommended_chunk_prefill_token_count = (chunk_size - max_decode_batch_size) // 2 + settings.sample_options = default_sample_options + settings.sched_metrics_port = args.sched_metrics_port + settings.gpu_only = args.memory_gpu_only + settings.use_self_defined_head_dim = False + settings.self_defined_head_dim = 576 + settings.full_kv_cache_on_each_gpu = True + settings.k_cache_on = True + settings.v_cache_on = True + + settings.kvc2_root_path = '/mnt/data/persist-kvc' + settings.kvc2_config_path = args.kvc2_config_dir + settings.memory_pool_size_GB = args.cpu_memory_size_GB + settings.evict_count = 40 + settings.kvc2_metrics_port = args.kvc2_metrics_port + settings.load_from_disk = False + settings.save_to_disk = True + + + settings.strategy_name = args.sched_strategy + + settings.auto_derive() + return settings + + + + + + diff --git a/archive/kt-sft/ktransformers/server/config/config.py b/archive/kt-sft/ktransformers/server/config/config.py new file mode 100644 index 0000000..0d7c17b --- /dev/null +++ b/archive/kt-sft/ktransformers/server/config/config.py @@ -0,0 +1,209 @@ +#!/usr/bin/env python +# coding=utf-8 +""" +Description : +Author : unicornchan +Date : 2024-06-11 16:35:42 +Version : 1.0.0 +LastEditors : WuHao +LastEditTime : 2024-08-12 06:31:14 +""" +import os +import shutil +import yaml +import psutil + +from ktransformers.server.config.singleton import Singleton +from typing import Optional + + +class Config(metaclass=Singleton): + """Singleton pattern Config class, used to get all configurations.""" + + CONFIG_FILE_NAME = "config.yaml" + + @staticmethod + def load() -> dict: + """load config file + + Returns: + dict: all configs + """ + base_path: str = os.path.dirname(os.path.dirname(os.path.dirname(__file__))) + config_yaml: str = os.path.join(base_path, "configs", Config.CONFIG_FILE_NAME) + + user_path: str = os.path.expanduser("~") + localstore_path: str = os.path.join(user_path, ".ktransformers") + kvc2_config_dir = os.path.join(localstore_path, "kvc2") + config_path: str = os.path.join(localstore_path, Config.CONFIG_FILE_NAME) + if not os.path.exists(config_yaml): + print(f"Can't find config file, {config_yaml}") + exit(-1) + if not os.path.exists(localstore_path): + os.mkdir(localstore_path) + if not os.path.exists(kvc2_config_dir): + os.mkdir(kvc2_config_dir) + if not os.path.exists(config_path): + shutil.copyfile(config_yaml, config_path) + with open(config_path, "r", encoding="utf-8") as fp: + config = yaml.safe_load(fp) + return config + + @staticmethod + def to_path(path: str) -> str: + """ + process file path + """ + base_path = os.path.dirname(os.path.dirname(os.path.dirname(__file__))) + real_path = path if os.path.isabs(path) else os.path.join(base_path, path) + return real_path + + def __init__(self): + cfg = Config.load() + self.base_path = os.path.dirname(os.path.dirname(os.path.dirname(__file__))) + self.user_path: str = os.path.expanduser("~") + self.localstore_path: str = os.path.join(self.user_path, ".ktransformers") + # log configs + self.log_dir = os.path.join(self.localstore_path, cfg["log"]["dir"]) + if not os.path.exists(self.log_dir): + os.mkdir(self.log_dir) + self.log_file = cfg["log"]["file"] + self.log_level = cfg["log"]["level"] + self.backup_count = cfg["log"]["backup_count"] + + self.kvc2_config_dir = os.path.join(self.localstore_path, "kvc2") + # server configs + self.server: dict = cfg.get("server", {}) + self.server_ip = self.server.get("ip", "0.0.0.0") + self.server_port = self.server.get("port", 9016) + self.api_key = self.server.get("api_key", "") + + # db configs + self.db_configs: dict = cfg.get("db", {}) + self.db_type = self.db_configs.get("type", "") + self.db_host = self.localstore_path + self.db_port = self.db_configs.get("port", "") + self.db_name = self.db_configs.get("database", "") + self.db_pool_size = self.db_configs.get("pool_size") + self.db_database = self.db_configs.get("database", "") + + # user config + self.user_config: dict = cfg.get("user", {}) + self.user_secret_key = self.user_config.get("secret_key", "") + self.user_algorithm = self.user_config.get("algorithm", "") + self.user_force_think = self.user_config.get("force_think", False) + + # model config + self.model: dict = cfg.get("model", {}) + self.backend_type: str = self.model.get("type", "transformers") + self.model_dir: str = self.model.get("path", "") + # to make sure it consistent with previous version + self.model_path: str = self.model_dir + self.model_name: str = self.model.get("name", "") + self.architectures: str = self.model.get("name", "") + self.model_device: str = self.model.get("device", "cuda:0") + self.gguf_path: Optional[str] = self.model.get("gguf_path", None) + self.use_cuda_graph = self.model.get("use_cuda_graph", True) + self.trust_remote_code = self.model.get("trust_remote_code", True) + # self.model_cache_lens = self.model.get("cache_lens") + self.optimize_config_path: Optional[str] = self.model.get( + "optimize_config_path", None + ) + + self.max_new_tokens = self.model.get("max_new_tokens", 2000) + self.json_mode = self.model.get("json_mode", False) + self.healing = self.model.get("healing", False) + self.ban_strings: Optional[list] = self.model.get("ban_strings", None) + self.gpu_split: Optional[str] = self.model.get("gpu_split", None) + self.length: Optional[int] = self.model.get("length", None) + self.rope_scale: Optional[float] = self.model.get("rope_scale", None) + self.rope_alpha: Optional[float] = self.model.get("rope_alpha", None) + self.no_flash_attn = self.model.get("no_flash_attn", False) + self.low_mem = self.model.get("low_mem", False) + self.experts_per_token: Optional[int] = self.model.get("experts_per_token", None) + self.load_q4 = self.model.get("load_q4", False) + self.fast_safetensors = self.model.get("fast_safetensors", False) + self.draft_model_dir: Optional[str] = self.model.get("draft_model_dir", None) + self.no_draft_scale = self.model.get("no_draft_scale", False) + self.modes = self.model.get("modes", False) + self.mode = self.model.get("mode", "llama") + self.username = self.model.get("username", "User") + self.botname = self.model.get("botname", "Chatbort") + self.system_prompt: Optional[str] = self.model.get("system_prompt", None) + self.temperature = self.model.get("temperature", 0.95) + self.smoothing_factor = self.model.get("smoothing_factor", 0.0) + self.dynamic_temperature: Optional[str] = self.model.get("dynamic_temperature", None) + self.top_k = self.model.get("top_k", 50) + self.top_p = self.model.get("top_p", 0.8) + self.top_a = self.model.get("top_a", 0.0) + self.skew = self.model.get("skew", 0.0) + self.typical = self.model.get("typical", 0.0) + self.repetition_penalty = self.model.get("repetition_penalty", 1.01) + self.frequency_penalty = self.model.get("frequency_penalty", 0.0) + self.presence_penalty = self.model.get("presence_penalty", 0.0) + self.response_chunk = self.model.get("response_chunk", 250) + self.no_code_formatting = self.model.get("no_code_formatting", False) + self.cache_8bit = self.model.get("cache_8bit", False) + self.cache_q4 = self.model.get("cache_q4", True) + self.ngram_decoding = self.model.get("ngram_decoding", False) + self.print_timings = self.model.get("print_timings", False) + self.amnesia = self.model.get("amnesia", False) + self.batch_size = self.model.get("batch_size", 1) + self.cache_lens = self.model.get("cache_lens", 4096) + self.device = self.model.get("device", "cuda:2") + + # web config + self.web: dict = cfg.get("web", {}) + self.web_cross_domain: bool = self.web.get("open_cross_domain", True) + self.mount_web: bool = self.web.get("mount", False) + + # ext + self.ext: dict = cfg.get("ext", {}) + self.cpu_infer = psutil.cpu_count(logical=False) - 3 + + # file config + self.local_store_configs: dict = cfg.get("local_store", {}) + self.file_upload_dir: str = os.path.join( + self.localstore_path, self.local_store_configs.get("file_upload_dir", "") + ) + self.assistant_store_dir: str = os.path.join( + self.localstore_path, self.local_store_configs.get("assistant_store_dir", "") + ) + + # long context config + self.long_context_config: dict = cfg.get("long_context", {}) + self.max_seq_len = self.long_context_config.get("max_seq_len", 32000) + self.block_size = self.long_context_config.get("block_size", 128) + self.local_windows_len = self.long_context_config.get("local_windows_len", 4096) + self.second_select_num = self.long_context_config.get("second_select_num", 32) + self.anchor_type = self.long_context_config.get("anchor_type", "DYNAMIC") + self.kv_type = self.long_context_config.get("kv_type", "FP16") + self.dense_layer_num = self.long_context_config.get("dense_layer_num", 2) + self.anchor_num = self.long_context_config.get("anchor_num", 1) + self.preselect_block = self.long_context_config.get("preselect_block", True) + self.head_select_mode = self.long_context_config.get("head_select_mode", "SHARED") + self.preselect_block_count = self.long_context_config.get("preselect_block_count", 32) + self.layer_step = self.long_context_config.get("layer_step", 1) + self.token_step = self.long_context_config.get("token_step", 100) + + # local chat + self.local_chat_config: dict = cfg.get("local_chat", {}) + self.prompt_file = self.local_chat_config.get("prompt_file", None) + + # asyncserver + self.sched_strategy = cfg["async_server"]["sched_strategy"] + self.sched_port = cfg["async_server"]["sched_port"] + self.sched_metrics_port = cfg["async_server"]["sched_metrics_port"] + self.kvc2_metrics_port = cfg["async_server"]["kvc2_metrics_port"] + self.max_batch_size = cfg["async_server"]["max_batch_size"] + self.page_size = cfg["attn"]["page_size"] + self.chunk_size = cfg["attn"]["chunk_size"] + self.memory_gpu_only = cfg["kvc2"]["gpu_only"] + self.cache_lens = ((self.cache_lens + self.page_size - 1) // self.page_size) * self.page_size + self.gpu_memory_size = 2*576*61*self.cache_lens + self.utilization_percentage = 1.0 #cfg["kvc2"]["utilization_percentage"] + self.cpu_memory_size_GB = cfg["kvc2"]["cpu_memory_size_GB"] + # only support 2 prefill task + self.max_prefill_batch_size = 2 + self.max_decode_batch_size = self.max_batch_size - self.max_prefill_batch_size + diff --git a/archive/kt-sft/ktransformers/server/config/log.py b/archive/kt-sft/ktransformers/server/config/log.py new file mode 100644 index 0000000..1d594fe --- /dev/null +++ b/archive/kt-sft/ktransformers/server/config/log.py @@ -0,0 +1,176 @@ +#!/usr/bin/env python +# coding=utf-8 +''' +Description : +Author : unicornchan +Date : 2024-06-12 02:48:39 +Version : 1.0.0 +LastEditors : chenxl +LastEditTime : 2024-07-27 01:55:50 +''' + +import codecs +import logging +import os +import re +import locale +from pathlib import Path +from logging.handlers import BaseRotatingHandler +import time +import colorlog + +from ktransformers.server.config.config import Config + + +class DailyRotatingFileHandler(BaseRotatingHandler): + """ + such as 'logging.TimeRotatingFileHandler', Additional features: + - support multiprocess + - support rotating daily + """ + + def __init__(self, filename, backupCount=0, encoding=None, delay=False, utc=False, **kwargs): # pylint: disable=unused-argument + self.backup_count = backupCount + self.utc = utc + self.suffix = "%Y-%m-%d" + self.base_log_path = Path(filename) + if not os.path.exists(self.base_log_path.parent): + os.makedirs(self.base_log_path.parent) + self.base_filename = self.base_log_path.name + self.current_filename = self._compute_fn() + self.current_log_path = self.base_log_path.with_name( + self.current_filename) + BaseRotatingHandler.__init__(self, filename, 'a', encoding, delay) + + # pylint: disable=unused-argument, invalid-name + def shouldRollover(self, record): + """ + Determine whether to rotate the log. If the log filename corresponding to the current + time is not consistent with the currently opened log filename, then it is necessary + to rotate the log + Args: + record: record is not used, as we are just comparing times, but it is needed so + the method signatures are the same + """ + if self.current_filename != self._compute_fn(): + return True + return False + + def doRollover(self): + """ + roll over + """ + # close last log file + if self.stream: + self.stream.close() + self.stream = None # type: ignore + + # gen new log file name + self.current_filename = self._compute_fn() + self.current_log_path = self.base_log_path.with_name( + self.current_filename) + + if not self.delay: + self.stream = self._open() # type: ignore + + self.delete_expired_files() + + def _compute_fn(self): + """ + gen log file name + """ + return self.base_filename + "." + time.strftime(self.suffix, time.localtime()) + + def _open(self): + """ + open a new log file, create soft link + """ + if self.encoding is None: + stream = open(str(self.current_log_path), self.mode, encoding=locale.getpreferredencoding()) + else: + stream = codecs.open(str(self.current_log_path), self.mode, self.encoding) + + if self.base_log_path.exists(): + try: + if not self.base_log_path.is_symlink() or os.readlink(self.base_log_path) != self.current_filename: + os.remove(self.base_log_path) + except OSError: + pass + + try: + os.symlink(self.current_filename, str(self.base_log_path)) + except OSError: + pass + return stream + + def delete_expired_files(self): + """ + delete expired files every day + """ + if self.backup_count <= 0: + return + + file_names = os.listdir(str(self.base_log_path.parent)) + result = [] + prefix = self.base_filename + "." + plen = len(prefix) + for file_name in file_names: + if file_name[:plen] == prefix: + suffix = file_name[plen:] + if re.match(r"^\d{4}-\d{2}-\d{2}(\.\w+)?$", suffix): + result.append(file_name) + if len(result) < self.backup_count: + result = [] + else: + result.sort() + result = result[:len(result) - self.backup_count] + + for file_name in result: + os.remove(str(self.base_log_path.with_name(file_name))) + + +class Logger(object): + """ + logger class + """ + level_relations = { + 'debug': logging.DEBUG, + 'info': logging.INFO, + 'warn': logging.WARNING, + 'error': logging.ERROR, + 'crit': logging.CRITICAL + } + + def __init__(self, level: str = 'info'): + fmt = '%(asctime)s %(levelname)s %(pathname)s[%(lineno)d] %(funcName)s: %(message)s' + cfg: Config = Config() + filename: str = os.path.join(cfg.log_dir, cfg.log_file) + backup_count: int = cfg.backup_count + th = DailyRotatingFileHandler(filename=filename, when='MIDNIGHT', backupCount=backup_count, encoding="utf-8") + th.setFormatter(logging.Formatter(fmt)) + + + color_fmt = ( + '%(log_color)s%(asctime)s %(levelname)s %(pathname)s[%(lineno)d]: %(message)s' + ) + color_formatter = colorlog.ColoredFormatter( + color_fmt, + log_colors={ + 'DEBUG': 'cyan', + 'INFO': 'green', + 'WARNING': 'yellow', + 'ERROR': 'red', + 'CRITICAL': 'bold_red' + } + ) + + sh = logging.StreamHandler() + sh.setFormatter(color_formatter) + + self.logger = logging.getLogger(filename) + self.logger.setLevel(self.level_relations.get(level)) # type: ignore + self.logger.addHandler(th) + self.logger.addHandler(sh) + + +logger = Logger(level=Config().log_level).logger diff --git a/archive/kt-sft/ktransformers/server/config/singleton.py b/archive/kt-sft/ktransformers/server/config/singleton.py new file mode 100644 index 0000000..1af5d6c --- /dev/null +++ b/archive/kt-sft/ktransformers/server/config/singleton.py @@ -0,0 +1,35 @@ +#!/usr/bin/env python +# coding=utf-8 +''' +Description : Implement singleton +Author : unicornchan +Date : 2024-06-11 17:08:36 +Version : 1.0.0 +LastEditors : chenxl +LastEditTime : 2024-07-27 01:55:56 +''' +import abc + +class Singleton(abc.ABCMeta, type): + """_summary_ + + Args: + abc.ABCMeta: Provide a mechanism for defining abstract methods and properties, + enforcing subclasses to implement these methods and properties. + type: Inherit from 'type' to make 'Singleton' a metaclass, + enabling the implementation of the Singleton + """ + _instances = {} + + def __call__(cls, *args, **kwds): + if cls not in cls._instances: + cls._instances[cls] = super(Singleton, cls).__call__(*args, **kwds) + return cls._instances[cls] + +class AbstractSingleton(abc.ABC, metaclass=Singleton): + """Provided an abstract Singleton base class, any class inheriting from + this base class will automatically become a Singleton class. + + Args: + abc.ABC: Abstract base class, it cannot be instantiated, only inherited. + """ diff --git a/archive/kt-sft/ktransformers/server/crud/__init__.py b/archive/kt-sft/ktransformers/server/crud/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/archive/kt-sft/ktransformers/server/crud/assistants/__init__.py b/archive/kt-sft/ktransformers/server/crud/assistants/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/archive/kt-sft/ktransformers/server/crud/assistants/assistants.py b/archive/kt-sft/ktransformers/server/crud/assistants/assistants.py new file mode 100644 index 0000000..2d942c9 --- /dev/null +++ b/archive/kt-sft/ktransformers/server/crud/assistants/assistants.py @@ -0,0 +1,66 @@ +from time import time +from typing import Optional,List +from uuid import uuid4 + +from ktransformers.server.models.assistants.assistants import Assistant +from ktransformers.server.schemas.assistants.assistants import AssistantCreate,AssistantObject,AssistantModify +from ktransformers.server.utils.sql_utils import SQLUtil +from ktransformers.server.config.log import logger +from ktransformers.server.schemas.base import Order + + +class AssistantDatabaseManager: + def __init__(self) -> None: + self.sql_util = SQLUtil() + + def create_assistant_object(self, assistant: AssistantCreate) -> AssistantObject: + assistant = AssistantObject( + **assistant.model_dump(mode='json'), + id=str(uuid4()), + object='assistant', + created_at=int(time()), + ) + return assistant + + def db_count_assistants(self) -> int: + with self.sql_util.get_db() as db: + return db.query(Assistant).count() + + def db_create_assistant(self, assistant: AssistantCreate): + ass_obj = self.create_assistant_object(assistant) + ass_obj.sync_db() + return ass_obj + + def db_list_assistants(self, limit: Optional[int], order: Order) -> List[AssistantObject]: + with self.sql_util.get_db() as db: + query = db.query(Assistant).order_by( + order.to_sqlalchemy_order()(Assistant.created_at)) + if limit is not None: + db_assistants = query.limit(limit) + else: + db_assistants = query.all() + return [AssistantObject.model_validate(a.__dict__) for a in db_assistants] + + def db_get_assistant_by_id(self, assistant_id: str) -> Optional[AssistantObject]: + with self.sql_util.get_db() as db: + db_assistant = db.query(Assistant).filter( + Assistant.id == assistant_id).first() + if db_assistant is None: + logger.debug(f"no assistant with id {str}") + return None + return AssistantObject.model_validate(db_assistant.__dict__) + + def db_update_assistant_by_id(self, assistant_id: str, assistant: AssistantModify): + with self.sql_util.get_db() as db: + db_assistant = db.query(Assistant).filter( + Assistant.id == assistant_id).first() + self.sql_util.db_update_commit_refresh(db, db_assistant, assistant) + return AssistantObject.model_validate(db_assistant.__dict__) + + def db_delete_assistant_by_id(self, assistant_id: str): + with self.sql_util.get_db() as db: + db_assistant = db.query(Assistant).filter( + Assistant.id == assistant_id).first() + db.delete(db_assistant) + db.commit() + diff --git a/archive/kt-sft/ktransformers/server/crud/assistants/messages.py b/archive/kt-sft/ktransformers/server/crud/assistants/messages.py new file mode 100644 index 0000000..779a6c8 --- /dev/null +++ b/archive/kt-sft/ktransformers/server/crud/assistants/messages.py @@ -0,0 +1,86 @@ +from time import time +from typing import Optional +from uuid import uuid4 + +from ktransformers.server.models.assistants.messages import Message +from ktransformers.server.schemas.assistants.messages import MessageCore, MessageCreate, MessageObject +from ktransformers.server.schemas.base import Order,ObjectID +from ktransformers.server.utils.sql_utils import SQLUtil + +class MessageDatabaseManager: + def __init__(self) -> None: + self.sql_util = SQLUtil() + + @staticmethod + def create_db_message_by_core(message: MessageCore): + message_dict = message.model_dump(mode="json") + return Message(**message_dict, id=str(uuid4()), created_at=int(time())) + + def create_db_message(self, message: MessageCreate): + return MessageDatabaseManager.create_db_message_by_core(message.to_core()) + + def db_add_message(self, message: Message): + with self.sql_util.get_db() as db: + db.add(message) + self.sql_util.db_add_commit_refresh(db, message) + + def db_create_message(self, thread_id: str, message: MessageCreate, status: MessageObject.Status): + db_message = self.create_db_message(message) + db_message.status = status.value + db_message.thread_id = thread_id + self.db_add_message(db_message) + return MessageObject.model_validate(db_message.__dict__) + + @staticmethod + def create_message_object(thread_id: ObjectID, run_id: ObjectID, message: MessageCreate): + core = message.to_core() + return MessageObject( + **core.model_dump(mode='json'), + id=str(uuid4()), + object='thread.message', + created_at=int(time()), + thread_id=thread_id, + run_id=run_id, + status=MessageObject.Status.in_progress, + ) + + def db_sync_message(self, message: MessageObject): + db_message = Message( + **message.model_dump(mode="json"), + ) + with self.sql_util.get_db() as db: + self.sql_util.db_merge_commit(db, db_message) + + def db_list_messages_of_thread( + self, thread_id: str, limit: Optional[int] = None, order: Order = Order.DESC): + + # logger.debug( + # f"list messages of: {thread_id}, limit {limit}, order {order}") + with self.sql_util.get_db() as db: + query = ( + db.query(Message) + .filter(Message.thread_id == thread_id) + .order_by(order.to_sqlalchemy_order()(Message.created_at)) + ) + if limit is not None: + messages = query.limit(limit) + else: + messages = query.all() + message_list = [MessageObject.model_validate(m.__dict__) for m in messages] + return message_list + + def db_get_message_by_id(self, thread_id: ObjectID, message_id: ObjectID) -> MessageObject: + with self.sql_util.get_db() as db: + message = db.query(Message).filter( + Message.id == message_id).first() + assert message.thread_id == thread_id + message_info = MessageObject.model_validate(message.__dict__) + return message_info + + def db_delete_message_by_id(self, thread_id: ObjectID, message_id: ObjectID): + with self.sql_util.get_db() as db: + message = db.query(Message).filter( + Message.id == message_id).first() + assert message.thread_id == thread_id + db.delete(message) + db.commit() diff --git a/archive/kt-sft/ktransformers/server/crud/assistants/runs.py b/archive/kt-sft/ktransformers/server/crud/assistants/runs.py new file mode 100644 index 0000000..bdb884d --- /dev/null +++ b/archive/kt-sft/ktransformers/server/crud/assistants/runs.py @@ -0,0 +1,50 @@ +from time import time +from uuid import uuid4 + +from ktransformers.server.models.assistants.runs import Run +from ktransformers.server.schemas.assistants.runs import RunCreate,RunObject +from ktransformers.server.schemas.base import ObjectID +from ktransformers.server.utils.sql_utils import SQLUtil + + +class RunsDatabaseManager: + def __init__(self) -> None: + self.sql_util = SQLUtil() + + def create_run_object(self, thread_id: ObjectID, run: RunCreate) -> RunObject: + run_obj = RunObject( + **run.model_dump(mode='json', exclude={"stream"}), + id=str(uuid4()), + object='run', + created_at=int(time()), + thread_id=thread_id, + status=RunObject.Status.queued, + ) + run_obj.set_compute_save(0) + return run_obj + + def db_create_run(self, thread_id: str, run: RunCreate): + db_run = Run( + **run.model_dump(mode="json", exclude={"stream"}), + id=str(uuid4()), + created_at=int(time()), + status="queued", + thread_id=thread_id, + ) + with self.sql_util.get_db() as db: + self.sql_util.db_add_commit_refresh(db, db_run) + run_obj = RunObject.model_validate(db_run.__dict__) + run_obj.set_compute_save(0) + return run_obj + + def db_sync_run(self, run: RunObject) -> None: + db_run = Run( + **run.model_dump(mode='json'), + ) + with self.sql_util.get_db() as db: + self.sql_util.db_merge_commit(db, db_run) + + def db_get_run(self, run_id: ObjectID) -> RunObject: + with self.sql_util.get_db() as db: + db_run = db.query(Run).filter(Run.id == run_id).first() + return RunObject.model_validate(db_run.__dict__) diff --git a/archive/kt-sft/ktransformers/server/crud/assistants/threads.py b/archive/kt-sft/ktransformers/server/crud/assistants/threads.py new file mode 100644 index 0000000..c87cb23 --- /dev/null +++ b/archive/kt-sft/ktransformers/server/crud/assistants/threads.py @@ -0,0 +1,93 @@ +from time import time +from typing import Optional,List +from uuid import uuid4 + +from ktransformers.server.models.assistants.messages import Message +from ktransformers.server.models.assistants.threads import Thread +from ktransformers.server.schemas.assistants.threads import ThreadCreate,ThreadObject +from ktransformers.server.schemas.base import ObjectID, Order +from ktransformers.server.schemas.conversation import ThreadPreview +from ktransformers.server.utils.sql_utils import SQLUtil +from ktransformers.server.crud.assistants.messages import MessageDatabaseManager +from ktransformers.server.config.log import logger +from ktransformers.server.crud.assistants.assistants import AssistantDatabaseManager + +class ThreadsDatabaseManager: + def __init__(self) -> None: + self.sql_util = SQLUtil() + self.message_manager = MessageDatabaseManager() + self.assistant_maanager = AssistantDatabaseManager() + + def db_create_thread(self, thread: ThreadCreate): + thread_id = str(uuid4()) + db_messages = [] + with self.sql_util.get_db() as db: + if thread.messages is not None: + logger.debug("Creating messages first for thread") + for message in thread.messages: + db_message: Message = MessageDatabaseManager.create_db_message_by_core( + message) + db_message.role = "user" + db_message.thread_id = thread_id + db.add(db_message) + db_messages.append(db_message) + + db_thread = Thread( + **thread.model_dump(exclude="messages"), + id=str(uuid4()), + created_at=int(time()), + messages=db_messages, + ) + + self.sql_util.db_add_commit_refresh(db, db_thread) + thread_obj = ThreadObject.model_validate(db_thread.__dict__) + + if 'assistant_id' in thread.meta_data: +# assistant = self.assistant_maanager.db_get_assistant_by_id(thread.meta_data['assistant_id'], db) + assistant = self.assistant_maanager.db_get_assistant_by_id(thread.meta_data['assistant_id']) + logger.info( + f'Append this related thread to assistant {assistant.id}') + assistant.append_related_threads([thread_obj.id]) + assistant.sync_db(db) + return thread_obj + + def db_get_thread_by_id(self, thread_id: ObjectID): + with self.sql_util.get_db() as db: + db_thread = db.query(Thread).filter(Thread.id == thread_id).first() + return ThreadObject.model_validate(db_thread.__dict__) + + def db_list_threads(self, limit: Optional[int], order: Order) -> List[ThreadObject]: + with self.sql_util.get_db() as db: + query = db.query(Thread).order_by(order.to_sqlalchemy_order()( + Thread.created_at)).filter(~Thread.meta_data.contains('assistant_id')) + + if limit is not None: + db_threads = query.limit(limit) + else: + db_threads = query.all() + + return [ThreadObject.model_validate(tool.__dict__) for tool in db_threads] + + def db_list_threads_preview(self, limit: Optional[int], order: Order) -> List[ThreadPreview]: + threads = self.db_list_threads(limit, order) + previews = [] + for thread in threads: + messages = self.message_manager.db_list_messages_of_thread( + thread.id, limit=2, order=Order.ASC) + if len(messages) == 2: + message = messages[0] + assistant = self.assistant_maanager.db_get_assistant_by_id( + messages[1].assistant_id) + else: + message = None + assistant = None + previews.append(ThreadPreview( + assistant=assistant, thread=thread, first_message=message)) + return previews + + def db_delete_thread_by_id(self, thread_id: ObjectID): + with self.sql_util.get_db() as db: + db_thread = db.query(Thread).filter(Thread.id == thread_id).first() + db.delete(db_thread) + # TODO delete related messages and runs and other stuff or just gc + db.commit() diff --git a/archive/kt-sft/ktransformers/server/exceptions.py b/archive/kt-sft/ktransformers/server/exceptions.py new file mode 100644 index 0000000..059e99a --- /dev/null +++ b/archive/kt-sft/ktransformers/server/exceptions.py @@ -0,0 +1,23 @@ +from fastapi import HTTPException, status + + +def db_exception(): + return HTTPException( + status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, + detail="DB Error", + ) + + +def not_implemented(what): + return HTTPException( + status_code=status.HTTP_501_NOT_IMPLEMENTED, + detail=f"{what} not implemented", + ) + + +def internal_server_error(what): + return HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=f"{what}") + + +def request_error(what): + return HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail=f"{what}") diff --git a/archive/kt-sft/ktransformers/server/main.py b/archive/kt-sft/ktransformers/server/main.py new file mode 100644 index 0000000..3341ee9 --- /dev/null +++ b/archive/kt-sft/ktransformers/server/main.py @@ -0,0 +1,122 @@ +import os +import re +from fastapi import FastAPI +from fastapi.staticfiles import StaticFiles +import uvicorn.logging +import uvicorn +import sys +project_dir = os.path.dirname(os.path.dirname(os.path.dirname(__file__))) +from fastapi.middleware.cors import CORSMiddleware +from ktransformers.server.args import ArgumentParser +from ktransformers.server.config.config import Config +from ktransformers.server.utils.create_interface import create_interface, GlobalInterface +from fastapi.openapi.utils import get_openapi +from fastapi import FastAPI +from fastapi.middleware.cors import CORSMiddleware +from ktransformers.server.api import router, post_db_creation_operations +from ktransformers.server.utils.sql_utils import Base, SQLUtil +from ktransformers.server.config.log import logger + + +def mount_app_routes(mount_app: FastAPI): + sql_util = SQLUtil() + logger.info("Creating SQL tables") + Base.metadata.create_all(bind=sql_util.sqlalchemy_engine) + post_db_creation_operations() + mount_app.include_router(router) + + +def create_app(): + cfg = Config() + if(hasattr(GlobalInterface.interface, "lifespan")): + app = FastAPI(lifespan=GlobalInterface.interface.lifespan) + else: + app = FastAPI() + if Config().web_cross_domain: + app.add_middleware( + CORSMiddleware, + allow_origins=["*"], + allow_credentials=True, + allow_methods=["*"], + allow_headers=["*"], + ) + mount_app_routes(app) + if cfg.mount_web: + mount_index_routes(app) + return app + + +def update_web_port(config_file: str): + ip_port_pattern = ( + r"(localhost|((25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)\.){3}(25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)):[0-9]{1,5}" + ) + with open(config_file, "r", encoding="utf-8") as f_cfg: + web_config = f_cfg.read() + ip_port = "localhost:" + str(Config().server_port) + new_web_config = re.sub(ip_port_pattern, ip_port, web_config) + with open(config_file, "w", encoding="utf-8") as f_cfg: + f_cfg.write(new_web_config) + + +def mount_index_routes(app: FastAPI): + project_dir = os.path.dirname(os.path.dirname(__file__)) + web_dir = os.path.join(project_dir, "website/dist") + web_config_file = os.path.join(web_dir, "config.js") + update_web_port(web_config_file) + if os.path.exists(web_dir): + app.mount("/web", StaticFiles(directory=web_dir), name="static") + else: + err_str = f"No website resources in {web_dir}, please complile the website by npm first" + logger.error(err_str) + print(err_str) + exit(1) + + +def run_api(app, host, port, **kwargs): + if kwargs.get("ssl_keyfile") and kwargs.get("ssl_certfile"): + uvicorn.run( + app, + host=host, + port=port, + ssl_keyfile=kwargs.get("ssl_keyfile"), + ssl_certfile=kwargs.get("ssl_certfile"), + ) + else: + uvicorn.run(app, host=host, port=port, log_level="debug") + + +def custom_openapi(app): + if app.openapi_schema: + return app.openapi_schema + openapi_schema = get_openapi( + title="ktransformers server", + version="1.0.0", + summary="This is a server that provides a RESTful API for ktransformers.", + description="We provided chat completion and openai assistant interfaces.", + routes=app.routes, + ) + openapi_schema["info"]["x-logo"] = {"url": "https://kvcache.ai/media/icon_1.png"} + app.openapi_schema = openapi_schema + return app.openapi_schema + + +def main(): + cfg = Config() + + arg_parser = ArgumentParser(cfg) + + args = arg_parser.parse_args() + create_interface(config=cfg, default_args=cfg) + app = create_app() + custom_openapi(app) + + run_api( + app=app, + host=args.host, + port=args.port, + ssl_keyfile=args.ssl_keyfile, + ssl_certfile=args.ssl_certfile, + ) + +if __name__ == "__main__": + main() diff --git a/archive/kt-sft/ktransformers/server/models/__init__.py b/archive/kt-sft/ktransformers/server/models/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/archive/kt-sft/ktransformers/server/models/assistants/__init__.py b/archive/kt-sft/ktransformers/server/models/assistants/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/archive/kt-sft/ktransformers/server/models/assistants/assistants.py b/archive/kt-sft/ktransformers/server/models/assistants/assistants.py new file mode 100644 index 0000000..58e4340 --- /dev/null +++ b/archive/kt-sft/ktransformers/server/models/assistants/assistants.py @@ -0,0 +1,29 @@ +from sqlalchemy import JSON, Column, Float, Integer, String, Text +from sqlalchemy.orm import relationship + +from ktransformers.server.utils.sql_utils import Base + + +class Assistant(Base): + __tablename__ = "assistants" + + id = Column(String, primary_key=True, index=True) + object = Column(String, default="assistant") + created_at = Column(Integer) + + name = Column(String, nullable=True) + description = Column(String, nullable=True) + model = Column(String) + instructions = Column(Text, nullable=True) + tools = Column(JSON) + tool_resources = Column(JSON) + temperature = Column(Float, nullable=True) + meta_data = Column(JSON, nullable=True) + top_p = Column(Float, nullable=True) + response_format = Column(JSON, default="auto") + + build_status = Column(JSON, nullable=True) + + runs = relationship("Run", back_populates="assistant") + + messages = relationship("Message", back_populates="assistant") diff --git a/archive/kt-sft/ktransformers/server/models/assistants/messages.py b/archive/kt-sft/ktransformers/server/models/assistants/messages.py new file mode 100644 index 0000000..c0e1ccb --- /dev/null +++ b/archive/kt-sft/ktransformers/server/models/assistants/messages.py @@ -0,0 +1,28 @@ +from sqlalchemy import JSON, Column, ForeignKey, Integer, String +from sqlalchemy.orm import relationship + +from ktransformers.server.utils.sql_utils import Base + + +class Message(Base): + __tablename__ = "messages" + + id = Column(String, primary_key=True, index=True) + object = Column(String, default="thread.message") + created_at = Column(Integer) + + thread_id = Column(String, ForeignKey("threads.id")) + status = Column(String, default="in_progress") + incomplete_details = Column(JSON, nullable=True) + completed_at = Column(Integer, nullable=True) + incomplete_at = Column(Integer, nullable=True) + role = Column(JSON) + content = Column(JSON) + assistant_id = Column(String, ForeignKey("assistants.id"), nullable=True) + run_id = Column(String, ForeignKey("runs.id"), nullable=True) + attachments = Column(JSON, nullable=True) + meta_data = Column(JSON, nullable=True) + + thread = relationship("Thread", back_populates="messages") + assistant = relationship("Assistant", back_populates="messages") + run = relationship("Run", back_populates="message") diff --git a/archive/kt-sft/ktransformers/server/models/assistants/run_steps.py b/archive/kt-sft/ktransformers/server/models/assistants/run_steps.py new file mode 100644 index 0000000..b241931 --- /dev/null +++ b/archive/kt-sft/ktransformers/server/models/assistants/run_steps.py @@ -0,0 +1,31 @@ +from sqlalchemy import JSON, Column, ForeignKey, Integer, String +from sqlalchemy.orm import relationship + +from ktransformers.server.utils.sql_utils import Base + + +class RunStep(Base): + __tablename__ = "run_steps" + # todo + id = Column(String, primary_key=True, index=True) + object = Column(String, default="thread.run.step") + created_at = Column(Integer) + + assistant_id = Column(String, ForeignKey("assistants.id")) + thread_id = Column(String, ForeignKey("threads.id")) + run_id = Column(String, ForeignKey("runs.id")) + type = Column(String) + status = Column(String) + step_details = Column(JSON) + last_error = Column(JSON, nullable=True) + expires_at = Column(Integer, nullable=True) + cancelled_at = Column(Integer, nullable=True) + failed_at = Column(Integer, nullable=True) + completed_at = Column(Integer, nullable=True) + + meta_data = Column(JSON, nullable=True) + usage = Column(JSON, nullable=True) + + assistant = relationship("Assistant", back_populates="run_steps") + thread = relationship("Thread", back_populates="run_steps") + run = relationship("Run", back_populates="run_steps") diff --git a/archive/kt-sft/ktransformers/server/models/assistants/runs.py b/archive/kt-sft/ktransformers/server/models/assistants/runs.py new file mode 100644 index 0000000..c002739 --- /dev/null +++ b/archive/kt-sft/ktransformers/server/models/assistants/runs.py @@ -0,0 +1,39 @@ +from sqlalchemy import JSON, Column, Float, ForeignKey, Integer, String, Text +from sqlalchemy.orm import relationship + +from ktransformers.server.utils.sql_utils import Base + + +class Run(Base): + __tablename__ = "runs" + + id = Column(String, primary_key=True, index=True) + object = Column(String, default="thread.run") + created_at = Column(Integer) + thread_id = Column(String, ForeignKey("threads.id")) + assistant_id = Column(String, ForeignKey("assistants.id")) + status = Column(String) + required_action = Column(JSON, nullable=True) + last_error = Column(JSON, nullable=True) + expires_at = Column(Integer, nullable=True) + started_at = Column(Integer, nullable=True) + cancelled_at = Column(Integer, nullable=True) + failed_at = Column(Integer, nullable=True) + completed_at = Column(Integer, nullable=True) + incomplete_details = Column(JSON, nullable=True) + # get from assistant + model = Column(String) + instructions = Column(Text, nullable=True) + tools = Column(JSON) + meta_data = Column(JSON, nullable=True) + usage = Column(JSON, nullable=True) + temperature = Column(Float, nullable=True) + top_p = Column(Float, nullable=True) + max_propmp_tokens = Column(Integer, nullable=True) + truncation_strategy = Column(JSON) + tool_choice = Column(JSON) + response_format = Column(JSON, default="auto") + + thread = relationship("Thread", back_populates="runs") + assistant = relationship("Assistant", back_populates="runs") + message = relationship("Message", back_populates="run") diff --git a/archive/kt-sft/ktransformers/server/models/assistants/threads.py b/archive/kt-sft/ktransformers/server/models/assistants/threads.py new file mode 100644 index 0000000..d5f161c --- /dev/null +++ b/archive/kt-sft/ktransformers/server/models/assistants/threads.py @@ -0,0 +1,18 @@ +from sqlalchemy import JSON, Column, Integer, String +from sqlalchemy.orm import relationship + +from ktransformers.server.utils.sql_utils import Base + + +class Thread(Base): + __tablename__ = "threads" + + id = Column(String, primary_key=True, index=True) + object = Column(String, default="thread") + created_at = Column(Integer) + + tool_resources = Column(JSON, nullable=True) + meta_data = Column(JSON, nullable=True) + + runs = relationship("Run", back_populates="thread") + messages = relationship("Message", back_populates="thread") diff --git a/archive/kt-sft/ktransformers/server/schemas/__init__.py b/archive/kt-sft/ktransformers/server/schemas/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/archive/kt-sft/ktransformers/server/schemas/assistants/__init__.py b/archive/kt-sft/ktransformers/server/schemas/assistants/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/archive/kt-sft/ktransformers/server/schemas/assistants/assistants.py b/archive/kt-sft/ktransformers/server/schemas/assistants/assistants.py new file mode 100644 index 0000000..24ea8d4 --- /dev/null +++ b/archive/kt-sft/ktransformers/server/schemas/assistants/assistants.py @@ -0,0 +1,202 @@ +from enum import Enum +from time import time +from typing import AsyncIterable, Callable, Dict, List, Optional, Union +from asyncio import Lock, Queue + +from fastapi import logger +from pydantic import BaseModel, Field, PrivateAttr, field_validator, model_validator +import torch + +from ktransformers.server.config.config import Config +from ktransformers.server.models.assistants.assistants import Assistant +from ktransformers.server.models.assistants.threads import Thread +from ktransformers.server.schemas.assistants.messages import Role +from ktransformers.server.schemas.assistants.runs import RunObject,RunStreamResponse,ObjectWithCreatedTime +from ktransformers.server.schemas.assistants.threads import ThreadObject +from ktransformers.server.schemas.base import Metadata,MetadataField,ObjectID +from ktransformers.server.schemas.assistants.tool import Tool,CodeInterpreter,FileSearch,RelatedThreads,FuntionTool,ToolResource,CodeInterpreterResource,FileSearchResource,RelatedThreadsResource,ToolType +from ktransformers.server.utils.sql_utils import SQLUtil + + +class AssistantBase(BaseModel): + name: Optional[str] = Field(None,description='The name of the assistant.') + description: Optional[str] = Field(None,description='The description of the assistant.') + instructions: Optional[str] = Field(None,description='Instructions which is added in front of the input of LLM') + tools: List[Tool] = Field([], max_length=128) + + @field_validator('tools', mode='before') + def validate_tools(cls, value): + re = [] + if not isinstance(value, list): + raise ValueError('Invalid type for tools') + + for tool in value: + if 'type' not in tool: + raise ValueError('Invalid type for tools') + if tool['type'] == 'code_interpreter': + re.append(CodeInterpreter(**tool)) + elif tool['type'] == 'file_search': + re.append(FileSearch(**tool)) + elif tool['type'] == 'related_threads': + re.append(RelatedThreads(**tool)) + elif tool['type'] == 'function': + re.append(FuntionTool(**tool)) + else: + raise ValueError('Invalid type for tools') + return re + + tool_resources: List[ToolResource] = Field([], max_length=128) + + @field_validator('tool_resources', mode='before') + def validate_tool_resources(cls, value): + re = [] + if not isinstance(value, list): + raise ValueError('Invalid type for tool resources') + + for tool_re in value: + if 'file_ids' in tool_re: + re.append(CodeInterpreterResource(**tool_re)) + elif 'vector_stores' in tool_re: + re.append(FileSearchResource(**tool_re)) + elif 'thread_ids' in tool_re: + re.append(RelatedThreadsResource(**tool_re)) + else: + raise ValueError('Invalid type for tool resources') + return re + + meta_data: Metadata = MetadataField + + @model_validator(mode='before') + def convert_meta_data(cls, values): + if 'meta_data' in values: + values['metadata'] = values['meta_data'] + return values + temperature: Optional[float] = Field(ge=0.0, le=2.0, default=1) + top_p: Optional[float] = Field(ge=0.0, le=1.0, default=1) + response_format: Union[str, Dict[str, str]] = "auto" + + +class AssistantCreate(AssistantBase): + model: str + + +class AssistantBuildStatus(BaseModel): + class Status(Enum): + not_build = "not_build" + in_queue = "in_queue" + parsing = "parsing" + prefilling = "prefilling" + dumping = "dumping" + completed = "completed" + paused = "paused" + + _lock: Lock = PrivateAttr(default_factory=Lock) + _queue: Optional[Queue] = PrivateAttr(None) + + status: Status = Field(default=Status.not_build) + total_file_count: int = Field(default=0) + parsed_file_count: int = Field(default=0) + + prefilling_current: int = Field(default=0) + prefilling_total: int = Field(default=0) + + build_started_time: Optional[int] = Field(default=None) + build_completed_time: Optional[int] = Field(default=None) + + # in megabytes + assistant_usage: int = Field(default=0, description='') + assistant_total_usage: int = Field(default=0) + disk_free_space: int = Field(default=0) + disk_total_space: int = Field(default=0) + + def to_stream_reply(self) -> str: + return f"event: assistant.build.status\ndata: {self.model_dump_json()}\n\n" + + +class AssistantObject(AssistantBase, ObjectWithCreatedTime): + model: Optional[str] = Field( + default=Config().model_name) + related_threads_objects: Optional[List] = Field(None, exclude=True) + _encoded_instruction: Optional[torch.Tensor] = PrivateAttr(default=None) + build_status: AssistantBuildStatus = Field(default=AssistantBuildStatus()) + + def as_api_response(self): + return self.model_dump(exclude={'build_status'}) + + def get_related_threads_ids(self) -> List[ObjectID]: + re = [] + for tool, tool_re in zip(self.tools, self.tool_resources): + if tool.type == ToolType.RELATED_THREADS: + re += tool_re.thread_ids or [] + return re + + def get_related_threads_objects(self) -> List: + # raise NotImplementedError # should be replaced + sql_utils = SQLUtil() + if self.related_threads_objects is None: + with sql_utils.get_db() as db: + db_threads = db.query(Thread).all() + self.related_threads_objects = [tool for tool in [ThreadObject.model_validate( + tool.__dict__) for tool in db_threads] if tool.is_related_threads and tool.meta_data['assistant_id'] == self.id] + # logger.debug( + # f'Found {len(self.related_threads_objects)} related threads') + return self.related_threads_objects + + def append_related_threads(self, thread_ids: List[ObjectID]): + # logger.debug(f'{self.tools} {self.tool_resources}') + for tool, tool_re in zip(self.tools, self.tool_resources): + if tool.type == ToolType.RELATED_THREADS: + tool_re.thread_ids += thread_ids + return + + self.tools.append(RelatedThreads(type=ToolType.RELATED_THREADS)) + self.tool_resources.append( + RelatedThreadsResource(thread_ids=thread_ids)) + + async def update_build_status(self, events: AsyncIterable) -> AsyncIterable: + async for event in events: + # logger.debug(event) + if isinstance(event, RunStreamResponse): + if event.event == RunObject.Status.completed: + self.build_status.status = AssistantBuildStatus.Status.completed + self.build_status.build_completed_time = int(time()) + self.sync_db() + yield self.build_status.model_copy() + elif isinstance(event, dict): + # logger.debug('dict') + if 'stage' in event: + if event['stage'] == 'prefill': + self.build_status.status = AssistantBuildStatus.Status.prefilling + self.build_status.prefilling_current = event['curr_progress'] + self.build_status.prefilling_total = event['max_progress'] + if event['stage'] == 'parse': + self.build_status.status = AssistantBuildStatus.Status.parsing + self.build_status.parsed_file_count = event['curr_progress'] + self.build_status.total_file_count = event['max_progress'] + yield self.build_status.model_copy() + + def get_build_status(self) -> AssistantBuildStatus: + return self.build_status + + + def sync_db(self)->None: + # raise NotImplementedError # should be replaced + sql_utils = SQLUtil() + db_assistant = Assistant( + **self.model_dump(mode='json'), + ) + with sql_utils.get_db() as db: + sql_utils.db_merge_commit(db, db_assistant) + + def get_encoded_instruction(self,encode_fn:Callable)->torch.Tensor: + if self._encoded_instruction is None: + logger.info(f'encoding assistant instruction: {self.instructions}') + self._encoded_instruction = encode_fn(self.instructions, Role.user) + return self._encoded_instruction + + +class AssistantModify(AssistantBase): + model: Optional[str] = None + + +# Non API Backend diff --git a/archive/kt-sft/ktransformers/server/schemas/assistants/messages.py b/archive/kt-sft/ktransformers/server/schemas/assistants/messages.py new file mode 100644 index 0000000..b65ca7c --- /dev/null +++ b/archive/kt-sft/ktransformers/server/schemas/assistants/messages.py @@ -0,0 +1,213 @@ +from enum import Enum +from typing import ForwardRef, List, Optional, Union,Callable + +import torch +from pydantic import BaseModel, PrivateAttr, model_validator + +from ktransformers.server.exceptions import not_implemented +from ktransformers.server.config.log import logger +from ktransformers.server.models.assistants.messages import Message +from ktransformers.server.schemas.base import Metadata, MetadataField, ObjectWithCreatedTime +from ktransformers.server.schemas.assistants.tool import Field,CodeInterpreter,FileSearch +from ktransformers.server.utils.sql_utils import SQLUtil + + +class IncompleteDetails(BaseModel): + reason: str + + +class ContentType(Enum): + image_file = "image_file" + image_url = "image_url" + text = "text" + + +class ContentObject(BaseModel): + type: ContentType + + +class ImageFile(BaseModel): + file_id: str + detail: str + + +class ImageFileObject(ContentObject): + image_file: ImageFile + + +class ImageUrl(BaseModel): + url: str + detail: str + + +class ImageUrlObject(ContentObject): + image_url: ImageUrl + + +class Annotation(BaseModel): + todo: str + + +class Text(BaseModel): + value: str + annotations: List[Annotation] = Field(default=[]) + + +class TextObject(ContentObject): + text: Text + delta_index: int = Field(default=0,exclude=True) + special_tokens_on: bool = Field(default=False,exclude=True) + last_two: str= Field(default='',exclude=True) + + def filter_append(self,text:str): + self.text.value+=text + self.delta_index+=1 + return True + + + +Content = Union[ImageFileObject, ImageUrlObject, TextObject] + + +class Attachment(BaseModel): + file_id: Optional[str] = Field(default=None) + tools: Optional[List[Union[CodeInterpreter, FileSearch]]] = Field(default=None) + + +class Role(Enum): + user = "user" + assistant = "assistant" + + def is_user(self)->bool: + return self == Role.user + + +class MessageCore(BaseModel): + role: Role + content: List[Content] + attachments: Optional[List[Attachment]] + meta_data: Metadata = MetadataField + @model_validator(mode='before') + @classmethod + def convert_meta_data(cls,values): + if 'meta_data' in values: + values['metadata'] = values['meta_data'] + return values + + +class MessageBase(MessageCore): + class Status(Enum): + created = "created" # only used for stream + in_progress = "in_progress" + incomplete = "incomplete" + completed = "completed" + thread_id: str + status: Status + incomplete_details: Optional[IncompleteDetails] = None + completed_at: Optional[int] = None + incomplete_at: Optional[int] = None + + assistant_id: Optional[str] = None + run_id: Optional[str] + + +MessageStreamResponse = ForwardRef('MessageStreamResponse') + +class MessageObject(MessageBase, ObjectWithCreatedTime): + _encoded_content: Optional[torch.Tensor] = PrivateAttr(default=None) + + + def get_text_content(self) -> str: + text_content = "" + for content in self.content: + if content.type == ContentType.text: + text_content += content.text.value + else: + raise not_implemented("Content other than text") + return text_content + + async def get_encoded_content(self,encode_fn:Callable): + if self._encoded_content is None: + logger.info(f'encoding {self.role.value} message({self.status.value}): {self.get_text_content()}') + self._encoded_content = encode_fn(self.get_text_content(),self.role) + + for f in self.get_attached_files(): + logger.info(f'encoding file: {f.filename}') + self._encoded_content = torch.cat([self._encoded_content, encode_fn(await f.get_str(),self.role)],dim=-1) + yield None + + yield self._encoded_content + + + def get_attached_files(self): + raise NotImplementedError # should be replaced + + + + def append_message_delta(self,text:str): + raise NotImplementedError # should be replaced + + def sync_db(self): + # raise NotImplementedError # should be replaced + sql_utils = SQLUtil() + db_message = Message( + **self.model_dump(mode="json"), + ) + with sql_utils.get_db() as db: + sql_utils.db_merge_commit(db, db_message) + + + def stream_response_with_event(self, event: MessageBase.Status) -> MessageStreamResponse: + match event: + case MessageObject.Status.created: + self.status = MessageObject.Status.in_progress + case _: + self.status = event + return MessageStreamResponse(message=self, event=event) + + +class MessageStreamResponse(BaseModel): + message: MessageObject + event: MessageObject.Status + + def to_stream_reply(self): + return f"event: thread.message.{self.event.value}\ndata: {self.message.model_dump_json()}\n\n" + + +class MessageCreate(BaseModel): + role: Role = Field(default=Role.user) + content: Union[str | List[Content]] + attachments: Optional[List[Attachment]] = None + meta_data: Metadata = MetadataField + @model_validator(mode='before') + @classmethod + def convert_meta_data(cls,values): + if 'meta_data' in values: + values['metadata'] = values['meta_data'] + return values + + def to_core(self) -> MessageCore: + # logger.debug(f"Converting message create to core {self.model_dump()}") + core = MessageCore( + role=self.role, + content=[], + attachments=self.attachments, + meta_data=self.meta_data, + ) + if isinstance(self.content, str): + core.content = [TextObject(type="text", text=Text(value=self.content, annotations=[]))] + elif isinstance(self.content, list): + core.content = self.content + else: + raise ValueError("Invalid content type") + return core + + +class MessageModify(BaseModel): + meta_data: Metadata = MetadataField + @model_validator(mode='before') + @classmethod + def convert_meta_data(cls,values): + if 'meta_data' in values: + values['metadata'] = values['meta_data'] + return values diff --git a/archive/kt-sft/ktransformers/server/schemas/assistants/runs.py b/archive/kt-sft/ktransformers/server/schemas/assistants/runs.py new file mode 100644 index 0000000..a76dd74 --- /dev/null +++ b/archive/kt-sft/ktransformers/server/schemas/assistants/runs.py @@ -0,0 +1,201 @@ +from enum import Enum +from typing import Dict, List, Optional, Union, ForwardRef + +from pydantic import BaseModel, Field, model_validator + +from ktransformers.server.models.assistants.runs import Run +from ktransformers.server.schemas.base import TODO, Metadata, MetadataField, ObjectWithCreatedTime +from ktransformers.server.schemas.assistants.threads import ThreadCreate +from ktransformers.server.schemas.assistants.tool import Tool, ToolResource +from ktransformers.server.utils.sql_utils import SQLUtil + + +class ToolCall(BaseModel): + id: str + type: str + function: TODO + + +class SubmitToolOutputs(BaseModel): + tool_calls: List[ToolCall] + + +class RequiredAction(BaseModel): + type: str + submit_tool_outputs: TODO + + +class LastError(BaseModel): + code: str + message: str + + +class IncompleteDetails(BaseModel): + reason: str + + +class Usage(BaseModel): + completion_tokens: int + prompt_tokens: int + total_tokens: int + + +class TruncationStrategy(BaseModel): + type: str = "auto" + last_message: Optional[int] + + +class ToolChoiceType(Enum): + none = "none" + auto = "auto" + required = "required" + + +class RunBase(BaseModel): + class Status(Enum): + created = "created" # only stream event will have this created status + queued = "queued" + in_progress = "in_progress" + requires_action = "requires_action" + cancelling = "cancelling" + cancelled = "cancelled" + failed = "failed" + completed = "completed" + expired = "expired" + + + thread_id: str + assistant_id: str + status: Status = Status.queued + required_action: Optional[RequiredAction] = Field(None) + last_error: Optional[LastError] = Field(None) + expires_at: Optional[int]= Field(None) + started_at: Optional[int] = Field(None) + cancelled_at: Optional[int] = Field(None) + failed_at: Optional[int] = Field(None) + completed_at: Optional[int] = Field(None) + incomplete_details: Optional[IncompleteDetails] = Field(None) + model: Optional[str] = Field(None) + instructions: Optional[str] = Field(None) + tools: Optional[List[Tool]] = Field([]) + meta_data: Metadata = MetadataField + @model_validator(mode='before') + @classmethod + def convert_meta_data(cls,values): + if 'meta_data' in values: + values['metadata'] = values['meta_data'] + return values + + def set_compute_save(self,save:int): + self.meta_data['compute_save'] = str(save) + + + usage: Optional[Usage] = Field(None) + temperature: Optional[float] = Field(None) + top_p: Optional[float]= Field(None) + max_propmp_tokens: Optional[int]= Field(None) + truncation_strategy: Optional[TruncationStrategy]= Field(None) + tool_choice: Optional[Union[ToolChoiceType, dict]]= Field(None) + response_format: Union[str, Dict[str, str]] = "auto" + + +RunStreamResponse = ForwardRef('RunStreamResponse') + +class RunObject(RunBase, ObjectWithCreatedTime): + def stream_response_with_event(self,event:RunBase.Status)->RunStreamResponse: + match event: + case RunBase.Status.created: + self.status = RunBase.Status.queued + case _: + self.status = event + return RunStreamResponse(run=self, event=event) + + + def sync_db(self): + # raise NotImplementedError # should be replaced in crud + sql_utils = SQLUtil() + db_run = Run( + **self.model_dump(mode='json'), + ) + with sql_utils.get_db() as db: + sql_utils.db_merge_commit(db, db_run) + + def create_message_creation_step(self): + raise NotImplementedError # should be replaced + + +class RunStreamResponse(BaseModel): + run: RunObject + event: RunObject.Status + def to_stream_reply(self): + return f"event: thread.run.{self.event.value}\ndata: {self.run.model_dump_json()}\n\n" + +class RunCreate(BaseModel): + assistant_id: str + model: Optional[str] = Field(default=None) + instructions: Optional[str] = Field(default=None) + # TODO: Add this + # additional_instructions: Optional[str] + # additional_messages: Optional[List[MessageCore]] + tools: List[Tool] = Field(default=[]) + meta_data: Metadata = MetadataField + @model_validator(mode='before') + @classmethod + def convert_meta_data(cls,values): + if 'meta_data' in values: + values['metadata'] = values['meta_data'] + return values + temperature: Optional[float] = Field(default=None) + top_p: Optional[float] = Field(default=None) + stream: Optional[bool] = Field(default=None) + max_propmp_tokens: Optional[int] = Field(default=None) + # TODO: Add this + # max_completion_tokens: Optional[int] + truncation_strategy: Optional[TruncationStrategy] = Field(default=None) + tool_choice: Optional[Union[ToolChoiceType, dict]] = Field(default=None) + response_format: Union[str, Dict[str, str]] = Field(default="auto") + + +class RunThreadCreate(BaseModel): + assistant_id: str + thread: Optional[ThreadCreate] + model: Optional[str] + instructions: Optional[str] + tools: List[Tool] + tool_resources: List[ToolResource] + meta_data: Metadata = MetadataField + @model_validator(mode='before') + @classmethod + def convert_meta_data(cls,values): + if 'meta_data' in values: + values['metadata'] = values['meta_data'] + return values + temperature: Optional[float] + top_p: Optional[float] + stream: Optional[bool] + max_propmp_tokens: Optional[int] + # TODO: Add this + # max_completion_tokens: Optional[int] + truncation_strategy: TruncationStrategy + tool_choice: Union[ToolChoiceType, dict] + response_format: Union[str, Dict[str, str]] = "auto" + + +class RunModify(BaseModel): + meta_data: Metadata = MetadataField + @model_validator(mode='before') + @classmethod + def convert_meta_data(cls,values): + if 'meta_data' in values: + values['metadata'] = values['meta_data'] + return values + + +class ToolOutput(BaseModel): + tool_call_id: Optional[str] + output: Optional[str] + + +class RunSubmit(BaseModel): + tool_outputs: List[ToolOutput] + stream: Optional[bool] diff --git a/archive/kt-sft/ktransformers/server/schemas/assistants/streaming.py b/archive/kt-sft/ktransformers/server/schemas/assistants/streaming.py new file mode 100644 index 0000000..0c3b1a7 --- /dev/null +++ b/archive/kt-sft/ktransformers/server/schemas/assistants/streaming.py @@ -0,0 +1,169 @@ +import asyncio +from typing import AsyncIterable, List, Union + +from fastapi import Request +from fastapi.responses import StreamingResponse +from pydantic import BaseModel + +from ktransformers.server.schemas.assistants.runs import RunStreamResponse +from ktransformers.server.schemas.endpoints.chat import ChatCompletionChunk +from ktransformers.server.config.log import logger +from ktransformers.server.schemas.base import Object +from ktransformers.server.schemas.assistants.messages import ContentType, ImageFileObject, ImageUrlObject, MessageObject, Text, TextObject + + +class TextObjectWithIndex(TextObject): + index: int + + +class ImageFileObjectWithIndex(ImageFileObject): + index: int + + +class ImageUrlObjectWithIndex(ImageUrlObject): + index: int + + +ContentWithIndex = Union[TextObjectWithIndex, + ImageFileObjectWithIndex, ImageUrlObjectWithIndex] + + +class MessageDeltaImpl(BaseModel): + # role: Optional[str] + content: List[ContentWithIndex] + + +class MessageDelta(Object): + delta: MessageDeltaImpl + + def to_stream_reply(self): + return f"event: thread.message.delta\ndata: {self.model_dump_json()}\n\n" + + +def text_delta(index: int, text: str): + return MessageDeltaImpl(content=[TextObjectWithIndex(index=index, type=ContentType.text, text=Text(value=text))]) + + +def append_message_delta(self: MessageObject, text: str): + + if len(self.content) == 0: + self.content.append(TextObject(type=ContentType.text, + text=Text(value=''), delta_index=0)) + + text_object: TextObject = self.content[0] + if text_object.filter_append(text): + return MessageDelta(id=self.id, object="thread.message.delta", delta=text_delta(text_object.delta_index, text)) + else: + return None + + +MessageObject.append_message_delta = append_message_delta + + +class RunStepDeltaImpl(BaseModel): + pass + + +class RunStepDelta(Object): + delta: RunStepDeltaImpl + + def to_stream_reply(self): + return f"event: thread.run.step.delta\ndata: {self.model_dump_json()}\n\n" + + +class Done(): + def to_stream_reply(self): + return f"data: [DONE]\n\n" + + +async def check_client_link(request: Request, async_events: AsyncIterable): + async for event in async_events: + if await request.is_disconnected(): + break + yield event + + +async def add_done(async_events: AsyncIterable): + async for event in async_events: + yield event + yield Done() + + +async def to_stream_reply(async_events: AsyncIterable): + async for event in async_events: + if isinstance(event, str): + yield event + else: + yield event.to_stream_reply() + + +async def filter_api_event(async_events: AsyncIterable): + async for event in async_events: + if isinstance(event, MessageDelta) or isinstance(event, RunStepDelta) or isinstance(event, RunStreamResponse) or isinstance(event, Done): + yield event + + +async def filter_chat_chunk(async_events: AsyncIterable): + async for event in async_events: + if isinstance(event, ChatCompletionChunk): + yield event + + +async def filter_by_types(async_events: AsyncIterable, types: List): + async for event in async_events: + for type in types: + if isinstance(event, type): + yield event + continue + + +def api_stream_response(request: Request, async_events: AsyncIterable): + return StreamingResponse(check_client_link(request, to_stream_reply(add_done(filter_api_event(async_events)))), media_type="text/event-stream") + + +def chat_stream_response(request: Request, async_events: AsyncIterable): + return StreamingResponse(check_client_link(request, to_stream_reply(add_done(filter_chat_chunk(async_events)))), media_type="text/event-stream") + + +def stream_response(request: Request, async_events: AsyncIterable): + return StreamingResponse(check_client_link(request, to_stream_reply(add_done(async_events))), media_type="text/event-stream") + + +def check_link_response(request: Request, async_events: AsyncIterable): + return StreamingResponse(check_client_link(request, async_events), media_type="text/event-stream") + + +def wrap_async_generator_into_queue(async_events: AsyncIterable) -> asyncio.Queue: + queue = asyncio.Queue() + + async def inner(): + # logger.debug('run inner') + async for event in async_events: + # logger.debug(f'put: {event}') + await queue.put(event) + await asyncio.sleep(0) + # logger.debug(f'put: None') + await queue.put(None) + asyncio.create_task(inner()) + return queue + + +async def unwrap_async_queue(queue: asyncio.Queue) -> AsyncIterable: + while True: + events = [await queue.get()] + events.extend([queue.get_nowait() for _ in range(queue.qsize())]) + + logger.debug(f'getting {len(events)} events') + for event in events: + if event is None: + break + yield event + + +async def unwrap_async_queue_slow(queue: asyncio.Queue) -> AsyncIterable: + while True: + event = await queue.get() + # logger.debug(f'unwrap_async_queue {event}') + if event is None: + break + yield event diff --git a/archive/kt-sft/ktransformers/server/schemas/assistants/threads.py b/archive/kt-sft/ktransformers/server/schemas/assistants/threads.py new file mode 100644 index 0000000..cf1c382 --- /dev/null +++ b/archive/kt-sft/ktransformers/server/schemas/assistants/threads.py @@ -0,0 +1,49 @@ +from enum import Enum +from typing import List +from typing_extensions import Self + +from pydantic import BaseModel, Field, model_validator + +from ktransformers.server.schemas.base import Metadata, MetadataField, ObjectWithCreatedTime +from ktransformers.server.schemas.assistants.tool import ToolResource +from ktransformers.server.schemas.assistants.messages import MessageCore + + +class ThreadBase(BaseModel): + meta_data: Metadata = MetadataField + @model_validator(mode='before') + @classmethod + def convert_meta_data(cls,values): + if 'meta_data' in values: + values['metadata'] = values['meta_data'] + return values + + tool_resources: List[ToolResource] = Field([], max_length=128) + + +class ThreadObject(ThreadBase, ObjectWithCreatedTime): + is_related_threads:bool = Field(False,exclude=True) + + @model_validator(mode='after') + def check_is_related_threads(self)->Self: + # logger.debug(f'check thread {self.id} is related thread? by {self}') + if 'assistant_id' in self.meta_data: + self.is_related_threads = True + return self + + class StreamEvent(Enum): + created = 'created' + + def to_stream_reply(self,event:StreamEvent): + return f"event: thread.{event.value}\ndata: {self.model_dump_json()}\n\n" + + +class ThreadCreate(ThreadBase): + messages: List[MessageCore] = Field(default=[]) + + +class ThreadModify(ThreadBase): + pass + + +# other than OpenAI API diff --git a/archive/kt-sft/ktransformers/server/schemas/assistants/tool.py b/archive/kt-sft/ktransformers/server/schemas/assistants/tool.py new file mode 100644 index 0000000..27fdcae --- /dev/null +++ b/archive/kt-sft/ktransformers/server/schemas/assistants/tool.py @@ -0,0 +1,54 @@ +from enum import Enum +from typing import List, Optional, Union + +from pydantic import BaseModel, Field + +from ktransformers.server.schemas.base import ObjectID + + +class ToolType(str, Enum): + CODE_INTERPRETER = "code_interpreter" + FILE_SEARCH = "file_search" + RELATED_THREADS = "related_threads" + FUNCTION = "function" + + +class ToolBase(BaseModel): + type: ToolType + + +class CodeInterpreter(ToolBase): + pass + + +class FileSearch(ToolBase): + pass + + +class RelatedThreads(ToolBase): + pass + + +class FuntionTool(ToolBase): + description: str + name: str + parameters: List[str] + + +Tool = Union[CodeInterpreter, FileSearch, RelatedThreads, FuntionTool] + + +class CodeInterpreterResource(BaseModel): + file_ids: Optional[List[str]] = Field(default_factory=list, max_length=20) + + +class FileSearchResource(BaseModel): + vector_store_ids: Optional[List[str]] = Field(default_factory=list, max_length=1) + vector_stores: Optional[List[str]] = Field(default_factory=list, max_length=1) + + +class RelatedThreadsResource(BaseModel): + thread_ids: List[ObjectID] = Field(default=[]) + + +ToolResource = Union[CodeInterpreterResource,FileSearchResource,RelatedThreadsResource] diff --git a/archive/kt-sft/ktransformers/server/schemas/base.py b/archive/kt-sft/ktransformers/server/schemas/base.py new file mode 100644 index 0000000..35058dc --- /dev/null +++ b/archive/kt-sft/ktransformers/server/schemas/base.py @@ -0,0 +1,46 @@ +from enum import Enum +from typing import Dict + +import sqlalchemy +from pydantic import BaseModel, ConfigDict, Field + +TODO = BaseModel + +ObjectID = str + + +class Object(BaseModel): + id: ObjectID + object: str + + model_config = ConfigDict(from_attributes=True) + + +# Pydantic Base Models +class ObjectWithCreatedTime(Object): + created_at: int + + + +class Order(str, Enum): + ASC = "asc" + DESC = "desc" + + def to_sqlalchemy_order(self): + match self: + case Order.ASC: + return sqlalchemy.asc + case Order.DESC: + return sqlalchemy.desc + + +Metadata = Dict[str, str] +MetadataField: Metadata = Field({},max_length=16, alias="metadata") + + +class DeleteResponse(Object): + deleted: bool = True + +class OperationResponse(BaseModel): + operation: str + status: str diff --git a/archive/kt-sft/ktransformers/server/schemas/conversation.py b/archive/kt-sft/ktransformers/server/schemas/conversation.py new file mode 100644 index 0000000..c64ea92 --- /dev/null +++ b/archive/kt-sft/ktransformers/server/schemas/conversation.py @@ -0,0 +1,12 @@ +from typing import Optional + +from pydantic import BaseModel + +from .assistants.assistants import AssistantObject +from .assistants.threads import ThreadObject +from .assistants.messages import MessageObject + +class ThreadPreview(BaseModel): + assistant: Optional[AssistantObject] = None + thread: ThreadObject + first_message: Optional[MessageObject] = None diff --git a/archive/kt-sft/ktransformers/server/schemas/endpoints/chat.py b/archive/kt-sft/ktransformers/server/schemas/endpoints/chat.py new file mode 100644 index 0000000..8f9d60f --- /dev/null +++ b/archive/kt-sft/ktransformers/server/schemas/endpoints/chat.py @@ -0,0 +1,100 @@ +from typing import List, Optional, Union, Dict, Any +from typing_extensions import Literal +from enum import Enum +from pydantic import BaseModel, Field +from ktransformers.server.config.config import Config +from ktransformers.server.schemas.base import Object + + +from openai.types.chat.chat_completion_chunk import Choice + +from uuid import uuid4 + +class CompletionUsage(BaseModel): + prompt_tokens: int + completion_tokens: int + total_tokens: int + prompt_tokens_details: Optional[Dict[str, Any]] = None + completion_tokens_details: Optional[Dict[str, Any]] = None + prefill_time: Optional[float] = None + decode_time: Optional[float] = None + +class Role(Enum): + system = 'system' + user = 'user' + assistant = 'assistant' + tool = 'tool' + function = 'function' + +class Message(BaseModel): + content: Optional[str] = None + role: Role + name: Optional[str] = None + tool_calls: Optional[List[Dict[str, Any]]] = {} + tool_call_id: Optional[str] = None + + def to_tokenizer_message(self): + message = {'role': self.role.value} + if self.content is not None: + message['content'] = self.content + if self.name is not None: + message['name'] = self.name + if self.tool_calls is not {}: + message['tool_calls'] = self.tool_calls + if self.tool_call_id is not None: + message['tool_call_id'] = self.tool_call_id + return message + +class FunctionParameters(BaseModel): + type: str = "object" + properties: Dict[str, Any] = {} + required: Optional[List[str]] = None + +class FunctionDefinition(BaseModel): + name: str + description: Optional[str] = None + parameters: FunctionParameters = Field(default_factory=FunctionParameters) + +class ToolFunction(BaseModel): + function: FunctionDefinition + +class Tool(BaseModel): + type: Literal["function"] + function: FunctionDefinition + +class ChatCompletionCreate(BaseModel): + messages: List[Message] + model: str + stream: bool = False + temperature: Optional[float] = Field(default=Config().temperature) + top_p: Optional[float] = Field(default=Config().top_p) + tools: Optional[List[Tool]] = None + tool_choice: Optional[Union[str, Dict[str, Any]]] = None + stream_options: Optional[Dict[str, Any]] = None + frequency_penalty: float = 0 + presence_penalty: float = 0 + max_tokens: Optional[int] = Field(default=None) + max_completion_tokens: Optional[int] = Field(default=None) + return_speed: Optional[bool] = Field(default=False) + def get_tokenizer_messages(self): + return [m.to_tokenizer_message() for m in self.messages] + +class ChatCompletionChunk(BaseModel): + id: str + choices: List[Choice] + created: int + model: str + object: Literal["chat.completion.chunk"] + service_tier: Optional[Literal["scale", "default"]] = None + system_fingerprint: Optional[str] = None + usage: Optional[CompletionUsage] = None + + def to_stream_reply(self): + return f"data: {self.model_dump_json()}\n\n" + +class RawUsage(BaseModel): + tokenize_time: float + prefill_time: float + decode_time: float + prefill_count: int + decode_count: int \ No newline at end of file diff --git a/archive/kt-sft/ktransformers/server/schemas/legacy/__init__.py b/archive/kt-sft/ktransformers/server/schemas/legacy/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/archive/kt-sft/ktransformers/server/schemas/legacy/completions.py b/archive/kt-sft/ktransformers/server/schemas/legacy/completions.py new file mode 100644 index 0000000..a506a97 --- /dev/null +++ b/archive/kt-sft/ktransformers/server/schemas/legacy/completions.py @@ -0,0 +1,51 @@ +from typing import List, Optional +from enum import Enum +from pydantic import BaseModel, Field +from ktransformers.server.config.config import Config +from ..base import Object + +class CompletionCreate(BaseModel): + model: str + prompt: str | List[str] + stream: bool = False + temperature: Optional[float] = Field(default=Config().temperature) + top_p: Optional[float] = Field(default=Config().top_p) + max_tokens: Optional[int] = Field(default=None) + max_completion_tokens: Optional[int] = Field(default=None) + + def get_tokenizer_messages(self): + if isinstance(self.prompt,List): + self.get_tokenizer_messages('\n'.join(self.prompt)) + return [{'content':self.prompt,'role':'user'}] + + +class FinishReason(Enum): + stop = 'stop' + length = 'length' + +class Choice(BaseModel): + index: int + text: str + logprobs: Optional[str] = None + finish_reason: FinishReason = None + + +class CompletionObject(Object): + created:int + choices: List[Choice] = [] + model:str = 'not implmented' + system_fingerprint:str = 'not implmented' + usage: Optional[str] = None + + def set_token(self,token:str): + if len(self.choices)==0: + self.choices.append(Choice(index=0,text='')) + self.choices[0].text = token + + def append_token(self,token:str): + if len(self.choices)==0: + self.choices.append(Choice(index=0,text='')) + self.choices[0].text += token + + def to_stream_reply(self): + return f"data:{self.model_dump_json()}\n\n" diff --git a/archive/kt-sft/ktransformers/server/utils/__init__.py b/archive/kt-sft/ktransformers/server/utils/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/archive/kt-sft/ktransformers/server/utils/create_interface.py b/archive/kt-sft/ktransformers/server/utils/create_interface.py new file mode 100644 index 0000000..992c831 --- /dev/null +++ b/archive/kt-sft/ktransformers/server/utils/create_interface.py @@ -0,0 +1,41 @@ +#!/usr/bin/env python +# coding=utf-8 +''' +Description : +Author : qiyuxinlin +Date : 2024-07-25 11:50:16 +Version : 1.0.0 +LastEditors : qiyuxinlin +LastEditTime : 2024-07-25 12:54:48 +Copyright (c) 2024 by KVCache.AI, All Rights Reserved. +''' +from ktransformers.server.config.config import Config +from ktransformers.server.backend.args import ConfigArgs +from ktransformers.server.backend.context_manager import ThreadContextManager +from ktransformers.server.backend.interfaces.exllamav2 import ExllamaInterface +from ktransformers.server.backend.interfaces.transformers import TransformersInterface +from ktransformers.server.backend.interfaces.ktransformers import KTransformersInterface + +def create_interface(config: Config, default_args: ConfigArgs): + if config.backend_type=='transformers': + from ktransformers.server.backend.interfaces.transformers import TransformersInterface as BackendInterface + elif config.backend_type == 'exllamav2': + from ktransformers.server.backend.interfaces.exllamav2 import ExllamaInterface as BackendInterface + elif config.backend_type == 'ktransformers': + from ktransformers.server.backend.interfaces.ktransformers import KTransformersInterface as BackendInterface + elif config.backend_type == 'balance_serve': + from ktransformers.server.backend.interfaces.balance_serve import BalanceServeInterface as BackendInterface + else: + raise NotImplementedError(f'{config.backend_type} not implemented') + GlobalInterface.interface = BackendInterface(default_args) + GlobalContextManager.context_manager = ThreadContextManager(GlobalInterface.interface) + +class GlobalContextManager: + context_manager: ThreadContextManager +class GlobalInterface: + interface: TransformersInterface | KTransformersInterface | ExllamaInterface + +def get_thread_context_manager() -> GlobalContextManager: + return GlobalContextManager.context_manager +def get_interface() -> GlobalInterface: + return GlobalInterface.interface \ No newline at end of file diff --git a/archive/kt-sft/ktransformers/server/utils/multi_timer.py b/archive/kt-sft/ktransformers/server/utils/multi_timer.py new file mode 100644 index 0000000..dbb54ca --- /dev/null +++ b/archive/kt-sft/ktransformers/server/utils/multi_timer.py @@ -0,0 +1,79 @@ +import time + + +def format_time(seconds): + units = [ + ("hours", 3600), + ("minutes", 60), + ("seconds", 1), + ("milliseconds", 1e-3), + ("microseconds", 1e-6), + ] + + for unit_name, unit_value in units: + if seconds >= unit_value: + time_value = seconds / unit_value + return f"{time_value:.2f} {unit_name}" + return "0 seconds" # Handle case for 0 seconds + + +class Profiler: + def __init__(self): + self.timers = {} + self.counters = {} + + def create_timer(self, name): + self.timers[name] = { + "start_time": None, + "elapsed_time": 0, + "running": False, + } + + def start_timer(self, name): + if name not in self.timers: + raise ValueError(f"Timer '{name}' does not exist.") + if self.timers[name]["running"]: + raise ValueError(f"Timer '{name}' is already running.") + self.timers[name]["start_time"] = time.time() + self.timers[name]["running"] = True + + def pause_timer(self, name): + if name not in self.timers: + raise ValueError(f"Timer '{name}' does not exist.") + if not self.timers[name]["running"]: + raise ValueError(f"Timer '{name}' is not running.") + self.timers[name]["elapsed_time"] += time.time() - self.timers[name]["start_time"] + self.timers[name]["running"] = False + + def get_timer_sec(self, name): + if name not in self.timers: + raise ValueError(f"Timer '{name}' does not exist.") + if self.timers[name]["running"]: + current_time = self.timers[name]["elapsed_time"] + (time.time() - self.timers[name]["start_time"]) + else: + current_time = self.timers[name]["elapsed_time"] + return current_time + + def get_all_timers(self): + all_timers = {} + for name in self.timers: + all_timers[name] = self.get_timer_sec(name) + return all_timers + + def report_timer_string(self, name): + return f"{name} elapsed time: {format_time(self.get_timer_sec(name))}" + + def create_and_start_timer(self, name): + self.create_timer(name) + self.start_timer(name) + + + # Counter + def inc(self,key:str,delta:int=1): + self.counters[key] = self.counters.get(key,0) + delta + + def set_counter(self,key:str,to=0): + self.counters[key] = to + + def get_counter(self,key:str): + return self.counters.get(key,0) diff --git a/archive/kt-sft/ktransformers/server/utils/sql_utils.py b/archive/kt-sft/ktransformers/server/utils/sql_utils.py new file mode 100644 index 0000000..68da000 --- /dev/null +++ b/archive/kt-sft/ktransformers/server/utils/sql_utils.py @@ -0,0 +1,128 @@ +#!/usr/bin/env python +# coding=utf-8 +''' +Description : +Author : chenxl +Date : 2024-06-12 09:12:58 +Version : 1.0.0 +LastEditors : chenxl +LastEditTime : 2024-07-27 01:56:04 +''' + +from urllib.parse import urlparse +import os +from contextlib import contextmanager +from sqlalchemy import create_engine +from sqlalchemy.orm import Session, sessionmaker, declarative_base + +from ktransformers.server.config.config import Config +from ktransformers.server.config.singleton import Singleton +from ktransformers.server.config.log import logger +from ktransformers.server.exceptions import db_exception + + +Base = declarative_base() + + +class SQLUtil(metaclass=Singleton): + """ + database connections init and management + """ + sqlalchemy_engine = None + session_local = None + + def __init__(self) -> None: + self.cfg: Config = Config() + if not self.sqlalchemy_engine: + SQLUtil.init_engine(self.cfg) + + @contextmanager + def get_db(self): + """ + After you finish using the session, it's crucial to close it. + """ + if not SQLUtil.sqlalchemy_engine: + SQLUtil.init_engine(self.cfg) + session = self.session_local() # type: ignore pylint: disable=not-callable + try: + yield session + finally: + session.close() + + @staticmethod + def init_engine(cfg: Config): + """ + initial engine and session maker Factory + """ + pool_size = cfg.db_pool_size + if SQLUtil.sqlalchemy_engine is None: + if cfg.db_type == "sqllite": + db_url = SQLUtil.create_sqllite_url(cfg) + else: + logger.error("Unsupported database type %s", cfg.db_type) + exit(-1) + SQLUtil.sqlalchemy_engine = create_engine( + db_url, connect_args={"check_same_thread": False}, pool_size=pool_size) + SQLUtil.session_local = sessionmaker( + autocommit=False, autoflush=False, bind=SQLUtil.sqlalchemy_engine) + + @staticmethod + def create_sqllite_url(cfg): + """ + create and validate SQLLite url + """ + path: str = cfg.db_host + database: str = cfg.db_database + absolute_path: str = os.path.join(path, database) + url = 'sqlite:///' + absolute_path + try: + result = urlparse(url) + if all([result.scheme, result.path, result.scheme == 'sqlite']): + return url + else: + logger.error("invalid sqllite url: %s", url) + exit(-1) + except ValueError: + logger.error("invalid sqllite url: %s", url) + exit(-1) + + def db_add_commit_refresh(self, session: Session, what): + """ + add data to database + """ + try: + session.add(what) + session.commit() + session.refresh(what) + except Exception as e: + logger.exception("db commit error with data %s", str(what.__dict__)) + ex = db_exception() + ex.detail = str(e) + session.rollback() + raise ex from e + + def db_merge_commit(self, session: Session, what): + try: + session.merge(what) + session.commit() + except Exception as e: + ex = db_exception() + ex.detail = str(e) + logger.exception("db merge commit error with data %s", str(what.__dict__)) + session.rollback() + raise ex from e + + def db_update_commit_refresh(self, session: Session, existing, what): + what = what.model_dump(mode="json") + try: + for key in what.keys(): + if what[key] is not None: + setattr(existing, key, what[key]) + session.commit() + session.refresh(existing) + except Exception as e: + ex = db_exception() + ex.detail = str(e) + logger.exception("db update commit refresh error with data %s", str(what.__dict__)) + session.rollback() + raise ex from e diff --git a/archive/kt-sft/ktransformers/sft/__init__.py b/archive/kt-sft/ktransformers/sft/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/archive/kt-sft/ktransformers/sft/flops_utils/__init__.py b/archive/kt-sft/ktransformers/sft/flops_utils/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/archive/kt-sft/ktransformers/sft/flops_utils/custom_profile.py b/archive/kt-sft/ktransformers/sft/flops_utils/custom_profile.py new file mode 100644 index 0000000..ca594c8 --- /dev/null +++ b/archive/kt-sft/ktransformers/sft/flops_utils/custom_profile.py @@ -0,0 +1,263 @@ +from distutils.version import LooseVersion + +from thop.vision.basic_hooks import * +from thop.rnn_hooks import * +from thop.utils import prGreen, prRed, prYellow +import sys, os + +project_dir = os.path.dirname(os.path.dirname(os.path.dirname(__file__))) +sys.path.insert(0, project_dir) + +from ktransformers.util.utils import prefill_and_generate + +# logger = logging.getLogger(__name__) +# logger.setLevel(logging.INFO) + + +if LooseVersion(torch.__version__) < LooseVersion("1.0.0"): + logging.warning( + "You are using an old version PyTorch {version}, which THOP does NOT support.".format( + version=torch.__version__ + ) + ) + +default_dtype = torch.float64 + +register_hooks = { + nn.ZeroPad2d: zero_ops, # padding does not involve any multiplication. + nn.Conv1d: count_convNd, + nn.Conv2d: count_convNd, + nn.Conv3d: count_convNd, + nn.ConvTranspose1d: count_convNd, + nn.ConvTranspose2d: count_convNd, + nn.ConvTranspose3d: count_convNd, + nn.BatchNorm1d: count_normalization, + nn.BatchNorm2d: count_normalization, + nn.BatchNorm3d: count_normalization, + nn.LayerNorm: count_normalization, + nn.InstanceNorm1d: count_normalization, + nn.InstanceNorm2d: count_normalization, + nn.InstanceNorm3d: count_normalization, + nn.PReLU: count_prelu, + nn.Softmax: count_softmax, + nn.ReLU: zero_ops, + nn.ReLU6: zero_ops, + nn.LeakyReLU: count_relu, + nn.MaxPool1d: zero_ops, + nn.MaxPool2d: zero_ops, + nn.MaxPool3d: zero_ops, + nn.AdaptiveMaxPool1d: zero_ops, + nn.AdaptiveMaxPool2d: zero_ops, + nn.AdaptiveMaxPool3d: zero_ops, + nn.AvgPool1d: count_avgpool, + nn.AvgPool2d: count_avgpool, + nn.AvgPool3d: count_avgpool, + nn.AdaptiveAvgPool1d: count_adap_avgpool, + nn.AdaptiveAvgPool2d: count_adap_avgpool, + nn.AdaptiveAvgPool3d: count_adap_avgpool, + nn.Linear: count_linear, + nn.Dropout: zero_ops, + nn.Upsample: count_upsample, + nn.UpsamplingBilinear2d: count_upsample, + nn.UpsamplingNearest2d: count_upsample, + nn.RNNCell: count_rnn_cell, + nn.GRUCell: count_gru_cell, + nn.LSTMCell: count_lstm_cell, + nn.RNN: count_rnn, + nn.GRU: count_gru, + nn.LSTM: count_lstm, + nn.Sequential: zero_ops, + nn.PixelShuffle: zero_ops, +} + +if LooseVersion(torch.__version__) >= LooseVersion("1.1.0"): + register_hooks.update({nn.SyncBatchNorm: count_normalization}) + + +def profile_origin(model, inputs, custom_ops=None, verbose=True, report_missing=False): + handler_collection = [] + types_collection = set() + if custom_ops is None: + custom_ops = {} + if report_missing: + verbose = True + + def add_hooks(m): + if len(list(m.children())) > 0: + return + + if hasattr(m, "total_ops") or hasattr(m, "total_params"): + logging.warning( + "Either .total_ops or .total_params is already defined in %s. " + "Be careful, it might change your code's behavior." % str(m) + ) + + m.register_buffer("total_ops", torch.zeros(1, dtype=default_dtype)) + m.register_buffer("total_params", torch.zeros(1, dtype=default_dtype)) + + for p in m.parameters(): + m.total_params += torch.DoubleTensor([p.numel()]) + + m_type = type(m) + + fn = None + if ( + m_type in custom_ops + ): # if defined both op maps, use custom_ops to overwrite. + fn = custom_ops[m_type] + if m_type not in types_collection and verbose: + print("[INFO] Customize rule %s() %s." % (fn.__qualname__, m_type)) + elif m_type in register_hooks: + fn = register_hooks[m_type] + if m_type not in types_collection and verbose: + print("[INFO] Register %s() for %s." % (fn.__qualname__, m_type)) + else: + if m_type not in types_collection and report_missing: + prRed( + "[WARN] Cannot find rule for %s. Treat it as zero Macs and zero Params." + % m_type + ) + + if fn is not None: + handler = m.register_forward_hook(fn) + handler_collection.append(handler) + types_collection.add(m_type) + + training = model.training + + model.eval() + model.apply(add_hooks) + + with torch.no_grad(): + model(*inputs) + + total_ops = 0 + total_params = 0 + for m in model.modules(): + if len(list(m.children())) > 0: # skip for non-leaf module + continue + total_ops += m.total_ops + total_params += m.total_params + + total_ops = total_ops.item() + total_params = total_params.item() + + # reset model to original status + model.train(training) + for handler in handler_collection: + handler.remove() + + # remove temporal buffers + for n, m in model.named_modules(): + if len(list(m.children())) > 0: + continue + if "total_ops" in m._buffers: + m._buffers.pop("total_ops") + if "total_params" in m._buffers: + m._buffers.pop("total_params") + + return total_ops, total_params + + +def custom_profile( + model: nn.Module, + inputs, + content, + tokenizer, + custom_ops=None, + verbose=True, + ret_layer_info=False, + report_missing=False, +): + handler_collection = {} + types_collection = set() + if custom_ops is None: + custom_ops = {} + if report_missing: + # overwrite `verbose` option when enable report_missing + verbose = True + + def add_hooks(m: nn.Module): + m.register_buffer("total_ops", torch.zeros(1, dtype=torch.float64)) + m.register_buffer("total_params", torch.zeros(1, dtype=torch.float64)) + + # for p in m.parameters(): + # m.total_params += torch.DoubleTensor([p.numel()]) + + m_type = type(m) + + fn = None + if m_type in custom_ops: + # if defined both op maps, use custom_ops to overwrite. + fn = custom_ops[m_type] + if m_type not in types_collection and verbose: + print("[INFO] Customize rule %s() %s." % (fn.__qualname__, m_type)) + elif m_type in register_hooks: + fn = register_hooks[m_type] + if m_type not in types_collection and verbose: + print("[INFO] Register %s() for %s." % (fn.__qualname__, m_type)) + else: + if m_type not in types_collection and report_missing: + prRed( + "[WARN] Cannot find rule for %s. Treat it as zero Macs and zero Params." + % m_type + ) + + if fn is not None: + handler_collection[m] = ( + m.register_forward_hook(fn), + m.register_forward_hook(count_parameters), + ) + types_collection.add(m_type) + + prev_training_status = model.training + + model.eval() + model.apply(add_hooks) + + messages = [{"role": "user", "content": content}] + input_tensor = tokenizer.apply_chat_template( + messages, add_generation_prompt=True, return_tensors="pt" + ) + + with torch.no_grad(): + # model(*inputs) + # TODO: model.model to deal with the PeftModelForCaualLM temp + simple_prefill_and_generate_for_test( + model.model, tokenizer, input_tensor.cuda(), max_new_tokens=1000, use_cuda_graph=False, mode = 'normal', force_think = False, chunk_prefill_size = 8192, + ) + + def dfs_count(module: nn.Module, prefix="\t") -> (int, int): + total_ops, total_params = module.total_ops.item(), 0 + ret_dict = {} + for n, m in module.named_children(): + # if not hasattr(m, "total_ops") and not hasattr(m, "total_params"): # and len(list(m.children())) > 0: + # m_ops, m_params = dfs_count(m, prefix=prefix + "\t") + # else: + # m_ops, m_params = m.total_ops, m.total_params + next_dict = {} + if m in handler_collection and not isinstance( + m, (nn.Sequential, nn.ModuleList) + ): + m_ops, m_params = m.total_ops.item(), m.total_params.item() + else: + m_ops, m_params, next_dict = dfs_count(m, prefix=prefix + "\t") + ret_dict[n] = (m_ops, m_params, next_dict) + total_ops += m_ops + total_params += m_params + # print(prefix, module._get_name(), (total_ops, total_params)) + return total_ops, total_params, ret_dict + + total_ops, total_params, ret_dict = dfs_count(model) + + # reset model to original status + model.train(prev_training_status) + for m, (op_handler, params_handler) in handler_collection.items(): + op_handler.remove() + params_handler.remove() + m._buffers.pop("total_ops") + m._buffers.pop("total_params") + + if ret_layer_info: + return total_ops, total_params, ret_dict + return total_ops, total_params diff --git a/archive/kt-sft/ktransformers/sft/flops_utils/lora_test_utils.py b/archive/kt-sft/ktransformers/sft/flops_utils/lora_test_utils.py new file mode 100644 index 0000000..627b86f --- /dev/null +++ b/archive/kt-sft/ktransformers/sft/flops_utils/lora_test_utils.py @@ -0,0 +1,615 @@ +from torch.profiler import profile, record_function, ProfilerActivity +import os +from transformers import TrainerCallback + +class ProfilerCallback(TrainerCallback): + def __init__(self, profiler): + self.profiler = profiler + + def on_step_end(self, args, state, control, **kwargs): + self.profiler.step() + +def _short(t): + return tuple(t.shape) if isinstance(t, torch.Tensor) else type(t) + +def install_shape_probes(model): + if os.environ.get("KT_DEBUG_MOE","0") != "1": + print("[KT_DEBUG_MOE] off"); return + + try: + acc = trainer.accelerator + cfg = getattr(acc, "dataloader_config", None) + if cfg is not None: + print("[ACCEL DL CONFIG]", + "split_batches=", getattr(cfg,"split_batches",None), + "dispatch_batches=", getattr(cfg,"dispatch_batches",None), + "even_batches=", getattr(cfg,"even_batches",None), + "use_seedable_sampler=", getattr(cfg,"use_seedable_sampler",None), + "non_blocking=", getattr(cfg,"non_blocking",None)) + except Exception as e: + print("[ACCEL DL CONFIG] ", e) + + try: + emb = model.base_model.model.model.embed_tokens + def _emb_pre(mod, inp): + x = inp[0] + if not hasattr(mod, "_dbg_once"): + print(f"[DBG] embed input_ids shape = {tuple(x.shape)} (expect B,S)") + mod._dbg_once = True + emb.register_forward_pre_hook(_emb_pre) + except Exception as e: + print("[DBG] embed hook failed:", e) + + try: + first_layer = model.base_model.model.model.layers[0] + _orig_fwd = first_layer.forward + def _wrap_fwd(self, *args, **kwargs): + hs = args[0] if args else kwargs.get("hidden_states") + if not hasattr(self, "_dbg_once_in"): + print(f"[DBG] L0.in hidden_states = {_short(hs)} (expect B,S,H)") + self._dbg_once_in = True + out = _orig_fwd(*args, **kwargs) + hs_out = out[0] if isinstance(out, (tuple, list)) else out + if not hasattr(self, "_dbg_once_out"): + print(f"[DBG] L0.out hidden_states = {_short(hs_out)}") + self._dbg_once_out = True + return out + first_layer.forward = MethodType(_wrap_fwd, first_layer) + except Exception as e: + print("[DBG] L0 wrap failed:", e) + + try: + moe_layer = None + for i, lyr in enumerate(model.base_model.model.model.layers): + if hasattr(lyr, "mlp"): + moe_layer = lyr.mlp + moe_idx = i + break + if moe_layer is not None: + _moe_orig = moe_layer.forward + def _moe_wrap(self, *args, **kwargs): + x = args[0] if args else kwargs.get("hidden_states") + if not hasattr(self, "_dbg_once"): + print(f"[DBG] MLP(in) @layer{moe_idx} hidden_states = {_short(x)}") + if isinstance(x, torch.Tensor) and x.dim() == 3: + B,S,H = x.shape + print(f"[DBG] tokens before flatten = B*S = {B}*{S} = {B*S}") + self._dbg_once = True + return _moe_orig(*args, **kwargs) + moe_layer.forward = MethodType(_moe_wrap, moe_layer) + else: + print("[DBG] no moe_layer found") + except Exception as e: + print("[DBG] moe wrap failed:", e) + + try: + from ktransformers.operators.experts import KTransformersExperts + def _experts_pre(mod, args): + if hasattr(mod, "_dbg_once"): return + try: + input_tensor, expert_ids, weights = args[:3] + print(f"[DBG] experts.in input_tensor={tuple(input_tensor.shape)} " + f"expert_ids={tuple(expert_ids.shape)} weights={tuple(weights.shape)}") + if input_tensor.dim()==2: + N = input_tensor.shape[0] + print(f"[DBG] N(input rows)={N}") + if expert_ids.dim()==2: + T,K = expert_ids.shape + print(f"[DBG] tokens(T)={T}, K={K}, T*K={T*K}") + mod._dbg_once = True + except Exception as e: + print("[DBG] experts hook parse err:", e) + count=0 + for name,m in model.named_modules(): + if isinstance(m, KTransformersExperts): + m.register_forward_pre_hook(_experts_pre); count+=1 + print(f"[KT_DEBUG_MOE] installed experts hook on {count} modules.") + except Exception as e: + print("[DBG] experts hook failed:", e) + +def inspect_device(model, write_file): + for name, module in model.named_modules(): + with open(write_file, 'a') as file: + file.write(f"Layer: {name}\n") + for param_name, param in module.named_parameters(recurse=False): + with open(write_file, 'a') as file: + file.write(f" Parameter '{param_name}' device: {param.device}\n") + for buffer_name, buffer in module.named_buffers(recurse=False): + with open(write_file, 'a') as file: + file.write(f" Buffer '{buffer_name}' device: {buffer.device}\n") + +def print_model_params(model): + # for layer_idx in range(len(model.model.orig_module.layers)): + for layer_idx in range(0, 3): + layer = model.model.orig_module.layers[layer_idx] + + print(f"\n================ Layer {layer_idx} Attention ================") + + q_proj = layer.self_attn.orig_module.q_proj.orig_module + print(f"\nq_proj.generate_linear.weight (shape: {q_proj.generate_linear.weight.shape})") + print(q_proj.generate_linear.weight.cpu()) + + # kv_a_proj = layer.self_attn.orig_module.kv_a_proj_with_mqa.orig_module + # print(f"\nkv_a_proj.weight (shape: {kv_a_proj.weight.shape})") + # print(kv_a_proj.weight.data[:3, :5].detach().cpu().numpy()) + + # o_proj = layer.self_attn.orig_module.o_proj.orig_module + # print(f"\no_proj.weight (shape: {o_proj.weight.shape})") + # print(o_proj.weight.data[:3, :5].detach().cpu().numpy()) + + # print(f"\n================ Layer {layer_idx} MLP/MoE ================") + + # if layer_idx == 0: + # mlp = layer.mlp + # for proj_type in ['gate_proj', 'up_proj', 'down_proj']: + # module = getattr(mlp, proj_type).orig_module + # print(f"\n{proj_type}.weight (shape: {module.weight.shape})") + # print(module.weight.data[:3, :5].detach().cpu().numpy()) + # else: + # moe = layer.mlp.orig_module + # print("\n[Shared Experts]") + # for proj_type in ['gate_proj', 'up_proj', 'down_proj']: + # module = getattr(moe.shared_experts, proj_type).orig_module + # print(f"\nshared_{proj_type}.weight (shape: {module.weight.shape})") + # print(module.weight.data[:3, :5].detach().cpu().numpy()) + + # print("\n[Experts]") + # for expert_idx in range(3): + # expert = moe.experts.orig_module[expert_idx] + # print(f"\nExpert {expert_idx}:") + # for proj_type in ['gate_proj', 'up_proj', 'down_proj']: + # module = getattr(expert, proj_type) + # print(f"{proj_type}.weight (shape: {module.weight.shape})") + # print(module.weight.data[:3, :5].detach().cpu().numpy()) + +def print_lora_params(model): + # for layer_idx in range(len(model.model.orig_module.layers)): + for layer_idx in range(0, 3): + layer = model.base_model.model.model.orig_module.layers[layer_idx] + # layer = model.model.orig_module.layers[layer_idx] + + q_proj_module = layer.self_attn.orig_module.q_proj.orig_module + + linear_weight = q_proj_module.generate_linear.weight + lora_A_weight = q_proj_module.lora_A["default"].weight + lora_B_weight = q_proj_module.lora_B["default"].weight + + print(f"\n=================== Layer {layer_idx} ===================") + + print("\nOriginal Linear (first row slice):") + print(linear_weight.cpu()) + + print("\nLora_A (first row slice):") + print(lora_A_weight.cpu()) + + print("\nLora_B (first row slice):") + print(lora_B_weight.cpu()) + +def print_grad_fn(grad_fn, indent=0): + """递归打印计算图节点""" + if grad_fn is None: + return + print(' ' * indent, f"Node: {str(grad_fn).split('(')[0]}") + print(' ' * indent, f" Metadata: {grad_fn.metadata}") + for child in getattr(grad_fn, 'next_functions', []): + if child[0] is not None: + print_grad_fn(child[0], indent + 2) + +def forward_hook(module, inputs, output): + if isinstance(output, (tuple, list)): + for i, o in enumerate(output): + if o is None: + print(f"{module.__class__.__name__} output index {i} is None") + else: + print(f"{module.__class__.__name__} output index {i}: requires_grad={o.requires_grad}, grad_fn={o.grad_fn}") + elif output is None: + print(f"{module.__class__.__name__} returned None") + else: + print(f"{module.__class__.__name__}: requires_grad={output.requires_grad}, grad_fn={output.grad_fn}") + +def check_moe_gradients(model): + moe_layer = model.base_model.model.model.orig_module.layers[1].mlp.orig_module + for name, param in moe_layer.named_parameters(): + if param.requires_grad and param.grad is not None: + grad_norm = torch.norm(param.grad) + print(f"MoE参数 {name} 梯度范数: {grad_norm}") + else: + print(f"MoE参数 {name} 无梯度") + +def disable_all_dropout(module): + for name, child in module.named_children(): + if isinstance(child, nn.Dropout): + child.p = 0 + child.inplace = False + disable_all_dropout(child) + +def verify_lora_layers(model): + for layer_path in target_layers: + module = model.get_submodule(layer_path) + orig_module = module.orig_module + + W = orig_module.weight.data # [576, 2048] -> [2048, 576] + lora_A = module.lora_A['default'].weight.data # [8, 2048] + lora_B = module.lora_B['default'].weight.data # [576, 8] + alpha_over_r = 32/8 # alpha=32, r=8 + + input_tensor = layer_data[layer_path]['input'] # [1, 512, 2048] + + try: + original_output = torch.matmul(input_tensor, W) # [1,512,2048] @ [2048,576] => [1,512,576] + except: + original_output = torch.matmul(input_tensor, W.T) # [1,512,2048] @ [2048,576] => [1,512,576] + + lora_effect = torch.matmul( + torch.matmul(input_tensor, lora_A.T), # [1,512,2048] @ [2048,8] => [1,512,8] + lora_B.T # [1,512,8] @ [8,576] => [1,512,576] + ) * alpha_over_r + + manual_output = original_output + lora_effect # [1,512,576] + + model_output = layer_data[layer_path]['output'] + + print(f"manual_output:{manual_output}") + print(f"model_output:{model_output}") + + if torch.allclose(manual_output, model_output, atol=1e-5): + print(f"{layer_path} 验证通过") + else: + print(f"{layer_path} 验证失败！最大误差：{torch.max(torch.abs(manual_output - model_output))}") + +def print_moe_stats(moe_layer: KExpertsTorch): + print(f"Total Params: {moe_layer.total_params/1e6:.2f}M") + + total_time = sum(moe_layer.times) + gflops = (moe_layer.total_flops / 1e9) / total_time if total_time !=0 else 0 + + print(f"Total Calls: {moe_layer.call_count}") + # print(f"Avg GFLOPS per Call: {gflops/moe_layer.call_count:.2f}") + print(f"Overall GFLOPS: {gflops:.2f}") + + if moe_layer.call_count > 0: + last_flops = moe_layer.flops_per_call[-1] + last_time = moe_layer.times[-1] + print(f"\nLast Call - FLOPs: {last_flops/1e9:.2f}G Time: {last_time*1000:.2f}ms " + f"GFLOPS: {(last_flops/1e9)/last_time:.2f}") + +def recursive_traverse(model, parent_name=''): + """ + 递归遍历模型，查找MoE层并调用print_moe_stats。 + """ + for name, module in model.named_children(): + full_name = f"{parent_name}.{name}" if parent_name else name + + if isinstance(module, KTransformersExperts): + print(f"Found MoE layer: {full_name}") + print_moe_stats(module.generate_experts) + + recursive_traverse(module, full_name) + +def log_step_state( + step: int, + inputs: dict, + loss: torch.Tensor, + model: nn.Module, + log_dir: str = "train_logs", +): + """ + 把当前 step 的输入 / loss / grad / param 保存到 log_dir/step_{step}.pt + """ + Path(log_dir).mkdir(parents=True, exist_ok=True) + + logged_inputs = { + k: v.detach().cpu() + for k, v in inputs.items() + if isinstance(v, torch.Tensor) + } + + loss_val = loss.detach().cpu() + + params, grads = {}, {} + for name, p in model.named_parameters(): + params[name] = p.detach().cpu() + grads[name] = p.grad.detach().cpu() if p.grad is not None else None + + torch.save( + { + "step": step, + "inputs": logged_inputs, + "loss": loss_val, + "params": params, + "grads": grads, + }, + f"{log_dir}/step_{step:08d}.pt", + ) + +def collect_gradients(model, input_ids): + torch.manual_seed(42) + + output = model(input_ids=input_ids) + + logits = output.logits + loss = logits.mean() + + model.zero_grad() + loss.backward() + + grads = [] + for name, param in model.named_parameters(): + if param.requires_grad and param.grad is not None: + grads.append(f"{name}: {param.grad.norm().item():.6f}") + + return grads + +def report_meta_tensors(model): + import torch, inspect + meta_modules = [] + for mod_name, mod in model.named_modules(): + metas = [] + for n, p in list(mod.named_parameters(recurse=False)): + if getattr(p, "is_meta", False) and p.is_meta: + metas.append(("param", n, tuple(p.shape))) + for n, b in list(mod.named_buffers(recurse=False)): + if getattr(b, "is_meta", False) and b.is_meta: + metas.append(("buffer", n, tuple(b.shape))) + if metas: + print(f"[META] {mod_name} ({type(mod).__name__}): {metas}") + meta_modules.append((mod_name, type(mod).__name__, metas)) + return meta_modules + +# def lora_and_load_adapter(model, tokenizer, sft_data_path, save_adapter_path, is_profiler=False): + # show some lora test + + ''' + # multi-gpu dataloader test + # _ = report_meta_tensors(model) + + # print("=== SAMPLE INSPECT ===") + # for i in range(2): + # summary = {} + # for k,v in ex.items(): + # if isinstance(v, list): + # if len(v)>0 and isinstance(v[0], list): + # summary[k] = f"list-of-lists len={len(v)} x len0={len(v[0])}" + # else: + # summary[k] = f"list len={len(v)}" + # elif torch.is_tensor(v): + # summary[k] = f"tensor shape={tuple(v.shape)}" + # else: + # summary[k] = str(type(v)) + # print(f"[SAMPLE {i}]", summary) + + # trainer.accelerator = Accelerator(device_placement=False) + # first_batch = next(iter(trainer.get_train_dataloader())) + # print("Batch keys:", list(first_batch.keys())) + + # acc = KAccelerator(device_placement=False) + # acc.state.device_ids = [0] + # acc.state.num_processes = 1 + # acc.state.num_gpus = 1 + # trainer.accelerator = acc + + # print("Accelerator device_ids:", trainer.accelerator.state.device_ids) + # print(f"type(trainer.model):{type(trainer.model)}") + # print(f"type(trainer.accelerator):{type(trainer.accelerator)}") + + + # print("-------------------------START TRAINING!!!-------------------------") + + # cfg = getattr(trainer.accelerator, "dataloader_config", None) + # print( + # "[ACCEL DL CONFIG]", + # "split_batches=", getattr(cfg, "split_batches", None), + # "dispatch_batches=", getattr(cfg, "dispatch_batches", None), + # "even_batches=", getattr(cfg, "even_batches", None), + # "use_seedable_sampler=", getattr(cfg, "use_seedable_sampler", None), + # "non_blocking=", getattr(cfg, "non_blocking", None), + # ) + # print("--------------------NEW DEBUG--------------------") + # install_shape_probes(trainer.model) # print some debug info about multi-gpu placement. + + # input_ids = torch.randint(0, 1000, (32, 128), device="cuda:0") + # gradients = collect_gradients(model, input_ids) + ''' + + # with open(f"/home/lpl/kt-sft/tmp/KSFTExpertsCPU_grads.txt", "w") as f: + # f.write("\n".join(gradients)) + # print(xx) + + # total_length = 0 + # valid_count = 0 + # for batch in tqdm(train_dataloader): + # input_ids = batch['input_ids'] + # # print(f"Token count per sample: {[len(ids) for ids in input_ids]}") + # for ids in input_ids: + # if not torch.equal(ids, torch.tensor([100001])): + # total_length += len(ids) + # valid_count += 1 + # # print(f"Input tensor: {input_ids}") + # # print(f"total_length:{total_length}") + # # break + + # if valid_count > 0: + # average_length = total_length / valid_count + # else: + + # print(xx) + + # from ktransformers.sft.flops_utils.custom_profile import custom_profile + + # for module in model.modules(): + # if not hasattr(module, 'total_ops'): + # module.register_buffer('total_ops', torch.zeros(1, dtype=torch.float64)) + # if not hasattr(module, 'total_params'): + # module.register_buffer('total_params', torch.zeros(1, dtype=torch.float64)) + + # # print(f"input:{input}") + # for inputs in tqdm(train_dataloader): + # # input_ids = batch['input_ids'] + # # del inputs['instruction'] + # # del inputs['input'] + # # del inputs['output'] + # # output = model(**inputs) + # model.eval() + # content = inputs['instruction'][0] + inputs['input'][0] + # # flops,params = custom_profile(model, inputs=inputs, content=content, tokenizer=tokenizer, custom_ops={YourModule: count_your_model}) + # # print('FLOPs = ' + str(flops / 1000 ** 3) + 'G') + # # print('Params = ' + str(params / 1000 ** 2) + 'M') + + # messages = [{"role": "user", "content": content}] + # input_tensor = tokenizer.apply_chat_template( + # messages, add_generation_prompt=True, return_tensors="pt" + # ) + # with torch.no_grad(): + # # model(*inputs) + # # model.model to deal with the PeftModelForCaualLM temp + # prefill_and_generate( + # model.model, tokenizer, input_tensor.cuda(), max_new_tokens=1000, use_cuda_graph=False, mode = 'normal', force_think = False, chunk_prefill_size = 8192, + # ) + # recursive_traverse(model) + + # output = model(input_ids=torch.tensor([[1,2,3]], dtype=torch.int32, device="cuda:0")) + # loss = output.logits.mean() + + # dot = make_dot(loss, params=dict(model.named_parameters())) + # dot.render("KT_compute_cpuinfer_moe_model_graph", format="svg") + + # with open("tmp/output_loss_KCPU.txt", "w") as file: + # file.write("Output (logits):\n") + # file.write("\n\nLoss:\n") + + # disable_all_dropout(model) + + # def print_dropout_status(module, prefix=""): + # for name, child in module.named_children(): + # if isinstance(child, nn.Dropout): + # print(f"{prefix}{name}: p={child.p}, training={child.training}") + # print_dropout_status(child, prefix + name + ".") + + # print_dropout_status(model) + + # for layer_path in target_layers: + # module = model.get_submodule(layer_path) + # hook = module.register_forward_hook( + # lambda m, i, o, ln=layer_path: record_layer_io(m, i, o, ln) + # ) + # hooks.append(hook) + + + # if is_profiler: + # profiler = profile( + # activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA], + # schedule=torch.profiler.schedule( + # ), + # on_trace_ready=torch.profiler.tensorboard_trace_handler('./logs'), + # record_shapes=False, + # with_stack=False + # ) + + # # profiler_args = { + # # "schedule": torch.profiler.schedule( + # # ) + # # } + + # trainer = KTrainer( + # model=model, + # train_dataset=train_dataset, + # data_collator=DataCollatorForSeq2Seq( + # tokenizer, pad_to_multiple_of=8, return_tensors="pt", padding=True + # ), + # callbacks=[ProfilerCallback(profiler)] + # ) + + # with profiler: + # trainer.train() + + # print("Training finished. Exporting profiler data...") + # with open("profiler_output.txt", "w") as f: + # f.write(profiler.key_averages().table(sort_by="cuda_time_total", row_limit=20)) + + # profiler.export_chrome_trace("trace.json") + + + + # verify_lora_layers(model) + + # model.save_pretrained(save_adapter_path) + + ''' + ----------------------- START: Lora Test ----------------------- + + + # for name, module in model.named_modules(): + # if "q_proj" in name or "kv_a_proj" in name or "o_proj" in name: + # print(name) + + # print_model_params(model) + + # model = KTransformersLinearLora() + + # inspect_device(model, '/home/yj/ktransformers/device1.txt') + # with open('/home/yj/ktransformers/device1.txt', 'a') as file: + # file.write(f"Base model device: {model.base_model.device}\n") + # file.write(f"LoRA adapter device: {model.lora_config['target_modules'].device}\n") + # print(f"Base model device: {model.base_model.device}") + # print(f"LoRA adapter device: {model.lora_config['target_modules'].device}") + + + # model = model.to('cuda') + + # for name, module in model.named_modules(): + # module.register_forward_hook(forward_hook) + + # for name, parms in model.named_parameters(): + # # parms.requires_grad = True + # print('-->name:', name) + # print('-->para:', parms) + # print('-->grad_requirs:',parms.requires_grad) + # print('-->grad_fn:',parms.grad_fn) + # print('-->grad_value:',parms.grad) + # print("===") + + # output = model(input_ids=torch.tensor([[1,2,3]], dtype=torch.int32, device="cuda:0")) + # loss = output.logits.mean() + + # dot = make_dot(loss, params=dict(model.named_parameters())) + # dot.render("KT_compute_graph", format="svg") + + # inspect_device(model, '/home/yj/ktransformers/device2.txt') + # with open('/home/yj/ktransformers/device2.txt', 'a') as file: + # file.write(f"Base model device: {model.base_model.device}\n") + # file.write(f"LoRA adapter device: {model.lora_config['target_modules'].device}\n") + # print(f"Base model device: {model.base_model.device}") + # print(f"LoRA adapter device: {model.lora_config['target_modules'].device}") + + # print_lora_params(model) + + # trainer = KTrainer( + # model=model, + # train_dataset=train_dataset, + # args=transformers.TrainingArguments( + # output_dir=save_adapter_path, + # per_device_train_batch_size=1, + # gradient_accumulation_steps=16, + # num_train_epochs=10, + # learning_rate=3e-4, + # fp16=False, + # logging_steps=10, + # save_steps=200, + # dataloader_drop_last=True, + # ddp_find_unused_parameters=False + # ), + # data_collator=DataCollatorForSeq2Seq( + # tokenizer, pad_to_multiple_of=8, return_tensors="pt", padding=True + # ), + # ) + + # model(input_ids=torch.tensor([[1,2,3]], dtype=torch.int32, device="cuda:0")) + + # trainer.train() + + # print_lora_params(model) + + # model = model.merge_and_unload() + ----------------------- END: Lora Test ----------------------- + + ''' \ No newline at end of file diff --git a/archive/kt-sft/ktransformers/sft/lora.py b/archive/kt-sft/ktransformers/sft/lora.py new file mode 100644 index 0000000..e862dba --- /dev/null +++ b/archive/kt-sft/ktransformers/sft/lora.py @@ -0,0 +1,497 @@ +from transformers import AutoTokenizer, DataCollatorForLanguageModeling +from transformers import Trainer, TrainingArguments +from transformers import Trainer +from transformers.training_args import OptimizerNames +from transformers.trainer_utils import seed_worker +from transformers.utils import ( + is_datasets_available, + is_sagemaker_mp_enabled, + is_torch_xpu_available, + is_torch_mlu_available, + is_torch_musa_available, + is_torch_npu_available, + is_torch_mps_available, + is_torch_hpu_available, + is_accelerate_available, + is_apex_available, + logging, +) +from packaging import version +import os +import inspect +import functools +from typing import Union, Any, Dict, List + +import torch +from torch.utils.data import DataLoader +import torch.nn as nn +from torch.utils.data import DataLoader, IterableDataset +from torch.utils.data import Dataset as TorchDataset + +from peft import LoraConfig, TaskType +from datasets import Dataset +from torchviz import make_dot +from tqdm import tqdm +import os, json +from pathlib import Path +from accelerate import Accelerator +if is_accelerate_available("0.28.0"): + from accelerate.utils import DataLoaderConfiguration +from accelerate import __version__ as accelerate_version +if version.parse(accelerate_version) > version.parse("1.3.0"): + from accelerate.utils import TorchTensorParallelPlugin +if is_sagemaker_mp_enabled(): + from transformers.trainer_utils import smp_forward_backward + +from ktransformers.sft.peft_utils.mapping import get_peft_model + +logger = logging.get_logger(__name__) + +class KAccelerator(Accelerator): + def __init__(self, *args, **kwargs): + kwargs.setdefault("device_placement", False) + super().__init__(*args, **kwargs) + + def prepare_model(self, model, *args, **kwargs): + return model + + def prepare(self, *args, **kwargs): + prepped = [] + for obj in args: + if isinstance(obj, nn.Module): + prepped.append(self.prepare_model(obj, **kwargs)) + else: + prepped.append(super().prepare(obj, **kwargs)) + return tuple(prepped) if len(prepped) > 1 else prepped[0] + +class KTrainer(Trainer): + def save_model(self, output_dir=None, _internal_call=False): + output_dir = output_dir or self.args.output_dir + os.makedirs(output_dir, exist_ok=True) + # only save LoRA adapter, including adapter_config.json + self.model.save_pretrained(output_dir) + + def _move_model_to_device(self, model, device): + print("[KTrainer] Due to the placement feature in KTransformers, skip moving model to", device) + return model + + def _wrap_model(self, model, training=True, dataloader=None): + self.model_wrapped = model + return model + + def create_accelerator_and_postprocess(self): + # We explicitly don't rely on the `Accelerator` to do gradient accumulation + grad_acc_kwargs = {} + if is_accelerate_available("0.28.0") and self.args.accelerator_config.gradient_accumulation_kwargs is not None: + grad_acc_kwargs = self.args.accelerator_config.gradient_accumulation_kwargs + + # check if num_steps is attempted to be passed in gradient_accumulation_kwargs + if "num_steps" in grad_acc_kwargs: + if self.args.gradient_accumulation_steps > 1: + # raise because we do not know which setting is intended. + raise ValueError( + "The `AcceleratorConfig`'s `num_steps` is set but `gradient_accumulation_steps` is greater than 1 in the passed `TrainingArguments`" + "If using the passed `AcceleratorConfig` is desired, do not set the `TrainingArguments` `gradient_accumulation_steps`." + ) + else: + self.args.gradient_accumulation_steps = grad_acc_kwargs["num_steps"] + + accelerator_config = self.args.accelerator_config.to_dict() + + if is_accelerate_available("0.28.0"): + # Extract dataloader config params from accelerator config + dataloader_params = ["split_batches", "dispatch_batches", "even_batches", "use_seedable_sampler"] + dataloader_config_dict = {param: accelerator_config.pop(param) for param in dataloader_params if param in accelerator_config} + if DataLoaderConfiguration is None: + raise ImportError("Your accelerate does not provide DataLoaderConfiguration but Trainer expects it.") + dataloader_config = DataLoaderConfiguration(**dataloader_config_dict) + if is_accelerate_available("1.1.0"): + dataloader_config.data_seed = self.args.data_seed + else: + dataloader_config = None + + non_blocking = accelerator_config.pop("non_blocking", False) + if not is_accelerate_available("0.30.0"): + if non_blocking: + raise ImportError( + "`non_blocking` is only supported in accelerate v0.30.0 and above. Please upgrade accelerate to use this feature." + ) + else: + if non_blocking and not self.args.dataloader_pin_memory: + logger.warning("`non_blocking` is enabled but `dataloader_pin_memory` is not. For best performance, enable both.") + if dataloader_config is not None: + dataloader_config.non_blocking = non_blocking + + accelerator_config.pop("gradient_accumulation_kwargs", None) + + args = { + "deepspeed_plugin": self.args.deepspeed_plugin, + "device_placement": False, + } + + if is_accelerate_available("0.28.0"): + args["dataloader_config"] = dataloader_config + else: + args.update(accelerator_config) + + if getattr(self.args, "tp_size", 1) > 1: + self.is_tp_enabled = True + if version.parse(accelerate_version) > version.parse("1.3.0") and TorchTensorParallelPlugin is not None: + args["torch_tp_plugin"] = TorchTensorParallelPlugin(tp_size=self.args.tp_size) + else: + raise ValueError("Requires accelerate>1.3.0 to use Tensor Parallelism.") + + self.accelerator = KAccelerator(**args) + + try: + self.accelerator.state.device_ids = [0] + self.accelerator.state.num_processes = 1 + self.accelerator.state.num_gpus = 1 + except Exception: + pass + + # some Trainer classes need to use `gather` instead of `gather_for_metrics`, thus we store a flag + self.gather_function = self.accelerator.gather_for_metrics + + if "use_gather_object" in inspect.signature(self.gather_function).parameters.keys(): + self.gather_function = functools.partial( + self.gather_function, use_gather_object=self.args.eval_use_gather_object + ) + + # deepspeed and accelerate flags covering both trainer args and accelerate launcher + self.is_deepspeed_enabled = getattr(self.accelerator.state, "deepspeed_plugin", None) is not None + self.is_fsdp_enabled = getattr(self.accelerator.state, "fsdp_plugin", None) is not None + self.is_tp_enabled = getattr(self.accelerator.state, "torch_tp_plugin", None) is not None + # post accelerator creation setup + if self.is_fsdp_enabled: + fsdp_plugin = self.accelerator.state.fsdp_plugin + for param in ["limit_all_gathers", "activation_checkpointing"]: + setattr(fsdp_plugin, param, self.args.fsdp_config.get(param, getattr(fsdp_plugin, param))) + if fsdp_plugin.activation_checkpointing and self.args.gradient_checkpointing: + raise ValueError( + "The activation_checkpointing in FSDP config and the gradient_checkpointing in training arg " + "can't be set to True simultaneously. Please use FSDP's activation_checkpointing logic " + "when using FSDP." + ) + + if self.is_deepspeed_enabled and getattr(self.args, "hf_deepspeed_config", None) is None: + self.propagate_args_to_deepspeed() + + # `save_only_model` can't be used with DeepSpeed/FSDP along with `load_best_model_at_end` + if ( + self.args.save_only_model + and (self.is_deepspeed_enabled or self.is_fsdp_enabled) + and self.args.load_best_model_at_end + ): + wrapper = "DeepSpeed" if self.is_deepspeed_enabled else "FSDP" + raise ValueError(f"{wrapper} can't be used with `save_only_model` along with `load_best_model_at_end`.") + + # `auto_find_batch_size` isn't supported yet with DeepSpeed Zero-3 + if ( + self.is_deepspeed_enabled + and self.accelerator.state.deepspeed_plugin.zero_stage == 3 + and self.args.auto_find_batch_size + ): + raise ValueError( + "`auto_find_batch_size` isn't supported yet with DeepSpeed Zero-3. Please consider using Zero-2, Zero-1, or FSDP" + ) + if ( + self.args.save_only_model + and self.is_fsdp_enabled + and "SHARDED_STATE_DICT" in str(self.accelerator.state.fsdp_plugin.state_dict_type) + ): + raise ValueError("save_only_model option is not compatible with FSDP state dict type 'SHARDED_STATE_DICT'") + + if dataloader_config is not None: + dataloader_config.split_batches = False + dataloader_config.dispatch_batches = False + dataloader_config.even_batches = False + + def get_train_dataloader(self) -> DataLoader: + """ + Returns the training DataLoader with per_device_train_batch_size + (no implicit multipliers by number of visible GPUs). + """ + if self.train_dataset is None: + raise ValueError("Trainer: training requires a train_dataset.") + + train_dataset = self.train_dataset + data_collator = self.data_collator + + if is_datasets_available(): + try: + import datasets + if isinstance(train_dataset, datasets.Dataset): + train_dataset = self._remove_unused_columns(train_dataset, description="training") + else: + data_collator = self._get_collator_with_removed_columns(data_collator, description="training") + except Exception: + data_collator = self._get_collator_with_removed_columns(data_collator, description="training") + else: + data_collator = self._get_collator_with_removed_columns(data_collator, description="training") + + dataloader_params = { + "batch_size": self.args.per_device_train_batch_size, + "collate_fn": data_collator, + "num_workers": self.args.dataloader_num_workers, + "pin_memory": self.args.dataloader_pin_memory, + "persistent_workers": self.args.dataloader_persistent_workers, + } + + if not isinstance(train_dataset, IterableDataset): + dataloader_params["sampler"] = self._get_train_sampler() + dataloader_params["drop_last"] = self.args.dataloader_drop_last + dataloader_params["worker_init_fn"] = seed_worker + if self.args.dataloader_num_workers > 0 and self.args.dataloader_prefetch_factor is not None: + dataloader_params["prefetch_factor"] = self.args.dataloader_prefetch_factor + + dl = DataLoader(train_dataset, **dataloader_params) + + try: + prepared = self.accelerator.prepare(dl, device_placement=[False]) + except TypeError: + prepared = self.accelerator.prepare(dl) + + return prepared + + def training_step( + self, + model: torch.nn.Module, + inputs: dict[str, Union[torch.Tensor, Any]], + num_items_in_batch=None + ) -> torch.Tensor: + model.train() + if hasattr(self.optimizer, "train") and callable(self.optimizer.train): + self.optimizer.train() + + inputs = self._prepare_inputs(inputs) + + if is_sagemaker_mp_enabled(): + loss_mb = smp_forward_backward(model, inputs, self.args.gradient_accumulation_steps) + return loss_mb.reduce_mean().detach().to(self.args.device) + + with self.compute_loss_context_manager(): + loss = self.compute_loss(model, inputs, num_items_in_batch=num_items_in_batch) + + del inputs + + if ( + self.args.torch_empty_cache_steps is not None + and self.state.global_step % self.args.torch_empty_cache_steps == 0 + ): + if is_torch_xpu_available(): + torch.xpu.empty_cache() + elif is_torch_mlu_available(): + torch.mlu.empty_cache() + elif is_torch_musa_available(): + torch.musa.empty_cache() + elif is_torch_npu_available(): + torch.npu.empty_cache() + elif is_torch_mps_available(min_version="2.0"): + torch.mps.empty_cache() + elif is_torch_hpu_available(): + logger.warning( + "`torch_empty_cache_steps` is set but HPU device/backend does not support empty_cache()." + ) + else: + torch.cuda.empty_cache() + + kwargs = {} + + if self.args.optim in [OptimizerNames.LOMO, OptimizerNames.ADALOMO]: + kwargs["learning_rate"] = self._get_learning_rate() + + if self.args.n_gpu > 1: + loss = loss.mean() + + if self.use_apex: + with amp.scale_loss(loss, self.optimizer) as scaled_loss: # type: ignore + scaled_loss.backward() + else: + if not self.model_accepts_loss_kwargs and self.compute_loss_func is None: + loss = loss / self.args.gradient_accumulation_steps + + if getattr(self.accelerator, "distributed_type", None) and \ + str(self.accelerator.distributed_type) == "DistributedType.DEEPSPEED": + kwargs["scale_wrt_gas"] = False + + self.accelerator.backward(loss, **kwargs) + + ret = loss.detach() + if ret.device != self.args.device: + ret = ret.to(self.args.device, non_blocking=True) + + if os.environ.get("KT_DBG_STEP", "0") == "1" and not hasattr(self, "_kt_dbg_once"): + try: + print(f"[KT-DBG] args.device={self.args.device} loss(before)={loss.device} loss(return)={ret.device}") + except Exception: + pass + self._kt_dbg_once = True + + return ret + +class SFTJsonListDataset(TorchDataset): + def __init__(self, path: str, tokenizer: AutoTokenizer, max_len: int = 512): + super().__init__() + with open(path, "r", encoding="utf-8") as f: + self.samples: List[Dict] = json.load(f) + self.tok = tokenizer + self.max_len = max_len + + @staticmethod + def build_example(ins: str, inp: str, out: str) -> Dict[str, str]: + ins = (ins or "").strip() + inp = (inp or "").strip() + out = (out or "").strip() + prompt = (ins + inp) if ins else inp + return {"prompt": prompt, "response": out} + + def __len__(self): + return len(self.samples) + + def __getitem__(self, idx: int): + rec = self.samples[idx] + eg = self.build_example(rec.get("instruction", ""), rec.get("input", ""), rec.get("output", "")) + + prompt_ids = self.tok( + eg["prompt"], + max_length=self.max_len, + truncation=True, + add_special_tokens=False, + )["input_ids"] + + response_ids = self.tok( + eg["response"], + max_length=self.max_len, + truncation=True, + add_special_tokens=False, + )["input_ids"] + + eos_id = self.tok.eos_token_id + input_ids = prompt_ids + response_ids + ([eos_id] if eos_id is not None else []) + input_ids = input_ids[: self.max_len] + + labels = [-100] * min(len(prompt_ids), self.max_len) + tail = input_ids[len(labels):] + labels = labels + tail + labels = labels[: self.max_len] + + attention_mask = [1] * len(input_ids) + + return { + "input_ids": torch.tensor(input_ids, dtype=torch.long), + "labels": torch.tensor(labels, dtype=torch.long), + "attention_mask": torch.tensor(attention_mask, dtype=torch.long), + } + +def lora_and_load_adapter(model, tokenizer, sft_data_path, save_adapter_path): + + Path(save_adapter_path).mkdir(parents=True, exist_ok=True) + + lora_config = LoraConfig( + task_type=TaskType.CAUSAL_LM, + target_modules=[ + "q_proj", # FOR DeepSeek-V2-Lite + "q_a_proj", # FOR DeepSeek-V3&R1 + "q_b_proj", + "kv_a_proj_with_mqa", + "kv_b_proj", + "o_proj", + "mlp.gate_proj", + "mlp.up_proj", + "mlp.down_proj", + "shared_experts.gate_proj", + "shared_experts.up_proj", + "shared_experts.down_proj", + ], + r=8, + lora_alpha=32, + lora_dropout=0.1, + ) + model = get_peft_model(model, lora_config) + model.print_trainable_parameters() + + train_dataset = SFTJsonListDataset(sft_data_path, tokenizer, max_len=512) + data_collator = DataCollatorForLanguageModeling(tokenizer, mlm=False) + + training_args = TrainingArguments( + output_dir=save_adapter_path, + per_device_train_batch_size=1, + gradient_accumulation_steps=16, + num_train_epochs=1, + # max_steps=30, # TODO: FOR TEST, will override any value given in num_train_epochs + learning_rate=1e-4, + fp16=False, + logging_steps=10, + save_steps=200, + dataloader_drop_last=True, + ddp_find_unused_parameters=False, + ) + + debug_path = os.path.join(save_adapter_path, "model_infra_debug.json") + with open(debug_path, "w", encoding="utf-8") as f: + json.dump({"model": str(model)}, f, ensure_ascii=False, indent=2) + + # output = model(input_ids=torch.tensor([[1,2,3]], dtype=torch.int32, device="cuda:0")) + # loss = output.logits.mean() + + # dot = make_dot(loss, params=dict(model.named_parameters())) + # dot.render("KT_compute_cpuinfer_moe_model_graph", format="svg") + + trainer = KTrainer( + model=model, + tokenizer=tokenizer, + args=training_args, + train_dataset=train_dataset, + data_collator=data_collator, + ) + model.config.use_cache = False + # model.gradient_checkpointing_enable() + # if hasattr(model, "enable_input_require_grads"): + # model.enable_input_require_grads() + + trainer.train() + +def inject_lora_layer(model, use_adapter_path): + + cfg_path = os.path.join(use_adapter_path, "adapter_config.json") + with open(cfg_path, "r", encoding="utf-8") as f: + data = json.load(f) + + task_type_str = (data.get("task_type") or "CAUSAL_LM").upper() + bias = data.get("bias", "none") + if bias in (None, False): + bias = "none" + if data.get("lora_bias") is True and bias == "none": + bias = "lora_only" + + tmods = data.get("target_modules") + if isinstance(tmods, str): + tmods = [m.strip() for m in tmods.split(",") if m.strip()] + + mts = data.get("modules_to_save", None) + if isinstance(mts, str): + mts = [m.strip() for m in mts.split(",") if m.strip()] + + rank_pattern = data.get("rank_pattern") or None + alpha_pattern = data.get("alpha_pattern") or None + + lora_config = LoraConfig( + r=data.get("r", 8), + lora_alpha=data.get("lora_alpha", 32), + lora_dropout=float(data.get("lora_dropout", 0.0)), + bias=bias, + task_type=TaskType[task_type_str], + target_modules=tmods, + modules_to_save=mts, + init_lora_weights=bool(data.get("init_lora_weights", True)), + inference_mode=bool(data.get("inference_mode", True)), + use_rslora=bool(data.get("use_rslora", False)), + use_dora=bool(data.get("use_dora", False)), + ) + print(f"lora_config:{lora_config.__dict__}") + + # model = inject_adapter_in_model(lora_config, model) + model = get_peft_model(model, lora_config) + model.config.use_cache = False + model.eval() \ No newline at end of file diff --git a/archive/kt-sft/ktransformers/sft/metrics.py b/archive/kt-sft/ktransformers/sft/metrics.py new file mode 100644 index 0000000..e531700 --- /dev/null +++ b/archive/kt-sft/ktransformers/sft/metrics.py @@ -0,0 +1,121 @@ +# Copyright 2025 HuggingFace Inc., THUDM, and the LlamaFactory team. +# +# This code is inspired by the HuggingFace's transformers library and the THUDM's ChatGLM implementation. +# https://github.com/huggingface/transformers/blob/v4.40.0/examples/pytorch/summarization/run_summarization.py +# https://github.com/THUDM/ChatGLM-6B/blob/main/ptuning/main.py +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from dataclasses import dataclass +from typing import TYPE_CHECKING, Optional + +import numpy as np +import torch +from transformers.utils import is_jieba_available, is_nltk_available + +from ktransformers.sft.metrics_utils.constants import IGNORE_INDEX +from ktransformers.sft.metrics_utils.misc import numpify +from ktransformers.sft.metrics_utils.packages import is_rouge_available + + +if TYPE_CHECKING: + from transformers import EvalPrediction, PreTrainedTokenizer + + +if is_jieba_available(): + import jieba # type: ignore + + +if is_nltk_available(): + from nltk.translate.bleu_score import SmoothingFunction, sentence_bleu # type: ignore + + +if is_rouge_available(): + from rouge_chinese import Rouge # type: ignore + + +def eval_logit_processor(logits: "torch.Tensor", labels: "torch.Tensor") -> "torch.Tensor": + r"""Compute the token with the largest likelihood to reduce memory footprint.""" + if isinstance(logits, (list, tuple)): + if logits[0].dim() == 3: # (batch_size, seq_len, vocab_size) + logits = logits[0] + else: # moe models have aux loss + logits = logits[1] + + if logits.dim() != 3: + raise ValueError("Cannot process the logits.") + + return torch.argmax(logits, dim=-1) + +@dataclass +class ComputeSimilarity: + r"""Compute text similarity scores and support `batch_eval_metrics`. + + Wraps the tokenizer into metric functions, used in CustomSeq2SeqTrainer. + """ + + tokenizer: "PreTrainedTokenizer" + + def _dump(self) -> Optional[dict[str, float]]: + result = None + if hasattr(self, "score_dict"): + result = {k: float(np.mean(v)) for k, v in self.score_dict.items()} + + self.score_dict = {"rouge-1": [], "rouge-2": [], "rouge-l": [], + "bleu-1": [], "bleu-2": [], "bleu-3": [], "bleu-4": []} + return result + + def __post_init__(self): + self._dump() + + def __call__(self, eval_preds: "EvalPrediction", compute_result: bool = True) -> Optional[dict[str, float]]: + preds, labels = numpify(eval_preds.predictions), numpify(eval_preds.label_ids) + + preds = np.where(preds != IGNORE_INDEX, preds, self.tokenizer.pad_token_id) + labels = np.where(labels != IGNORE_INDEX, labels, self.tokenizer.pad_token_id) + + decoded_preds = self.tokenizer.batch_decode(preds, skip_special_tokens=True) + decoded_labels = self.tokenizer.batch_decode(labels, skip_special_tokens=True) + + for pred, label in zip(decoded_preds, decoded_labels): + hypothesis = list(jieba.cut(pred)) + reference = list(jieba.cut(label)) + + if len(" ".join(hypothesis).split()) == 0 or len(" ".join(reference).split()) == 0: + result = {"rouge-1": {"f": 0.0}, "rouge-2": {"f": 0.0}, "rouge-l": {"f": 0.0}} + else: + rouge = Rouge() + scores = rouge.get_scores(" ".join(hypothesis), " ".join(reference)) + result = scores[0] + + refs = [reference] + hyp = hypothesis + smooth = SmoothingFunction().method3 + bleu1 = sentence_bleu(refs, hyp, weights=(1.0, 0.0, 0.0, 0.0), smoothing_function=smooth) + bleu2 = sentence_bleu(refs, hyp, weights=(0.5, 0.5, 0.0, 0.0), smoothing_function=smooth) + bleu3 = sentence_bleu(refs, hyp, weights=(1/3, 1/3, 1/3, 0.0), smoothing_function=smooth) + bleu4 = sentence_bleu(refs, hyp, weights=(0.25, 0.25, 0.25, 0.25), smoothing_function=smooth) + + for k, v in result.items(): + self.score_dict[k].append(round(v["f"] * 100, 4)) + + # bleu_score = sentence_bleu([list(label)], list(pred), smoothing_function=SmoothingFunction().method3) + # self.score_dict["bleu-4"].append(round(bleu_score * 100, 4)) + + self.score_dict["bleu-1"].append(round(bleu1 * 100, 4)) + self.score_dict["bleu-2"].append(round(bleu2 * 100, 4)) + self.score_dict["bleu-3"].append(round(bleu3 * 100, 4)) + self.score_dict["bleu-4"].append(round(bleu4 * 100, 4)) + + if compute_result: + return self._dump() \ No newline at end of file diff --git a/archive/kt-sft/ktransformers/sft/metrics_utils/__init__.py b/archive/kt-sft/ktransformers/sft/metrics_utils/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/archive/kt-sft/ktransformers/sft/metrics_utils/constants.py b/archive/kt-sft/ktransformers/sft/metrics_utils/constants.py new file mode 100644 index 0000000..f252a3a --- /dev/null +++ b/archive/kt-sft/ktransformers/sft/metrics_utils/constants.py @@ -0,0 +1,3464 @@ +# Copyright 2025 the LlamaFactory team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import os +from collections import OrderedDict, defaultdict +from enum import Enum, unique +from typing import Optional + +from peft.utils import SAFETENSORS_WEIGHTS_NAME as SAFE_ADAPTER_WEIGHTS_NAME +from peft.utils import WEIGHTS_NAME as ADAPTER_WEIGHTS_NAME +from transformers.utils import SAFE_WEIGHTS_INDEX_NAME, SAFE_WEIGHTS_NAME, WEIGHTS_INDEX_NAME, WEIGHTS_NAME + + +AUDIO_PLACEHOLDER = os.getenv("AUDIO_PLACEHOLDER", "