diff --git a/.gemini/config.yaml b/.gemini/config.yaml new file mode 100644 index 00000000000..2499d3f0951 --- /dev/null +++ b/.gemini/config.yaml @@ -0,0 +1,6 @@ +# https://developers.google.com/gemini-code-assist/docs/customize-gemini-behavior-github +have_fun: false # Just review the code +code_review: + comment_severity_threshold: HIGH # Reduce quantity of comments + pull_request_opened: + summary: false # Don't summarize the PR in a separate comment diff --git a/.github/workflows/accuracy_test.yaml b/.github/workflows/accuracy_test.yaml index 7140f262f75..044c5dcfd00 100644 --- a/.github/workflows/accuracy_test.yaml +++ b/.github/workflows/accuracy_test.yaml @@ -70,6 +70,8 @@ jobs: runner: linux-aarch64-a2-1 - model_name: Qwen3-30B-A3B runner: linux-aarch64-a2-2 + - model_name: DeepSeek-V2-Lite + runner: linux-aarch64-a2-2 fail-fast: false name: ${{ matrix.model_name }} accuracy @@ -200,9 +202,8 @@ jobs: markdown_name="${model_base_name}" echo "markdown_name=$markdown_name" >> $GITHUB_OUTPUT mkdir -p ./benchmarks/accuracy - pytest -sv ./tests/e2e/singlecard/models/test_lm_eval_correctness.py \ - --config ./tests/e2e/singlecard/models/configs/${{ matrix.model_name }}.yaml \ - --report_output ./benchmarks/accuracy/${model_base_name}.md + pytest -sv ./tests/e2e/models/test_lm_eval_correctness.py \ + --config ./tests/e2e/models/configs/${{ matrix.model_name }}.yaml - name: Generate step summary if: ${{ always() }} @@ -225,14 +226,14 @@ jobs: outputs: model_name: ${{ steps.set_output.outputs.model_name }} - + vllm_ascend_version: ${{ env.GHA_VLLM_ASCEND_VERSION }} + create_pr: runs-on: ubuntu-latest needs: accuracy_tests if: ${{ github.event_name == 'workflow_dispatch' && github.event.inputs.vllm-ascend-version == 'latest' }} env: UPSTREAM_REPO: vllm-project/vllm-ascend - steps: - name: Checkout repository uses: actions/checkout@v4 @@ -257,10 +258,10 @@ jobs: TIMESTAMP=$(date +%Y%m%d%H%M%S) BRANCH_NAME="auto-pr/accuracy-report-${TIMESTAMP}" echo "BRANCH_NAME=${BRANCH_NAME}" >> $GITHUB_ENV - git checkout -B "${BRANCH_NAME}" upstream/${{ github.event.inputs.vllm-ascend-version }} + git checkout -B "${BRANCH_NAME}" upstream/main - name: Download only current run reports - uses: actions/download-artifact@v4 + uses: actions/download-artifact@v5 with: path: ./docs/source/developer_guide/evaluation/accuracy_report pattern: report-* @@ -298,7 +299,7 @@ jobs: GITHUB_TOKEN: ${{ secrets.PAT_TOKEN }} run: | git add ./docs/source/developer_guide/evaluation/accuracy_report/*.md - git commit -s -m "[Doc] Update accuracy reports for ${{ github.event.inputs.vllm-ascend-version }}" + git commit -s -m "[Doc] Update accuracy reports for ${{ needs.accuracy_tests.outputs.vllm_ascend_version }}" git push -f origin "${{ env.BRANCH_NAME }}" - name: Create PR in upstream via API @@ -310,9 +311,9 @@ jobs: owner: 'vllm-project', repo: 'vllm-ascend', head: `vllm-ascend-ci:${{ env.BRANCH_NAME }}`, - base: '${{ github.event.inputs.vllm-ascend-version }}', - title: `[Doc] Update accuracy reports for ${{ github.event.inputs.vllm-ascend-version }}`, - body: `The accuracy results running on NPU Altlas A2 have changed, updating reports for: All models (Qwen/Qwen3-30B-A3B, Qwen2.5-VL-7B-Instruct, Qwen3-8B-Base) + base: 'main', + title: `[Doc] Update accuracy reports for ${{ needs.accuracy_tests.outputs.vllm_ascend_version }}`, + body: `The accuracy results running on NPU Altlas A2 have changed, updating reports for: All models (Qwen3-30B-A3B, Qwen2.5-VL-7B-Instruct, Qwen3-8B-Base, DeepSeek-V2-Lite) - [Workflow run][1] diff --git a/.github/workflows/vllm_ascend_test.yaml b/.github/workflows/vllm_ascend_test.yaml index d46b4a9fd47..0c0deed9a07 100644 --- a/.github/workflows/vllm_ascend_test.yaml +++ b/.github/workflows/vllm_ascend_test.yaml @@ -185,6 +185,9 @@ jobs: run: | pip install -r requirements-dev.txt pip install -v -e . + if [[ "${{ matrix.vllm_version }}" == "v0.10.0" ]]; then + pip install "transformers<4.54.0" + fi - name: Run e2e test env: @@ -211,8 +214,7 @@ jobs: --ignore=tests/e2e/singlecard/test_embedding.py \ --ignore=tests/e2e/singlecard/spec_decode_v1/test_v1_mtp_correctness.py \ --ignore=tests/e2e/singlecard/spec_decode_v1/test_v1_spec_decode.py \ - --ignore=tests/e2e/singlecard/test_offline_inference_310p.py \ - --ignore=tests/e2e/singlecard/models/test_lm_eval_correctness.py + --ignore=tests/e2e/singlecard/test_offline_inference_310p.py e2e-2-cards: needs: [e2e] if: ${{ needs.e2e.result == 'success' }} @@ -268,6 +270,9 @@ jobs: run: | pip install -r requirements-dev.txt pip install -v -e . + if [[ "${{ matrix.vllm_version }}" == "v0.10.0" ]]; then + pip install "transformers<4.54.0" + fi - name: Run vllm-project/vllm-ascend test env: diff --git a/.github/workflows/vllm_ascend_test_long_term.yaml b/.github/workflows/vllm_ascend_test_long_term.yaml deleted file mode 100644 index 0dfa7e30944..00000000000 --- a/.github/workflows/vllm_ascend_test_long_term.yaml +++ /dev/null @@ -1,102 +0,0 @@ -# -# Copyright (c) 2025 Huawei Technologies Co., Ltd. All Rights Reserved. -# This file is a part of the vllm-ascend project. -# -# 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. -# -name: 'e2e test / long-term-test' - -on: - schedule: - # Runs at 23:00 UTC (7:00 AM Beijing) every day - - cron: '0 23 * * *' - pull_request: - types: [ labeled ] - -# Bash shells do not use ~/.profile or ~/.bashrc so these shells need to be explicitly -# declared as "shell: bash -el {0}" on steps that need to be properly activated. -# It's used to activate ascend-toolkit environment variables. -defaults: - run: - shell: bash -el {0} - -# only cancel in-progress runs of the same workflow -concurrency: - group: ${{ github.workflow }}-${{ github.ref }} - cancel-in-progress: true - -jobs: - long-term-test: - # long-term-test will be triggered when tag 'long-term-test' & 'ready-for-test' or schedule job - if: ${{ contains(github.event.pull_request.labels.*.name, 'long-term-test') && contains(github.event.pull_request.labels.*.name, 'ready-for-test') || github.event_name == 'schedule' }} - strategy: - max-parallel: 2 - matrix: - os: [linux-aarch64-a2-1, linux-aarch64-a2-2] - vllm_version: [main, v0.10.0] - name: vLLM Ascend long term test - runs-on: ${{ matrix.os }} - container: - image: swr.cn-southwest-2.myhuaweicloud.com/base_image/ascend-ci/cann:8.2.rc1-910b-ubuntu22.04-py3.11 - env: - VLLM_LOGGING_LEVEL: ERROR - VLLM_USE_MODELSCOPE: True - steps: - - name: Check npu and CANN info - run: | - npu-smi info - cat /usr/local/Ascend/ascend-toolkit/latest/"$(uname -i)"-linux/ascend_toolkit_install.info - - - name: Config mirrors - run: | - sed -Ei 's@(ports|archive).ubuntu.com@cache-service.nginx-pypi-cache.svc.cluster.local:8081@g' /etc/apt/sources.list - pip config set global.index-url http://cache-service.nginx-pypi-cache.svc.cluster.local/pypi/simple - pip config set global.trusted-host cache-service.nginx-pypi-cache.svc.cluster.local - apt-get update -y - apt install git -y - - - name: Checkout vllm-project/vllm-ascend repo - uses: actions/checkout@v4 - - - name: Install system dependencies - run: | - apt-get -y install `cat packages.txt` - apt-get -y install gcc g++ cmake libnuma-dev - - - name: Checkout vllm-project/vllm repo - uses: actions/checkout@v4 - with: - repository: vllm-project/vllm - ref: ${{ matrix.vllm_version }} - path: ./vllm-empty - - - name: Install vllm-project/vllm from source - working-directory: ./vllm-empty - run: | - VLLM_TARGET_DEVICE=empty pip install -e . - - - name: Install vllm-project/vllm-ascend - env: - PIP_EXTRA_INDEX_URL: https://mirrors.huaweicloud.com/ascend/repos/pypi - run: | - pip install -r requirements-dev.txt - pip install -v -e . - - - name: Run vllm-project/vllm-ascend long term test - run: | - if [[ "${{ matrix.os }}" == "linux-arm64-npu-1" ]]; then - pytest -sv tests/e2e/long_term/accuracy/accuracy_singlecard.py - else - # accuracy test multi card - pytest -sv tests/e2e/long_term/accuracy/accuracy_multicard.py - fi diff --git a/codecov.yml b/codecov.yml index 933ced8a444..3bf401b0e5d 100644 --- a/codecov.yml +++ b/codecov.yml @@ -17,12 +17,10 @@ coverage: status: - # non-voting, new code must be fully tested + # Patch coverage is mandatory and must be >= 80% patch: default: - target: 100% - # non-voting - informational: true + target: 80% # non-voting project: default: diff --git a/csrc/torch_binding.cpp b/csrc/torch_binding.cpp index f2a0d1f5de6..8bdc4b5606c 100644 --- a/csrc/torch_binding.cpp +++ b/csrc/torch_binding.cpp @@ -27,6 +27,17 @@ namespace vllm_ascend { +AscendType get_dtype_from_torch(at::ScalarType scalarType) +{ + if (scalarType == at::ScalarType::Float) { + return AscendType::FP32; + } else if (scalarType == at::ScalarType::BFloat16) { + return AscendType::BF16; + } else { + return AscendType::FP16; + } +} + std::tuple rotary_embedding(at::Tensor &positions, at::Tensor &query, at::Tensor &key, int64_t head_size, at::Tensor &cos_sin_cache, bool is_neox) { diff --git a/csrc/torch_binding_meta.cpp b/csrc/torch_binding_meta.cpp new file mode 100644 index 00000000000..1f9464c3ae7 --- /dev/null +++ b/csrc/torch_binding_meta.cpp @@ -0,0 +1,86 @@ +#include +#include +#include +#include +#include +#include +#include "utils.h" +/* + * How to write a meta implementation for a custom operator (meta kernel): + * + * Meta implementations are used for shape and dtype inference, tracing, and export. + * They do NOT perform any real computation or allocate device memory. + * Instead, they return empty tensors with the correct shapes, dtypes, and device types. + * + * Steps to write a meta implementation: + * 1. The function signature should match the operator's schema, but only use the arguments + * necessary to infer output shapes and dtypes. + * 2. Use input tensor shapes, dtypes, and any relevant arguments to compute the output shapes. + * 3. Return empty tensors (e.g., at::empty_symint, at::empty_like) with the correct shape and dtype. + * 4. Do NOT perform any real computation or data movement. + * 5. Register the meta implementation with the "Meta" dispatch key using TORCH_LIBRARY_IMPL or similar. + * + * Example: + * std::tuple my_op_meta( + * at::Tensor &input, int64_t some_param) { + * // Infer output shape based on input and parameters + * auto out_shape = ...; + * at::Tensor out = at::empty_symint(out_shape, input.options()); + * // Return empty tensor(s) with correct shape/dtype + * return {out, ...}; + * } + * + * See below for real examples. + */ + +namespace vllm_ascend { +namespace meta { + +std::tuple rotary_embedding_meta( + at::Tensor &positions, + at::Tensor &query, + at::Tensor &key, + int64_t head_size, + at::Tensor &cos_sin_cache, + bool is_neox) { + auto num_tokens = positions.sym_numel(); + auto query_hidden_size = query.sym_numel() / num_tokens; + auto key_hidden_size = key.sym_numel() / num_tokens; + + auto num_heads = query_hidden_size / head_size; + auto num_kv_heads = key_hidden_size / head_size; + at::Tensor query_dst = at::empty_symint({num_tokens, num_heads, head_size}, query.options()); + at::Tensor key_dst = at::empty_symint({num_tokens, num_kv_heads, head_size}, key.options()); + + return {query_dst, key_dst}; +} + +std::tuple get_masked_input_and_mask_meta( + at::Tensor &input, + const int64_t org_vocab_start_index, + const int64_t org_vocab_end_index, + const int64_t num_org_vocab_padding, + const int64_t added_vocab_start_index, + const int64_t added_vocab_end_index) { + + at::Tensor masked_input = at::empty_like(input); + at::Tensor mask = at::empty_like(input, input.options().dtype(at::kBool)); + + return {masked_input, mask}; +} + + +} // namespace meta +} // namespace vllm_ascend + +namespace { + // Register the meta implementations of the custom kernels for symbolic tracing, this will also + // the custom kernel been captured into aclgraph + TORCH_LIBRARY_IMPL_EXPAND(_C, Meta, ops) { + // Rotary embedding meta implementation + ops.impl("rotary_embedding", &vllm_ascend::meta::rotary_embedding_meta); + // Masked input and mask meta implementation + ops.impl("get_masked_input_and_mask", &vllm_ascend::meta::get_masked_input_and_mask_meta); + +} +} \ No newline at end of file diff --git a/csrc/utils.h b/csrc/utils.h index e94ad2d8447..74481e1b14e 100644 --- a/csrc/utils.h +++ b/csrc/utils.h @@ -29,15 +29,3 @@ } -namespace vllm_ascend { -AscendType get_dtype_from_torch(at::ScalarType scalarType) -{ - if (scalarType == at::ScalarType::Float) { - return AscendType::FP32; - } else if (scalarType == at::ScalarType::BFloat16) { - return AscendType::BF16; - } else { - return AscendType::FP16; - } -} -} // namespace vllm_ascend diff --git a/docs/source/faqs.md b/docs/source/faqs.md index 81d22f26b0b..4250fd0b3b1 100644 --- a/docs/source/faqs.md +++ b/docs/source/faqs.md @@ -10,7 +10,7 @@ ### 1. What devices are currently supported? -Currently, **ONLY** Atlas A2 series(Ascend-cann-kernels-910b)，Atlas A2 series(Atlas-A3-cann-kernels) and Atlas 300I(Ascend-cann-kernels-310p) series are supported: +Currently, **ONLY** Atlas A2 series(Ascend-cann-kernels-910b)，Atlas A3 series(Atlas-A3-cann-kernels) and Atlas 300I(Ascend-cann-kernels-310p) series are supported: - Atlas A2 Training series (Atlas 800T A2, Atlas 900 A2 PoD, Atlas 200T A2 Box16, Atlas 300T A2) - Atlas 800I A2 Inference series (Atlas 800I A2) @@ -36,6 +36,33 @@ TAG=v0.7.3rc2 docker pull m.daocloud.io/quay.io/ascend/vllm-ascend:$TAG ``` +#### Load Docker Images for offline environment +If you want to use container image for offline environments (no internet connection), you need to download container image in a environment with internet access: + +**Exporting Docker images:** + +```{code-block} bash + :substitutions: +# Pull the image on a machine with internet access +TAG=|vllm_ascend_version| +docker pull quay.io/ascend/vllm-ascend:$TAG + +# Export the image to a tar file and compress to tar.gz +docker save quay.io/ascend/vllm-ascend:$TAG | gzip > vllm-ascend-$TAG.tar.gz +``` + +**Importing Docker images in environment without internet access:** + +```{code-block} bash + :substitutions: +# Transfer the tar/tar.gz file to the offline environment and load it +TAG=|vllm_ascend_version| +docker load -i vllm-ascend-$TAG.tar.gz + +# Verify the image is loaded +docker images | grep vllm-ascend +``` + ### 3. What models does vllm-ascend supports? Find more details [here](https://vllm-ascend.readthedocs.io/en/latest/user_guide/support_matrix/supported_models.html). @@ -161,10 +188,10 @@ for output in outputs: 2. Set the following enveriments parameters: ```bash -export LCCL_DETERMINISTIC = 1 -export HCCL_DETERMINISTIC = 1 -export ATB_MATMUL_SHUFFLE_K_ENABLE = 0 -export ATB_LLM_LCOC_ENABLE = 0 +export LCCL_DETERMINISTIC=1 +export HCCL_DETERMINISTIC=true +export ATB_MATMUL_SHUFFLE_K_ENABLE=0 +export ATB_LLM_LCOC_ENABLE=0 ``` ### 19. How to fix the error "ImportError: Please install vllm[audio] for audio support" for Qwen2.5-Omni model？ diff --git a/docs/source/user_guide/configuration/additional_config.md b/docs/source/user_guide/configuration/additional_config.md index df01430df1d..75d01494641 100644 --- a/docs/source/user_guide/configuration/additional_config.md +++ b/docs/source/user_guide/configuration/additional_config.md @@ -32,6 +32,7 @@ The following table lists the additional configuration options available in vLLM | `expert_map_path` | str | `None` | When using expert load balancing for the MOE model, an expert map path needs to be passed in. | | `chunked_prefill_for_mla` | bool | `False` | Whether to enable the fused operator-like chunked_prefill. | | `kv_cache_dtype` | str | `None` | When using the kv cache quantization method, kv cache dtype needs to be set, currently only int8 is supported. | +| `enable_shared_expert_dp` | bool | `True` | When the shared expert in DP, it has better performance but consumes more memory. When the memory is sensitive, this switch can be turned off manually. | The details of each config option are as follows: diff --git a/pyproject.toml b/pyproject.toml index e394895dec5..1a140ce879f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -19,8 +19,6 @@ requires = [ "msgpack", "quart", "numba", - # Remove after https://github.com/vllm-project/vllm-ascend/issues/2034 - "transformers<4.54.0", ] build-backend = "setuptools.build_meta" diff --git a/requirements.txt b/requirements.txt index 6384149ac05..7808e852594 100644 --- a/requirements.txt +++ b/requirements.txt @@ -13,8 +13,6 @@ setuptools-scm>=8 torch>=2.7.1 torchvision wheel -# Remove after https://github.com/vllm-project/vllm-ascend/issues/2034 -transformers<4.54.0 # requirements for disaggregated prefill msgpack diff --git a/tests/e2e/long_term/accuracy/accuracy_multicard.py b/tests/e2e/long_term/accuracy/accuracy_multicard.py deleted file mode 100644 index 4479c4bf992..00000000000 --- a/tests/e2e/long_term/accuracy/accuracy_multicard.py +++ /dev/null @@ -1,167 +0,0 @@ -# -# Copyright (c) 2025 Huawei Technologies Co., Ltd. All Rights Reserved. -# Copyright 2023 The vLLM 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. -# This file is a part of the vllm-ascend project. -# Adapted from vllm-project/blob/main/tests/entrypoints/llm/test_accuracy.py -# -import gc -import multiprocessing -import sys -from multiprocessing import Queue - -import lm_eval -import pytest -import torch - -SERVER_HOST = "127.0.0.1" -SERVER_PORT = 8000 -HEALTH_URL = f"http://{SERVER_HOST}:{SERVER_PORT}/health" -COMPLETIONS_URL = f"http://{SERVER_HOST}:{SERVER_PORT}/v1/completions" - -# pre-trained model path on Hugging Face. -# Qwen/Qwen2.5-0.5B-Instruct: accuracy test for DP. -# Qwen/Qwen3-30B-A3B: accuracy test for EP and DP. -# deepseek-ai/DeepSeek-V2-Lite: accuracy test for TP. -MODEL_NAME = ["Qwen/Qwen3-30B-A3B", "deepseek-ai/DeepSeek-V2-Lite"] - -# Benchmark configuration mapping models to evaluation tasks: -# - Text model: GSM8K (grade school math reasoning) -# - Vision-language model: MMMU Art & Design validation (multimodal understanding) -TASK = { - "Qwen/Qwen2.5-0.5B-Instruct": "gsm8k", - "Qwen/Qwen3-30B-A3B": "gsm8k", - "deepseek-ai/DeepSeek-V2-Lite": "gsm8k" -} -# Answer validation requiring format consistency. -FILTER = { - "Qwen/Qwen2.5-0.5B-Instruct": "exact_match,strict-match", - "Qwen/Qwen3-30B-A3B": "exact_match,strict-match", - "deepseek-ai/DeepSeek-V2-Lite": "exact_match,strict-match" -} -# 3% relative tolerance for numerical accuracy. -RTOL = 0.03 -# Baseline accuracy after VLLM optimization. -EXPECTED_VALUE = { - "Qwen/Qwen2.5-0.5B-Instruct": 0.316, - "Qwen/Qwen3-30B-A3B": 0.888, - "deepseek-ai/DeepSeek-V2-Lite": 0.375 -} -# Maximum context length configuration for each model. -MAX_MODEL_LEN = { - "Qwen/Qwen2.5-0.5B-Instruct": 4096, - "Qwen/Qwen3-30B-A3B": 4096, - "deepseek-ai/DeepSeek-V2-Lite": 4096 -} -# Model types distinguishing text-only and vision-language models. -MODEL_TYPE = { - "Qwen/Qwen2.5-0.5B-Instruct": "vllm", - "Qwen/Qwen3-30B-A3B": "vllm", - "deepseek-ai/DeepSeek-V2-Lite": "vllm" -} -# wrap prompts in a chat-style template. -APPLY_CHAT_TEMPLATE = { - "Qwen/Qwen2.5-0.5B-Instruct": False, - "Qwen/Qwen3-30B-A3B": False, - "deepseek-ai/DeepSeek-V2-Lite": False -} -# Few-shot examples handling as multi-turn dialogues. -FEWSHOT_AS_MULTITURN = { - "Qwen/Qwen2.5-0.5B-Instruct": False, - "Qwen/Qwen3-30B-A3B": False, - "deepseek-ai/DeepSeek-V2-Lite": False -} -# MORE_ARGS extra CLI args per model -MORE_ARGS = { - "Qwen/Qwen2.5-0.5B-Instruct": - None, - "Qwen/Qwen3-30B-A3B": - "tensor_parallel_size=2,enable_expert_parallel=True,enforce_eager=True", - "deepseek-ai/DeepSeek-V2-Lite": - "tensor_parallel_size=2,trust_remote_code=True,enforce_eager=True" -} - -multiprocessing.set_start_method("spawn", force=True) - - -def run_test(queue, model, max_model_len, model_type, more_args): - try: - if model_type == "vllm-vlm": - model_args = (f"pretrained={model},max_model_len={max_model_len}," - "dtype=auto,max_images=2") - else: - model_args = (f"pretrained={model},max_model_len={max_model_len}," - "dtype=auto") - if more_args is not None: - model_args = f"{model_args},{more_args}" - results = lm_eval.simple_evaluate( - model=model_type, - model_args=model_args, - tasks=TASK[model], - batch_size="auto", - apply_chat_template=APPLY_CHAT_TEMPLATE[model], - fewshot_as_multiturn=FEWSHOT_AS_MULTITURN[model], - ) - result = results["results"][TASK[model]][FILTER[model]] - print("result:", result) - queue.put(result) - except Exception as e: - error_msg = f"{type(e).__name__}: {str(e)}" - queue.put(error_msg) - sys.exit(1) - finally: - gc.collect() - torch.npu.empty_cache() - - -@pytest.mark.parametrize("model", MODEL_NAME) -def test_lm_eval_accuracy(monkeypatch: pytest.MonkeyPatch, model): - with monkeypatch.context(): - result_queue: Queue[float] = multiprocessing.Queue() - p = multiprocessing.Process(target=run_test, - args=(result_queue, model, - MAX_MODEL_LEN[model], - MODEL_TYPE[model], MORE_ARGS[model])) - p.start() - p.join() - result = result_queue.get() - print(result) - assert (EXPECTED_VALUE[model] - RTOL < result < EXPECTED_VALUE[model] + RTOL), \ - f"Expected: {EXPECTED_VALUE[model]}±{RTOL} | Measured: {result}" - - -DP_DENSCE_MODEL = ["Qwen/Qwen2.5-0.5B-Instruct"] -DP_MOE_MOEDL = ["Qwen/Qwen3-30B-A3B"] - -DP_MORE_ARGS = { - "Qwen/Qwen2.5-0.5B-Instruct": - "tensor_parallel_size=2,data_parallel_size=2", - "Qwen/Qwen3-30B-A3B": - "tensor_parallel_size=2,data_parallel_size=2,enable_expert_parallel=True,max_model_len=1024,enforce_eager=True", -} - - -@pytest.mark.parametrize("model", DP_DENSCE_MODEL) -def test_lm_eval_accuracy_dp(model): - result_queue: Queue[float] = multiprocessing.Queue() - p = multiprocessing.Process(target=run_test, - args=(result_queue, model, - MAX_MODEL_LEN[model], MODEL_TYPE[model], - DP_MORE_ARGS[model])) - p.start() - p.join() - result = result_queue.get() - print(result) - assert (EXPECTED_VALUE[model] - RTOL < result < EXPECTED_VALUE[model] + RTOL), \ - f"Expected: {EXPECTED_VALUE[model]}±{RTOL} | Measured: {result}" diff --git a/tests/e2e/long_term/accuracy/accuracy_singlecard.py b/tests/e2e/long_term/accuracy/accuracy_singlecard.py deleted file mode 100644 index 2860dd56e7c..00000000000 --- a/tests/e2e/long_term/accuracy/accuracy_singlecard.py +++ /dev/null @@ -1,115 +0,0 @@ -# -# Copyright (c) 2025 Huawei Technologies Co., Ltd. All Rights Reserved. -# Copyright 2023 The vLLM 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. -# This file is a part of the vllm-ascend project. -# Adapted from vllm-project/blob/main/tests/entrypoints/llm/test_accuracy.py -# - -import gc -import multiprocessing -import sys -from multiprocessing import Queue - -import lm_eval -import pytest -import torch - -# pre-trained model path on Hugging Face. -MODEL_NAME = ["Qwen/Qwen2.5-0.5B-Instruct", "Qwen/Qwen2.5-VL-3B-Instruct"] -# Benchmark configuration mapping models to evaluation tasks: -# - Text model: GSM8K (grade school math reasoning) -# - Vision-language model: MMMU Art & Design validation (multimodal understanding) -TASK = { - "Qwen/Qwen2.5-0.5B-Instruct": "gsm8k", - "Qwen/Qwen2.5-VL-3B-Instruct": "mmmu_val_art_and_design" -} -# Answer validation requiring format consistency. -FILTER = { - "Qwen/Qwen2.5-0.5B-Instruct": "exact_match,strict-match", - "Qwen/Qwen2.5-VL-3B-Instruct": "acc,none" -} -# 3% relative tolerance for numerical accuracy. -RTOL = 0.03 -# Baseline accuracy after VLLM optimization. -EXPECTED_VALUE = { - "Qwen/Qwen2.5-0.5B-Instruct": 0.316, - "Qwen/Qwen2.5-VL-3B-Instruct": 0.566 -} -# Maximum context length configuration for each model. -MAX_MODEL_LEN = { - "Qwen/Qwen2.5-0.5B-Instruct": 4096, - "Qwen/Qwen2.5-VL-3B-Instruct": 8192 -} -# Model types distinguishing text-only and vision-language models. -MODEL_TYPE = { - "Qwen/Qwen2.5-0.5B-Instruct": "vllm", - "Qwen/Qwen2.5-VL-3B-Instruct": "vllm-vlm" -} -# wrap prompts in a chat-style template. -APPLY_CHAT_TEMPLATE = {"vllm": False, "vllm-vlm": True} -# Few-shot examples handling as multi-turn dialogues. -FEWSHOT_AS_MULTITURN = {"vllm": False, "vllm-vlm": True} -# batch_size -BATCH_SIZE = { - "Qwen/Qwen2.5-0.5B-Instruct": "auto", - "Qwen/Qwen2.5-VL-3B-Instruct": 1 -} - -multiprocessing.set_start_method("spawn", force=True) - - -def run_test(queue, model, max_model_len, model_type): - try: - if model_type == "vllm-vlm": - model_args = (f"pretrained={model},max_model_len={max_model_len}," - "tensor_parallel_size=1,dtype=auto,max_images=2") - else: - model_args = (f"pretrained={model},max_model_len={max_model_len}," - "tensor_parallel_size=1,dtype=auto") - results = lm_eval.simple_evaluate( - model=model_type, - model_args=model_args, - tasks=TASK[model], - batch_size=BATCH_SIZE[model], - apply_chat_template=APPLY_CHAT_TEMPLATE[model_type], - fewshot_as_multiturn=FEWSHOT_AS_MULTITURN[model_type], - ) - result = results["results"][TASK[model]][FILTER[model]] - print("result:", result) - queue.put(result) - except Exception as e: - queue.put(e) - sys.exit(1) - finally: - gc.collect() - torch.npu.empty_cache() - - -@pytest.mark.parametrize("model", MODEL_NAME) -def test_lm_eval_accuracy(monkeypatch: pytest.MonkeyPatch, model): - with monkeypatch.context(): - result_queue: Queue[float] = multiprocessing.Queue() - p = multiprocessing.Process(target=run_test, - args=(result_queue, model, - MAX_MODEL_LEN[model], - MODEL_TYPE[model])) - p.start() - p.join() - result = result_queue.get() - if isinstance(result, Exception): - pytest.fail(f"Subprocess failed with exception: {str(result)}") - print(result) - assert (EXPECTED_VALUE[model] - RTOL < result < EXPECTED_VALUE[model] + RTOL), \ - f"Expected: {EXPECTED_VALUE[model]}±{RTOL} | Measured: {result}" diff --git a/tests/e2e/models/configs/DeepSeek-V2-Lite.yaml b/tests/e2e/models/configs/DeepSeek-V2-Lite.yaml new file mode 100644 index 00000000000..7df0544d636 --- /dev/null +++ b/tests/e2e/models/configs/DeepSeek-V2-Lite.yaml @@ -0,0 +1,13 @@ +model_name: "deepseek-ai/DeepSeek-V2-Lite" +tasks: +- name: "gsm8k" + metrics: + - name: "exact_match,strict-match" + value: 0.375 + - name: "exact_match,flexible-extract" + value: 0.375 +tensor_parallel_size: 2 +apply_chat_template: False +fewshot_as_multiturn: False +trust_remote_code: True +enforce_eager: True diff --git a/tests/e2e/singlecard/models/configs/Qwen2.5-VL-7B-Instruct.yaml b/tests/e2e/models/configs/Qwen2.5-VL-7B-Instruct.yaml similarity index 100% rename from tests/e2e/singlecard/models/configs/Qwen2.5-VL-7B-Instruct.yaml rename to tests/e2e/models/configs/Qwen2.5-VL-7B-Instruct.yaml diff --git a/tests/e2e/singlecard/models/configs/Qwen3-30B-A3B.yaml b/tests/e2e/models/configs/Qwen3-30B-A3B.yaml similarity index 100% rename from tests/e2e/singlecard/models/configs/Qwen3-30B-A3B.yaml rename to tests/e2e/models/configs/Qwen3-30B-A3B.yaml diff --git a/tests/e2e/singlecard/models/configs/Qwen3-8B-Base.yaml b/tests/e2e/models/configs/Qwen3-8B-Base.yaml similarity index 100% rename from tests/e2e/singlecard/models/configs/Qwen3-8B-Base.yaml rename to tests/e2e/models/configs/Qwen3-8B-Base.yaml diff --git a/tests/e2e/singlecard/models/configs/accuracy.txt b/tests/e2e/models/configs/accuracy.txt similarity index 100% rename from tests/e2e/singlecard/models/configs/accuracy.txt rename to tests/e2e/models/configs/accuracy.txt diff --git a/tests/e2e/singlecard/models/conftest.py b/tests/e2e/models/conftest.py similarity index 53% rename from tests/e2e/singlecard/models/conftest.py rename to tests/e2e/models/conftest.py index 2b25c1a9294..a75659f4f4e 100644 --- a/tests/e2e/singlecard/models/conftest.py +++ b/tests/e2e/models/conftest.py @@ -21,14 +21,14 @@ def pytest_addoption(parser): parser.addoption( "--config", action="store", - default="./tests/e2e/singlecard/models/configs/Qwen3-8B-Base.yaml", + default="./tests/e2e/models/configs/Qwen3-8B-Base.yaml", help="Path to the model config YAML file", ) parser.addoption( - "--report_output", + "--report-dir", action="store", - default="./benchmarks/accuracy/Qwen3-8B-Base.md", - help="Path to the report output file", + default="./benchmarks/accuracy", + help="Directory to store report files", ) @@ -49,25 +49,24 @@ def config(pytestconfig): @pytest.fixture(scope="session") -def report_output(pytestconfig): - return pytestconfig.getoption("--report_output") +def report_dir(pytestconfig): + return pytestconfig.getoption("report_dir") def pytest_generate_tests(metafunc): if "config_filename" in metafunc.fixturenames: - # If config specified, use the --config directly - single_config = metafunc.config.getoption("--config") - if single_config: - metafunc.parametrize("config_filename", - [Path(single_config).resolve()]) - return - # Otherwise, check --config-list-file - rel_path = metafunc.config.getoption("--config-list-file") - config_list_file = Path(rel_path).resolve() - config_dir = config_list_file.parent - with open(config_list_file, encoding="utf-8") as f: - configs = [ - config_dir / line.strip() for line in f - if line.strip() and not line.startswith("#") - ] - metafunc.parametrize("config_filename", configs) + + if metafunc.config.getoption("--config-list-file"): + rel_path = metafunc.config.getoption("--config-list-file") + config_list_file = Path(rel_path).resolve() + config_dir = config_list_file.parent + with open(config_list_file, encoding="utf-8") as f: + configs = [ + config_dir / line.strip() for line in f + if line.strip() and not line.startswith("#") + ] + metafunc.parametrize("config_filename", configs) + else: + single_config = metafunc.config.getoption("--config") + config_path = Path(single_config).resolve() + metafunc.parametrize("config_filename", [config_path]) diff --git a/tests/e2e/singlecard/models/report_template.md b/tests/e2e/models/report_template.md similarity index 100% rename from tests/e2e/singlecard/models/report_template.md rename to tests/e2e/models/report_template.md diff --git a/tests/e2e/singlecard/models/test_lm_eval_correctness.py b/tests/e2e/models/test_lm_eval_correctness.py similarity index 94% rename from tests/e2e/singlecard/models/test_lm_eval_correctness.py rename to tests/e2e/models/test_lm_eval_correctness.py index 3453a057121..567d3de70fe 100644 --- a/tests/e2e/singlecard/models/test_lm_eval_correctness.py +++ b/tests/e2e/models/test_lm_eval_correctness.py @@ -48,7 +48,7 @@ def build_model_args(eval_config, tp_size): } for s in [ "max_images", "gpu_memory_utilization", "enable_expert_parallel", - "tensor_parallel_size" + "tensor_parallel_size", "enforce_eager" ]: val = eval_config.get(s, None) if val is not None: @@ -60,8 +60,7 @@ def build_model_args(eval_config, tp_size): return model_args -def generate_report(tp_size, eval_config, report_data, report_output, - env_config): +def generate_report(tp_size, eval_config, report_data, report_dir, env_config): env = Environment(loader=FileSystemLoader(TEST_DIR)) template = env.get_template("report_template.md") model_args = build_model_args(eval_config, tp_size) @@ -85,12 +84,14 @@ def generate_report(tp_size, eval_config, report_data, report_output, num_fewshot=eval_config.get("num_fewshot", "N/A"), rows=report_data["rows"]) + report_output = os.path.join( + report_dir, f"{os.path.basename(eval_config['model_name'])}.md") os.makedirs(os.path.dirname(report_output), exist_ok=True) with open(report_output, 'w', encoding='utf-8') as f: f.write(report_content) -def test_lm_eval_correctness_param(config_filename, tp_size, report_output, +def test_lm_eval_correctness_param(config_filename, tp_size, report_dir, env_config): eval_config = yaml.safe_load(config_filename.read_text(encoding="utf-8")) model_args = build_model_args(eval_config, tp_size) @@ -143,6 +144,5 @@ def test_lm_eval_correctness_param(config_filename, tp_size, report_output, metric_name.replace(',', '_stderr,') if metric_name == "acc,none" else metric_name.replace(',', '_stderr,')] }) - generate_report(tp_size, eval_config, report_data, report_output, - env_config) + generate_report(tp_size, eval_config, report_data, report_dir, env_config) assert success diff --git a/tests/e2e/multicard/moe/test_moe_comm.py b/tests/e2e/multicard/moe/test_moe_comm.py new file mode 100644 index 00000000000..b1de5e680f9 --- /dev/null +++ b/tests/e2e/multicard/moe/test_moe_comm.py @@ -0,0 +1,153 @@ +# Copyright (c) 2025 Huawei Technologies Co., Ltd. All Rights Reserved. +# Copyright 2023 The vLLM 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. +# This file is a part of the vllm-ascend project. + +from types import SimpleNamespace + +import pytest +import torch +from transformers import PretrainedConfig +from vllm import forward_context + +from vllm_ascend.distributed import moe_comm_method +from vllm_ascend.distributed.moe_comm_method import (AllGatherCommImpl, + NativeAllGatherCommImpl) + + +@pytest.mark.parametrize("num_tokens", [16, 128]) +@pytest.mark.parametrize("hidden_size", [64, 128]) +@pytest.mark.parametrize("global_num_experts", [8, 16]) +@pytest.mark.parametrize("top_k_num", [2, 4]) +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) +@pytest.mark.parametrize("num_local_experts", [4, 8]) +@pytest.mark.parametrize("ep_rank", [0, 1]) +def test_all_gather_comm_impl( + num_tokens, + hidden_size, + global_num_experts, + top_k_num, + dtype, + num_local_experts, + ep_rank, +): + """ + Tests the AllGatherCommImpl against the NativeAllGatherCommImpl. + + This test compares the outputs of the NPU-optimized AllGatherCommImpl + with a native PyTorch implementation (NativeAllGatherCommImpl) to ensure + correctness across various configurations. + """ + if top_k_num > global_num_experts: + pytest.skip("top_k_num cannot be greater than global_num_experts") + if num_local_experts > global_num_experts: + pytest.skip( + "num_local_experts cannot be greater than global_num_experts") + + device = torch.device("npu") + hf_config = PretrainedConfig( + num_experts_per_tok=top_k_num, + num_experts=global_num_experts, + ) + + # Instantiate implementations + native_impl = NativeAllGatherCommImpl(device, dtype, hf_config) + + all_gather_impl = AllGatherCommImpl(device, dtype, hf_config) + + # TODO: Find out if this is the correct way to mock the forward context and ep group + # Mock get_forward_context to return an object with moe_comm_method + forward_context._forward_context = SimpleNamespace( + moe_comm_method=all_gather_impl) + # Mock get_ep_group to return a fake group with the specified ep_rank + fake_ep_group = SimpleNamespace(rank_in_group=ep_rank) + moe_comm_method.get_ep_group = lambda: fake_ep_group + + # --- Input Data --- + hidden_states = torch.randn(num_tokens, + hidden_size, + device=device, + dtype=dtype) + topk_ids = torch.randint(0, + global_num_experts, (num_tokens, top_k_num), + device=device, + dtype=torch.int32) + topk_weights = torch.rand(num_tokens, top_k_num, device=device).to(dtype) + topk_weights = torch.nn.functional.softmax(topk_weights, dim=1) + + num_experts = global_num_experts + + expert_map = None + if num_local_experts < global_num_experts: + # Create a map where some experts are local and some are not + expert_map = torch.full((global_num_experts, ), -1, device=device) + expert_map[ep_rank * num_local_experts:(ep_rank + 1) * + num_local_experts] = torch.arange(num_local_experts, + device=device) + num_experts = num_local_experts + + # --- Run Native Implementation (Golden Reference) --- + native_hidden_states_out = hidden_states.clone() + ( + native_permuted_hidden, + native_expert_tokens, + _, + ) = native_impl._pre_process(hidden_states, topk_ids, topk_weights, + expert_map, num_experts) + # Simulate MLP output + native_mlp_output = torch.randn_like(native_permuted_hidden) + native_impl._post_process(native_mlp_output, native_hidden_states_out) + + # --- Run AllGather Implementation --- + all_gather_hidden_states_out = hidden_states.clone() + ( + all_gather_permuted_hidden, + all_gather_expert_tokens, + _, + ) = torch.ops.vllm.moe_comm_pre_process(hidden_states, topk_ids, + topk_weights, expert_map, + num_experts) + + # Use the same simulated MLP output for a fair comparison + all_gather_mlp_output = native_mlp_output.clone() + + torch.ops.vllm.moe_comm_post_process(all_gather_mlp_output, + all_gather_hidden_states_out) + + # --- Assertions --- + # Define tolerance based on dtype + atol = 1e-3 if dtype == torch.float16 else 1e-2 + rtol = 1e-3 if dtype == torch.float16 else 1e-2 + + # 1. Compare expert_tokens from pre_process + assert torch.allclose(native_expert_tokens.to( + all_gather_expert_tokens.device), + all_gather_expert_tokens, + atol=atol, + rtol=rtol), "Expert tokens do not match." + + # 2. Compare permuted_hidden_states from pre_process + num_valid_tokens = native_expert_tokens.sum() + assert torch.allclose(native_permuted_hidden[:num_valid_tokens].to( + all_gather_permuted_hidden.device), + all_gather_permuted_hidden[:num_valid_tokens], + atol=atol, + rtol=rtol), "Permuted hidden states do not match." + + # 3. Compare final hidden_states from post_process + assert torch.allclose(native_hidden_states_out.to( + all_gather_hidden_states_out.device), + all_gather_hidden_states_out, + atol=atol, + rtol=rtol), "Final hidden states do not match." diff --git a/tests/e2e/multicard/test_external_launcher.py b/tests/e2e/multicard/test_external_launcher.py index c5eecab81c4..24c66bfcb4c 100644 --- a/tests/e2e/multicard/test_external_launcher.py +++ b/tests/e2e/multicard/test_external_launcher.py @@ -24,11 +24,14 @@ import subprocess import sys from pathlib import Path +from unittest.mock import patch import pytest +import torch_npu MODELS = ["Qwen/Qwen3-0.6B"] MOE_MODELS = ["Qwen/Qwen3-30B-A3B"] +DEVICE_NAME = torch_npu.npu.get_device_name(0)[:10] @pytest.mark.parametrize("model", MODELS) @@ -147,3 +150,38 @@ def test_external_launcher_and_sleepmode(): assert "Generated text:" in output assert "Sleep and wake up successfully!!" in output assert proc.returncode == 0 + + +@pytest.mark.skipif( + DEVICE_NAME != "Ascend910B", + reason="This test is only for Ascend910B devices.", +) +@pytest.mark.parametrize("model", MODELS) +@patch.dict(os.environ, {"VLLM_ASCEND_ENABLE_MATMUL_ALLREDUCE": "1"}) +def test_mm_allreduce(model): + script = Path( + __file__ + ).parent.parent.parent.parent / "examples" / "offline_external_launcher.py" + env = os.environ.copy() + cmd = [ + sys.executable, + str(script), + "--model", + model, + "--trust-remote-code", + ] + + print(f"Running subprocess: {' '.join(cmd)}") + proc = subprocess.run( + cmd, + env=env, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + timeout=600, + ) + + output = proc.stdout.decode() + print(output) + + assert "Generated text:" in output + assert proc.returncode == 0 diff --git a/tests/e2e/singlecard/ops/test_rotary_embedding.py b/tests/e2e/singlecard/ops/test_rotary_embedding.py index a3504a88b24..c750f010e75 100644 --- a/tests/e2e/singlecard/ops/test_rotary_embedding.py +++ b/tests/e2e/singlecard/ops/test_rotary_embedding.py @@ -17,11 +17,12 @@ # Only Neox style true scenario is supported for now IS_NEOX_STYLE = [True] DTYPES = [torch.half] -HEAD_SIZES = [64, 96, 128, 256] +HEAD_SIZES = [64, 64, 96, 128, 256] ROTARY_DIMS = [None, 32] # None means rotary dim == head size NUM_HEADS = [17] # Arbitrary values for testing BATCH_SIZES = [5] # Arbitrary values for testing SEQ_LENS = [11, 4096] # Arbitrary values for testing +NUM_TOKENS = [10, 21] SEEDS = [0] DEVICES = [f"npu:{0}"] # Set tolerance to 1 for quant ops @@ -198,3 +199,146 @@ def test_rotary_embedding_quant_with_leading_dim( ref_key, atol=DEFAULT_ATOL, rtol=DEFAULT_RTOL) + + +class ModelwithRotaryEmbedding(nn.Module): + + def __init__( + self, + hidden_size: int, + num_heads: int, + head_size: int, + rotary_dim: int, + max_position_embeddings: int, + base: int, + is_neox_style: bool, + dtype: torch.dtype, + ) -> None: + super().__init__() + self.qkv_proj = nn.Linear(hidden_size, num_heads * head_size * 3) + self.rope = RotaryEmbedding( + head_size=head_size, + rotary_dim=rotary_dim, + max_position_embeddings=max_position_embeddings, + base=base, + is_neox_style=is_neox_style, + dtype=dtype, + ) + self.o_proj = nn.Linear(num_heads * head_size, hidden_size) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + offsets: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + # we simulated a simple attention layer to test if it can be seamlessly captured into aclgraph + qkv = self.qkv_proj(hidden_states) + q, k, v = qkv.chunk(3, dim=-1) + query, key = torch.ops._C.rotary_embedding( + positions, + q, + k, + self.rope.head_size, + self.rope.cos_sin_cache, + self.rope.is_neox_style, + ) + query = query.view(q.shape) + key = key.view(k.shape) + o = self.o_proj(query) + return o + + +# The first graph seems will have some accuracy issue when directly run pytest on the ops folder, +# add a warmup graph replay for workaround +ACL_GRPAH_FIRST_RUN = True + + +@pytest.mark.parametrize("is_neox_style", IS_NEOX_STYLE) +@pytest.mark.parametrize("num_tokens", BATCH_SIZES) +@pytest.mark.parametrize("num_heads", NUM_HEADS) +@pytest.mark.parametrize("head_size", HEAD_SIZES) +@pytest.mark.parametrize("rotary_dim", ROTARY_DIMS) +@pytest.mark.parametrize("dtype", DTYPES) +@pytest.mark.parametrize("seed", SEEDS) +@pytest.mark.parametrize("device", DEVICES) +@torch.inference_mode() +def test_capture_rotary_embedding_in_aclgraph( + is_neox_style: bool, + num_tokens: int, + num_heads: int, + head_size: int, + rotary_dim: int, + dtype: torch.dtype, + seed: int, + device: str, + max_position_embeddings: int = 8192, + base: int = 10000, +): + """Test if the rotary embedding can be captured in aclgraph.""" + torch.manual_seed(seed) + torch.set_default_device(device) + if rotary_dim is None: + rotary_dim = head_size + model = ModelwithRotaryEmbedding( + hidden_size=num_heads * head_size, + num_heads=num_heads, + head_size=head_size, + rotary_dim=rotary_dim, + max_position_embeddings=max_position_embeddings, + base=base, + is_neox_style=is_neox_style, + dtype=dtype, + ) + + def custom_op_checking_backend(gm: torch.fx.GraphModule, example_input): + # Validate if the rotary_embedding custom kernel is indeed inside the graph by + # string match + graph = str(gm.graph) + assert "_C.rotary_embedding" in graph + return gm + + static_positions = torch.randint(0, max_position_embeddings, + (num_tokens, )) + static_hidden_states = torch.randn(num_tokens, + num_heads * head_size, + dtype=dtype, + device="npu") + compiled_model = torch.compile(model, backend=custom_op_checking_backend) + stream = torch.npu.Stream() + stream.wait_stream(torch.npu.current_stream()) + with torch.npu.stream(stream): + # warmup the fx graph before capture + for i in range(3): + static_output = compiled_model(static_positions, + static_hidden_states, + offsets=None) + stream.wait_stream(torch.npu.current_stream()) + + aclgraph = torch.npu.NPUGraph() + + with torch.npu.graph(aclgraph): + # Capture the model in aclgraph. + static_output = compiled_model(static_positions, static_hidden_states) + # Capture the model in aclgraph. + random_filled_positions = torch.randint(0, + max_position_embeddings, + (num_tokens, ), + device="npu") + random_filled_hidden_states = torch.randn(num_tokens, + num_heads * head_size, + dtype=dtype, + device="npu") + static_positions.copy_(random_filled_positions) + static_hidden_states.copy_(random_filled_hidden_states) + + aclgraph.replay() + global ACL_GRPAH_FIRST_RUN + if ACL_GRPAH_FIRST_RUN: + ACL_GRPAH_FIRST_RUN = False + return + output_reference = model(static_positions, static_hidden_states) + torch.testing.assert_close(static_output, + output_reference, + atol=DEFAULT_ATOL, + rtol=DEFAULT_RTOL) diff --git a/tests/e2e/singlecard/sample/test_rejection_sampler.py b/tests/e2e/singlecard/sample/test_rejection_sampler.py index 123e7c20c08..2a331202861 100644 --- a/tests/e2e/singlecard/sample/test_rejection_sampler.py +++ b/tests/e2e/singlecard/sample/test_rejection_sampler.py @@ -77,8 +77,9 @@ def test_perfect_match(rejection_sampler): metadata = create_sampling_metadata(all_greedy=True) logits = create_logits_tensor(output_tokens) - bonus_token_tensor = torch.tensor([output_tokens[0][-1]], - device=logits.device) + bonus_token_tensor = torch.tensor([[output_tokens[0][-1]]], + device=logits.device, + dtype=torch.int32) spec_decode_metadata = SpecDecodeMetadata.make_dummy(spec_tokens, device=logits.device) @@ -102,8 +103,9 @@ def test_early_mismatch(rejection_sampler): metadata = create_sampling_metadata(all_greedy=True) logits = create_logits_tensor(output_tokens) - bonus_token_tensor = torch.tensor([output_tokens[0][-1]], - device=logits.device) + bonus_token_tensor = torch.tensor([[output_tokens[0][-1]]], + device=logits.device, + dtype=torch.int32) spec_decode_metadata = SpecDecodeMetadata.make_dummy(spec_tokens, device=logits.device) @@ -131,7 +133,9 @@ def test_multiple_sequences(rejection_sampler): metadata = create_sampling_metadata(all_greedy=True) logits = create_logits_tensor(output_tokens) bonus_token_tensor = torch.tensor( - [output_tokens[0][-1], output_tokens[1][-1]], device=logits.device) + [output_tokens[0][-1], output_tokens[1][-1]], + device=logits.device, + dtype=torch.int32).unsqueeze(1) spec_decode_metadata = SpecDecodeMetadata.make_dummy(spec_tokens, device=logits.device) @@ -155,8 +159,9 @@ def test_single_token_sequence(rejection_sampler): metadata = create_sampling_metadata(all_greedy=True) logits = create_logits_tensor(output_tokens) - bonus_token_tensor = torch.tensor([output_tokens[0][-1]], - device=logits.device) + bonus_token_tensor = torch.tensor([[output_tokens[0][-1]]], + device=logits.device, + dtype=torch.int32) spec_decode_metadata = SpecDecodeMetadata.make_dummy(spec_tokens, device=logits.device) @@ -178,8 +183,9 @@ def test_empty_sequence(rejection_sampler): metadata = create_sampling_metadata(all_greedy=True) logits = create_logits_tensor(output_tokens) - bonus_token_tensor = torch.tensor([output_tokens[0][-1]], - device=logits.device) + bonus_token_tensor = torch.tensor([[output_tokens[0][-1]]], + device=logits.device, + dtype=torch.int32) spec_decode_metadata = SpecDecodeMetadata.make_dummy(spec_tokens, device=logits.device) @@ -203,7 +209,9 @@ def test_multiple_mismatches(rejection_sampler): metadata = create_sampling_metadata(all_greedy=True) logits = create_logits_tensor(output_tokens) bonus_token_tensor = torch.tensor( - [output_tokens[0][-1], output_tokens[1][-1]], device=logits.device) + [output_tokens[0][-1], output_tokens[1][-1]], + device=logits.device, + dtype=torch.int32).unsqueeze(1) spec_decode_metadata = SpecDecodeMetadata.make_dummy(spec_tokens, device=logits.device) @@ -237,7 +245,8 @@ def test_parametrized_cases(rejection_sampler, spec_tokens, output_tokens, metadata = create_sampling_metadata(all_greedy=True) logits = create_logits_tensor(output_tokens) bonus_token_tensor = torch.tensor([tokens[-1] for tokens in output_tokens], - device=logits.device) + device=logits.device, + dtype=torch.int32).unsqueeze(1) spec_decode_metadata = SpecDecodeMetadata.make_dummy(spec_tokens, device=logits.device) diff --git a/tests/e2e/singlecard/spec_decode_v1/test_v1_spec_decode.py b/tests/e2e/singlecard/spec_decode_v1/test_v1_spec_decode.py index 56fa6cc6392..c7b173a6e38 100644 --- a/tests/e2e/singlecard/spec_decode_v1/test_v1_spec_decode.py +++ b/tests/e2e/singlecard/spec_decode_v1/test_v1_spec_decode.py @@ -101,6 +101,7 @@ def test_ngram_correctness( del spec_llm +@pytest.mark.skipif(True, reason="oom in CI, fix me") @pytest.mark.parametrize("use_eagle3", [False, True], ids=["eagle", "eagle3"]) def test_eagle_correctness( test_prompts: list[list[dict[str, Any]]], diff --git a/tests/ut/attention/test_mla_v1.py b/tests/ut/attention/test_mla_v1.py index 2ecc3f7bd74..497b7b53abc 100644 --- a/tests/ut/attention/test_mla_v1.py +++ b/tests/ut/attention/test_mla_v1.py @@ -664,6 +664,7 @@ def test_rope_single(self, mock_rope): def test_forward_decode_without_graph(self, mock_page_attention_mla, mock_up_proj): self.impl.running_in_graph = False + self.impl.running_chunkprefilll_with_torchair = False num_tokens = 100 num_blocks = 256 block_size = 4 @@ -690,3 +691,40 @@ def test_forward_decode_without_graph(self, mock_page_attention_mla, self.assertEqual(result.shape[2], self.impl.v_head_dim) mock_up_proj.assert_called_once() mock_page_attention_mla.assert_called_once() + + @patch("vllm_ascend.attention.mla_v1.AscendMLAImpl._forward_prefill") + @patch("torch_npu._npu_reshape_and_cache") + def test_forward_without_graph(self, _, mock_forward_prefill): + self.impl.running_in_graph = False + self.impl.torchair_graph_enabled = False + + num_tokens = 100 + num_blocks = 256 + block_size = 4 + rotary_emb_return_value = (torch.randn(num_tokens, 16, + self.impl.kv_lora_rank), + torch.randn(0, 1, self.impl.kv_lora_rank)) + self.impl.rotary_emb.side_effect = lambda *args, **kwargs: rotary_emb_return_value + self.impl.o_proj.side_effect = lambda *args, **kwargs: torch.randn( + 1, num_blocks, 128) + + hidden_states_or_q_c = torch.randn(num_tokens, self.impl.q_lora_rank) + hidden_states_or_kv_c_normed = torch.randn(num_tokens, + self.impl.kv_lora_rank) + k_pe = torch.randn(num_tokens, self.impl.qk_rope_head_dim) + kv_cache = (torch.randn(num_blocks, block_size, self.impl.num_heads, + self.impl.kv_lora_rank), + torch.randn(num_blocks, block_size, self.impl.num_heads, + self.impl.qk_rope_head_dim)) + output = torch.randn(num_tokens, self.impl.num_heads, + self.impl.v_head_dim) + + metadata = MagicMock() + metadata.num_decodes = 0 + metadata.num_prefills = num_tokens + mock_forward_prefill.return_value = torch.randn( + 0, self.impl.num_heads * self.impl.v_head_dim) + result = self.impl.forward(None, hidden_states_or_q_c, + hidden_states_or_kv_c_normed, k_pe, + kv_cache, metadata, output, False) + self.assertEqual(result.shape[0], num_tokens) diff --git a/tests/ut/distributed/test_communicator.py b/tests/ut/distributed/test_communicator.py new file mode 100644 index 00000000000..880cb246ea7 --- /dev/null +++ b/tests/ut/distributed/test_communicator.py @@ -0,0 +1,155 @@ +import unittest +from unittest.mock import MagicMock, Mock, patch + +import torch +import torch.distributed as dist + +from vllm_ascend.distributed.communicator import NPUCommunicator + + +class TestNPUCommunicator(unittest.TestCase): + + @patch("vllm.config.get_current_vllm_config", return_value=None) + @patch("torch.npu.current_device", return_value=MagicMock()) + @patch("torch.npu.set_device", return_value=MagicMock()) + @patch("torch.distributed.get_process_group_ranks", + return_value={ + 0: 0, + 1: 1 + }) + @patch("torch.distributed.get_group_rank", return_value={0: 0, 1: 1}) + @patch("torch.distributed.is_initialized", return_value=True) + @patch("torch.distributed.get_rank", return_value=1) + @patch("torch.distributed.is_initialized", return_value=True) + @patch("torch.distributed.get_backend", return_value="hccl") + @patch("torch.distributed.get_rank", return_value=1) + @patch("torch.distributed.get_world_size", return_value=2) + @patch("torch.distributed.get_process_group_ranks", return_value=[0, 1]) + @patch("torch.npu.device") + def test_all_to_all_with_sizes(self, *_): + + def patched_all_to_all(output_tensor_list, + input_tensor_list, + group=None, + async_op=False): + output_tensor_list[:] = ([ + torch.tensor([10, 20]), + torch.tensor([50, 60]) + ]) + + torch.distributed.all_to_all = patched_all_to_all + + scatter_sizes = [2, 2] + gather_sizes = [2, 2] + input_ = torch.tensor([10, 20, 30, 40]) + + comm = NPUCommunicator(cpu_group=dist.group.WORLD) + + output = comm.all_to_all(input_, + scatter_sizes=scatter_sizes, + gather_sizes=gather_sizes) + + assert output.tolist() == [10, 20, 50, 60] + + @patch("vllm.config.get_current_vllm_config", return_value=None) + @patch("torch.npu.current_device", return_value=MagicMock()) + @patch("torch.npu.set_device", return_value=MagicMock()) + @patch("torch.distributed.get_process_group_ranks", + return_value={ + 0: 0, + 1: 1 + }) + @patch("torch.distributed.get_group_rank", return_value={0: 0, 1: 1}) + @patch("torch.distributed.is_initialized", return_value=True) + @patch("torch.distributed.get_rank", return_value=1) + @patch("torch.distributed.is_initialized", return_value=True) + @patch("torch.distributed.get_backend", return_value="hccl") + @patch("torch.distributed.get_rank", return_value=1) + @patch("torch.distributed.get_world_size", return_value=2) + @patch("torch.distributed.get_process_group_ranks", return_value=[0, 1]) + @patch("torch.npu.device") + def test_all_to_all_without_sizes(self, *_): + + def patched_all_to_all(output_tensor_list, + input_tensor_list, + group=None, + async_op=False): + output_tensor_list[:] = ([ + torch.tensor([[10, 20]]), + torch.tensor([[50, 60]]) + ]) + + torch.distributed.all_to_all = patched_all_to_all + + input_ = torch.tensor([[10, 20], [30, 40]]) + + comm = NPUCommunicator(cpu_group=dist.group.WORLD) + output = comm.all_to_all(input_, scatter_dim=0, gather_dim=0) + + assert output.tolist() == [[10, 20], [50, 60]] + + @patch("vllm.config.get_current_vllm_config", return_value=None) + @patch("torch.npu.current_device", return_value=MagicMock()) + @patch("torch.npu.set_device", return_value=MagicMock()) + @patch("torch.distributed.get_process_group_ranks", + return_value={ + 0: 0, + 1: 1 + }) + @patch("torch.distributed.get_group_rank", return_value={0: 0, 1: 1}) + @patch("torch.distributed.is_initialized", return_value=True) + @patch("torch.distributed.get_rank", return_value=1) + @patch("torch.distributed.is_initialized", return_value=True) + @patch("torch.distributed.get_backend", return_value="hccl") + @patch("torch.distributed.get_rank", return_value=1) + @patch("torch.distributed.get_world_size", return_value=2) + @patch("torch.distributed.get_process_group_ranks", return_value=[0, 1]) + @patch("torch.npu.device") + def test_dispatch(self, *_): + comm = NPUCommunicator(cpu_group=dist.group.WORLD) + comm.all2all_manager = Mock() + hidden_states = torch.randn(2, 4, 8) + router_logits = torch.randn(2, 4, 2) + + mock_dispatch_result = (torch.randn(2, 4, 8), torch.randn(2, 4, 2)) + comm.all2all_manager.dispatch.return_value = mock_dispatch_result + + result_hidden, result_logits = comm.dispatch(hidden_states, + router_logits) + + assert torch.allclose(result_hidden, mock_dispatch_result[0]) + assert torch.allclose(result_logits, mock_dispatch_result[1]) + + comm.all2all_manager.dispatch.assert_called_once_with( + hidden_states, router_logits) + + @patch("vllm.config.get_current_vllm_config", return_value=None) + @patch("torch.npu.current_device", return_value=MagicMock()) + @patch("torch.npu.set_device", return_value=MagicMock()) + @patch("torch.distributed.get_process_group_ranks", + return_value={ + 0: 0, + 1: 1 + }) + @patch("torch.distributed.get_group_rank", return_value={0: 0, 1: 1}) + @patch("torch.distributed.is_initialized", return_value=True) + @patch("torch.distributed.get_rank", return_value=1) + @patch("torch.distributed.is_initialized", return_value=True) + @patch("torch.distributed.get_backend", return_value="hccl") + @patch("torch.distributed.get_rank", return_value=1) + @patch("torch.distributed.get_world_size", return_value=2) + @patch("torch.distributed.get_process_group_ranks", return_value=[0, 1]) + @patch("torch.npu.device") + def test_combine(self, *_): + comm = NPUCommunicator(cpu_group=dist.group.WORLD) + comm.all2all_manager = Mock() + hidden_states = torch.randn(2, 4, 8) + + mock_combine_result = torch.randn(2, 4, 8) + comm.all2all_manager.combine.return_value = mock_combine_result + + result = comm.combine(hidden_states) + + assert torch.allclose(result, mock_combine_result) + + comm.all2all_manager.combine.assert_called_once_with(hidden_states) diff --git a/tests/ut/ops/test_fused_ops.py b/tests/ut/ops/test_fused_ops.py index 6c89f6fc1d8..8c16ec4c2f7 100644 --- a/tests/ut/ops/test_fused_ops.py +++ b/tests/ut/ops/test_fused_ops.py @@ -20,6 +20,7 @@ import torch.nn as nn import torch_npu from pytest_mock import MockerFixture +from vllm.model_executor.layers.fused_moe import FusedMoEMethodBase from vllm_ascend.ascend_forward_context import _get_fused_moe_state from vllm_ascend.ops.fused_moe import (AscendFusedMoE, @@ -59,6 +60,7 @@ def mock_dist_env(mocker: MockerFixture): patch('vllm_ascend.ops.fused_moe.get_tp_group', return_value=mock_dp_and_tp_group(mocker)), \ patch('vllm.distributed.parallel_state.get_tp_group', return_value=mock_dp_and_tp_group(mocker)), \ patch('vllm_ascend.ops.fused_moe.get_dp_group', return_value=mock_dp_and_tp_group(mocker)), \ + patch('vllm.model_executor.layers.fused_moe.layer.get_dp_group', return_value=mock_dp_and_tp_group(mocker)), \ patch('torch.distributed.all_gather', return_value=MagicMock(return_value=torch.randn(10,32))), \ patch('torch.distributed.all_to_all_single', return_value=torch.randn(8, 32)), \ patch('vllm_ascend.ops.fused_moe.tensor_model_parallel_all_reduce', @@ -112,7 +114,7 @@ def mock_moe_env(mocker: MockerFixture): torch.randn(16, 2) )), \ patch("torch_npu.npu_grouped_matmul", return_value=( - (torch.randn(8, 2), torch.randn(8, 2)) + [torch.randn(16, 2)] )), \ patch("torch_npu.npu_swiglu", return_value=( torch.randn(16, 2) @@ -180,6 +182,18 @@ def __init__(self, shared_experts, num_tokens): self.apply = MagicMock(return_value=(torch.randn(num_tokens, 32))) +class MockFusedMoEMethod(FusedMoEMethodBase): + + def create_weights(self, layer: torch.nn.Module, num_experts: int, + hidden_size: int, intermediate_size_per_partition: int, + params_dtype: torch.dtype, **extra_weight_attrs): + pass + + def apply(self, hidden_states: torch.Tensor, + expert_weights: torch.Tensor) -> torch.Tensor: + pass + + class TestAscendFusedMoe: def test_init_no_quant(self, mock_dist_env, default_moe_config): @@ -213,7 +227,7 @@ def test_init_no_quant(self, mock_dist_env, default_moe_config): def test_init_with_quant(self, mock_dist_env, default_moe_config): mock_quant_config = MagicMock() - mock_quant_method = MagicMock() + mock_quant_method = MockFusedMoEMethod() mock_quant_config.get_quant_method.return_value = mock_quant_method moe = AscendFusedMoE(**default_moe_config, diff --git a/tests/ut/sample/test_rejection_sampler.py b/tests/ut/sample/test_rejection_sampler.py index b6aaf868c5f..adbf376dd79 100644 --- a/tests/ut/sample/test_rejection_sampler.py +++ b/tests/ut/sample/test_rejection_sampler.py @@ -32,11 +32,12 @@ class TestAscendRejectionSampler(TestBase): def test_rejection_greedy_sample_pytorch(self): """Test greedy rejection sampling: stop when draft doesn't match, otherwise append bonus token""" batch_size = 2 - max_spec_len = 3 + max_spec_len = 2 output_token_ids = torch.full((batch_size, max_spec_len + 1), PLACEHOLDER_TOKEN_ID) cu_num_draft_tokens = torch.tensor([2, 4]) + num_draft_tokens = [2, 2] draft_token_ids = torch.tensor([10, 11, 20, 21]) target_argmax = torch.tensor([10, 99, 20, 22]) bonus_token_ids = torch.tensor([[100], [200]]) @@ -49,8 +50,9 @@ def test_rejection_greedy_sample_pytorch(self): draft_token_ids, target_argmax, bonus_token_ids, - is_greedy, + num_draft_tokens, max_spec_len, + is_greedy, ) assert output_token_ids[0, 0].item() == 10 diff --git a/vllm_ascend/ascend_config.py b/vllm_ascend/ascend_config.py index 659f4415f77..777ff9ffac4 100644 --- a/vllm_ascend/ascend_config.py +++ b/vllm_ascend/ascend_config.py @@ -47,6 +47,9 @@ def __init__(self, vllm_config): self.expert_map_path = additional_config.get("expert_map_path", None) self.chunked_prefill_for_mla = additional_config.get( "chunked_prefill_for_mla", False) + self.enable_shared_expert_dp = additional_config.get( + "enable_shared_expert_dp", True + ) and not self.torchair_graph_config.enabled and vllm_config.parallel_config.enable_expert_parallel class TorchairGraphConfig: @@ -166,6 +169,10 @@ def check_ascend_config(vllm_config, enforce_eager): raise NotImplementedError( "Torchair graph mode only works with following model types:" f"{TORCHAIR_MODEL_LIST}.") + if ascend_config.enable_shared_expert_dp: + logger.warning( + "enable_shared_expert_dp is not supported for torchair graph mode currently, " + "it has been disabled automatically.") # aclgraph case else: # aclgraph doesn't work with deepseek model and only qwen model is well tested. diff --git a/vllm_ascend/ascend_forward_context.py b/vllm_ascend/ascend_forward_context.py index c86253472ff..c045ad6306e 100644 --- a/vllm_ascend/ascend_forward_context.py +++ b/vllm_ascend/ascend_forward_context.py @@ -5,11 +5,12 @@ import torch from vllm.config import VllmConfig -from vllm.distributed import get_dp_group, get_ep_group, get_tp_group +from vllm.distributed import (get_dp_group, get_ep_group, + get_tensor_model_parallel_world_size) from vllm.forward_context import get_forward_context, set_forward_context import vllm_ascend.envs as envs -from vllm_ascend.platform import NPUPlatform +from vllm_ascend.distributed.moe_comm_method import MoECommMethod class FusedMoEState(Enum): @@ -54,6 +55,8 @@ def set_ascend_forward_context( num_tokens_across_dp: Optional[torch.Tensor] = None, with_prefill: bool = True, in_profile_run: bool = False, + reserved_mc2_mask: Optional[torch.Tensor] = None, + moe_comm_method: Optional[MoECommMethod] = None, num_actual_tokens: Optional[int] = None, ): """A context manager that stores the current forward context, @@ -66,6 +69,7 @@ def set_ascend_forward_context( num_tokens=num_tokens, num_tokens_across_dp=num_tokens_across_dp): forward_context = get_forward_context() + forward_context.moe_comm_method = moe_comm_method forward_context.with_prefill = with_prefill ep_size = (get_ep_group().world_size if vllm_config.parallel_config.enable_expert_parallel else 1) @@ -97,16 +101,17 @@ def set_ascend_forward_context( if num_tokens is not None: if num_actual_tokens is None: num_actual_tokens = num_tokens - tp_world_size = get_tp_group().world_size + tp_world_size = get_tensor_model_parallel_world_size() # NOTE: token num which need to pad to when mc2 forward_context.padded_num_tokens = math.ceil( max_tokens_across_dp / tp_world_size) * tp_world_size - mc2_mask = torch.zeros(forward_context.padded_num_tokens, - dtype=torch.bool, - device=NPUPlatform.device_type) - mc2_mask[:num_actual_tokens] = True - forward_context.mc2_mask = mc2_mask + if reserved_mc2_mask is not None: + mc2_mask = reserved_mc2_mask[:forward_context. + padded_num_tokens] + mc2_mask[:num_actual_tokens] = True + mc2_mask[num_actual_tokens:] = False + forward_context.mc2_mask = mc2_mask try: yield diff --git a/vllm_ascend/attention/mla_v1.py b/vllm_ascend/attention/mla_v1.py index a8f8ae82332..e7dccf33ab1 100644 --- a/vllm_ascend/attention/mla_v1.py +++ b/vllm_ascend/attention/mla_v1.py @@ -621,6 +621,7 @@ def __init__( ascend_config = get_ascend_config() self.torchair_graph_enabled = ascend_config.torchair_graph_config.enabled self.enable_kv_nz = ascend_config.torchair_graph_config.enable_kv_nz + self.enable_shared_expert_dp = ascend_config.enable_shared_expert_dp # Adapt torch air graph mode with spec decoding. speculative_config = get_current_vllm_config().speculative_config @@ -635,6 +636,8 @@ def _v_up_proj_and_o_proj(self, x, enable_multistream_mla: bool = False): x = torch.bmm(x, self.W_UV) # Convert from (N, B, V) to (B, N * V) x = x.transpose(0, 1).reshape(-1, self.num_heads * self.v_head_dim) + if hasattr(self, "running_in_graph") and not self.running_in_graph: + return x MAX_O_PROJ_PREFETCH_SIZE = 16 * 1024 * 1024 # 16MB npu_prefetch(self.o_proj.weight, x, @@ -905,14 +908,7 @@ def _forward_prefill( ] and not ascend_config.chunked_prefill_for_mla: attn_output = attn_output_torch - current_ms_metadata = get_multistream_comm_context() - if current_ms_metadata is None: - return self.o_proj(attn_output, is_prefill=True)[0] - else: - current_ms_metadata.before_comm_event.record() - with torch.npu.stream(current_ms_metadata.comm_stream): - current_ms_metadata.before_comm_event.wait() - return self.o_proj(attn_output, is_prefill=True)[0] + return attn_output def exec_kv( self, @@ -998,7 +994,7 @@ def _forward_decode( decode_meta = attn_metadata.decode assert decode_meta is not None num_tokens = q_nope.size(0) - if self.running_in_graph: + if self.running_in_graph or self.running_chunkprefilll_with_torchair: # shape of knope/k_pe for npu graph mode should be: # [num_blocks, num_kv_heads, block_size, self.kv_lora_rank/self.qk_rope_head_dim] block_size = kv_c_and_k_pe_cache[0].shape[1] @@ -1112,6 +1108,7 @@ def forward( self.running_in_graph = self.torchair_graph_enabled and attn_metadata.attn_state in [ AscendAttentionState.DecodeOnly, AscendAttentionState.SpecDecoding ] + self.running_chunkprefilll_with_torchair = self.torchair_graph_enabled and attn_metadata.attn_state == AscendAttentionState.ChunkedPrefill num_actual_toks = attn_metadata.num_actual_tokens if k_pe is None and not self.running_in_graph: kv_c, k_pe = self.kv_a_proj_with_mqa( @@ -1148,18 +1145,25 @@ def forward( if has_decode: decode_k_nope = None assert attn_metadata.decode is not None - if self.running_in_graph: + if self.running_in_graph or self.running_chunkprefilll_with_torchair: cos = attn_metadata.decode.cos sin = attn_metadata.decode.sin - with npu_stream_switch("mla_secondary", - 0, - enabled=enable_multistream_mla): - npu_wait_tensor(hidden_states_or_kv_c_normed, - ckq, - enabled=enable_multistream_mla) + if self.running_chunkprefilll_with_torchair: + decode_hs = ( + hidden_states_or_kv_c_normed[:num_decode_tokens]) + slots = attn_metadata.slot_mapping[:num_decode_tokens] decode_k_pe, decode_k_nope, decode_kv = self.exec_kv( - hidden_states_or_kv_c_normed, cos, sin, kv_cache, - attn_metadata.slot_mapping) + decode_hs, cos, sin, kv_cache, slots) + else: + with npu_stream_switch("mla_secondary", + 0, + enabled=enable_multistream_mla): + npu_wait_tensor(hidden_states_or_kv_c_normed, + ckq, + enabled=enable_multistream_mla) + decode_k_pe, decode_k_nope, decode_kv = self.exec_kv( + hidden_states_or_kv_c_normed, cos, sin, kv_cache, + attn_metadata.slot_mapping) # Without explicitly controlling the order, IndexByTensor operations # would be placed after `matmul W_KV_T` hindering the overlapping of # KvRmsNormRopeCache and SingleRope. @@ -1183,6 +1187,8 @@ def forward( decode_k_pe, enabled=enable_multistream_mla) decode_q_pe = self.rope_single(decode_q_pe, cos, sin) + elif self.running_chunkprefilll_with_torchair: + decode_q_pe = self.rope_single(decode_q_pe, cos, sin) else: decode_q_pe[...], decode_k_pe[...] = self.rotary_emb( attn_metadata.decode.input_positions, @@ -1221,16 +1227,15 @@ def forward( kv_cache ) > 1, "the number of kv cache should be greater than 1, namely (nope_cache and rope_cache)" if self.torchair_graph_enabled: - if kv_cache[0].numel( - ) > 0 and attn_metadata.attn_state == AscendAttentionState.PrefillNoCache: + if kv_cache[0].numel() > 0 and has_prefill: slots = attn_metadata.slot_mapping # NOTE: Separate the kv cache in advance to avoid OOM or other issues - torch_npu._npu_reshape_and_cache(key=kv_c_normed.view( - num_tokens, self.num_kv_heads, -1), - value=prefill_k_pe, - key_cache=kv_cache[0], - value_cache=kv_cache[1], - slot_indices=slots) + torch_npu._npu_reshape_and_cache( + key=kv_c_normed.view(num_tokens, self.num_kv_heads, -1), + value=prefill_k_pe, + key_cache=kv_cache[0], + value_cache=kv_cache[1], + slot_indices=slots[num_decode_tokens:]) else: kv_c_normed = kv_c_normed.view( [num_actual_toks, self.num_kv_heads, -1]) @@ -1240,6 +1245,12 @@ def forward( key_cache=kv_cache[0], value_cache=kv_cache[1], slot_indices=attn_metadata.slot_mapping) + if not self.running_in_graph: + o_proj_input_shape = (num_actual_toks, + self.num_heads * self.v_head_dim) + o_proj_input = torch.empty(o_proj_input_shape, + dtype=hidden_states_or_q_c.dtype, + device=hidden_states_or_q_c.device) if has_prefill: # FIX: aicore move should be also placed on the comm stream in dbo, # otherwise it may affect the accuracy @@ -1250,11 +1261,12 @@ def forward( attn_metadata) current_ms_metadata = get_multistream_comm_context() if current_ms_metadata is not None: + current_ms_metadata.before_comm_event.record() with torch.npu.stream(current_ms_metadata.comm_stream): - output[num_decode_tokens:] = output_prefill - current_ms_metadata.after_comm_event.record() + current_ms_metadata.before_comm_event.wait() + o_proj_input[num_decode_tokens:] = output_prefill else: - output[num_decode_tokens:] = output_prefill + o_proj_input[num_decode_tokens:] = output_prefill if has_decode: if self.running_in_graph: @@ -1271,9 +1283,32 @@ def forward( current_ms_metadata = get_multistream_comm_context() if current_ms_metadata is not None: with torch.npu.stream(current_ms_metadata.comm_stream): - output[:num_decode_tokens] = output_decode - current_ms_metadata.after_comm_event.record() + o_proj_input[:num_decode_tokens] = output_decode else: - output[:num_decode_tokens] = output_decode + o_proj_input[:num_decode_tokens] = output_decode + current_ms_metadata = get_multistream_comm_context() + MAX_O_PROJ_PREFETCH_SIZE = 16 * 1024 * 1024 # 16MB + if current_ms_metadata is None: + npu_prefetch(self.o_proj.weight, + o_proj_input, + max_size=MAX_O_PROJ_PREFETCH_SIZE, + enabled=enable_multistream_mla) + + output[...] = self.o_proj( + o_proj_input, + is_prefill=True, + is_force_scatter=self.enable_shared_expert_dp)[0] + else: + with torch.npu.stream(current_ms_metadata.comm_stream): + npu_prefetch(self.o_proj.weight, + o_proj_input, + max_size=MAX_O_PROJ_PREFETCH_SIZE, + enabled=enable_multistream_mla) + output[...] = self.o_proj( + o_proj_input, + is_prefill=True, + is_force_scatter=self.enable_shared_expert_dp)[0] + current_ms_metadata.after_comm_event.record() + del o_proj_input return output_padded diff --git a/vllm_ascend/distributed/moe_comm_method.py b/vllm_ascend/distributed/moe_comm_method.py new file mode 100644 index 00000000000..f347ab06cb4 --- /dev/null +++ b/vllm_ascend/distributed/moe_comm_method.py @@ -0,0 +1,449 @@ +from abc import ABC, abstractmethod + +import torch +import torch_npu +from transformers.configuration_utils import PretrainedConfig +from vllm.distributed.parallel_state import get_ep_group, get_tp_group +from vllm.forward_context import ForwardContext, get_forward_context +from vllm.utils import direct_register_custom_op + +from vllm_ascend.distributed.parallel_state import get_mc2_group +from vllm_ascend.utils import AscendSocVersion, get_ascend_soc_version + + +class MoECommMethod(ABC): + """Base class for MoE communication methods.""" + + def __init__( + self, + device: torch.device, + dtype: torch.dtype, + hf_config: PretrainedConfig, + ): + self.device = device + self.dtype = dtype + self.top_k_num = getattr(hf_config, "num_experts_per_tok", 0) + # global_num_experts may be called num_experts or n_routed_experts in different models. + possible_keys = ["num_experts", "n_routed_experts"] + for key in possible_keys: + if hasattr(hf_config, key): + self.global_num_experts = getattr(hf_config, key) + break + else: + self.global_num_experts = 0 + + @abstractmethod + def _pre_process( + self, + hidden_states: torch.Tensor, + topk_ids: torch.Tensor, + topk_weights: torch.Tensor, + expert_map: torch.Tensor, + num_experts: int, + ) -> tuple[torch.Tensor, torch.Tensor, int]: + """Pre-process before MLP. + + Args: + hidden_states (torch.Tensor): Tensor of shape (num_tokens, hidden_size) + topk_ids (torch.Tensor): Tensor of shape (num_tokens, top_k_num) + topk_weights (torch.Tensor): Tensor of shape (num_tokens, top_k_num) + expert_map (torch.Tensor): Tensor of shape (global_num_experts, ) + Mapping from global expert IDs to local expert IDs. + num_experts (int): Number of local experts (experts on this device). + + Returns: + tuple[torch.Tensor, torch.Tensor, int]: Return a tuple containing: + - permuted_hidden_states (torch.Tensor): Tensor of shape + (num_tokens * top_k_num, hidden_size) after permuting + hidden_states based on topk_ids. + - expert_tokens (torch.Tensor): Tensor of shape (num_experts, ) + Number of tokens assigned to each expert. + - group_list_type (int): Type of group list, 0 for `cumsum` + and 1 for `count`. This is mainly for `npu_grouped_matmul` + to determine how to handle the output. + Raises: + NotImplementedError: If the method is not implemented in the subclass. + """ + pass + + @abstractmethod + def _post_process(self, mlp_output: torch.Tensor, + hidden_states: torch.Tensor) -> None: + """Post-process after MLP. + + Args: + mlp_output (torch.Tensor): Tensor of shape + (num_tokens * top_k_num, hidden_size) after MLP. + hidden_states (torch.Tensor): Tensor of shape + (num_tokens, hidden_size) to be updated with the final output. + """ + pass + + +class DummyCommImpl(MoECommMethod): + + def _pre_process( + self, + hidden_states: torch.Tensor, + topk_ids: torch.Tensor, + topk_weights: torch.Tensor, + expert_map: torch.Tensor, + num_experts: int, + ) -> tuple[torch.Tensor, torch.Tensor, int]: + """Dummy implementation, see moe_comm_pre_process_fake for details.""" + return moe_comm_pre_process_fake(hidden_states, topk_ids, topk_weights, + expert_map, num_experts) + + def _post_process(self, mlp_output: torch.Tensor, + hidden_states: torch.Tensor) -> None: + """Dummy implementation that does nothing.""" + pass + + +class NativeAllGatherCommImpl(MoECommMethod): + """This implementation should be compatible with all scenarios. + + Note that this implementation purely consists of native PyTorch ops + and does not use any NPU-specific ops. So the performance may not be optimal. + But it is a good fallback for scenarios where NPU-specific ops are not available. + """ + + def _pre_process( + self, + hidden_states: torch.Tensor, + topk_ids: torch.Tensor, + topk_weights: torch.Tensor, + expert_map: torch.Tensor, + num_experts: int, + ) -> tuple[torch.Tensor, torch.Tensor, int]: + num_tokens = hidden_states.shape[0] + + # Generate token indices and flatten + token_indices = torch.arange(num_tokens, + device=self.device, + dtype=torch.int64) + token_indices = (token_indices.unsqueeze(1).expand( + -1, self.top_k_num).reshape(-1)) + + # Flatten token-to-expert mappings and map to local experts + weights_flat = topk_weights.view(-1) + experts_flat = topk_ids.view(-1) + local_experts_flat = (expert_map[experts_flat] + if expert_map is not None else experts_flat) + + # Filter valid token-expert pairs + mask = local_experts_flat != -1 + # FIXME: npu_grouped_matmul output random values at [num_valid_tokens:, ...] + # So we need to filter out invalid tokens by zeroing their weights. + # This is a workaround and should be removed after the issue is fixed + filtered_weights = torch.where(mask, weights_flat, + torch.zeros_like(weights_flat)).to( + self.dtype) + filtered_experts = torch.where( + mask, + local_experts_flat, + torch.full_like(local_experts_flat, num_experts), + ).to(topk_ids.dtype) + + # Sort by local expert IDs + sort_indices = torch.argsort(filtered_experts.view(torch.float32)) + self.sorted_token_indices = token_indices[sort_indices] + self.sorted_weights = filtered_weights[sort_indices] + + # Compute token counts with minlength of num_experts + # This is equivalent to but faster than: + # >>> token_counts = torch.bincount(filtered_experts, minlength=num_experts)[:-1] + token_counts = torch.zeros(num_experts + 1, + device=self.device, + dtype=torch.int64) + ones = torch.ones_like(filtered_experts, dtype=torch.int64) + token_counts.scatter_add_(0, filtered_experts.to(torch.int64), ones) + expert_tokens = token_counts[:num_experts] + + # Rearrange hidden_states + permuted_hidden_states = hidden_states[self.sorted_token_indices] + + group_list_type = 1 # `count` mode + + return permuted_hidden_states, expert_tokens, group_list_type + + def _post_process(self, mlp_output: torch.Tensor, + hidden_states: torch.Tensor) -> None: + mlp_output = mlp_output * self.sorted_weights.unsqueeze(1) + + final_hidden_states = torch.zeros_like(hidden_states) + final_hidden_states.index_add_(0, self.sorted_token_indices, + mlp_output) + + hidden_states[:] = final_hidden_states + + +class AllGatherCommImpl(MoECommMethod): + """This implementation is the same as NativeAllGatherCommImpl, + but uses NPU-specific ops for better performance. + + This implementation should be compatible with all scenarios, and + thus it is the default implementation for MoE communication methods. + It uses `torch_npu.npu_moe_init_routing_v2` for pre-processing + and `torch_npu.npu_moe_token_unpermute` for post-processing + to handle the token-to-expert mapping and communication efficiently. + + NOTE(Yizhou): TBH, it is really weird that we were supposed to use + `torch_npu.npu_moe_init_routing_v2` and `torch_npu.npu_moe_finalize_routing` + or `torch_npu.npu_moe_token_permute` and `torch_npu.npu_moe_token_unpermute` + for pre-processing and post-processing, respectively. + But `npu_moe_finalize_routing` will lead to accuracy issues so we have to + use `torch_npu.npu_moe_token_unpermute` instead. + This is a workaround and should be removed after the issue is fixed. + """ + + def _pre_process( + self, + hidden_states: torch.Tensor, + topk_ids: torch.Tensor, + topk_weights: torch.Tensor, + expert_map: torch.Tensor, # noqa: F841 + num_experts: int, + ) -> tuple[torch.Tensor, torch.Tensor, int]: + num_tokens = hidden_states.shape[0] + + self.topk_weights = topk_weights + self.topk_ids = topk_ids + + first_expert_idx = 0 + if expert_map is not None: + # FIXME: npu_grouped_matmul output random values at [num_valid_tokens:, ...] + # So we need to filter out invalid tokens by zeroing their weights. + # This is a workaround and should be removed after the issue is fixed + mask = expert_map[topk_ids] != -1 + # NOTE: This is equivalent to self.topk_weights[~mask] = 0.0, + # but ~mask will dispatch to aclnnNonzeroV2, which is not supported in ACL Graph + self.topk_weights = torch.where(mask, topk_weights, 0.0) + + first_expert_idx = get_ep_group().rank_in_group * num_experts + last_expert_idx = first_expert_idx + num_experts + + permuted_hidden_states, expanded_row_idx, expert_tokens, _ = ( + torch_npu.npu_moe_init_routing_v2( + hidden_states, + topk_ids, + active_num=num_tokens * self.top_k_num, + expert_num=self.global_num_experts, + expert_tokens_num_type=1, # Only support `count` mode now + expert_tokens_num_flag=True, # Output `expert_tokens` + active_expert_range=[first_expert_idx, last_expert_idx], + quant_mode=-1, + )) + self.expanded_row_idx = expanded_row_idx + permuted_hidden_states = permuted_hidden_states + + group_list_type = 1 # `count` mode + + return permuted_hidden_states, expert_tokens, group_list_type + + def _post_process(self, mlp_output: torch.Tensor, + hidden_states: torch.Tensor) -> None: + hidden_states[:] = torch_npu.npu_moe_token_unpermute( + permuted_tokens=mlp_output, + sorted_indices=self.expanded_row_idx, + probs=self.topk_weights) + + +class MC2CommImpl(MoECommMethod): + """This implementation is for the scenarios listed below: + 1. `enable_expert_parallel=True`. + 2. `npu_moe_distribute_dispatch` and `npu_moe_distribute_combine` are available. + 3. `enable_expert_parallel=False` is not supported. + + This implementation uses the MC2 communication method, which is optimized for + Communication and Computation parallelism on Ascend devices. + """ + + def __init__( + self, + device: torch.device, + dtype: torch.dtype, + hf_config: PretrainedConfig, + ): + super().__init__(device, dtype, hf_config) + + # Shared communication configurations + ep_group = get_mc2_group() + self.ep_rank_id = ep_group.rank_in_group + self.ep_world_size = ep_group.world_size + self.tp_world_size = get_tp_group().world_size + + device_group = ep_group.device_group + local_rank = torch.distributed.get_rank(group=device_group) + backend = device_group._get_backend(torch.device("npu")) + self.moe_all_to_all_group_name = backend.get_hccl_comm_name(local_rank) + + # Feature flags + self.enable_dispatch_v2 = hasattr(torch_npu, + "npu_moe_distribute_dispatch_v2") + self.is_ascend_a3 = get_ascend_soc_version() == AscendSocVersion.A3 + self.need_extra_args = self.is_ascend_a3 # or is_torchair + + # Intermediate tensors to be passed from pre_process to post_process + self.topk_ids = None + self.topk_weights = None + self.mc2_mask = None + self.assist_info_for_combine = None + self.ep_recv_counts = None + self.tp_recv_counts = None + + def _pre_process( + self, + hidden_states: torch.Tensor, + topk_ids: torch.Tensor, + topk_weights: torch.Tensor, + expert_map: torch.Tensor, + num_experts: int, + ) -> tuple[torch.Tensor, torch.Tensor, int]: + # Store tensors needed for post_process + self.topk_ids = topk_ids + self.topk_weights = topk_weights.to(torch.float32) + self.mc2_mask = get_forward_context().mc2_mask + + dispatch_kwargs = { + "x": hidden_states, + "expert_ids": self.topk_ids, + "expert_shard_type": 0, + "shared_expert_rank_num": 0, + "moe_expert_num": self.global_num_experts, + "global_bs": 0, + "scales": None, + "quant_mode": 0, + "group_ep": self.moe_all_to_all_group_name, + "ep_world_size": self.ep_world_size, + "ep_rank_id": self.ep_rank_id, + } + + if self.need_extra_args: + dispatch_kwargs.update({ + "group_tp": self.moe_all_to_all_group_name, + "tp_world_size": 1, + "tp_rank_id": 0, + }) + if self.is_ascend_a3 and self.enable_dispatch_v2: + dispatch_kwargs.update({ + "x_active_mask": self.mc2_mask, + }) + + dispatch = torch_npu.npu_moe_distribute_dispatch_v2 if self.enable_dispatch_v2 else torch_npu.npu_moe_distribute_dispatch + + ( + permuted_hidden_states, + _, # dynamic_scale is not used + self.assist_info_for_combine, + expert_tokens, + self.ep_recv_counts, + self.tp_recv_counts, + ) = dispatch(**dispatch_kwargs)[:6] + + group_list_type = 1 + + return permuted_hidden_states, expert_tokens, group_list_type + + def _post_process(self, mlp_output: torch.Tensor, + hidden_states: torch.Tensor) -> None: + combine_kwargs = { + "expand_x": mlp_output, + "expert_ids": self.topk_ids, + "expert_scales": self.topk_weights, + "expert_shard_type": 0, + "shared_expert_rank_num": 0, + "moe_expert_num": self.global_num_experts, + "global_bs": 0, + "ep_send_counts": self.ep_recv_counts, + "group_ep": self.moe_all_to_all_group_name, + "ep_world_size": self.ep_world_size, + "ep_rank_id": self.ep_rank_id, + } + + if self.enable_dispatch_v2: + combine_kwargs[ + "assist_info_for_combine"] = self.assist_info_for_combine + else: + combine_kwargs["expand_idx"] = self.assist_info_for_combine + + if self.need_extra_args: + combine_kwargs.update({ + "tp_send_counts": self.tp_recv_counts, + "group_tp": self.moe_all_to_all_group_name, + "tp_world_size": 1, + "tp_rank_id": 0, + }) + if self.is_ascend_a3 and self.enable_dispatch_v2: + combine_kwargs.update({ + "x_active_mask": self.mc2_mask, + }) + + combine = torch_npu.npu_moe_distribute_combine_v2 if self.enable_dispatch_v2 else torch_npu.npu_moe_distribute_combine + + hidden_states[:] = combine(**combine_kwargs) + + +def moe_comm_pre_process( + hidden_states: torch.Tensor, + topk_ids: torch.Tensor, + topk_weights: torch.Tensor, + expert_map: torch.Tensor, + num_experts: int, +) -> tuple[torch.Tensor, torch.Tensor, int]: + """This function is a wrapper for the pre_process method of the + MoECommMethod instance stored in the ForwardContext. So it can be + used as a custom op in the vllm framework. + """ + forward_context: ForwardContext = get_forward_context() + self = forward_context.moe_comm_method + return self._pre_process(hidden_states, topk_ids, topk_weights, expert_map, + num_experts) + + +def moe_comm_pre_process_fake( + hidden_states: torch.Tensor, + topk_ids: torch.Tensor, + topk_weights: torch.Tensor, + expert_map: torch.Tensor, + num_experts: int, +) -> tuple[torch.Tensor, torch.Tensor, int]: + """This is a fake implementation of the pre_process method. + torch.compile will use this implementation to generate FX graph. + """ + top_k_num = topk_ids.shape[1] + permuted_hidden_states = hidden_states.repeat_interleave(top_k_num, dim=0) + expert_tokens = torch.zeros((num_experts, ), + dtype=torch.int64, + device=hidden_states.device) + group_list_type = 0 + return permuted_hidden_states, expert_tokens, group_list_type + + +def moe_comm_post_process(mlp_output: torch.Tensor, + hidden_states: torch.Tensor) -> None: + """This function is a wrapper for the post_process method of the + MoECommMethod instance stored in the ForwardContext. So it can be + used as a custom op in the vllm framework. + """ + forward_context: ForwardContext = get_forward_context() + self = forward_context.moe_comm_method + self._post_process(mlp_output, hidden_states) + return + + +direct_register_custom_op( + op_name="moe_comm_pre_process", + op_func=moe_comm_pre_process, + mutates_args=[], + fake_impl=moe_comm_pre_process_fake, + dispatch_key="PrivateUse1", +) + +direct_register_custom_op( + op_name="moe_comm_post_process", + op_func=moe_comm_post_process, + mutates_args=["hidden_states"], + fake_impl=lambda x, y: None, # No-op for fake implementation + dispatch_key="PrivateUse1", +) diff --git a/vllm_ascend/meta_registration.py b/vllm_ascend/meta_registration.py new file mode 100644 index 00000000000..600b5e74803 --- /dev/null +++ b/vllm_ascend/meta_registration.py @@ -0,0 +1,86 @@ +import torch +from torch.library import Library + +# This file provides a template and registration utilities for writing "meta" implementations +# of custom operators in Python for the vllm_ascend project. +# +# We offer two ways to implement meta implementations for custom ops: +# 1. Python meta implementation (as shown in this file): Write a Python function that +# takes the same arguments as your operator and returns empty tensors with the correct +# shapes and dtypes. This is useful for rapid prototyping and for ops that are only +# used in Python. +# 2. C++ meta implementation: You can also implement the meta function in C++ for better +# performance or to match the C++ op logic more closely. See `torch_binding_meta.cpp` +# for examples of C++ meta implementations and how to register them. +# +# Both approaches enable tracing, export, and shape inference in PyTorch and vLLM, which +# is essential for supporting `torch.compile` and aclgraph. + +# How to add a new meta implementation in Python: +# ------------------------------------- +# 1. Write a Python function that takes the same arguments as your operator, and returns +# empty tensors (using torch.empty_like, torch.empty, etc.) with the correct shapes and dtypes. +# Do NOT perform any real computation or allocate device memory. +# +# 2. Register your meta function using `register_meta_if_necessary`, providing: +# - The namespace (usually "_C" for custom ops) +# - The operator name (as registered in C++) +# - The Python meta function +# - (Optional) The overload name, if your op has overloads +# +# 3. The registration utility will check if a meta implementation already exists for your op, +# and only register if necessary. This avoids duplicate registrations. +# +# 4. Example meta implementations are provided below for rotary_embedding and get_masked_input_and_mask. +# +# 5. When developing new custom ops, always provide a meta implementation to enable tracing, +# export, and shape inference in PyTorch and vLLM to enable the capture of `torch.compile` +# and aclgraph. +# +# For more details, see: https://pytorch.org/docs/stable/notes/extending.html#meta-tensors + +lib = Library("_C", "IMPL") + + +def register_meta_if_necessary(ns: str, op_name: str, fn, overload: str = ""): + if overload != "": + op_name = op_name + "." + overload + schema_to_find = ns + "::" + op_name + meta_impl_list = torch._C._dispatch_get_registrations_for_dispatch_key( + "Meta") + if schema_to_find in meta_impl_list: + return + lib.impl(op_name, fn, "Meta") + + +def rotary_embedding_meta(positions: torch.Tensor, query: torch.Tensor, + key: torch.Tensor, head_size: int, + cos_sin_cache: torch.Tensor, is_neox: bool): + + num_tokens = positions.numel() + query_hidden_size = query.numel() // num_tokens + key_hidden_size = key.numel() // num_tokens + num_heads = query_hidden_size // head_size + num_kv_heads = key_hidden_size // head_size + + query_dst = torch.empty_like(query).view(num_tokens, num_heads, head_size) + key_dst = torch.empty_like(key).view(num_tokens, num_kv_heads, head_size) + return query_dst, key_dst + + +def get_masked_input_and_mask_meta(input: torch.Tensor, + org_vocab_start_index: int, + org_vocab_end_index: int, + num_org_vocab_padding: int, + added_vocab_start_index: int, + added_vocab_end_index: int): + + masked_input = torch.empty_like(input) + mask = torch.empty_like(input).to(torch.bool) + + return masked_input, mask + + +register_meta_if_necessary("_C", "rotary_embedding", rotary_embedding_meta) +register_meta_if_necessary("_C", "get_masked_input_and_mask", + get_masked_input_and_mask_meta) diff --git a/vllm_ascend/models/deepseek_v2.py b/vllm_ascend/models/deepseek_v2.py index ce051c4d846..0e4cf83374f 100644 --- a/vllm_ascend/models/deepseek_v2.py +++ b/vllm_ascend/models/deepseek_v2.py @@ -141,7 +141,8 @@ class CustomDeepseekV2RowParallelLinearReplaceAllreduce(RowParallelLinear): def forward( self, input_, - is_prefill=True + is_prefill=True, + is_force_scatter=False ) -> Union[torch.Tensor, tuple[torch.Tensor, Optional[nn.Parameter]]]: if self.input_is_parallel: input_parallel = input_ @@ -160,7 +161,13 @@ def forward( input_parallel, bias=bias_) if self.reduce_results and self.tp_size > 1: - if not is_prefill and output_parallel.shape[0] % self.tp_size == 0: + num_tokens = output_parallel.shape[0] + if is_force_scatter and num_tokens % self.tp_size: + output_parallel = nn.functional.pad( + output_parallel, (0, 0, 0, -num_tokens % self.tp_size)) + if is_force_scatter or (not is_prefill + and output_parallel.shape[0] % self.tp_size + == 0): output = tensor_model_parallel_reduce_scatter(output_parallel, dim=0) else: @@ -180,7 +187,8 @@ class CustomDeepseekV2RowParallelLinear(RowParallelLinear): def forward( self, input_, - is_prefill=True + is_prefill=True, + is_force_scatter=False ) -> Union[torch.Tensor, tuple[torch.Tensor, Optional[nn.Parameter]]]: if self.input_is_parallel: input_parallel = input_ @@ -347,13 +355,15 @@ def __init__( reduce_results = not self.all_reduce_merge intermediate_size = (config.moe_intermediate_size * config.n_shared_experts) + enable_shared_expert_dp = ascend_config.enable_shared_expert_dp self.shared_experts = CustomDeepseekV2MLP( hidden_size=config.hidden_size, intermediate_size=intermediate_size, hidden_act=config.hidden_act, quant_config=quant_config, reduce_results=reduce_results, - force_replicate=self.enable_multistream_moe, + force_replicate=self.enable_multistream_moe + or enable_shared_expert_dp, prefix=f"{prefix}.shared_experts", ) else: @@ -447,9 +457,11 @@ def __init__( self.kv_lora_rank = kv_lora_rank self.num_heads = num_heads - tp_size = get_tensor_model_parallel_world_size() - assert num_heads % tp_size == 0 - self.num_local_heads = num_heads // tp_size + self.tp_size = get_tensor_model_parallel_world_size() + assert num_heads % self.tp_size == 0 + self.num_local_heads = num_heads // self.tp_size + self.layers = config.num_hidden_layers + self.first_k_dense_replace = config.first_k_dense_replace self.scaling = self.qk_head_dim**-0.5 self.rope_theta = rope_theta @@ -462,6 +474,7 @@ def __init__( self.torchair_graph_enabled = ascend_config.torchair_graph_config.enabled self.enable_multistream_mla = \ ascend_config.torchair_graph_config.enable_multistream_mla + self.enable_shared_expert_dp = ascend_config.enable_shared_expert_dp if self.q_lora_rank is not None: self.q_a_proj = ReplicatedLinear(self.hidden_size, @@ -501,8 +514,9 @@ def __init__( prefix=f"{prefix}.kv_b_proj") if (config.n_routed_experts is not None and self.debug_layer_idx >= config.first_k_dense_replace - and self.debug_layer_idx % config.moe_layer_freq == 0 and - ascend_config.torchair_graph_config.enable_multistream_moe): + and self.debug_layer_idx % config.moe_layer_freq == 0 + and (ascend_config.torchair_graph_config.enable_multistream_moe + or self.enable_shared_expert_dp)): self.o_proj = CustomDeepseekV2RowParallelLinearReplaceAllreduce( self.num_heads * self.v_head_dim, self.hidden_size, @@ -596,13 +610,27 @@ def forward( output = output.view(-1, output_shape[-1]) return output else: - kv_c, k_pe = self.kv_a_proj_with_mqa(hidden_states)[0].split( + kv_no_split = self.kv_a_proj_with_mqa(hidden_states)[0] + if self.enable_shared_expert_dp and self.debug_layer_idx > self.first_k_dense_replace and self.debug_layer_idx < self.layers: + hidden_states_or_q_c = get_tp_group().all_gather( + hidden_states_or_q_c, 0) + kv_no_split = get_tp_group().all_gather(kv_no_split, 0) + + kv_c, k_pe = kv_no_split.split( [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1) kv_c_normed = self.kv_a_layernorm(kv_c.contiguous()) + if not self.enable_shared_expert_dp or self.debug_layer_idx < self.first_k_dense_replace: + output_shape = hidden_states.shape + else: + num_tokens = hidden_states_or_q_c.shape[0] + rows = num_tokens // self.tp_size + if num_tokens % self.tp_size: + rows += 1 + output_shape = (rows, hidden_states.shape[1]) return self.mla_attn(hidden_states_or_q_c, kv_c_normed, k_pe, - output_shape=hidden_states.shape) + output_shape=output_shape) class CustomDeepseekV2DecoderLayer(DeepseekV2DecoderLayer): @@ -677,6 +705,8 @@ def __init__( eps=config.rms_norm_eps) self.routed_scaling_factor = config.routed_scaling_factor self.first_k_dense_replace = config.first_k_dense_replace + self.tp_group = get_tp_group().device_group + self.enable_shared_expert_dp = ascend_config.enable_shared_expert_dp def forward( self, @@ -731,6 +761,18 @@ def forward( # first layer. residual *= 1. / self.routed_scaling_factor + tp_size = get_tensor_model_parallel_world_size() + if self.enable_shared_expert_dp and ( + self.layer_idx == self.first_k_dense_replace + or self.layer_idx == self.layers) and tp_size > 1: + num_tokens, _ = residual.shape + if num_tokens % tp_size: + residual = nn.functional.pad(residual, + (0, 0, 0, -num_tokens % tp_size)) + chunk_residual = torch.tensor_split(residual, tp_size, dim=0) + tp_rank = get_tensor_model_parallel_rank() + residual = chunk_residual[tp_rank] + # Fully Connected hidden_states, residual = self.post_attention_layernorm( hidden_states, residual) @@ -756,6 +798,22 @@ def forward( dim=0) residual = tensor_model_parallel_all_gather(residual, dim=0) + # for last layer of main model and mtp layer. + if self.enable_shared_expert_dp and self.layer_idx >= ( + self.layers - 1) and tp_size > 1: + hidden_states = get_tp_group().all_gather(hidden_states, 0) + residual = get_tp_group().all_gather(residual, 0) + + attn_metadata = get_forward_context().attn_metadata + if attn_metadata is not None: + num_tokens = attn_metadata.num_actual_tokens + else: + num_tokens = hidden_states.shape[0] + + if num_tokens < hidden_states.shape[0]: + hidden_states = hidden_states[:num_tokens] + residual = residual[:num_tokens] + return hidden_states, residual diff --git a/vllm_ascend/ops/common_fused_moe.py b/vllm_ascend/ops/common_fused_moe.py index eeb8ec32237..b97aef7de11 100644 --- a/vllm_ascend/ops/common_fused_moe.py +++ b/vllm_ascend/ops/common_fused_moe.py @@ -19,12 +19,13 @@ import torch from vllm.config import CompilationLevel, get_current_vllm_config +from vllm.forward_context import get_forward_context from vllm.model_executor.layers.fused_moe.layer import \ UnquantizedFusedMoEMethod from vllm_ascend.ascend_config import get_ascend_config -from vllm_ascend.ops.fused_moe import (fused_experts, fused_experts_moge, - select_experts) +from vllm_ascend.ops.fused_moe import (fused_experts_moge, select_experts, + unified_fused_experts) from vllm_ascend.utils import is_310p original_unquantized_fused_moe_init_func = UnquantizedFusedMoEMethod.__init__ @@ -95,20 +96,18 @@ def forward_oot( expert_map=expert_map, apply_router_weight_on_input=apply_router_weight_on_input) - # If use aclgraph, we need to set max_num_tokens to make - # the input shape of `npu_moe_init_routing` fixed - max_num_tokens = self.max_num_batched_tokens if self.use_aclgraph else None + moe_comm_method = get_forward_context().moe_comm_method - return fused_experts( + return unified_fused_experts( hidden_states=x, w1=layer.w13_weight, w2=layer.w2_weight, topk_weights=topk_weights, topk_ids=topk_ids, - top_k=top_k, + global_num_experts=global_num_experts, expert_map=expert_map, - apply_router_weight_on_input=apply_router_weight_on_input, - max_num_tokens=max_num_tokens) + moe_comm_method=moe_comm_method, + ) UnquantizedFusedMoEMethod.__init__ = unquantized_fused_moe_init_func diff --git a/vllm_ascend/ops/fused_moe.py b/vllm_ascend/ops/fused_moe.py index f35fb105758..aeb75cfa0df 100644 --- a/vllm_ascend/ops/fused_moe.py +++ b/vllm_ascend/ops/fused_moe.py @@ -43,6 +43,7 @@ from vllm_ascend.ascend_forward_context import FusedMoEState from vllm_ascend.distributed.communication_op import \ data_parallel_reduce_scatter +from vllm_ascend.distributed.moe_comm_method import MoECommMethod from vllm_ascend.distributed.parallel_state import get_mc2_group from vllm_ascend.ops.expert_load_balancer import ExpertLoadBalancer from vllm_ascend.ops.moe_dispatcher.token_dispatcher import ( @@ -57,6 +58,62 @@ MOE_ALL2ALL_BUFFER: bool = envs_ascend.MOE_ALL2ALL_BUFFER +def unified_fused_experts( + hidden_states: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + activation: str = "silu", + apply_router_weight_on_input: bool = False, + use_int8_w8a8: bool = False, + use_int4_w4a8: bool = False, + global_num_experts: Optional[int] = None, + expert_map: Optional[torch.Tensor] = None, + w1_scale: Optional[torch.Tensor] = None, + w2_scale: Optional[torch.Tensor] = None, + w1_scale_bias: torch.Tensor = None, + w2_scale_bias: torch.Tensor = None, + moe_comm_method: Optional[MoECommMethod] = None, + # For TorchAir graph + is_torchair: bool = False, + # For Cube/Vector parallel + shared_experts: Optional[Any] = None, + quantized_x_for_share: Optional[Any] = None, + dynamic_scale_for_share: Optional[Any] = None, + # For load balance + log2phy: torch.Tensor = None, + global_redundant_expert_num: int = 0, +) -> torch.Tensor: + # Check constraints + assert hidden_states.shape[1] == w1.shape[2], ( + f"Hidden size mismatch {hidden_states.shape[1]} != {w1.shape[2]}") + + assert topk_weights.shape == topk_ids.shape, "topk shape mismatch" + assert hidden_states.is_contiguous(), "Hidden_states must be contiguous" + assert w1.stride(-1) == 1, "Stride of last dimension must be 1" + assert w2.stride(-1) == 1, "Stride of last dimension must be 1" + assert hidden_states.dtype in [ + torch.float32, torch.float16, torch.bfloat16 + ] + assert moe_comm_method is not None, "Missing communication context" + + num_experts = w1.shape[0] + + permuted_hidden_states, expert_tokens, group_list_type = torch.ops.vllm.moe_comm_pre_process( + hidden_states, topk_ids, topk_weights, expert_map, num_experts) + mlp_output = apply_mlp( + permuted_hidden_states, + w1, + w2, + expert_tokens, + group_list_type=group_list_type, + ) + torch.ops.vllm.moe_comm_post_process(mlp_output, hidden_states) + + return hidden_states + + def process_topk_ids(topk_ids: torch.Tensor, expert_num: int, ep_size: int, max_row_per_ep_rank: int, num_tokens: int, top_k: int) -> tuple[torch.Tensor, torch.Tensor]: @@ -205,11 +262,9 @@ def fused_experts_with_mc2( group_list_type=1, group_type=0, group_list=group_list, - ) + )[0] - # TODO: Remove this in the future. - gate_up_out = torch.cat(gate_up_out_list, dim=0) - gate_up_out = torch_npu.npu_swiglu(gate_up_out) + gate_up_out = torch_npu.npu_swiglu(gate_up_out_list) w2 = w2.transpose(1, 2) down_out_list = torch_npu.npu_grouped_matmul( @@ -219,9 +274,7 @@ def fused_experts_with_mc2( group_list_type=1, group_type=0, group_list=group_list, - ) - - down_out_list = torch.cat(down_out_list, dim=0) + )[0] # moeCombine kwargs_mc2 = { @@ -312,9 +365,8 @@ def apply_mlp( group_list_type=group_list_type, group_type=0, group_list=group_list, - ) + )[0] - hidden_states = torch.cat(hidden_states, dim=0) hidden_states = torch_npu.npu_swiglu(hidden_states) w2 = w2.transpose(1, 2) @@ -325,9 +377,8 @@ def apply_mlp( group_list_type=group_list_type, group_type=0, group_list=group_list, - ) + )[0] - hidden_states = torch.cat(hidden_states, dim=0) return hidden_states @@ -417,23 +468,19 @@ def fused_experts_with_all2all( group_list_type=0, group_type=0, group_list=expert_tokens, - ) + )[0] - # TODO: Remove this in the future. - hidden_states = torch.cat(gate_up_out_list, dim=0) - hidden_states = torch_npu.npu_swiglu(hidden_states) + hidden_states = torch_npu.npu_swiglu(gate_up_out_list) w2 = w2.transpose(1, 2) - down_out_list = torch_npu.npu_grouped_matmul( + hidden_states = torch_npu.npu_grouped_matmul( x=[hidden_states], weight=[w2], split_item=2, group_list_type=0, group_type=0, group_list=expert_tokens, - ) - - hidden_states = torch.cat(down_out_list, dim=0) + )[0] if expert_map is not None: resorted_idx = torch.argsort(sorted_idx) @@ -823,11 +870,9 @@ def fused_experts( group_list_type=0, group_type=0, group_list=expert_tokens, - ) + )[0] - # TODO: Remove this in the future. - gate_up_out = torch.cat(gate_up_out_list, dim=0) - gate_up_out = torch_npu.npu_swiglu(gate_up_out) + gate_up_out = torch_npu.npu_swiglu(gate_up_out_list) w2 = w2.transpose(1, 2) down_out_list = torch_npu.npu_grouped_matmul( @@ -837,9 +882,7 @@ def fused_experts( group_list_type=0, group_type=0, group_list=expert_tokens, - ) - - down_out_list = torch.cat(down_out_list, dim=0) + )[0] if expert_map is not None: weighted_down_out = down_out_list * sorted_weights.unsqueeze(1) @@ -1195,8 +1238,27 @@ def __init__( ): # TODO: This could not initialize FusedMoE baseclass, # fixme and make __init__() of AscendFusedMoE more clear - super(FusedMoE, self).__init__() - + super().__init__( + num_experts=num_experts, + top_k=top_k, + hidden_size=hidden_size, + intermediate_size=intermediate_size, + params_dtype=params_dtype, + reduce_results=reduce_results, + renormalize=renormalize, + use_grouped_topk=use_grouped_topk, + num_expert_group=num_expert_group, + topk_group=topk_group, + quant_config=quant_config, + tp_size=tp_size, + ep_size=ep_size, + dp_size=dp_size, + prefix=prefix, + custom_routing_function=custom_routing_function, + scoring_func=scoring_func, + e_score_correction_bias=e_score_correction_bias, + activation=activation, + ) AscendFusedMoE.moe_counter += 1 self.moe_instance_id = AscendFusedMoE.moe_counter @@ -1263,6 +1325,7 @@ def __init__( self.enable_multistream_moe = \ ascend_config.torchair_graph_config.enable_multistream_moe and \ self.torchair_graph_enabled + self.enable_shared_expert_dp = ascend_config.enable_shared_expert_dp if self.scoring_func != "softmax" and not self.use_grouped_topk: raise ValueError("Only softmax scoring function is supported for " @@ -1403,22 +1466,24 @@ def forward(self, else: # TODO: Determine if we can remove the padding padding_size = tp_size - if num_tokens < padding_size: + if num_tokens < padding_size and not self.enable_shared_expert_dp: hidden_states = nn.functional.pad( hidden_states, (0, 0, 0, padding_size - num_tokens)) router_logits = nn.functional.pad( router_logits, (0, 0, 0, padding_size - num_tokens)) if tp_size > 1: - chunk_hidden_states = torch.tensor_split(hidden_states, - tp_size, - dim=0) - chunk_router_logits = torch.tensor_split(router_logits, - tp_size, - dim=0) - chunk_mc2_mask = torch.tensor_split(mc2_mask, tp_size, dim=0) tp_rank = get_tensor_model_parallel_rank() - hidden_states = chunk_hidden_states[tp_rank] - router_logits = chunk_router_logits[tp_rank] + if not self.enable_shared_expert_dp: + chunk_hidden_states = torch.tensor_split(hidden_states, + tp_size, + dim=0) + chunk_router_logits = torch.tensor_split(router_logits, + tp_size, + dim=0) + hidden_states = chunk_hidden_states[tp_rank] + router_logits = chunk_router_logits[tp_rank] + + chunk_mc2_mask = torch.tensor_split(mc2_mask, tp_size, dim=0) mc2_mask = chunk_mc2_mask[tp_rank] if self.dp_size > 1: @@ -1485,7 +1550,7 @@ def forward(self, if (fused_moe_state not in [ FusedMoEState.AllGather, FusedMoEState.AllGatherEP, FusedMoEState.NaiveMulticast - ] and not replace_allreduce): + ] and not replace_allreduce and not self.enable_shared_expert_dp): if tp_size > 1: dist.all_gather(list(chunk_hidden_states), e_hidden_states, self.tp_group) @@ -1495,7 +1560,7 @@ def forward(self, final_hidden_states = e_hidden_states if num_tokens < padding_size: final_hidden_states = final_hidden_states[:num_tokens] - elif self.dp_size > 1: + elif self.dp_size > 1 and not self.enable_shared_expert_dp: if fused_moe_state == FusedMoEState.NaiveMulticast: start = 0 if self.dp_rank == 0 else cu_tokens_across_dp_cpu[ self.dp_rank - 1] diff --git a/vllm_ascend/patch/worker/patch_common/patch_linear.py b/vllm_ascend/patch/worker/patch_common/patch_linear.py index f5fbcecb770..57cc4e0b58a 100644 --- a/vllm_ascend/patch/worker/patch_common/patch_linear.py +++ b/vllm_ascend/patch/worker/patch_common/patch_linear.py @@ -25,6 +25,7 @@ from vllm.distributed import (get_tensor_model_parallel_rank, split_tensor_along_last_dim) from vllm.distributed.parallel_state import get_tp_group +from vllm.logger import logger from vllm.model_executor.layers.linear import RowParallelLinear from vllm_ascend import envs @@ -142,4 +143,5 @@ def calc_output(self, input_parallel: torch.Tensor) -> torch.Tensor: if envs.VLLM_ASCEND_ENABLE_MATMUL_ALLREDUCE: + logger.info("AscendRowParallelLinear: Matmul all-reduce is enabled. ") vllm.model_executor.layers.linear.RowParallelLinear = AscendRowParallelLinear diff --git a/vllm_ascend/platform.py b/vllm_ascend/platform.py index eb7ea8276cc..f101ccdc7a3 100644 --- a/vllm_ascend/platform.py +++ b/vllm_ascend/platform.py @@ -205,8 +205,15 @@ def check_and_update_config(cls, vllm_config: VllmConfig) -> None: register_ascend_customop() @classmethod - def get_attn_backend_cls(cls, selected_backend, head_size, dtype, - kv_cache_dtype, block_size, use_v1, use_mla): + def get_attn_backend_cls(cls, + selected_backend, + head_size, + dtype, + kv_cache_dtype, + block_size, + use_v1, + use_mla, + has_sink=False): if not use_v1: raise ValueError("vLLM Ascend does not support V0 engine.") diff --git a/vllm_ascend/sample/rejection_sampler.py b/vllm_ascend/sample/rejection_sampler.py index 832f0179dd7..e0d770df26e 100644 --- a/vllm_ascend/sample/rejection_sampler.py +++ b/vllm_ascend/sample/rejection_sampler.py @@ -147,16 +147,25 @@ def rejection_sample( if not sampling_metadata.all_random: # Rejection sampling for greedy sampling requests. target_argmax = target_probs.argmax(dim=-1) - rejection_greedy_sample_pytorch( - output_token_ids, - cu_num_draft_tokens, - draft_token_ids, - target_argmax, - bonus_token_ids, - is_greedy, - max_spec_len, - # num_warps=1, - ) + if min(num_draft_tokens) == 1 and max( + num_draft_tokens) == 1 and sampling_metadata.all_greedy: + rejection_greedy_sample_spec_len_1_pytorch( + output_token_ids, + draft_token_ids, + target_argmax, + bonus_token_ids, + ) + else: + rejection_greedy_sample_pytorch( + output_token_ids, + cu_num_draft_tokens, + draft_token_ids, + target_argmax, + bonus_token_ids, + num_draft_tokens, + max_spec_len, + is_greedy, + ) if sampling_metadata.all_greedy: return output_token_ids @@ -284,47 +293,89 @@ def sample_recovered_tokens( return recovered_token_ids -def rejection_greedy_sample_pytorch( - output_token_ids, # [batch_size, max_spec_len + 1] - cu_num_draft_tokens, # [batch_size] - draft_token_ids, # [num_tokens] - target_argmax, # [num_tokens] - bonus_token_ids, # [batch_size] - is_greedy=None, # [batch_size] or None - max_spec_len=None, +def rejection_greedy_sample_spec_len_1_pytorch( + output_token_ids, # [batch_size, 2] + draft_token_ids, # [num_tokens] + target_argmax, # [num_tokens] + bonus_token_ids, # [batch_size] ): - batch_size = output_token_ids.shape[0] - - if is_greedy is None: - is_greedy = torch.ones(batch_size, - dtype=torch.bool, - device=output_token_ids.device) - - for req_idx in range(batch_size): - if not is_greedy[req_idx]: - continue - - if req_idx == 0: - start_idx = 0 - else: - start_idx = cu_num_draft_tokens[req_idx - 1].item() - end_idx = cu_num_draft_tokens[req_idx].item() - num_draft_tokens = end_idx - start_idx - - rejected = False - for pos in range(num_draft_tokens): - if not rejected: - draft_token_id = draft_token_ids[start_idx + pos].item() - target_argmax_id = target_argmax[start_idx + pos].item() - - output_token_ids[req_idx, pos] = target_argmax_id + batch_size = output_token_ids.size(0) + num_tokens = draft_token_ids.size(0) + assert batch_size == num_tokens + accept_req_mask = draft_token_ids == target_argmax + output_token_ids[:, 0] = target_argmax + bonus_token_ids = bonus_token_ids.squeeze(1) + output_token_ids[accept_req_mask, 1] = bonus_token_ids[accept_req_mask] - if draft_token_id != target_argmax_id: - rejected = True - if not rejected: - bonus_token_id = bonus_token_ids[req_idx].item() - output_token_ids[req_idx, num_draft_tokens] = bonus_token_id +def rejection_greedy_sample_pytorch( + output_token_ids, # [batch_size, max_spec_len + 1] + cu_num_draft_tokens, # [batch_size] + draft_token_ids, # [num_tokens] + target_argmax, # [num_tokens] + bonus_token_ids, # [batch_size] + draft_tokens_per_req, # [batch_size], list + max_spec_len, + is_greedy=None, # [batch_size] or None +): + batch_size = output_token_ids.size(0) + num_tokens = draft_token_ids.size(0) + device = output_token_ids.device + draft_tokens_per_req = torch.tensor(draft_tokens_per_req).to( + device, non_blocking=True) + if is_greedy is None: + is_greedy = torch.ones(batch_size, dtype=torch.bool, device=device) + + start_indices = cu_num_draft_tokens - draft_tokens_per_req + req_ids = torch.arange(batch_size, device=device) + token_req_ids = torch.repeat_interleave(req_ids, draft_tokens_per_req) + token_positions = torch.arange( + num_tokens, device=device) - start_indices[token_req_ids] + + # Find the first mismatch position of each request. + mismatch_global = (draft_token_ids != target_argmax) + if max_spec_len == 0: + first_mismatch_pos_per_req = torch.zeros(batch_size, + dtype=torch.long, + device=device) + else: + # [bs, max_spec_len] + pos_matrix = torch.full((batch_size, max_spec_len), + -1, + dtype=torch.long, + device=device) + pos_matrix[token_req_ids, token_positions] = token_positions + mismatch_matrix = torch.full((batch_size, max_spec_len), + False, + dtype=torch.bool, + device=device) + mismatch_matrix[token_req_ids, token_positions] = mismatch_global + mismatch_positions = torch.where(mismatch_matrix, pos_matrix, + max_spec_len * 2) + first_mismatch_pos_per_req, _ = torch.min(mismatch_positions, dim=1) + no_mismatch_mask = (first_mismatch_pos_per_req == max_spec_len * 2) + first_mismatch_pos_per_req[no_mismatch_mask] = draft_tokens_per_req[ + no_mismatch_mask] + + # Copy matched target tokens into output. + copy_len = torch.minimum(first_mismatch_pos_per_req + 1, + draft_tokens_per_req) + copy_indices = torch.arange(max_spec_len + 1, + device=device).expand(batch_size, -1) + copy_mask = copy_indices < copy_len.unsqueeze(1) + greedy_mask = is_greedy.unsqueeze(1) + final_copy_mask = copy_mask & greedy_mask + global_idx = start_indices.unsqueeze(1) + copy_indices + output_token_ids[final_copy_mask] = target_argmax[ + global_idx[final_copy_mask]].to(output_token_ids.dtype) + # Fill bonus token. + needs_bonus = is_greedy & (first_mismatch_pos_per_req + >= draft_tokens_per_req) + if torch.any(needs_bonus): + bonus_rows = torch.where(needs_bonus)[0] + bonus_cols = draft_tokens_per_req[bonus_rows] + bonus_token_ids = bonus_token_ids.squeeze(1) + output_token_ids[bonus_rows, bonus_cols] = bonus_token_ids[bonus_rows] def rejection_random_sample_pytorch( diff --git a/vllm_ascend/torchair/torchair_model_runner.py b/vllm_ascend/torchair/torchair_model_runner.py index 200167438b4..f42f83d1583 100644 --- a/vllm_ascend/torchair/torchair_model_runner.py +++ b/vllm_ascend/torchair/torchair_model_runner.py @@ -17,9 +17,19 @@ # Adapted from vllm-project/vllm/vllm/worker/gpu_model_runner.py # +from typing import Optional + import torch +import torch_npu from vllm.config import VllmConfig +from vllm.forward_context import get_forward_context +from vllm.logger import logger +from vllm_ascend.platform import NPUPlatform +from vllm_ascend.torchair.utils import (check_torchair_cache_exist, + write_kv_cache_bytes_to_file) +from vllm_ascend.utils import (ACL_FORMAT_FRACTAL_ND, ACL_FORMAT_FRACTAL_NZ, + maybe_converting_weight_acl_format) from vllm_ascend.worker.model_runner_v1 import NPUModelRunner @@ -27,3 +37,135 @@ class NPUTorchairModelRunner(NPUModelRunner): def __init__(self, vllm_config: VllmConfig, device: torch.device): super().__init__(vllm_config, device) + + def _get_forward_metadata_across_dp_and_pad( + self, num_tokens: int, with_prefill: bool, enable_dbo: bool + ) -> tuple[int, Optional[torch.Tensor], bool, bool]: + """Override from NPUModelRunner to pad num_tokens""" + if self.dp_size == 1: + if not with_prefill: + maybe_padded_num_tokens = self.select_torchair_padded_batch_size( + num_tokens) + return maybe_padded_num_tokens, None, with_prefill, enable_dbo + return num_tokens, None, with_prefill, enable_dbo + + num_tokens_across_dp, with_prefill, enable_dbo = self._get_forward_metadata_across_dp( + num_tokens, with_prefill, enable_dbo) + + if not with_prefill: + max_num_token = num_tokens_across_dp.max().item() + maybe_padded_num_tokens = self.select_torchair_padded_batch_size( + max_num_token) + num_tokens_across_dp = torch.full((self.dp_size, ), + maybe_padded_num_tokens, + dtype=torch.int32, + device="cpu") + else: + maybe_padded_num_tokens = num_tokens + + return maybe_padded_num_tokens, num_tokens_across_dp, with_prefill, enable_dbo + + def _build_attention_metadata(self, with_prefill, num_reqs, skip_attn): + # NOTE: If torchair graph mode and not with_prefill, + # we can't skip_attn, it will cause graph recompile. + if not with_prefill: + attn_metadata = self.attn_metadata_builder.build_torchair_graph_dummy( + num_reqs=num_reqs, num_actual_tokens=1) + else: + attn_metadata = super()._build_attention_metadata( + with_prefill, num_reqs, skip_attn) + return attn_metadata + + def _generate_dummy_run_hidden_states(self, with_prefill, + is_torchair_compile, input_ids, + positions, attn_metadata, num_tokens, + intermediate_tensors, inputs_embeds): + + if not with_prefill: + # Only mark static while compiling + if is_torchair_compile: + torch._dynamo.mark_static(input_ids) + torch._dynamo.mark_static(positions) + torch._dynamo.mark_static(attn_metadata.decode.block_table) + torch._dynamo.mark_static(attn_metadata.decode.input_positions) + torch._dynamo.mark_static(get_forward_context().mc2_mask) + if hasattr(attn_metadata.decode, "sin"): + torch._dynamo.mark_static(attn_metadata.decode.sin) + torch._dynamo.mark_static(attn_metadata.decode.cos) + torch._dynamo.mark_static(attn_metadata.slot_mapping) + if self.speculative_config: + torch._dynamo.mark_static(attn_metadata.decode.attn_mask) + for kv in self.kv_caches: + assert isinstance(kv, tuple), "kv_cache must be a tuple" + torch._dynamo.mark_static(kv[0]) + torch._dynamo.mark_static(kv[1]) + + maybe_converting_weight_acl_format(self.model, + ACL_FORMAT_FRACTAL_NZ) + + compiled_model = self._get_torchair_lazy_compiled_model(num_tokens) + model_kwargs = {} + model_kwargs["kv_caches"] = self.kv_caches + model_kwargs["attn_metadata"] = attn_metadata + hidden_states = compiled_model( + input_ids=input_ids, + positions=positions, + intermediate_tensors=intermediate_tensors, + inputs_embeds=None, + **model_kwargs, + ) + else: + hidden_states = super()._generate_dummy_run_hidden_states( + with_prefill, is_torchair_compile, input_ids, positions, + attn_metadata, num_tokens, intermediate_tensors, inputs_embeds) + return hidden_states + + def _convert_torch_format(self, kv_cache): + kv_cache = torch_npu.npu_format_cast(kv_cache, ACL_FORMAT_FRACTAL_ND) + return kv_cache + + def _compile_torchair_graph(self, torchair_graph_batch_sizes) -> None: + # Trigger torchair graph capture for specific shapes. + # Capture the large shapes first so that the smaller shapes + # can reuse the memory pool allocated for the large shapes. + for idx, num_tokens in enumerate(reversed(torchair_graph_batch_sizes)): + for _ in range(self.vllm_config.compilation_config. + cudagraph_num_of_warmups): + self._dummy_run(num_tokens, is_torchair_compile=True) + self._dummy_run(num_tokens, is_torchair_compile=True) + logger.info("Batchsize %d is compiled successfully: %d/%d.", + num_tokens, idx + 1, len(torchair_graph_batch_sizes)) + + def _capture_model(self): + """Override from NPUModelRunner to use torchair graph capture.""" + # TODO(NeverRaR): Calling graph_capture(device=self.device) in + # torchair graph capture can cause some issues, so now we just + # temporarily split the codepath for the two different graph patterns. + torchair_graph_batch_sizes = self.torchair_graph_batch_sizes + graph_num = len(torchair_graph_batch_sizes) + + if self.use_cached_npu_graph and not check_torchair_cache_exist(): + # If caching is enabled but does not exist, we will compile the model twice. The first + # time is used to generate the cache, and the second time is used to load the cache to + # skip the overhead caused by Dynamo guard mechanism. + logger.info( + "Use cached npu graph but cache doesn't exist! Now we compile graph to genetate torchair cache, this usually takes %.1f~%.1f mins.", + 0.5 * graph_num, 1.5 * graph_num) + self._compile_torchair_graph(torchair_graph_batch_sizes) + NPUPlatform.synchronize() + torch._dynamo.reset() + self.torchair_compiled_models.clear() + if self.use_cached_npu_graph: + logger.info( + "Loading torchair graph cache, this usually takes %.1f~%.1f mins.", + 0.3 * graph_num, 0.5 * graph_num) + self._compile_torchair_graph(torchair_graph_batch_sizes) + else: + logger.info( + "Capturing torchair graph, this usually takes %.1f~%.1f mins.", + 0.5 * graph_num, 1.5 * graph_num) + self._compile_torchair_graph(torchair_graph_batch_sizes) + + if self.new_kv_cache_bytes > 0: + write_kv_cache_bytes_to_file(torch.distributed.get_rank(), + self.new_kv_cache_bytes) diff --git a/vllm_ascend/utils.py b/vllm_ascend/utils.py index ee620b4bb99..7c0f77f4f81 100644 --- a/vllm_ascend/utils.py +++ b/vllm_ascend/utils.py @@ -214,8 +214,12 @@ def enable_custom_op(): if _CUSTOM_OP_ENABLED is not None: return _CUSTOM_OP_ENABLED try: + # isort: off # register custom ops into torch_library here import vllm_ascend.vllm_ascend_C # type: ignore # noqa: F401 + # register the meta implementation for custom kernel if necessary + import vllm_ascend.meta_registration # type: ignore # noqa: F401 + # isort: on _CUSTOM_OP_ENABLED = True except ImportError: _CUSTOM_OP_ENABLED = False diff --git a/vllm_ascend/worker/model_runner_v1.py b/vllm_ascend/worker/model_runner_v1.py index d7944b8d74c..9891a029f07 100644 --- a/vllm_ascend/worker/model_runner_v1.py +++ b/vllm_ascend/worker/model_runner_v1.py @@ -26,7 +26,7 @@ import weakref from contextlib import contextmanager, nullcontext from dataclasses import dataclass -from typing import TYPE_CHECKING, Dict, List, Optional, Union, cast +from typing import TYPE_CHECKING, Dict, List, Optional, Type, Union, cast import numpy as np import numpy.typing as npt @@ -43,7 +43,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1 import KVConnectorBase_V1 from vllm.distributed.parallel_state import (get_dp_group, get_pp_group, get_tp_group) -from vllm.forward_context import get_forward_context +from vllm.forward_context import DPMetadata, get_forward_context from vllm.logger import logger from vllm.model_executor.layers.fused_moe import FusedMoE from vllm.model_executor.layers.rotary_embedding import MRotaryEmbedding @@ -79,11 +79,12 @@ AscendMetadata) from vllm_ascend.attention.attention_v1_torchair import AscendTorchairMetadata from vllm_ascend.attention.mla_v1 import AscendMLAMetadata +from vllm_ascend.distributed.moe_comm_method import (AllGatherCommImpl, + DummyCommImpl, + MoECommMethod) from vllm_ascend.multistream.ms_split import compute_split_seq_index from vllm_ascend.platform import NPUPlatform from vllm_ascend.sample.rejection_sampler import AscendRejectionSampler -from vllm_ascend.torchair.utils import (check_torchair_cache_exist, - write_kv_cache_bytes_to_file) from vllm_ascend.utils import (ACL_FORMAT_FRACTAL_ND, ACL_FORMAT_FRACTAL_NZ, ProfileExecuteDuration, is_310p, maybe_converting_weight_acl_format, @@ -110,6 +111,9 @@ if is_310p(): torch_npu.npu.set_compile_mode(jit_compile=False) + ACL_FORMAT = ACL_FORMAT_FRACTAL_NZ +else: + ACL_FORMAT = ACL_FORMAT_FRACTAL_ND @dataclass @@ -334,7 +338,8 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): self.use_aclgraph = (self.vllm_config.compilation_config.level == CompilationLevel.PIECEWISE - and not self.model_config.enforce_eager) + and not self.model_config.enforce_eager and + not ascend_config.torchair_graph_config.enabled) self.aclgraph_batch_sizes = list( reversed( self.vllm_config.compilation_config.cudagraph_capture_sizes)) @@ -374,6 +379,14 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): self.is_kv_producer = vllm_config.kv_transfer_config.is_kv_producer self.is_kv_consumer = vllm_config.kv_transfer_config.is_kv_consumer + self.reserved_mc2_mask = torch.zeros( + 512, + dtype=torch.bool, + device=self.device, + ) + + self.moe_comm_method = AllGatherCommImpl + def check_batch_sizes_consistency(self) -> None: if not dist.is_initialized(): return @@ -640,26 +653,11 @@ def _get_forward_metadata_across_dp_and_pad( self, num_tokens: int, with_prefill: bool, enable_dbo: bool ) -> tuple[int, Optional[torch.Tensor], bool, bool]: if self.dp_size == 1: - if self.torchair_graph_enabled and not with_prefill: - maybe_padded_num_tokens = self.select_torchair_padded_batch_size( - num_tokens) - return maybe_padded_num_tokens, None, with_prefill, enable_dbo return num_tokens, None, with_prefill, enable_dbo - maybe_padded_num_tokens = num_tokens num_tokens_across_dp, with_prefill, enable_dbo = self._get_forward_metadata_across_dp( num_tokens, with_prefill, enable_dbo) - - if self.torchair_graph_enabled and not with_prefill: - max_num_token = num_tokens_across_dp.max().item() - maybe_padded_num_tokens = self.select_torchair_padded_batch_size( - max_num_token) - num_tokens_across_dp = torch.full((self.dp_size, ), - maybe_padded_num_tokens, - dtype=torch.int32, - device="cpu") - - return maybe_padded_num_tokens, num_tokens_across_dp, with_prefill, enable_dbo + return num_tokens, num_tokens_across_dp, with_prefill, enable_dbo def _check_dbo_is_valid(self, query_lens: torch.Tensor, attn_state: AscendAttentionState, @@ -844,7 +842,7 @@ def get_supported_tasks(self) -> "tuple[SupportedTask, ...]": def _make_attention_mask(self, seq_lens, query_lens, position, attn_state) -> torch.Tensor: # Chunk Prefill situation. - if attn_state == AscendAttentionState.ChunkedPrefill: + if attn_state == AscendAttentionState.ChunkedPrefill and not self.vllm_config.model_config.use_mla: return self.attn_mask_builder.get_splitfuse_attn_mask( seq_lens, query_lens, position, self.dtype, self.device) # Prefill without cache situation. @@ -1017,6 +1015,32 @@ def _gather_mm_embeddings( mm_embeds.append(mm_embeds_item) return mm_embeds + def get_dp_padding(self, + num_tokens: int) -> tuple[int, Optional[torch.Tensor]]: + """This implementation is derived from vLLM's `GPUModelRunner.get_dp_padding`. + Please note that vLLM may refactor or modify this function over time, + at present, we are using the version introduced in PR #18935. + """ + dp_size = self.vllm_config.parallel_config.data_parallel_size + dp_rank = self.vllm_config.parallel_config.data_parallel_rank + + # For DP: Don't pad when setting enforce_eager. + # This lets us set enforce_eager on the prefiller in a P/D setup and + # still use ACL graphs (enabled by this padding) on the decoder. + + if dp_size == 1 or self.vllm_config.model_config.enforce_eager: + # Early exit. + return 0, None + + num_tokens_across_dp = DPMetadata.num_tokens_across_dp( + num_tokens, dp_size, dp_rank) + max_tokens_across_dp_cpu = torch.max(num_tokens_across_dp).item() + num_tokens_after_padding = torch.tensor([max_tokens_across_dp_cpu] * + dp_size, + device="cpu", + dtype=torch.int32) + return max_tokens_across_dp_cpu - num_tokens, num_tokens_after_padding + def _process_reqs( self, scheduler_output: "SchedulerOutput", @@ -1039,6 +1063,11 @@ def _process_reqs( # Eager mode. num_input_tokens = total_num_scheduled_tokens + # Padding for DP + num_pad, num_tokens_across_dp_native = self.get_dp_padding( + num_input_tokens) + num_input_tokens += num_pad + modified_batch = self.attn_metadata_builder.reorder_batch( self.input_batch, scheduler_output) if modified_batch: @@ -1264,13 +1293,26 @@ def _process_reqs( for k, v in self.intermediate_tensors.items() }) + moe_comm_method = self.moe_comm_method + + # NOTE: Currently this padding logic is really messy, + # MC2 may not be available in eager mode + # TODO: Unify the padding logic between TorchAir and ACL Graph ASAP + if self.use_aclgraph: + num_tokens_across_dp = num_tokens_across_dp_native + else: + num_input_tokens = padded_num_tokens_across_dp + # Run forward pass with set_ascend_forward_context( attn_metadata, self.vllm_config, - num_tokens=padded_num_tokens_across_dp, + num_tokens=num_input_tokens, num_tokens_across_dp=num_tokens_across_dp, with_prefill=with_prefill, + reserved_mc2_mask=self.reserved_mc2_mask, + moe_comm_method=moe_comm_method(self.device, self.dtype, + self.model_config.hf_config), num_actual_tokens=total_num_scheduled_tokens): with ProfileExecuteDuration().capture_async("forward"): self.maybe_setup_kv_connector(scheduler_output) @@ -1605,9 +1647,12 @@ def execute_model( intermediate_tensors)) kv_connector_output = None if not vllm_version_is("0.10.0"): - kv_connector_output = KVConnectorOutput( - finished_sending=finished_sending, - finished_recving=finished_recving) + if finished_sending is not None and finished_recving is not None: + kv_connector_output = KVConnectorOutput( + finished_sending=finished_sending, + finished_recving=finished_recving) + else: + kv_connector_output = None finished_sending = None finished_recving = None with ProfileExecuteDuration().capture_async("post process"): @@ -1844,6 +1889,31 @@ def get_finished_kv_transfer( scheduler_output.finished_req_ids) return None, None + def _build_attention_metadata(self, with_prefill, num_reqs, skip_attn): + if skip_attn: + attn_metadata = None + else: + # TODO(zzzzwwjj): when aclgraph and full graph mode, we need build attn_metadata + attn_metadata = None + return attn_metadata + + def _generate_dummy_run_hidden_states(self, with_prefill, + is_torchair_compile, input_ids, + positions, attn_metadata, num_tokens, + intermediate_tensors, inputs_embeds): + maybe_converting_weight_acl_format(self.model, ACL_FORMAT_FRACTAL_ND) + hidden_states = self.model(input_ids=input_ids, + positions=positions, + intermediate_tensors=intermediate_tensors, + inputs_embeds=inputs_embeds) + if self.use_aux_hidden_state_outputs: + hidden_states, _ = hidden_states + else: + hidden_states = hidden_states + if self.use_spec_decode and isinstance(self.drafter, EagleProposer): + self.drafter.dummy_run(num_tokens) + return hidden_states + @torch.inference_mode() def _dummy_run( self, @@ -1851,6 +1921,7 @@ def _dummy_run( skip_attn: bool = True, with_prefill: bool = False, is_torchair_compile: bool = False, + moe_comm_method: Type[MoECommMethod] = DummyCommImpl, ) -> torch.Tensor: # Padding for DP (num_tokens, num_tokens_across_dp, with_prefill, @@ -1880,20 +1951,11 @@ def _dummy_run( if self.is_kv_producer: with_prefill = True - # NOTE: If torchair graph mode and not with_prefill, - # we can't skip_attn, it will cause graph recompile. - if self.torchair_graph_enabled and not with_prefill: - attn_metadata = self.attn_metadata_builder.build_torchair_graph_dummy( - num_reqs=num_reqs, num_actual_tokens=1) - elif skip_attn: - attn_metadata = None - else: - # TODO(zzzzwwjj): when aclgraph and full graph mode, we need build attn_metadata - attn_metadata = None + attn_metadata = self._build_attention_metadata(with_prefill, num_reqs, + skip_attn) with self.maybe_dummy_run_with_lora(self.lora_config, num_scheduled_tokens): - model = self.model if self.is_multimodal_model: input_ids = None inputs_embeds = self.inputs_embeds[:num_tokens] @@ -1927,63 +1989,15 @@ def _dummy_run( num_tokens_across_dp=num_tokens_across_dp, with_prefill=with_prefill, in_profile_run=self.in_profile_run, + reserved_mc2_mask=self.reserved_mc2_mask, + moe_comm_method=moe_comm_method( + self.device, self.dtype, self.model_config.hf_config), num_actual_tokens=0, ): - model_kwargs = {} - if self.torchair_graph_enabled and not with_prefill: - # Only mark static while compiling - if is_torchair_compile: - torch._dynamo.mark_static(input_ids) - torch._dynamo.mark_static(positions) - torch._dynamo.mark_static( - attn_metadata.decode.block_table) - torch._dynamo.mark_static( - attn_metadata.decode.input_positions) - torch._dynamo.mark_static( - get_forward_context().mc2_mask) - if hasattr(attn_metadata.decode, "sin"): - torch._dynamo.mark_static(attn_metadata.decode.sin) - torch._dynamo.mark_static(attn_metadata.decode.cos) - torch._dynamo.mark_static(attn_metadata.slot_mapping) - if self.speculative_config: - torch._dynamo.mark_static( - attn_metadata.decode.attn_mask) - for kv in self.kv_caches: - assert isinstance( - kv, tuple), "kv_cache must be a tuple" - torch._dynamo.mark_static(kv[0]) - torch._dynamo.mark_static(kv[1]) - - maybe_converting_weight_acl_format(self.model, - ACL_FORMAT_FRACTAL_NZ) - - compiled_model = self._get_torchair_lazy_compiled_model( - num_tokens) - model_kwargs["kv_caches"] = self.kv_caches - model_kwargs["attn_metadata"] = attn_metadata - hidden_states = compiled_model( - input_ids=input_ids, - positions=positions, - intermediate_tensors=intermediate_tensors, - inputs_embeds=None, - **model_kwargs, - ) - else: - maybe_converting_weight_acl_format(self.model, - ACL_FORMAT_FRACTAL_ND) - - hidden_states = model( - input_ids=input_ids, - positions=positions, - intermediate_tensors=intermediate_tensors, - inputs_embeds=inputs_embeds) - if self.use_aux_hidden_state_outputs: - hidden_states, _ = hidden_states - else: - hidden_states = hidden_states - if self.use_spec_decode and isinstance( - self.drafter, EagleProposer): - self.drafter.dummy_run(num_tokens) + hidden_states = self._generate_dummy_run_hidden_states( + with_prefill, is_torchair_compile, input_ids, positions, + attn_metadata, num_tokens, intermediate_tensors, + inputs_embeds) if self.speculative_config and self.speculative_config.method == "deepseek_mtp": assert isinstance(self.drafter, MtpProposer) self.drafter.dummy_run( @@ -2094,8 +2108,8 @@ def load_model(self) -> None: if isinstance(module, (MergedColumnParallelLinear, QKVParallelLinear, RowParallelLinear)): - module.weight.data = torch_npu.npu_format_cast( - module.weight.data, ACL_FORMAT_FRACTAL_NZ) + module.weight.data = self._convert_torch_format( + module.weight.data) if self.drafter: logger.info("Loading drafter model...") if isinstance(self.drafter, EagleProposer): @@ -2180,6 +2194,10 @@ def _get_torchair_lazy_compiled_model(self, batch_size: int): ge_cache=False) return self.torchair_compiled_models[batch_size] + def _convert_torch_format(self, tensor): + tensor = torch_npu.npu_format_cast(tensor, ACL_FORMAT) + return tensor + def initialize_kv_cache(self, kv_cache_config: KVCacheConfig) -> None: """ Initialize KV cache based on `kv_cache_config`. @@ -2188,9 +2206,6 @@ def initialize_kv_cache(self, kv_cache_config: KVCacheConfig) -> None: cache size of each layer """ self.kv_cache_config = kv_cache_config - import torch_npu - acl_format = ACL_FORMAT_FRACTAL_NZ if is_310p( - ) and not self.torchair_graph_enabled else ACL_FORMAT_FRACTAL_ND kv_caches: Dict[str, torch.Tensor] = {} def align_memory(tensor: torch.Tensor, alignment: int) -> torch.Tensor: @@ -2249,7 +2264,6 @@ def align_memory(tensor: torch.Tensor, alignment: int) -> torch.Tensor: kv_cache_spec.head_size) dtype = kv_cache_spec.dtype if self.model_config.is_deepseek_mla: - num_blocks, block_size, num_kv_heads, head_size = kv_cache_shape rope_dim = self.model_config.hf_text_config.qk_rope_head_dim nope_dim = head_size - rope_dim @@ -2265,10 +2279,8 @@ def align_memory(tensor: torch.Tensor, alignment: int) -> torch.Tensor: nope_cache = torch.zeros(nope_cache_shape, dtype=dtype, device=self.device) - rope_cache = torch_npu.npu_format_cast( - rope_cache, acl_format) - nope_cache = torch_npu.npu_format_cast( - nope_cache, acl_format) + rope_cache = self._convert_torch_format(rope_cache) + nope_cache = self._convert_torch_format(nope_cache) else: # In order to transfer kv cache through the reigster_memory api from llmdatadist, the memory @@ -2306,8 +2318,7 @@ def align_memory(tensor: torch.Tensor, alignment: int) -> torch.Tensor: kv_cache = torch.zeros(cache_shape, dtype=dtype, device=self.device) - kv_cache = torch_npu.npu_format_cast( - kv_cache, acl_format) + kv_cache = self._convert_torch_format(kv_cache) else: cache_size = math.prod(cache_shape) cache_size_aligned = cache_size + alignment @@ -2370,67 +2381,35 @@ def get_kv_cache_spec(self) -> dict[str, KVCacheSpec]: return kv_cache_spec - def _compile_torchair_graph(self, torchair_graph_batch_sizes) -> None: - # Trigger torchair graph capture for specific shapes. + def _capture_model(self): + if not self.use_aclgraph: + logger.info("Skipping NPU graph capture for eager mode.") + return + # Trigger ACL graph capture for specific shapes. # Capture the large shapes first so that the smaller shapes # can reuse the memory pool allocated for the large shapes. - for idx, num_tokens in enumerate(reversed(torchair_graph_batch_sizes)): - for _ in range(self.vllm_config.compilation_config. - cudagraph_num_of_warmups): - self._dummy_run(num_tokens, is_torchair_compile=True) - self._dummy_run(num_tokens, is_torchair_compile=True) - logger.info("Batchsize %d is compiled successfully: %d/%d.", - num_tokens, idx + 1, len(torchair_graph_batch_sizes)) + with graph_capture(device=self.device): + skip_attn = not self.vllm_config.compilation_config.full_cuda_graph + for num_tokens in reversed(self.aclgraph_batch_sizes): + for _ in range(self.vllm_config.compilation_config. + cudagraph_num_of_warmups): + self._dummy_run( + num_tokens, + skip_attn=skip_attn, + moe_comm_method=self.moe_comm_method, + ) + self._dummy_run( + num_tokens, + skip_attn=skip_attn, + moe_comm_method=self.moe_comm_method, + ) def capture_model(self) -> None: start_time = time.perf_counter() start_free_npu_memory = torch.npu.mem_get_info()[0] - # TODO(NeverRaR): Calling graph_capture(device=self.device) in - # torchair graph capture can cause some issues, so now we just - # temporarily split the codepath for the two different graph patterns. - if self.torchair_graph_enabled: - torchair_graph_batch_sizes = self.torchair_graph_batch_sizes - graph_num = len(torchair_graph_batch_sizes) - - if self.use_cached_npu_graph and not check_torchair_cache_exist(): - # If caching is enabled but does not exist, we will compile the model twice. The first - # time is used to generate the cache, and the second time is used to load the cache to - # skip the overhead caused by Dynamo guard mechanism. - logger.info( - "Use cached npu graph but cache doesn't exist! Now we compile graph to genetate torchair cache, this usually takes %.1f~%.1f mins.", - 0.5 * graph_num, 1.5 * graph_num) - self._compile_torchair_graph(torchair_graph_batch_sizes) - NPUPlatform.synchronize() - torch._dynamo.reset() - self.torchair_compiled_models.clear() - if self.use_cached_npu_graph: - logger.info( - "Loading torchair graph cache, this usually takes %.1f~%.1f mins.", - 0.3 * graph_num, 0.5 * graph_num) - self._compile_torchair_graph(torchair_graph_batch_sizes) - else: - logger.info( - "Capturing torchair graph, this usually takes %.1f~%.1f mins.", - 0.5 * graph_num, 1.5 * graph_num) - self._compile_torchair_graph(torchair_graph_batch_sizes) - - if self.new_kv_cache_bytes > 0: - write_kv_cache_bytes_to_file(torch.distributed.get_rank(), - self.new_kv_cache_bytes) - elif self.use_aclgraph: - # Trigger ACL graph capture for specific shapes. - # Capture the large shapes first so that the smaller shapes - # can reuse the memory pool allocated for the large shapes. - # TODO(zzzzwwjj): Check dummy_run with ACL Graph and full graph mode - with graph_capture(device=self.device): - for num_tokens in reversed(self.aclgraph_batch_sizes): - for _ in range(self.vllm_config.compilation_config. - cudagraph_num_of_warmups): - self._dummy_run(num_tokens) - self._dummy_run(num_tokens) - else: - logger.info("Skipping NPU graph capture for eager mode.") - return + + self._capture_model() + end_time = time.perf_counter() end_free_npu_memory = torch.npu.mem_get_info()[0] elapsed_time = end_time - start_time