From b2598c3271af2d55a948e29799e90bc0cfcda188 Mon Sep 17 00:00:00 2001 From: Ronald1995 Date: Thu, 7 Aug 2025 17:19:23 +0800 Subject: [PATCH 01/22] enable mm allreduce test (#2192) ### What this PR does / why we need it? This PR is to add e2e test for using npu_mm_all_reduce_base fusion kernel. ### Does this PR introduce _any_ user-facing change? no ### How was this patch tested? not involved - vLLM version: v0.10.0 - vLLM main: https://github.com/vllm-project/vllm/commit/5d5d419ca6aa55034eef0144f24e66789b486cb5 Signed-off-by: Ronald1995 --- tests/e2e/multicard/test_external_launcher.py | 38 +++++++++++++++++++ .../patch/worker/patch_common/patch_linear.py | 2 + 2 files changed, 40 insertions(+) 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/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 From dceef080b140305841a98d71cd8cfdeccbd390af Mon Sep 17 00:00:00 2001 From: huangxialu Date: Thu, 7 Aug 2025 17:20:19 +0800 Subject: [PATCH 02/22] [main] remove torch.cat and replace it by List[0] (#2153) ### What this PR does / why we need it? torch_npu.npu_grouped_matmul: https://www.hiascend.com/document/detail/zh/Pytorch/710/apiref/torchnpuCustomsapi/context/torch_npu-npu_grouped_matmul.md According to the document, when `split_item` is 2 or 3, `torch_npu.npu_grouped_matmul` will return a list which has one element. Therefore, the `torch.cat` after `torch_npu.npu_grouped_matmul` is unnecessary. ### Does this PR introduce _any_ user-facing change? not involved ### How was this patch tested? ut and e2e covered: `tests/ut/ops/test_fused_ops.py`, `tests/e2e/singlecard/ops/test_fused_moe.py` **performance**: (qwen3 30B, 2k->20k) base: Total Token throughput (tok/s): 667.76 remove cat: Total Token throughput (tok/s): 680.82 - vLLM version: v0.10.0 - vLLM main: https://github.com/vllm-project/vllm/commit/fa00c5d75bc63c87f5822f839db1342f19e4acc8 Signed-off-by: huangxialu --- tests/ut/ops/test_fused_ops.py | 2 +- vllm_ascend/ops/fused_moe.py | 38 +++++++++++----------------------- 2 files changed, 13 insertions(+), 27 deletions(-) diff --git a/tests/ut/ops/test_fused_ops.py b/tests/ut/ops/test_fused_ops.py index 6c89f6fc1d8..46192f61684 100644 --- a/tests/ut/ops/test_fused_ops.py +++ b/tests/ut/ops/test_fused_ops.py @@ -112,7 +112,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) diff --git a/vllm_ascend/ops/fused_moe.py b/vllm_ascend/ops/fused_moe.py index f35fb105758..f3c14bfe8f6 100644 --- a/vllm_ascend/ops/fused_moe.py +++ b/vllm_ascend/ops/fused_moe.py @@ -205,11 +205,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 +217,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 +308,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 +320,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 +411,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 +813,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 +825,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) From ad1083761f06737eb1ad2424dcd74aab47b559cb Mon Sep 17 00:00:00 2001 From: Mengqing Cao Date: Fri, 8 Aug 2025 10:20:23 +0800 Subject: [PATCH 03/22] [CI][Quickfix] Fix AscendFusedMoE init error (#2268) ### What this PR does / why we need it? Fix AscendFusedMoE init error. Use `super().__init__()` instead of `super(FusedMoE, self).__init__()` to ensure the member variables in base class could be called by the children class ### Does this PR introduce _any_ user-facing change? N/A ### How was this patch tested? CI passed with new existing test. - vLLM version: v0.10.0 - vLLM main: https://github.com/vllm-project/vllm/commit/766bc8162cb37ad32605eee051d4f049ec325926 --------- Signed-off-by: MengqingCao --- tests/ut/ops/test_fused_ops.py | 16 +++++++++++++++- vllm_ascend/ops/fused_moe.py | 23 +++++++++++++++++++++-- 2 files changed, 36 insertions(+), 3 deletions(-) diff --git a/tests/ut/ops/test_fused_ops.py b/tests/ut/ops/test_fused_ops.py index 46192f61684..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', @@ -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/vllm_ascend/ops/fused_moe.py b/vllm_ascend/ops/fused_moe.py index f3c14bfe8f6..625146d592d 100644 --- a/vllm_ascend/ops/fused_moe.py +++ b/vllm_ascend/ops/fused_moe.py @@ -1181,8 +1181,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 From 0bd5ff52997691221f61b0504dc722458536af06 Mon Sep 17 00:00:00 2001 From: Icey <1790571317@qq.com> Date: Fri, 8 Aug 2025 11:09:16 +0800 Subject: [PATCH 04/22] Fix accuracy test config and add DeepSeek-V2-Lite test (#2261) ### What this PR does / why we need it? This PR fix accuracy test related to https://github.com/vllm-project/vllm-ascend/pull/2073, users can now perform accuracy tests on multiple models simultaneously and generate different report files by running: ```bash cd ~/vllm-ascend pytest -sv ./tests/e2e/models/test_lm_eval_correctness.py \ --config-list-file ./tests/e2e/models/configs/accuracy.txt ``` ### Does this PR introduce _any_ user-facing change? no ### How was this patch tested? image - vLLM version: v0.10.0 - vLLM main: https://github.com/vllm-project/vllm/commit/766bc8162cb37ad32605eee051d4f049ec325926 --------- Signed-off-by: Icey <1790571317@qq.com> --- .github/workflows/accuracy_test.yaml | 9 +- .github/workflows/vllm_ascend_test.yaml | 3 +- .../workflows/vllm_ascend_test_long_term.yaml | 102 ----------- .../long_term/accuracy/accuracy_multicard.py | 167 ------------------ .../long_term/accuracy/accuracy_singlecard.py | 115 ------------ .../e2e/models/configs/DeepSeek-V2-Lite.yaml | 13 ++ .../configs/Qwen2.5-VL-7B-Instruct.yaml | 0 .../models/configs/Qwen3-30B-A3B.yaml | 0 .../models/configs/Qwen3-8B-Base.yaml | 0 .../models/configs/accuracy.txt | 0 tests/e2e/{singlecard => }/models/conftest.py | 43 +++-- .../models/report_template.md | 0 .../models/test_lm_eval_correctness.py | 12 +- 13 files changed, 46 insertions(+), 418 deletions(-) delete mode 100644 .github/workflows/vllm_ascend_test_long_term.yaml delete mode 100644 tests/e2e/long_term/accuracy/accuracy_multicard.py delete mode 100644 tests/e2e/long_term/accuracy/accuracy_singlecard.py create mode 100644 tests/e2e/models/configs/DeepSeek-V2-Lite.yaml rename tests/e2e/{singlecard => }/models/configs/Qwen2.5-VL-7B-Instruct.yaml (100%) rename tests/e2e/{singlecard => }/models/configs/Qwen3-30B-A3B.yaml (100%) rename tests/e2e/{singlecard => }/models/configs/Qwen3-8B-Base.yaml (100%) rename tests/e2e/{singlecard => }/models/configs/accuracy.txt (100%) rename tests/e2e/{singlecard => }/models/conftest.py (53%) rename tests/e2e/{singlecard => }/models/report_template.md (100%) rename tests/e2e/{singlecard => }/models/test_lm_eval_correctness.py (94%) diff --git a/.github/workflows/accuracy_test.yaml b/.github/workflows/accuracy_test.yaml index 7140f262f75..74296c6fb6d 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() }} @@ -312,7 +313,7 @@ jobs: 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) + 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..178eac8dbce 100644 --- a/.github/workflows/vllm_ascend_test.yaml +++ b/.github/workflows/vllm_ascend_test.yaml @@ -211,8 +211,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' }} 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/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 From 3e65c406b8844b36782e1fb617788e2a582942ac Mon Sep 17 00:00:00 2001 From: Icey <1790571317@qq.com> Date: Fri, 8 Aug 2025 14:12:11 +0800 Subject: [PATCH 05/22] Fix accuracy test create PR (#2274) ### What this PR does / why we need it? Fix create PR of accuracy test ### Does this PR introduce _any_ user-facing change? no ### How was this patch tested? Local testing: https://github.com/nv-action/vllm-benchmarks/pull/87 - vLLM version: v0.10.0 - vLLM main: https://github.com/vllm-project/vllm/commit/099c0464637f330f8ea38b07fe0694717c16d815 --------- Signed-off-by: Icey <1790571317@qq.com> --- .github/workflows/accuracy_test.yaml | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/.github/workflows/accuracy_test.yaml b/.github/workflows/accuracy_test.yaml index 74296c6fb6d..044d157e7fa 100644 --- a/.github/workflows/accuracy_test.yaml +++ b/.github/workflows/accuracy_test.yaml @@ -226,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 @@ -258,7 +258,7 @@ 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 @@ -299,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 @@ -311,8 +311,8 @@ 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 }}`, + 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] From ee6f79c44a85720fd3564edafd3a7b2ac4adbf82 Mon Sep 17 00:00:00 2001 From: yangqinghao-cmss Date: Sat, 9 Aug 2025 08:26:04 +0800 Subject: [PATCH 06/22] Add ut for test_communicator.py (#2293) ### What this PR does / why we need it? Add ut for test_communicator.py - vLLM version: v0.10.0 - vLLM main: https://github.com/vllm-project/vllm/commit/e5ebeeba531755a78f68413e88a23d061404f3e3 Signed-off-by: yangqinghao-cmss --- tests/ut/distributed/test_communicator.py | 155 ++++++++++++++++++++++ 1 file changed, 155 insertions(+) create mode 100644 tests/ut/distributed/test_communicator.py 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) From 9260910c8dee96b7fee4382723d682999c918584 Mon Sep 17 00:00:00 2001 From: wangxiyuan Date: Mon, 11 Aug 2025 11:22:32 +0800 Subject: [PATCH 07/22] [CI] Fix broken CI (#2302) 1. disable test_eagle_ccorrectness test, we'll reopen it once oom error fixed. 2. drop transformers version limit for main, since vLLM rely on >=4.55.0, see: https://github.com/vllm-project/vllm/commit/65552b476b1c475ef433995d2699bb27428693b3 3. fix kv_connector_output bug, see: https://github.com/vllm-project/vllm/commit/796bae07c59716b7b61d57343826bfbeabdd01bb - vLLM version: v0.10.0 - vLLM main: https://github.com/vllm-project/vllm/commit/d1af8b7be9c5ad9d2926ce215771e9cd7279147b Signed-off-by: wangxiyuan --- .github/workflows/vllm_ascend_test.yaml | 6 ++++++ pyproject.toml | 2 -- requirements.txt | 2 -- .../e2e/singlecard/spec_decode_v1/test_v1_spec_decode.py | 1 + vllm_ascend/worker/model_runner_v1.py | 9 ++++++--- 5 files changed, 13 insertions(+), 7 deletions(-) diff --git a/.github/workflows/vllm_ascend_test.yaml b/.github/workflows/vllm_ascend_test.yaml index 178eac8dbce..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: @@ -267,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/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/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/vllm_ascend/worker/model_runner_v1.py b/vllm_ascend/worker/model_runner_v1.py index d7944b8d74c..ba1657c611b 100644 --- a/vllm_ascend/worker/model_runner_v1.py +++ b/vllm_ascend/worker/model_runner_v1.py @@ -1605,9 +1605,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"): From 1ab15414bb291e5669b5211ce3d66879f61670fc Mon Sep 17 00:00:00 2001 From: wangxiyuan Date: Mon, 11 Aug 2025 14:06:49 +0800 Subject: [PATCH 08/22] [2/N][Refactor] torchair model runner refactor (#2204) There is lot of torchair code in model runner leading the code hard for maintenance. We'll create new torchair_model_runner to split torchair related logic. Following the workflow #2203 What's this PR do: move `torchair` related logic into `_get_forward_metadata_across_dp` and override it in torchair model runner - vLLM version: v0.10.0 - vLLM main: https://github.com/vllm-project/vllm/commit/1b9902806915040ac9b3029f2ab7522ec505afc3 Signed-off-by: wangxiyuan --- vllm_ascend/torchair/torchair_model_runner.py | 28 +++++++++++++++++++ vllm_ascend/worker/model_runner_v1.py | 17 +---------- 2 files changed, 29 insertions(+), 16 deletions(-) diff --git a/vllm_ascend/torchair/torchair_model_runner.py b/vllm_ascend/torchair/torchair_model_runner.py index 200167438b4..845e88ee638 100644 --- a/vllm_ascend/torchair/torchair_model_runner.py +++ b/vllm_ascend/torchair/torchair_model_runner.py @@ -17,6 +17,8 @@ # Adapted from vllm-project/vllm/vllm/worker/gpu_model_runner.py # +from typing import Optional + import torch from vllm.config import VllmConfig @@ -27,3 +29,29 @@ 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]: + 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 diff --git a/vllm_ascend/worker/model_runner_v1.py b/vllm_ascend/worker/model_runner_v1.py index ba1657c611b..d3a29852bd6 100644 --- a/vllm_ascend/worker/model_runner_v1.py +++ b/vllm_ascend/worker/model_runner_v1.py @@ -640,26 +640,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, From c0f0b708137838c1387b4f98eeaeda383d06f0d9 Mon Sep 17 00:00:00 2001 From: Pleaplusone Date: Mon, 11 Aug 2025 15:59:42 +0800 Subject: [PATCH 09/22] [core] Support capture custom ops into aclgraph (#2113) ### What this PR does / why we need it? Thanks to the PR https://github.com/vllm-project/vllm-ascend/pull/426 make vllm-ascend support the aclgraph inference to reduce the host overhead. However, the capability of aclgraph strongly relies on the functionality provided by `torch.compile`, which is the key feature supported in torch 2.x . Therefore, capture custom op into aclgraph is only possible when it can be recognize and captured by `torch.compile`. In this PR, we register the meta implementation of current custom ops to enable the fx graph capture. And by doing that, insert those custom ops into aclgraph become a natural thing to the ascend runtime. ### Does this PR introduce _any_ user-facing change? No user face change. ### How was this patch tested? Tested in unittest, we will integrate the `rotary_embedding` op into a small custom model and use `torch.compile` and aclgraph to capture and replay it to verify its functionality. - vLLM version: v0.10.0 - vLLM main: https://github.com/vllm-project/vllm/commit/1b9902806915040ac9b3029f2ab7522ec505afc3 --------- Signed-off-by: ganyi --- csrc/torch_binding.cpp | 11 ++ csrc/torch_binding_meta.cpp | 86 +++++++++++ csrc/utils.h | 12 -- .../singlecard/ops/test_rotary_embedding.py | 146 +++++++++++++++++- vllm_ascend/meta_registration.py | 86 +++++++++++ vllm_ascend/utils.py | 4 + 6 files changed, 332 insertions(+), 13 deletions(-) create mode 100644 csrc/torch_binding_meta.cpp create mode 100644 vllm_ascend/meta_registration.py 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/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/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/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 From ca274001b0a576e77ba30e5b7329a019e653f0f9 Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Mon, 11 Aug 2025 16:02:12 +0800 Subject: [PATCH 10/22] Bump actions/download-artifact from 4 to 5 (#2311) Bumps [actions/download-artifact](https://github.com/actions/download-artifact) from 4 to 5. - vLLM version: v0.10.0 - vLLM main: https://github.com/vllm-project/vllm/commit/ebf7605b0dd58ff5d572d1918e52ca732025eee0 Signed-off-by: dependabot[bot] Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com> --- .github/workflows/accuracy_test.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/accuracy_test.yaml b/.github/workflows/accuracy_test.yaml index 044d157e7fa..044c5dcfd00 100644 --- a/.github/workflows/accuracy_test.yaml +++ b/.github/workflows/accuracy_test.yaml @@ -261,7 +261,7 @@ jobs: 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-* From 29aaba5f845c276892c962dab21e4f29681e478a Mon Sep 17 00:00:00 2001 From: whx <56632993+whx-sjtu@users.noreply.github.com> Date: Mon, 11 Aug 2025 17:37:49 +0800 Subject: [PATCH 11/22] [Perf][MTP] Optimize reject sampler in greedy situation. (#2137) This PR port optimization in PR #2002 to main and makes it cleaner. - vLLM version: v0.10.0 - vLLM main: https://github.com/vllm-project/vllm/commit/afa5b7ca0b417abadfa85e32f28969b72e58a885 --------- Signed-off-by: whx-sjtu <2952154980@qq.com> --- .../sample/test_rejection_sampler.py | 31 ++-- tests/ut/sample/test_rejection_sampler.py | 6 +- vllm_ascend/sample/rejection_sampler.py | 147 ++++++++++++------ 3 files changed, 123 insertions(+), 61 deletions(-) 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/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/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( From 881e36d6a93ffd06939c19b2c6f2fd584451507c Mon Sep 17 00:00:00 2001 From: wangxiyuan Date: Mon, 11 Aug 2025 18:03:19 +0800 Subject: [PATCH 12/22] [3/N][Refactor] torchair model runner refactor (#2207) There is lot of torchair code in model runner leading the code hard for maintenance. We'll create new torchair_model_runner to split torchair related logic. Following the workflow #2203, this is the first PR. What's this PR do: create common function `_build_attention_metadata` and `_generate_dummy_run_hidden_states` for dummy_run - vLLM version: v0.10.0 - vLLM main: https://github.com/vllm-project/vllm/commit/ebf7605b0dd58ff5d572d1918e52ca732025eee0 Signed-off-by: wangxiyuan --- vllm_ascend/torchair/torchair_model_runner.py | 58 +++++++++++ vllm_ascend/worker/model_runner_v1.py | 97 ++++++------------- 2 files changed, 89 insertions(+), 66 deletions(-) diff --git a/vllm_ascend/torchair/torchair_model_runner.py b/vllm_ascend/torchair/torchair_model_runner.py index 845e88ee638..4dd62acc525 100644 --- a/vllm_ascend/torchair/torchair_model_runner.py +++ b/vllm_ascend/torchair/torchair_model_runner.py @@ -21,7 +21,10 @@ import torch from vllm.config import VllmConfig +from vllm.forward_context import get_forward_context +from vllm_ascend.utils import (ACL_FORMAT_FRACTAL_NZ, + maybe_converting_weight_acl_format) from vllm_ascend.worker.model_runner_v1 import NPUModelRunner @@ -55,3 +58,58 @@ def _get_forward_metadata_across_dp_and_pad( 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 diff --git a/vllm_ascend/worker/model_runner_v1.py b/vllm_ascend/worker/model_runner_v1.py index d3a29852bd6..a0fe9e0802c 100644 --- a/vllm_ascend/worker/model_runner_v1.py +++ b/vllm_ascend/worker/model_runner_v1.py @@ -1832,6 +1832,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, @@ -1868,20 +1893,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] @@ -1917,61 +1933,10 @@ def _dummy_run( in_profile_run=self.in_profile_run, 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( From eb43a475f429192e7509e85e28b1c65d5097f373 Mon Sep 17 00:00:00 2001 From: zhenghaojiang Date: Mon, 11 Aug 2025 19:58:59 +0800 Subject: [PATCH 13/22] [Feat] chunkprefill mla support torchair graph (#1772) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit chunkprefill mla only support eager mode now，we want to optimaze it by support torchair graph, the idea is simple, when all the request is running in decode, use torchair graph to deal with it, else when chunkprefill or prefill only, use the eager mode - vLLM version: v0.10.0 - vLLM main: https://github.com/vllm-project/vllm/commit/ebf7605b0dd58ff5d572d1918e52ca732025eee0 Signed-off-by: haojiangzheng Co-authored-by: haojiangzheng --- tests/ut/attention/test_mla_v1.py | 1 + vllm_ascend/attention/mla_v1.py | 45 ++++++++++++++++++------------- 2 files changed, 28 insertions(+), 18 deletions(-) diff --git a/tests/ut/attention/test_mla_v1.py b/tests/ut/attention/test_mla_v1.py index 2ecc3f7bd74..652cff3bf42 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 diff --git a/vllm_ascend/attention/mla_v1.py b/vllm_ascend/attention/mla_v1.py index a8f8ae82332..48713fc385d 100644 --- a/vllm_ascend/attention/mla_v1.py +++ b/vllm_ascend/attention/mla_v1.py @@ -998,7 +998,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 +1112,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 +1149,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 +1191,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 +1231,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]) From c8b0f5f7998d06d26f2e203f2a6e8beb1e3f9521 Mon Sep 17 00:00:00 2001 From: wangxiyuan Date: Mon, 11 Aug 2025 21:39:24 +0800 Subject: [PATCH 14/22] [4/N][Refactor] torchair model runner refactor (#2208) There is lot of torchair code in model runner leading the code hard for maintenance. We'll create new torchair_model_runner to split torchair related logic. Following the workflow #2203, this is the first PR. What's this PR do: create common function `_convert_torch_foramt` for initialize_kv_cache - vLLM version: v0.10.0 - vLLM main: https://github.com/vllm-project/vllm/commit/14a5d903ab826b723a24a2d89631006394de76a1 Signed-off-by: wangxiyuan --- vllm_ascend/torchair/torchair_model_runner.py | 7 +++++- vllm_ascend/worker/model_runner_v1.py | 24 +++++++++---------- 2 files changed, 18 insertions(+), 13 deletions(-) diff --git a/vllm_ascend/torchair/torchair_model_runner.py b/vllm_ascend/torchair/torchair_model_runner.py index 4dd62acc525..b3b8ecbe903 100644 --- a/vllm_ascend/torchair/torchair_model_runner.py +++ b/vllm_ascend/torchair/torchair_model_runner.py @@ -20,10 +20,11 @@ from typing import Optional import torch +import torch_npu from vllm.config import VllmConfig from vllm.forward_context import get_forward_context -from vllm_ascend.utils import (ACL_FORMAT_FRACTAL_NZ, +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 @@ -113,3 +114,7 @@ def _generate_dummy_run_hidden_states(self, with_prefill, 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 diff --git a/vllm_ascend/worker/model_runner_v1.py b/vllm_ascend/worker/model_runner_v1.py index a0fe9e0802c..ae1cff3feed 100644 --- a/vllm_ascend/worker/model_runner_v1.py +++ b/vllm_ascend/worker/model_runner_v1.py @@ -110,6 +110,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 @@ -2047,8 +2050,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): @@ -2133,6 +2136,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`. @@ -2141,9 +2148,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: @@ -2202,7 +2206,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 @@ -2218,10 +2221,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 @@ -2259,8 +2260,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 From 9c6d108330574176f79eea52f989ea6049336af8 Mon Sep 17 00:00:00 2001 From: jack Date: Mon, 11 Aug 2025 22:21:29 +0800 Subject: [PATCH 15/22] Configure Gemini (#2298) ### What this PR does / why we need it? This PR requests Gemini AI to review PRs. ### Does this PR introduce _any_ user-facing change? No ### How was this patch tested? NA - vLLM version: v0.10.0 - vLLM main: https://github.com/vllm-project/vllm/commit/14a5d903ab826b723a24a2d89631006394de76a1 Signed-off-by: QwertyJack <7554089+QwertyJack@users.noreply.github.com> --- .gemini/config.yaml | 6 ++++++ 1 file changed, 6 insertions(+) create mode 100644 .gemini/config.yaml 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 From 81817908caca6c22206a52c9247d93de9628d4db Mon Sep 17 00:00:00 2001 From: Ronald1995 Date: Tue, 12 Aug 2025 08:05:01 +0800 Subject: [PATCH 16/22] ut: add ci guard for ut coverage (#2317) ### What this PR does / why we need it? add ci guard for ut coverage, if ut coverage of patch pr is below 80%, the ci will failed/ ### Does this PR introduce _any_ user-facing change? not involved ### How was this patch tested? not involved - vLLM version: v0.10.0 - vLLM main: https://github.com/vllm-project/vllm/commit/458e74eb907f96069e6d8a4f3c9f457001fef2ea --------- Signed-off-by: Ronald1995 --- codecov.yml | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) 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: From dc585f148acc91916916556ffd37fda708657cd3 Mon Sep 17 00:00:00 2001 From: Wang Kunpeng <1289706727@qq.com> Date: Tue, 12 Aug 2025 14:12:12 +0800 Subject: [PATCH 17/22] [main][prefill optimization] Optimize parallel strategies to reduce communication overhead (#2198) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ### What this PR does / why we need it? 1.Shared Expert Sharding Strategy Update: Switched from TP-aligned to pure DP for shared experts, enabling more efficient execution. 2.O_Proj AllReduce → ReduceScatter: Reduced communication overhead by using ReduceScatter, made possible by pure DP sharding. 3.AllGather Postponed: Delayed to after QKV down projection to reduce synchronization impact during prefill. ### How was this patch tested? Adding ut case in `tests/ut/attention/test_mla_v1.py` #### How to run use parameter `--additional_config='{"enable_shared_expert_dp": true}'` ##### a.How to run eager mode eg: python -m vllm.entrypoints.openai.api_server --model=/model_path --trust-remote-code -tp 8 -dp 2 --enable_expert_parallel --port 8002 --max-model-len 5120 --max-num-batched-tokens 16384 --enforce-eager --disable-log-requests --additional_config='{"ascend_scheduler_config":{"enabled":true},"enable_shared_expert_dp": true,"chunked_prefill_for_mla":true}' ##### b.How to run graph mode eg: python -m vllm.entrypoints.openai.api_server --model=/model_path --trust-remote-code -tp 8 -dp 2 --enable_expert_parallel --port 8002 --max-model-len 5120 --max-num-batched-tokens 16384 --disable-log-requests --additional_config='{"ascend_scheduler_config":{"enabled":true},"enable_shared_expert_dp": true,"chunked_prefill_for_mla":true,"torchair_graph_config":{"enabled":true}}' - vLLM version: v0.10.0 - vLLM main: https://github.com/vllm-project/vllm/commit/9edd1db02bc6dce6da503503a373657f3466a78b --------- Signed-off-by: Wang Kunpeng <1289706727@qq.com> Signed-off-by: SlightwindSec Co-authored-by: SlightwindSec --- .../configuration/additional_config.md | 1 + tests/ut/attention/test_mla_v1.py | 37 +++++++++ vllm_ascend/ascend_config.py | 7 ++ vllm_ascend/attention/mla_v1.py | 54 +++++++++---- vllm_ascend/models/deepseek_v2.py | 80 ++++++++++++++++--- vllm_ascend/ops/fused_moe.py | 27 ++++--- 6 files changed, 169 insertions(+), 37 deletions(-) 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/tests/ut/attention/test_mla_v1.py b/tests/ut/attention/test_mla_v1.py index 652cff3bf42..497b7b53abc 100644 --- a/tests/ut/attention/test_mla_v1.py +++ b/tests/ut/attention/test_mla_v1.py @@ -691,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/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/attention/mla_v1.py b/vllm_ascend/attention/mla_v1.py index 48713fc385d..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, @@ -1249,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 @@ -1259,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: @@ -1280,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/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/fused_moe.py b/vllm_ascend/ops/fused_moe.py index 625146d592d..aec6e72264a 100644 --- a/vllm_ascend/ops/fused_moe.py +++ b/vllm_ascend/ops/fused_moe.py @@ -1268,6 +1268,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 " @@ -1408,22 +1409,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: @@ -1490,7 +1493,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) @@ -1500,7 +1503,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] From 49ec6c98b7b4cb0929cec45bce361280a9768a45 Mon Sep 17 00:00:00 2001 From: Mengqing Cao Date: Tue, 12 Aug 2025 14:12:53 +0800 Subject: [PATCH 18/22] [Doc] Update faq (#2334) ### What this PR does / why we need it? - update determinitic calculation - update support device ### Does this PR introduce _any_ user-facing change? - Users should update ray and protobuf when using ray as distributed backend - Users should change to use `export HCCL_DETERMINISTIC=true` when enabling determinitic calculation ### How was this patch tested? N/A - vLLM version: v0.10.0 - vLLM main: https://github.com/vllm-project/vllm/commit/ea1292ad3ee724e44b3dfec2a26778cd614729f9 Signed-off-by: MengqingCao --- docs/source/faqs.md | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/docs/source/faqs.md b/docs/source/faqs.md index 81d22f26b0b..26384f81f0e 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) @@ -161,10 +161,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？ From 1a70564e7c1ecf5cb16a65276eecc14f179cfc4c Mon Sep 17 00:00:00 2001 From: wangxiyuan Date: Tue, 12 Aug 2025 14:24:50 +0800 Subject: [PATCH 19/22] [5/N][Refactor] torchair model runner refactor (#2216) There is lot of torchair code in model runner leading the code hard for maintenance. We'll create new torchair_model_runner to split torchair related logic. Following the workflow #2203 What's this PR do: create common function `_capture_model` for capture_model - vLLM version: v0.10.0 - vLLM main: https://github.com/vllm-project/vllm/commit/1891a265d316217f9c1e552cf7c380ef5bf1eec1 Signed-off-by: wangxiyuan --- vllm_ascend/torchair/torchair_model_runner.py | 51 +++++++++++++ vllm_ascend/worker/model_runner_v1.py | 72 ++++--------------- 2 files changed, 66 insertions(+), 57 deletions(-) diff --git a/vllm_ascend/torchair/torchair_model_runner.py b/vllm_ascend/torchair/torchair_model_runner.py index b3b8ecbe903..f42f83d1583 100644 --- a/vllm_ascend/torchair/torchair_model_runner.py +++ b/vllm_ascend/torchair/torchair_model_runner.py @@ -23,7 +23,11 @@ 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 @@ -37,6 +41,7 @@ def __init__(self, vllm_config: VllmConfig, device: torch.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( @@ -118,3 +123,49 @@ def _generate_dummy_run_hidden_states(self, with_prefill, 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/worker/model_runner_v1.py b/vllm_ascend/worker/model_runner_v1.py index ae1cff3feed..594649c6d45 100644 --- a/vllm_ascend/worker/model_runner_v1.py +++ b/vllm_ascend/worker/model_runner_v1.py @@ -82,8 +82,6 @@ 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, @@ -2323,67 +2321,27 @@ 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)) + # 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) 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 From 992271b027770257915386d06a715eba1dbc9c12 Mon Sep 17 00:00:00 2001 From: yiz-liu <136800916+yiz-liu@users.noreply.github.com> Date: Tue, 12 Aug 2025 21:10:20 +0800 Subject: [PATCH 20/22] [1/N][Feat] Support MoE models with ACL Graph and refactor MoE communication logic (#2125) ### What this PR does / why we need it? This PR refactors the MoE (Mixture of Experts) communication logic by introducing a strategy pattern. It defines an abstract base class, `MoECommMethod`, which encapsulates different communication strategies for MoE layers. By decoupling the MoE implementation from any single communication method, this change makes it simpler to add, replace, or optimize communication strategies in the future. Plan / Roadmap 1. Introduce `MoECommMethod`, implement `AllGatherImpl`, and adapt ACL Graph handling to cover all scenarios (this PR). 2. Implement `MC2CommImpl` and `AllToAllCommImpl` to optimize performance in specific scenarios. 3. Enable W8A8 / Int8 models to use `unified_fused_experts`. Other notes * Data-parallel (DP) communication currently does not work with vLLM's dispatch/combine mechanisms; an alternative approach is required to resolve this incompatibility. - vLLM version: v0.10.0 - vLLM main: https://github.com/vllm-project/vllm/commit/f7ad6a1eb3deb9ca70a6bce3705dbd16cf9d8b28 --------- Signed-off-by: Yizhou Liu --- tests/e2e/multicard/moe/test_moe_comm.py | 153 +++++++ vllm_ascend/ascend_forward_context.py | 21 +- vllm_ascend/distributed/moe_comm_method.py | 449 +++++++++++++++++++++ vllm_ascend/ops/common_fused_moe.py | 17 +- vllm_ascend/ops/fused_moe.py | 57 +++ vllm_ascend/platform.py | 11 +- vllm_ascend/worker/model_runner_v1.py | 82 +++- 7 files changed, 764 insertions(+), 26 deletions(-) create mode 100644 tests/e2e/multicard/moe/test_moe_comm.py create mode 100644 vllm_ascend/distributed/moe_comm_method.py 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/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/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/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 aec6e72264a..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]: 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/worker/model_runner_v1.py b/vllm_ascend/worker/model_runner_v1.py index 594649c6d45..3aeabc6e18e 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,6 +79,9 @@ 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 @@ -335,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)) @@ -375,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 @@ -1003,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", @@ -1025,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: @@ -1250,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) @@ -1865,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, @@ -1932,6 +1989,9 @@ 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, ): hidden_states = self._generate_dummy_run_hidden_states( @@ -2328,13 +2388,21 @@ def _capture_model(self): # 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): + 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) - self._dummy_run(num_tokens) + 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() From 8bfd16a145ae715ac1633d6b3aae6474559869a5 Mon Sep 17 00:00:00 2001 From: jack Date: Wed, 13 Aug 2025 16:00:43 +0800 Subject: [PATCH 21/22] [Doc] Add container image save/load FAQ for offline environments (#2347) ### What this PR does / why we need it? Add Docker export/import guide for air-gapped environments ### Does this PR introduce _any_ user-facing change? No ### How was this patch tested? NA - vLLM version: v0.10.0 - vLLM main: https://github.com/vllm-project/vllm/commit/d16aa3dae446d93f870a2e51b240e18a01cac294 Signed-off-by: QwertyJack <7554089+QwertyJack@users.noreply.github.com> --- docs/source/faqs.md | 27 +++++++++++++++++++++++++++ 1 file changed, 27 insertions(+) diff --git a/docs/source/faqs.md b/docs/source/faqs.md index 26384f81f0e..4250fd0b3b1 100644 --- a/docs/source/faqs.md +++ b/docs/source/faqs.md @@ -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). From 0f7492d18e572f53cce67615aa714ad09090f859 Mon Sep 17 00:00:00 2001 From: zhenghaojiang Date: Wed, 13 Aug 2025 17:15:59 +0800 Subject: [PATCH 22/22] [Bugfix] fix the oom when chunkprefill with long context like 64k (#2319) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The attn mask was declared in the mla.py，we don't need the splitfuse mask when mla chunkprefill, and this mask will cause memory problem when long context like 64k or 128k - vLLM version: v0.10.0 - vLLM main: https://github.com/vllm-project/vllm/commit/14a5d903ab826b723a24a2d89631006394de76a1 --------- Signed-off-by: haojiangzheng --- vllm_ascend/worker/model_runner_v1.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vllm_ascend/worker/model_runner_v1.py b/vllm_ascend/worker/model_runner_v1.py index 3aeabc6e18e..9891a029f07 100644 --- a/vllm_ascend/worker/model_runner_v1.py +++ b/vllm_ascend/worker/model_runner_v1.py @@ -842,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.