EngineX/xc-llm-ascend
3
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

upgrade to vllm 0.11.2 (#4400)

Browse Source 
Bump vLLM version to v0.11.2

What's broken and changed by vLLM:
1. structured_output is broken by
https://github.com/vllm-project/vllm/pull/26866
2. get_mrope_input_positions is broken by
https://github.com/vllm-project/vllm/pull/28399
3. graph mode is broken by
https://github.com/vllm-project/vllm/pull/25110 we'll upgrade torch to
2.8 to fix the problem later
4. embedding is broken by
https://github.com/vllm-project/vllm/pull/27583
5. `get_attn_backend_cls` and attention backend is broken are broken by
https://github.com/vllm-project/vllm/pull/28534
6. spec decode is broken by
https://github.com/vllm-project/vllm/pull/28771
7. sp feature is broken by
https://github.com/vllm-project/vllm/pull/27126
8. mtp is broken by https://github.com/vllm-project/vllm/pull/27922
9. lora is broken by https://github.com/vllm-project/vllm/pull/21068
10. execute_model is broken by
https://github.com/vllm-project/vllm/pull/26866
11. `VLLM_DISABLE_SHARED_EXPERTS_STREAM` env is broken by
https://github.com/vllm-project/vllm/pull/28159
12. kv cahe is broken by https://github.com/vllm-project/vllm/pull/27753
13. dp is broken by https://github.com/vllm-project/vllm/pull/25110

 
What's broken and changed by ourself:
1. qwen vl is broken by https://github.com/vllm-project/vllm/pull/28455
We'll remove model files in the future to avoid this kind of error
2. Engine core is broken by
https://github.com/vllm-project/vllm/pull/23691 We'll remove the patch
file in the future.
3. Ascend scheduler is broken by
https://github.com/vllm-project/vllm/pull/28733 We'll remove ascend
scheudler later.
4. qwen3-next is broken by
https://github.com/vllm-project/vllm/pull/28083 We'll remove model files
in the future to avoid this kind of error
5. qwen vl is broken by https://github.com/vllm-project/vllm/pull/27764.
We'll remove model files in the future

Known issue:
1. ray doesn't work 
2. the accuracy of qwen3-next is not correct
3. qwen3-vl is broken
4. prefix cache+ ascend scheduler + deepseek v2 lite is broken.

Co-authored-by: MengqingCao <cmq0113@163.com>
Co-authored-by: hfadzxy <starmoon_zhang@163.com>
Co-authored-by: leo-pony <nengjunma@outlook.com>
Co-authored-by: 22dimensions <waitingwind@foxmail.com>
Co-authored-by: shen-shanshan <467638484@qq.com>


- vLLM version: v0.11.2

---------

Signed-off-by: wangxiyuan <wangxiyuan1007@gmail.com>
Signed-off-by: MengqingCao <cmq0113@163.com>
Signed-off-by: hfadzxy <starmoon_zhang@163.com>
Signed-off-by: leo-pony <nengjunma@outlook.com>
Co-authored-by: MengqingCao <cmq0113@163.com>
Co-authored-by: hfadzxy <starmoon_zhang@163.com>
Co-authored-by: leo-pony <nengjunma@outlook.com>
This commit is contained in:
wangxiyuan
2025-11-26 11:48:58 +08:00
committed by GitHub
parent d5f77f14d0
commit bc69d7cfe1
54 changed files with 744 additions and 437 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
.github/workflows/_e2e_nightly_multi_node.yaml vendored
View File

@@ -32,7 +32,7 @@ on:
description: how many pods will be pulled up via lws.yaml, indicates number of nodes we need
vllm_version:
required: false
default: "2918c1b49c88c29783c86f78d2c4221cb9622379"
default: "v0.11.2"
type: string
description: vllm version to use
vllm_ascend_remote_url:

2
.github/workflows/format_pr_body.yaml vendored
View File

@@ -36,7 +36,7 @@ jobs:
- name: Get vLLM version
run: |
VLLM_COMMIT=2918c1b49c88c29783c86f78d2c4221cb9622379
VLLM_COMMIT=v0.11.2
echo "VLLM_COMMIT=https://github.com/vllm-project/vllm/commit/$VLLM_COMMIT" >> $GITHUB_ENV
- name: Checkout repository

2
.github/workflows/nightly_benchmarks.yaml vendored
View File

@@ -51,7 +51,7 @@ jobs:
strategy:
matrix:
include:
- vllm_branch: 2918c1b49c88c29783c86f78d2c4221cb9622379
- vllm_branch: v0.11.2
vllm_ascend_branch: main
max-parallel: 1
container:

11
.github/workflows/vllm_ascend_test.yaml vendored
View File

@@ -42,7 +42,7 @@ jobs:
lint:
uses: ./.github/workflows/pre-commit.yml
with:
vllm: 2918c1b49c88c29783c86f78d2c4221cb9622379
vllm: v0.11.2
changes:
runs-on: ubuntu-latest
outputs:
@@ -83,7 +83,7 @@ jobs:
VLLM_USE_MODELSCOPE: True
strategy:
matrix:
vllm_version: [2918c1b49c88c29783c86f78d2c4221cb9622379]
vllm_version: [v0.11.2]
steps:
- name: Install packages
run: |
@@ -121,7 +121,10 @@ jobs:
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/Ascend/ascend-toolkit/latest/x86_64-linux/devlib
pytest -sv --cov --cov-report=xml:unittests-coverage.xml tests/ut \
--ignore tests/ut/torchair/models/test_torchair_deepseek_mtp.py \
--ignore tests/ut/torchair/models/test_torchair_deepseek_v2.py
--ignore tests/ut/torchair/models/test_torchair_deepseek_v2.py \
--ignore tests/ut/models/test_qwen2_vl.py \
--ignore tests/ut/models/test_qwen2_5_vl.py \
--ignore tests/ut/models/test_qwen2_5_vl_without_padding.py
- name: Upload coverage to Codecov
# only upload coverage when commits merged
@@ -138,7 +141,7 @@ jobs:
name: e2e-light
strategy:
matrix:
vllm_version: [2918c1b49c88c29783c86f78d2c4221cb9622379]
vllm_version: [v0.11.2]
# Note (yikun): If CI resource are limited we can split job into two chain jobs
needs: [lint, changes]
# only trigger e2e test after lint passed and the change is e2e related with pull request.

2
.github/workflows/vllm_ascend_test_full.yaml vendored
View File

@@ -69,7 +69,7 @@ jobs:
name: e2e-full
strategy:
matrix:
vllm_version: [2918c1b49c88c29783c86f78d2c4221cb9622379]
vllm_version: [v0.11.2]
needs: [changes]
if: ${{ needs.changes.outputs.e2e_tracker == 'true' }}
uses: ./.github/workflows/_e2e_test.yaml

4
.github/workflows/vllm_ascend_test_nightly_a2.yaml vendored
View File

@@ -86,7 +86,7 @@ jobs:
tests: tests/e2e/nightly/ops
uses: ./.github/workflows/_e2e_nightly_single_node.yaml
with:
vllm: 2918c1b49c88c29783c86f78d2c4221cb9622379
vllm: v0.11.2
runner: ${{ matrix.test_config.os }}
tests: ${{ matrix.test_config.tests }}
image: 'swr.cn-southwest-2.myhuaweicloud.com/base_image/ascend-ci/vllm-ascend:nightly-a2'
@@ -125,7 +125,7 @@ jobs:
- Qwen3-Next-80B-A3B-Instruct
uses: ./.github/workflows/_e2e_nightly_single_node_models.yaml
with:
vllm: 2918c1b49c88c29783c86f78d2c4221cb9622379
vllm: v0.11.2
runner: ${{ matrix.test_config.os }}
model_list: ${{ toJson(matrix.test_config.model_list) }}
image: swr.cn-southwest-2.myhuaweicloud.com/base_image/ascend-ci/cann:8.2.rc1-910b-ubuntu22.04-py3.11

2
.github/workflows/vllm_ascend_test_nightly_a3.yaml vendored
View File

@@ -136,7 +136,7 @@ jobs:
tests: tests/e2e/nightly/models/test_deepseek_v3_2_exp_w8a8.py
uses: ./.github/workflows/_e2e_nightly_single_node.yaml
with:
vllm: 2918c1b49c88c29783c86f78d2c4221cb9622379
vllm: v0.11.2
runner: ${{ matrix.test_config.os }}
image: 'swr.cn-southwest-2.myhuaweicloud.com/base_image/ascend-ci/vllm-ascend:nightly-a3'
tests: ${{ matrix.test_config.tests }}

2
.github/workflows/vllm_ascend_test_report.yaml vendored
View File

@@ -72,7 +72,7 @@ jobs:
- DeepSeek-V2-Lite
uses: ./.github/workflows/_e2e_nightly_single_node_models.yaml
with:
vllm: 2918c1b49c88c29783c86f78d2c4221cb9622379
vllm: v0.11.2
runner: ${{ matrix.runner }}
image: swr.cn-southwest-2.myhuaweicloud.com/base_image/ascend-ci/cann:8.3.rc1-910b-ubuntu22.04-py3.11
model_list: ${{ toJson(matrix.model_list) }}

6
Dockerfile
View File

@@ -46,10 +46,8 @@ RUN pip config set global.index-url ${PIP_INDEX_URL}
# Install vLLM
ARG VLLM_REPO=https://github.com/vllm-project/vllm.git
ARG VLLM_TAG=2918c1b49c88c29783c86f78d2c4221cb9622379
# Revert this change once VLLM_TAG is specified to branch or tag
# RUN git clone --depth 1 $VLLM_REPO --branch $VLLM_TAG /vllm-workspace/vllm
RUN git clone $VLLM_REPO /vllm-workspace/vllm && (cd /vllm-workspace/vllm && git checkout $VLLM_TAG)
ARG VLLM_TAG=v0.11.2
RUN git clone --depth 1 $VLLM_REPO --branch $VLLM_TAG /vllm-workspace/vllm
# In x86, triton will be installed by vllm. But in Ascend, triton doesn't work correctly. we need to uninstall it.
RUN VLLM_TARGET_DEVICE="empty" python3 -m pip install -v -e /vllm-workspace/vllm/[audio] --extra-index https://download.pytorch.org/whl/cpu/ && \
python3 -m pip uninstall -y triton && \

6
Dockerfile.310p
View File

@@ -37,10 +37,8 @@ RUN pip config set global.index-url ${PIP_INDEX_URL}
# Install vLLM
ARG VLLM_REPO=https://github.com/vllm-project/vllm.git
ARG VLLM_TAG=2918c1b49c88c29783c86f78d2c4221cb9622379
# Revert this change once VLLM_TAG is specified to branch or tag
# RUN git clone --depth 1 $VLLM_REPO --branch $VLLM_TAG /vllm-workspace/vllm
RUN git clone $VLLM_REPO /vllm-workspace/vllm && (cd /vllm-workspace/vllm && git checkout $VLLM_TAG)
ARG VLLM_TAG=v0.11.2
RUN git clone --depth 1 $VLLM_REPO --branch $VLLM_TAG /vllm-workspace/vllm
# In x86, triton will be installed by vllm. But in Ascend, triton doesn't work correctly. we need to uninstall it.
RUN VLLM_TARGET_DEVICE="empty" python3 -m pip install -v -e /vllm-workspace/vllm/[audio] --extra-index https://download.pytorch.org/whl/cpu/ && \
python3 -m pip uninstall -y triton && \

6
Dockerfile.310p.openEuler
View File

@@ -34,10 +34,8 @@ COPY . /vllm-workspace/vllm-ascend/
# Install vLLM
ARG VLLM_REPO=https://github.com/vllm-project/vllm.git
ARG VLLM_TAG=2918c1b49c88c29783c86f78d2c4221cb9622379
# Revert this change once VLLM_TAG is specified to branch or tag
# RUN git clone --depth 1 $VLLM_REPO --branch $VLLM_TAG /vllm-workspace/vllm
RUN git clone $VLLM_REPO /vllm-workspace/vllm && (cd /vllm-workspace/vllm && git checkout $VLLM_TAG)
ARG VLLM_TAG=v0.11.2
RUN git clone --depth 1 $VLLM_REPO --branch $VLLM_TAG /vllm-workspace/vllm
# In x86, triton will be installed by vllm. But in Ascend, triton doesn't work correctly. we need to uninstall it.
RUN VLLM_TARGET_DEVICE="empty" python3 -m pip install -e /vllm-workspace/vllm/[audio] --extra-index https://download.pytorch.org/whl/cpu/ && \
python3 -m pip uninstall -y triton && \

6
Dockerfile.a3
View File

@@ -45,10 +45,8 @@ RUN apt-get update -y && \
# Install vLLM
ARG VLLM_REPO=https://github.com/vllm-project/vllm.git
ARG VLLM_TAG=2918c1b49c88c29783c86f78d2c4221cb9622379
# Revert this change once VLLM_TAG is specified to branch or tag
# RUN git clone --depth 1 $VLLM_REPO --branch $VLLM_TAG /vllm-workspace/vllm
RUN git clone $VLLM_REPO /vllm-workspace/vllm && (cd /vllm-workspace/vllm && git checkout $VLLM_TAG)
ARG VLLM_TAG=v0.11.2
RUN git clone --depth 1 $VLLM_REPO --branch $VLLM_TAG /vllm-workspace/vllm
# In x86, triton will be installed by vllm. But in Ascend, triton doesn't work correctly. we need to uninstall it.
RUN VLLM_TARGET_DEVICE="empty" python3 -m pip install -v -e /vllm-workspace/vllm/[audio] --extra-index https://download.pytorch.org/whl/cpu/ && \
python3 -m pip uninstall -y triton && \

6
Dockerfile.a3.openEuler
View File

@@ -48,10 +48,8 @@ RUN yum update -y && \
# Install vLLM
ARG VLLM_REPO=https://github.com/vllm-project/vllm.git
ARG VLLM_TAG=2918c1b49c88c29783c86f78d2c4221cb9622379
# Revert this change once VLLM_TAG is specified to branch or tag
# RUN git clone --depth 1 $VLLM_REPO --branch $VLLM_TAG /vllm-workspace/vllm
RUN git clone $VLLM_REPO /vllm-workspace/vllm && (cd /vllm-workspace/vllm && git checkout $VLLM_TAG)
ARG VLLM_TAG=v0.11.2
RUN git clone --depth 1 $VLLM_REPO --branch $VLLM_TAG /vllm-workspace/vllm
# In x86, triton will be installed by vllm. But in Ascend, triton doesn't work correctly. we need to uninstall it.
RUN VLLM_TARGET_DEVICE="empty" python3 -m pip install -e /vllm-workspace/vllm/[audio] --extra-index https://download.pytorch.org/whl/cpu/ && \
python3 -m pip uninstall -y triton && \

6
Dockerfile.openEuler
View File

@@ -48,10 +48,8 @@ RUN yum update -y && \
# Install vLLM
ARG VLLM_REPO=https://github.com/vllm-project/vllm.git
ARG VLLM_TAG=2918c1b49c88c29783c86f78d2c4221cb9622379
# Revert this change once VLLM_TAG is specified to branch or tag
# RUN git clone --depth 1 $VLLM_REPO --branch $VLLM_TAG /vllm-workspace/vllm
RUN git clone $VLLM_REPO /vllm-workspace/vllm && (cd /vllm-workspace/vllm && git checkout $VLLM_TAG)
ARG VLLM_TAG=v0.11.2
RUN git clone --depth 1 $VLLM_REPO --branch $VLLM_TAG /vllm-workspace/vllm
# In x86, triton will be installed by vllm. But in Ascend, triton doesn't work correctly. we need to uninstall it.
RUN VLLM_TARGET_DEVICE="empty" python3 -m pip install -e /vllm-workspace/vllm/[audio] --extra-index https://download.pytorch.org/whl/cpu/ && \
python3 -m pip uninstall -y triton && \

2
docs/source/community/versioning_policy.md
View File

@@ -43,7 +43,7 @@ The table below is the release compatibility matrix for vLLM Ascend release.
For main branch of vLLM Ascend, we usually make it compatible with the latest vLLM release and a newer commit hash of vLLM. Please note that this table is usually updated. Please check it regularly.
| vLLM Ascend | vLLM | Python | Stable CANN | PyTorch/torch_npu |
|-------------|--------------|------------------|-------------|--------------------|
| main | v0.11.0/2918c1b49c88c29783c86f78d2c4221cb9622379 | >= 3.10, < 3.12 | 8.3.RC1 | 2.7.1 / 2.7.1 |
| main | v0.11.2 | >= 3.10, < 3.12 | 8.3.RC1 | 2.7.1 / 2.7.1 |
## Release cadence

2
docs/source/conf.py
View File

@@ -77,7 +77,7 @@ myst_substitutions = {
# CANN image tag
'cann_image_tag': "8.3.rc2-910b-ubuntu22.04-py3.11",
# vllm version in ci
'ci_vllm_version': 'v0.11.0',
'ci_vllm_version': 'v0.11.2',
}
# For cross-file header anchors

1
tests/e2e/multicard/test_prefix_caching.py
View File

@@ -85,6 +85,7 @@ def test_prefix_cache_with_v1_scheduler(model: str, max_tokens: int) -> None:
)
@pytest.mark.skip(reason="Fix me, the accuracy is not correct")
@pytest.mark.parametrize("model", MODELS)
@pytest.mark.parametrize("max_tokens", [50])
def test_prefix_cache_with_ascend_scheduler(model: str,

2
tests/e2e/multicard/test_qwen3_next.py
View File

@@ -24,6 +24,7 @@ Run `pytest tests/e2e/multicard/test_qwen3_next.py`.
import os
from unittest.mock import patch
import pytest
from modelscope import snapshot_download # type: ignore
from tests.e2e.conftest import VllmRunner
@@ -63,6 +64,7 @@ def test_models_distributed_Qwen3_NEXT_TP4_FULL_DECODE_ONLY():
del vllm_model
@pytest.mark.skip(reason="Fix me, the accuracy is not correct")
def test_models_distributed_Qwen3_NEXT_MTP_TP4_SIMILARITY():
example_prompts = [
"Hello, my name is",

10
tests/e2e/singlecard/test_ascend_scheduler.py
View File

@@ -18,9 +18,8 @@ def test_concurrent_partial_prefill(enforce_eager):
},
},
max_num_seqs=3,
max_num_batched_tokens=2048,
max_num_batched_tokens=8192,
enforce_eager=enforce_eager,
max_model_len=2048,
gpu_memory_utilization=0.7) as vllm_model:
outputs = vllm_model.model.generate(["Hello my name is Robert and I"] *
3)
@@ -38,9 +37,8 @@ def test_prefix_cache_stats_is_recorded(enforce_eager):
},
},
max_num_seqs=3,
max_num_batched_tokens=2048,
max_num_batched_tokens=8192,
enforce_eager=enforce_eager,
max_model_len=2048,
gpu_memory_utilization=0.7) as vllm_model:
# 17 tokens will make sure first 16 tokens are cached in a block
input_tokens = {"prompt_token_ids": [101] * 129}
@@ -51,7 +49,7 @@ def test_prefix_cache_stats_is_recorded(enforce_eager):
@pytest.mark.parametrize("max_tokens",
[4]) # cannot align results when max_tokens > 4
@pytest.mark.parametrize("chunked_prefill_token_size", [16])
@pytest.mark.parametrize("chunked_prefill_token_size", [2048])
def test_chunked_prefill_with_ascend_scheduler(
max_tokens: int, chunked_prefill_token_size: int) -> None:
example_prompts = [
@@ -93,7 +91,7 @@ def test_chunked_prefill_with_ascend_scheduler(
@pytest.mark.parametrize("max_tokens",
[4]) # cannot align results when max_tokens > 4
@pytest.mark.parametrize("chunked_prefill_token_size", [16])
@pytest.mark.parametrize("chunked_prefill_token_size", [2048])
def test_chunked_prefill_with_scheduler_dynamic_batch(
max_tokens: int, chunked_prefill_token_size: int) -> None:
example_prompts = [

2
tests/ut/attention/test_attention_v1.py
View File

@@ -14,7 +14,7 @@ from vllm_ascend.attention.utils import AscendCommonAttentionMetadata
class TestAscendAttentionBackend(TestBase):
def test_get_name(self):
self.assertEqual(AscendAttentionBackend.get_name(), "ASCEND")
self.assertEqual(AscendAttentionBackend.get_name(), "CUSTOM")
def test_get_impl_cls(self):
self.assertEqual(AscendAttentionBackend.get_impl_cls(),

19
tests/ut/compilation/test_acl_graph.py
View File

@@ -107,8 +107,7 @@ class TestACLGraphWrapper(TestBase):
wrapper = ACLGraphWrapper(runnable=self.mock_runnable,
vllm_config=self.mock_vllm_config,
runtime_mode=CUDAGraphMode.FULL,
graph_pool=self.mock_graph_pool)
runtime_mode=CUDAGraphMode.FULL)
self.assertEqual(wrapper.runnable, self.mock_runnable)
self.assertEqual(wrapper.vllm_config, self.mock_vllm_config)
@@ -130,7 +129,6 @@ class TestACLGraphWrapper(TestBase):
runnable=self.mock_runnable,
vllm_config=self.mock_vllm_config,
runtime_mode=CUDAGraphMode.FULL,
graph_pool=self.mock_graph_pool,
cudagraph_options=self.mock_cudagraph_options)
self.assertEqual(wrapper.runnable, self.mock_runnable)
@@ -152,8 +150,7 @@ class TestACLGraphWrapper(TestBase):
with self.assertRaises(AssertionError):
ACLGraphWrapper(runnable=self.mock_runnable,
vllm_config=self.mock_vllm_config,
runtime_mode=CUDAGraphMode.NONE,
graph_pool=self.mock_graph_pool)
runtime_mode=CUDAGraphMode.NONE)
@patch('vllm_ascend.compilation.acl_graph.get_forward_context')
@patch('vllm_ascend.compilation.acl_graph.current_platform')
@@ -171,7 +168,6 @@ class TestACLGraphWrapper(TestBase):
runnable=self.mock_runnable,
vllm_config=self.mock_vllm_config,
runtime_mode=CUDAGraphMode.FULL,
graph_pool=self.mock_graph_pool,
cudagraph_options=self.mock_cudagraph_options)
result = wrapper("arg1", "arg2")
@@ -196,7 +192,6 @@ class TestACLGraphWrapper(TestBase):
runnable=self.mock_runnable,
vllm_config=self.mock_vllm_config,
runtime_mode=CUDAGraphMode.FULL,
graph_pool=self.mock_graph_pool,
cudagraph_options=self.mock_cudagraph_options)
result = wrapper("arg1", "arg2")
@@ -247,7 +242,6 @@ class TestACLGraphWrapper(TestBase):
runnable=self.mock_runnable,
vllm_config=self.mock_vllm_config,
runtime_mode=CUDAGraphMode.FULL,
graph_pool=self.mock_graph_pool,
cudagraph_options=self.mock_cudagraph_options)
# Create a real torch tensor for the test, not a mock
@@ -319,7 +313,6 @@ class TestACLGraphWrapper(TestBase):
runnable=self.mock_runnable,
vllm_config=self.mock_vllm_config,
runtime_mode=CUDAGraphMode.FULL,
graph_pool=self.mock_graph_pool,
cudagraph_options=self.mock_cudagraph_options)
# Create a real torch tensor for the test, not a mock
@@ -392,7 +385,6 @@ class TestACLGraphWrapper(TestBase):
runnable=self.mock_runnable,
vllm_config=self.mock_vllm_config,
runtime_mode=CUDAGraphMode.FULL,
graph_pool=self.mock_graph_pool,
cudagraph_options=self.mock_cudagraph_options)
# First call to capture the graph
@@ -447,7 +439,6 @@ class TestACLGraphWrapper(TestBase):
runnable=self.mock_runnable,
vllm_config=self.mock_vllm_config,
runtime_mode=CUDAGraphMode.FULL,
graph_pool=self.mock_graph_pool,
cudagraph_options=self.mock_cudagraph_options)
# First call to capture the graph
@@ -518,7 +509,6 @@ class TestACLGraphWrapper(TestBase):
runnable=self.mock_runnable,
vllm_config=self.mock_vllm_config,
runtime_mode=CUDAGraphMode.FULL,
graph_pool=self.mock_graph_pool,
cudagraph_options=self.mock_cudagraph_options)
# Create a real torch tensor for the test, not a mock
@@ -588,7 +578,6 @@ class TestACLGraphWrapper(TestBase):
runnable=self.mock_runnable,
vllm_config=self.mock_vllm_config,
runtime_mode=CUDAGraphMode.FULL,
graph_pool=self.mock_graph_pool,
cudagraph_options=self.mock_cudagraph_options)
# Create a real torch tensor for the test, not a mock
@@ -659,7 +648,6 @@ class TestACLGraphWrapper(TestBase):
runnable=self.mock_runnable,
vllm_config=self.mock_vllm_config,
runtime_mode=CUDAGraphMode.FULL,
graph_pool=self.mock_graph_pool,
cudagraph_options=self.mock_cudagraph_options)
# Create a real torch tensor for the test, not a mock
@@ -680,7 +668,6 @@ class TestACLGraphWrapper(TestBase):
runnable=mock_runnable,
vllm_config=self.mock_vllm_config,
runtime_mode=CUDAGraphMode.FULL,
graph_pool=self.mock_graph_pool,
cudagraph_options=self.mock_cudagraph_options)
# Should be able to access attributes of the runnable
@@ -699,7 +686,6 @@ class TestACLGraphWrapper(TestBase):
runnable=mock_runnable,
vllm_config=self.mock_vllm_config,
runtime_mode=CUDAGraphMode.FULL,
graph_pool=self.mock_graph_pool,
cudagraph_options=self.mock_cudagraph_options)
# Should raise AttributeError for non-existent attributes
@@ -715,7 +701,6 @@ class TestACLGraphWrapper(TestBase):
runnable=self.mock_runnable,
vllm_config=self.mock_vllm_config,
runtime_mode=CUDAGraphMode.FULL,
graph_pool=self.mock_graph_pool,
cudagraph_options=self.mock_cudagraph_options)
unwrapped = wrapper.unwrap()

25
tests/ut/core/test_schedule_config.py
View File

@@ -48,7 +48,7 @@ class TestAscendSchedulerConfig(TestBase):
enable_chunked_prefill=False,
policy="fcfs",
scheduler_cls="vllm_ascend.core.scheduler.AscendScheduler",
max_num_batched_tokens=2048,
max_num_batched_tokens=8192,
max_model_len=2048,
max_long_partial_prefills=1,
long_prefill_token_threshold=512,
@@ -58,8 +58,8 @@ class TestAscendSchedulerConfig(TestBase):
self.assertEqual(ascend_config.policy, "fcfs")
self.assertEqual(ascend_config.scheduler_cls,
"vllm_ascend.core.scheduler.AscendScheduler")
self.assertEqual(ascend_config.max_num_batched_tokens, 2048)
self.assertEqual(ascend_config.encoder_cache_size, 2048)
self.assertEqual(ascend_config.max_num_batched_tokens, 8192)
self.assertEqual(ascend_config.encoder_cache_size, 8192)
self.assertEqual(ascend_config.max_long_partial_prefills, 1)
self.assertEqual(ascend_config.long_prefill_token_threshold, 512)
@@ -69,7 +69,7 @@ class TestAscendSchedulerConfig(TestBase):
self.basic_scheduler_config,
AscendSchedulerConfig(
policy="custom_policy",
max_num_batched_tokens=2048,
max_num_batched_tokens=8192,
max_model_len=2048,
),
)
@@ -86,7 +86,8 @@ class TestAscendSchedulerConfig(TestBase):
def test_valid_config_with_multimodal(self):
config = AscendSchedulerConfig.initialize_from_config(
SchedulerConfig(is_multimodal_model=True), {})
SchedulerConfig(is_multimodal_model=True,
max_num_batched_tokens=8192), {})
self.assertTrue(config.is_multimodal_model)
def test_valid_config_with_chunked_prefill(self):
@@ -94,12 +95,12 @@ class TestAscendSchedulerConfig(TestBase):
self.basic_scheduler_config,
AscendSchedulerConfig(
enable_chunked_prefill=True,
max_num_batched_tokens=2048,
max_model_len=4096,
max_num_batched_tokens=8192,
max_model_len=8192,
),
)
self.assertEqual(ascend_config.max_num_batched_tokens, 2048)
self.assertEqual(ascend_config.max_model_len, 4096)
self.assertEqual(ascend_config.max_num_batched_tokens, 8192)
self.assertEqual(ascend_config.max_model_len, 8192)
self.assertTrue(ascend_config.enable_chunked_prefill)
def test_invalid_config_without_chunked_prefill(self):
@@ -109,7 +110,7 @@ class TestAscendSchedulerConfig(TestBase):
AscendSchedulerConfig(
enable_chunked_prefill=False,
max_num_batched_tokens=2048,
max_model_len=4096,
max_model_len=8192,
),
)
self.assertIn(
@@ -117,7 +118,7 @@ class TestAscendSchedulerConfig(TestBase):
str(context.exception),
)
self.assertIn("max_num_batched_tokens (2048)", str(context.exception))
self.assertIn("max_model_len (4096)", str(context.exception))
self.assertIn("max_model_len (8192)", str(context.exception))
def test_initialize_from_config_with_pd_transfer(self):
ascend_config = AscendSchedulerConfig.initialize_from_config(
@@ -125,7 +126,7 @@ class TestAscendSchedulerConfig(TestBase):
AscendSchedulerConfig(
enable_pd_transfer=True,
decode_max_num_seqs=48,
max_num_batched_tokens=4096,
max_num_batched_tokens=8192,
max_model_len=4096,
),
)

104
tests/ut/core/test_scheduler.py
View File

@@ -3,6 +3,7 @@
from typing import Any, Dict, List, Optional, Tuple
from unittest.mock import MagicMock, patch
import numpy as np
import torch
from vllm.config import (CacheConfig, KVTransferConfig, ModelConfig,
SchedulerConfig, SpeculativeConfig, VllmConfig)
@@ -80,7 +81,10 @@ def make_output(scheduler):
req.request_id: i
for i, req in enumerate(scheduler.running)
}
sampled_token_ids = [[1000]] * len(scheduler.running)
sampled_token_ids = [
np.array([1000], dtype=np.int64) for _ in scheduler.running
]
logprobs = None
modelrunner_output = ModelRunnerOutput(
@@ -361,16 +365,15 @@ class TestAscendScheduler(TestBase):
},
num_common_prefix_blocks=0,
finished_req_ids=set(),
free_encoder_mm_hashes=[],
structured_output_request_ids={},
grammar_bitmask=None)
free_encoder_mm_hashes=[])
model_output = ModelRunnerOutput(
req_ids=[req.request_id for req in requests],
req_id_to_index={
req.request_id: i
for i, req in enumerate(requests)
},
sampled_token_ids=[[EOS_TOKEN_ID], [10, 11]
sampled_token_ids=[np.array([EOS_TOKEN_ID]),
np.array([10, 11])
], # First request hits EOS, second continues
logprobs=None,
prompt_logprobs_dict={},
@@ -414,17 +417,16 @@ class TestAscendScheduler(TestBase):
},
num_common_prefix_blocks=0,
finished_req_ids=set(),
free_encoder_mm_hashes=[],
structured_output_request_ids={},
grammar_bitmask=None)
free_encoder_mm_hashes=[])
model_output = ModelRunnerOutput(
req_ids=[req.request_id for req in requests],
req_id_to_index={
req.request_id: i
for i, req in enumerate(requests)
},
sampled_token_ids=[[10, 42, 12],
[13, 14]], # First request hits stop token
sampled_token_ids=[np.array([10, 42, 12]),
np.array([13, 14])
], # First request hits stop token
logprobs=None,
prompt_logprobs_dict={},
pooler_output=[])
@@ -466,17 +468,16 @@ class TestAscendScheduler(TestBase):
},
num_common_prefix_blocks=0,
finished_req_ids=set(),
free_encoder_mm_hashes=[],
structured_output_request_ids={},
grammar_bitmask=None)
free_encoder_mm_hashes=[])
model_output = ModelRunnerOutput(
req_ids=[req.request_id for req in requests],
req_id_to_index={
req.request_id: i
for i, req in enumerate(requests)
},
sampled_token_ids=[[10, 11, 12],
[13]], # First request exceeds max_tokens
sampled_token_ids=[np.array([10, 11, 12]),
np.array([13])
], # First request exceeds max_tokens
logprobs=None,
prompt_logprobs_dict={},
pooler_output=[])
@@ -511,13 +512,11 @@ class TestAscendScheduler(TestBase):
},
num_common_prefix_blocks=0,
finished_req_ids=set(),
free_encoder_mm_hashes=[],
structured_output_request_ids={},
grammar_bitmask=None)
free_encoder_mm_hashes=[])
model_output = ModelRunnerOutput(
req_ids=[requests[0].request_id],
req_id_to_index={requests[0].request_id: 0},
sampled_token_ids=[[EOS_TOKEN_ID, 10, 11]],
sampled_token_ids=[np.array([EOS_TOKEN_ID, 10, 11])],
logprobs=None,
prompt_logprobs_dict={},
pooler_output=[])
@@ -574,7 +573,7 @@ class TestAscendScheduler(TestBase):
model_runner_output = ModelRunnerOutput(
req_ids=[requests[0].request_id],
req_id_to_index={requests[0].request_id: 0},
sampled_token_ids=[[0]],
sampled_token_ids=[np.array([0], dtype=np.int64)],
logprobs=None,
prompt_logprobs_dict={},
pooler_output=[])
@@ -590,7 +589,7 @@ class TestAscendScheduler(TestBase):
model_runner_output = ModelRunnerOutput(
req_ids=[requests[1].request_id],
req_id_to_index={requests[1].request_id: 0},
sampled_token_ids=[[0]],
sampled_token_ids=[np.array([0], dtype=np.int64)],
logprobs=None,
prompt_logprobs_dict={},
pooler_output=[])
@@ -608,10 +607,12 @@ class TestAscendScheduler(TestBase):
spec_tokens_list: List[List[List[int]]] = [[[1, 2, 3]], [[1, 2, 3]],
[[1, 2], [3]], [[1]], [[]],
[[1, 2, 3], [4, 5, 6]]]
output_tokens_list: List[List[List[int]]] = [[[1, 2, 3, 4]], [[1, 5]],
[[1, 2, 5], [3, 4]],
[[1, 2]], [[5]],
[[1, 2, 7], [4, 8]]]
output_tokens_list: List[List[List[int]]] = [
[np.array([1, 2, 3, 4])], [np.array([1, 5])],
[np.array([1, 2, 5]), np.array([3, 4])], [np.array([1, 2])],
[np.array([5])], [np.array([1, 2, 7]),
np.array([4, 8])]
]
expected_list: List[Tuple[int, int,
int, List[int]]] = [(1, 3, 3, [1, 1, 1]),
(1, 3, 1, [1, 0, 0]),
@@ -649,7 +650,9 @@ class TestAscendScheduler(TestBase):
model_runner_output = ModelRunnerOutput(
req_ids=req_ids,
req_id_to_index=req_to_index,
sampled_token_ids=[[0] for _ in range(len(requests))],
sampled_token_ids=[
np.array([0]) for _ in range(len(requests))
],
logprobs=None,
prompt_logprobs_dict={},
pooler_output=[])
@@ -1054,16 +1057,15 @@ class TestSchedulerDynamicBatch(TestBase):
},
num_common_prefix_blocks=0,
finished_req_ids=set(),
free_encoder_mm_hashes=[],
structured_output_request_ids={},
grammar_bitmask=None)
free_encoder_mm_hashes=[])
model_output = ModelRunnerOutput(
req_ids=[req.request_id for req in requests],
req_id_to_index={
req.request_id: i
for i, req in enumerate(requests)
},
sampled_token_ids=[[EOS_TOKEN_ID], [10, 11]
sampled_token_ids=[np.array([EOS_TOKEN_ID]),
np.array([10, 11])
], # First request hits EOS, second continues
logprobs=None,
prompt_logprobs_dict={},
@@ -1107,17 +1109,16 @@ class TestSchedulerDynamicBatch(TestBase):
},
num_common_prefix_blocks=0,
finished_req_ids=set(),
free_encoder_mm_hashes=[],
structured_output_request_ids={},
grammar_bitmask=None)
free_encoder_mm_hashes=[])
model_output = ModelRunnerOutput(
req_ids=[req.request_id for req in requests],
req_id_to_index={
req.request_id: i
for i, req in enumerate(requests)
},
sampled_token_ids=[[10, 42, 12],
[13, 14]], # First request hits stop token
sampled_token_ids=[np.array([10, 42, 12]),
np.array([13, 14])
], # First request hits stop token
logprobs=None,
prompt_logprobs_dict={},
pooler_output=[])
@@ -1159,17 +1160,16 @@ class TestSchedulerDynamicBatch(TestBase):
},
num_common_prefix_blocks=0,
finished_req_ids=set(),
free_encoder_mm_hashes=[],
structured_output_request_ids={},
grammar_bitmask=None)
free_encoder_mm_hashes=[])
model_output = ModelRunnerOutput(
req_ids=[req.request_id for req in requests],
req_id_to_index={
req.request_id: i
for i, req in enumerate(requests)
},
sampled_token_ids=[[10, 11, 12],
[13]], # First request exceeds max_tokens
sampled_token_ids=[np.array([10, 11, 12]),
np.array([13])
], # First request exceeds max_tokens
logprobs=None,
prompt_logprobs_dict={},
pooler_output=[])
@@ -1204,13 +1204,11 @@ class TestSchedulerDynamicBatch(TestBase):
},
num_common_prefix_blocks=0,
finished_req_ids=set(),
free_encoder_mm_hashes=[],
structured_output_request_ids={},
grammar_bitmask=None)
free_encoder_mm_hashes=[])
model_output = ModelRunnerOutput(
req_ids=[requests[0].request_id],
req_id_to_index={requests[0].request_id: 0},
sampled_token_ids=[[EOS_TOKEN_ID, 10, 11]],
sampled_token_ids=[np.array([EOS_TOKEN_ID, 10, 11])],
logprobs=None,
prompt_logprobs_dict={},
pooler_output=[])
@@ -1267,7 +1265,7 @@ class TestSchedulerDynamicBatch(TestBase):
model_runner_output = ModelRunnerOutput(
req_ids=[requests[0].request_id],
req_id_to_index={requests[0].request_id: 0},
sampled_token_ids=[[0]],
sampled_token_ids=[np.array([0])],
logprobs=None,
prompt_logprobs_dict={},
pooler_output=[])
@@ -1283,7 +1281,7 @@ class TestSchedulerDynamicBatch(TestBase):
model_runner_output = ModelRunnerOutput(
req_ids=[requests[1].request_id],
req_id_to_index={requests[1].request_id: 0},
sampled_token_ids=[[0]],
sampled_token_ids=[np.array([0])],
logprobs=None,
prompt_logprobs_dict={},
pooler_output=[])
@@ -1301,10 +1299,12 @@ class TestSchedulerDynamicBatch(TestBase):
spec_tokens_list: List[List[List[int]]] = [[[1, 2, 3]], [[1, 2, 3]],
[[1, 2], [3]], [[1]], [[]],
[[1, 2, 3], [4, 5, 6]]]
output_tokens_list: List[List[List[int]]] = [[[1, 2, 3, 4]], [[1, 5]],
[[1, 2, 5], [3, 4]],
[[1, 2]], [[5]],
[[1, 2, 7], [4, 8]]]
output_tokens_list: List[List[List[int]]] = [
[np.array([1, 2, 3, 4])], [np.array([1, 5])],
[np.array([1, 2, 5]), np.array([3, 4])], [np.array([1, 2])],
[np.array([5])], [np.array([1, 2, 7]),
np.array([4, 8])]
]
expected_list: List[Tuple[int, int,
int, List[int]]] = [(1, 3, 3, [1, 1, 1]),
(1, 3, 1, [1, 0, 0]),
@@ -1342,7 +1342,9 @@ class TestSchedulerDynamicBatch(TestBase):
model_runner_output = ModelRunnerOutput(
req_ids=req_ids,
req_id_to_index=req_to_index,
sampled_token_ids=[[0] for _ in range(len(requests))],
sampled_token_ids=[
np.array([0]) for _ in range(len(requests))
],
logprobs=None,
prompt_logprobs_dict={},
pooler_output=[])

3
tests/ut/kv_connector/utils.py
View File

@@ -6,6 +6,7 @@
import os
from typing import Any, Optional
import numpy as np
import torch
from vllm import SamplingParams
from vllm.config import (CacheConfig, DeviceConfig, KVTransferConfig,
@@ -188,7 +189,7 @@ def create_model_runner_output(
# Make sampled tokens.
sampled_token = EOS_TOKEN_ID if use_eos else 0
sampled_token_ids = [[sampled_token] for _ in req_ids]
sampled_token_ids = [np.array([sampled_token]) for _ in req_ids]
# Make output data structure.
extra_args = {}

4
tests/ut/test_platform.py
View File

@@ -549,7 +549,6 @@ class TestNPUPlatform(TestBase):
dtype="float16",
kv_cache_dtype="float16",
block_size=64,
use_v1=True,
#use_sfa=False,
use_mla=True,
)
@@ -570,7 +569,6 @@ class TestNPUPlatform(TestBase):
dtype="float16",
kv_cache_dtype="float16",
block_size=64,
use_v1=True,
#use_sfa=False,
use_mla=True,
)
@@ -592,7 +590,6 @@ class TestNPUPlatform(TestBase):
dtype="float16",
kv_cache_dtype="float16",
block_size=64,
use_v1=True,
#use_sfa=False,
use_mla=False,
)
@@ -614,7 +611,6 @@ class TestNPUPlatform(TestBase):
dtype="float16",
kv_cache_dtype="float16",
block_size=64,
use_v1=True,
#use_sfa=False,
use_mla=False,
)

4
tests/ut/torchair/test_torchair_worker.py
View File

@@ -57,6 +57,8 @@ class TestNPUTorchairWorker(TestBase):
worker.model_config = MagicMock()
worker.model_config.seed = 42
worker.vllm_config = MagicMock()
worker.parallel_config = MagicMock()
worker.parallel_config.local_world_size = 0
result = worker._init_device()
@@ -89,6 +91,8 @@ class TestNPUTorchairWorker(TestBase):
worker.model_config = MagicMock()
worker.model_config.seed = 42
worker.vllm_config = MagicMock()
worker.parallel_config = MagicMock()
worker.parallel_config.local_world_size = 0
result = worker._init_device()

2
tests/ut/worker/test_worker_v1.py
View File

@@ -327,6 +327,8 @@ class TestNPUWorker(TestBase):
worker = NPUWorker()
worker.local_rank = 1
worker.model_config = MagicMock()
worker.parallel_config = MagicMock()
worker.parallel_config.local_world_size = 0
worker.model_config.seed = 42
# Test _init_device

6
vllm_ascend/attention/attention_v1.py
View File

@@ -56,13 +56,17 @@ if prefill_context_parallel_enable():
# isort: on
from vllm.attention.backends.registry import (AttentionBackendEnum,
register_backend)
@register_backend(AttentionBackendEnum.CUSTOM, "ASCEND")
class AscendAttentionBackend(AttentionBackend):
accept_output_buffer: bool = True
@staticmethod
def get_name() -> str:
return "ASCEND"
return "CUSTOM"
@staticmethod
def get_impl_cls() -> Type["AscendAttentionBackendImpl"]:

5
vllm_ascend/compilation/acl_graph.py
View File

@@ -62,11 +62,9 @@ class ACLGraphWrapper:
runnable: Callable,
vllm_config: VllmConfig,
runtime_mode: CUDAGraphMode,
graph_pool: Any = None,
cudagraph_options: Optional[CUDAGraphOptions] = None):
self.runnable = runnable
self.vllm_config = vllm_config
self.graph_pool = graph_pool
self.runtime_mode = runtime_mode
self.compilation_config = vllm_config.compilation_config
@@ -76,8 +74,7 @@ class ACLGraphWrapper:
# assert runtime_mode is not NONE(no aclgraph), otherwise, we don't
# need to initialize a ACLGraphWrapper.
assert self.runtime_mode != CUDAGraphMode.NONE
if self.graph_pool is None:
self.graph_pool = current_platform.get_global_graph_pool()
self.graph_pool = current_platform.get_global_graph_pool()
if cudagraph_options is None:
cudagraph_options = CUDAGraphOptions()

7
vllm_ascend/core/recompute_scheduler.py
View File

@@ -92,7 +92,7 @@ class RecomputeScheduler(SchedulerInterface):
self.max_num_running_reqs = self.scheduler_config.max_num_seqs
self.max_num_scheduled_tokens = \
self.scheduler_config.max_num_batched_tokens
self.max_model_len = self.scheduler_config.max_model_len
self.max_model_len = self.vllm_config.model_config.max_model_len
self.enable_kv_cache_events = (
self.kv_events_config is not None
and self.kv_events_config.enable_kv_cache_events)
@@ -928,8 +928,9 @@ class RecomputeScheduler(SchedulerInterface):
continue
req_index = model_runner_output.req_id_to_index[req_id]
generated_token_ids = sampled_token_ids[
req_index] if sampled_token_ids else []
generated_token_ids: list[int] = (
sampled_token_ids[req_index].tolist()
if sampled_token_ids else [])
scheduled_spec_token_ids = (
scheduler_output.scheduled_spec_decode_tokens.get(req_id))

11
vllm_ascend/core/scheduler.py
View File

@@ -219,7 +219,8 @@ class AscendScheduler(Scheduler):
# Schedule encoder inputs.
if request.has_encoder_inputs:
(encoder_inputs_to_schedule, num_new_tokens,
new_encoder_budget) = self._try_schedule_encoder_inputs(
new_encoder_budget,
_) = self._try_schedule_encoder_inputs(
request, num_computed_tokens, num_new_tokens,
encoder_budget)
if num_new_tokens == 0 or len(
@@ -464,7 +465,6 @@ class AscendScheduler(Scheduler):
num_scheduled_tokens, scheduled_spec_decode_tokens,
req_to_new_blocks)
scheduled_cached_reqs = cached_reqs_data
scheduler_output = SchedulerOutput(
scheduled_new_reqs=new_reqs_data,
scheduled_cached_reqs=scheduled_cached_reqs,
@@ -480,10 +480,7 @@ class AscendScheduler(Scheduler):
finished_req_ids=self.finished_req_ids, # type: ignore
free_encoder_mm_hashes=self.encoder_cache_manager.
get_freed_mm_hashes(),
structured_output_request_ids={},
grammar_bitmask=None,
)
# NOTE(Kuntai): this function is designed for multiple purposes:
# 1. Plan the KV cache store
# 2. Wrap up all the KV cache load / save ops into an opaque object
@@ -539,10 +536,10 @@ class AscendScheduler(Scheduler):
def _get_prompt_limit(self, request: Request) -> int:
if (self.scheduler_config.chunked_prefill_enabled
and not self.scheduler_config.is_multi_step):
prompt_limit = self.scheduler_config.max_model_len
prompt_limit = self.vllm_config.model_config.max_model_len
else:
prompt_limit = min(
self.scheduler_config.max_model_len,
self.vllm_config.model_config.max_model_len,
self.scheduler_config.max_num_batched_tokens,
)

13
vllm_ascend/core/scheduler_dynamic_batch.py
View File

@@ -35,7 +35,7 @@ from vllm.v1.structured_output import StructuredOutputManager
class BudgetRefiner:
"""This budget refiner can make dynamic adjustment to the token budget
"""This budget refiner can make dynamic adjustment to the token budget
in the chunked prefill scheduling strategy."""
def __init__(self, default_budget, slo_limit=-1) -> None:
@@ -416,8 +416,8 @@ class SchedulerDynamicBatch(Scheduler):
# Schedule encoder inputs.
if request.has_encoder_inputs:
(encoder_inputs_to_schedule, num_new_tokens,
new_encoder_compute_budget
) = self._try_schedule_encoder_inputs(
new_encoder_compute_budget,
_) = self._try_schedule_encoder_inputs(
request, num_computed_tokens, num_new_tokens,
encoder_compute_budget)
if num_new_tokens == 0:
@@ -549,11 +549,6 @@ class SchedulerDynamicBatch(Scheduler):
scheduled_spec_decode_tokens,
req_to_new_blocks,
)
scheduled_requests = (scheduled_new_reqs + scheduled_running_reqs +
scheduled_resumed_reqs)
structured_output_request_ids, grammar_bitmask = (
self.get_grammar_bitmask(scheduled_requests,
scheduled_spec_decode_tokens))
scheduler_output = SchedulerOutput(
scheduled_new_reqs=new_reqs_data,
scheduled_cached_reqs=cached_reqs_data,
@@ -569,8 +564,6 @@ class SchedulerDynamicBatch(Scheduler):
finished_req_ids=self.finished_req_ids,
free_encoder_mm_hashes=self.encoder_cache_manager.
get_freed_mm_hashes(),
structured_output_request_ids=structured_output_request_ids,
grammar_bitmask=grammar_bitmask,
)
# NOTE(Kuntai): this function is designed for multiple purposes:

23
vllm_ascend/models/qwen2_5_vl.py
View File

@@ -40,7 +40,6 @@ from vllm.model_executor.models.qwen2_5_vl import (
Qwen2_5_VLDummyInputsBuilder, Qwen2_5_VLForConditionalGeneration,
Qwen2_5_VLMultiModalProcessor, Qwen2_5_VLProcessingInfo)
from vllm.model_executor.models.utils import maybe_prefix
from vllm.model_executor.models.vision import conv3d_to_linear_weight
from vllm.multimodal import MULTIMODAL_REGISTRY
from vllm_ascend.ascend_forward_context import set_ascend_forward_context
@@ -144,8 +143,14 @@ class AscendQwen2_5_VisionBlock(Qwen2_5_VisionBlock):
quant_config: Optional[QuantizationConfig] = None,
prefix: str = "",
) -> None:
super().__init__(dim, num_heads, mlp_hidden_dim, act_fn, norm_layer,
quant_config, prefix)
super().__init__(dim=dim,
num_heads=num_heads,
mlp_hidden_dim=mlp_hidden_dim,
act_fn=act_fn,
norm_layer=norm_layer,
quant_config=quant_config,
prefix=prefix)
self.attn = AscendQwen2_5_VisionAttention(embed_dim=dim,
num_heads=num_heads,
projection_size=dim,
@@ -161,14 +166,6 @@ class AscendQwen2_5_VisionBlock(Qwen2_5_VisionBlock):
return x
class AscendQwen2_5_VisionPatchEmbed(Qwen2_5_VisionPatchEmbed):
def forward(self, x: torch.Tensor) -> torch.Tensor:
x = x.matmul(
self.proj.weight.data.view(self.hidden_size, -1).transpose(0, 1))
return x
class AscendQwen2_5_VisionRotaryEmbedding(Qwen2_5_VisionRotaryEmbedding):
def __init__(self, dim: int, theta: float = 10000.0) -> None:
@@ -195,7 +192,7 @@ class AscendQwen2_5_VisionTransformer(Qwen2_5_VisionTransformer):
head_dim = self.hidden_size // self.num_heads
self.rotary_pos_emb = AscendQwen2_5_VisionRotaryEmbedding(head_dim //
2)
self.patch_embed = AscendQwen2_5_VisionPatchEmbed(
self.patch_embed = Qwen2_5_VisionPatchEmbed(
patch_size=vision_config.patch_size,
temporal_patch_size=vision_config.temporal_patch_size,
in_channels=vision_config.in_channels,
@@ -357,8 +354,6 @@ class AscendQwen2_5_VisionTransformer(Qwen2_5_VisionTransformer):
params_dict = dict(self.named_parameters(remove_duplicate=False))
loaded_params: Set[str] = set()
for name, loaded_weight in weights:
if name.endswith("patch_embed.proj.weight"):
loaded_weight = conv3d_to_linear_weight(loaded_weight)
for (param_name, weight_name, shard_id) in stacked_params_mapping:
if weight_name not in name:
continue

20
vllm_ascend/models/qwen2_5_vl_without_padding.py
View File

@@ -146,8 +146,13 @@ class AscendQwen2_5_VisionBlock_Without_Padding(Qwen2_5_VisionBlock):
norm_layer: Optional[Callable[[int], nn.Module]] = None,
quant_config: Optional[QuantizationConfig] = None,
prefix: str = "") -> None:
super().__init__(dim, num_heads, mlp_hidden_dim, act_fn, norm_layer,
quant_config, prefix)
super().__init__(dim=dim,
num_heads=num_heads,
mlp_hidden_dim=mlp_hidden_dim,
act_fn=act_fn,
norm_layer=norm_layer,
quant_config=quant_config,
prefix=prefix)
self.attn = AscendQwen2_5_VisionAttention_Without_Padding(
embed_dim=dim,
num_heads=num_heads,
@@ -367,8 +372,15 @@ class AscendQwen3_VisionBlock(Qwen3_VisionBlock):
prefix: str = "",
use_data_parallel: bool = False,
) -> None:
super().__init__(dim, num_heads, mlp_hidden_dim, act_fn, norm_layer,
quant_config, prefix, use_data_parallel)
super().__init__(dim=dim,
num_heads=num_heads,
mlp_hidden_dim=mlp_hidden_dim,
act_fn=act_fn,
norm_layer=norm_layer,
quant_config=quant_config,
prefix=prefix,
use_data_parallel=use_data_parallel)
self.attn = AscendQwen2_5_VisionAttention_Without_Padding(
embed_dim=dim,
num_heads=num_heads,

322
vllm_ascend/models/qwen3_next.py
View File

@@ -16,7 +16,7 @@ from vllm.config import (CacheConfig, ModelConfig, SpeculativeConfig,
from vllm.distributed import (divide, get_tensor_model_parallel_rank,
get_tensor_model_parallel_world_size)
from vllm.forward_context import get_forward_context
from vllm.model_executor.layers.fla.ops import RMSNormGated, chunk
from vllm.model_executor.layers.fla.ops import chunk
from vllm.model_executor.layers.fla.ops.fused_recurrent import \
fused_recurrent_gated_delta_rule
from vllm.model_executor.layers.fused_moe import FusedMoE
@@ -24,6 +24,7 @@ from vllm.model_executor.layers.fused_moe import FusedMoE
# yapf: disable
from vllm.model_executor.layers.layernorm import \
GemmaRMSNorm as Qwen3NextRMSNorm
from vllm.model_executor.layers.layernorm import RMSNormGated
# yapf: enable
from vllm.model_executor.layers.linear import (ColumnParallelLinear,
MergedColumnParallelLinear,
@@ -179,6 +180,83 @@ class CustomQwen3NextGatedDeltaNet(Qwen3NextGatedDeltaNet, MambaBase):
raise ValueError(f"Duplicate layer name: {prefix}")
compilation_config.static_forward_context[prefix] = self
def forward(
self,
hidden_states: torch.Tensor,
output: torch.Tensor,
):
"""
Forward pass with three parts:
1. Input projection
2. Core attention (custom op)
3. Output projection
"""
num_tokens = hidden_states.size(0)
# ============================================================
# Part 1: Input Projection
# ============================================================
forward_context = get_forward_context()
attn_metadata: AttentionMetadata = forward_context.attn_metadata
if attn_metadata is None:
# V1 profile run
return
assert isinstance(attn_metadata, dict)
attn_metadata = attn_metadata[self.prefix]
assert isinstance(attn_metadata, GDNAttentionMetadata)
num_actual_tokens = (attn_metadata.num_prefill_tokens +
attn_metadata.num_decode_tokens +
attn_metadata.num_spec_decode_tokens)
# 1. Set up dimensions for reshapes later
projected_states, _ = self.in_proj(hidden_states[:num_actual_tokens])
projected_states_qkvz, projected_states_ba = torch.split(
projected_states,
[
self.projection_size_qkvz // self.tp_size,
self.projection_size_ba // self.tp_size
],
dim=-1,
)
query, key, value, z, b, a = self.fix_query_key_value_ordering(
projected_states_qkvz, projected_states_ba)
query, key, value = map(lambda x: rearrange(x, 'l p d -> l (p d)'),
(query, key, value))
mixed_qkv = torch.cat((query, key, value), dim=-1)
# ============================================================
# Part 2: Core Attention (Custom Op)
# ============================================================
core_attn_out = torch.zeros(
(num_tokens, self.num_v_heads // self.tp_size, self.head_v_dim),
dtype=hidden_states.dtype,
device=hidden_states.device,
)
torch.ops.vllm.gdn_attention_core(
mixed_qkv,
b,
a,
core_attn_out,
self.prefix,
)
# ============================================================
# Part 3: Output Projection
# ============================================================
z_shape_og = z.shape
# Reshape input data into 2D tensor
core_attn_out = core_attn_out.reshape(-1, core_attn_out.shape[-1])
z = z.reshape(-1, z.shape[-1])
core_attn_out = self.norm(core_attn_out, z)
core_attn_out = core_attn_out.reshape(z_shape_og)
core_attn_out = rearrange(core_attn_out, "... h d -> ... (h d)")
output[:num_tokens], _ = self.out_proj(core_attn_out)
def _forward(
self,
hidden_states: torch.Tensor,
@@ -437,6 +515,248 @@ class CustomQwen3NextGatedDeltaNet(Qwen3NextGatedDeltaNet, MambaBase):
output[:num_actual_tokens], _ = self.out_proj(core_attn_out)
def _forward_core(
self,
mixed_qkv: torch.Tensor,
b: torch.Tensor,
a: torch.Tensor,
core_attn_out: torch.Tensor,
):
"""
Core attention computation (called by custom op).
"""
forward_context = get_forward_context()
attn_metadata: AttentionMetadata = forward_context.attn_metadata
if attn_metadata is None:
# V1 profile run
return
assert isinstance(attn_metadata, dict)
attn_metadata = attn_metadata[self.prefix]
assert isinstance(attn_metadata, GDNAttentionMetadata)
has_initial_state = attn_metadata.has_initial_state
spec_query_start_loc = attn_metadata.spec_query_start_loc
non_spec_query_start_loc = attn_metadata.non_spec_query_start_loc
spec_sequence_masks = attn_metadata.spec_sequence_masks
spec_token_indx = attn_metadata.spec_token_indx
non_spec_token_indx = attn_metadata.non_spec_token_indx
spec_state_indices_tensor = attn_metadata.spec_state_indices_tensor # noqa: E501
non_spec_state_indices_tensor = attn_metadata.non_spec_state_indices_tensor # noqa: E501
self_kv_cache = self.kv_cache[forward_context.virtual_engine]
conv_state = self_kv_cache[0].transpose(-1, -2)
ssm_state = self_kv_cache[1]
num_actual_tokens = (attn_metadata.num_prefill_tokens +
attn_metadata.num_decode_tokens +
attn_metadata.num_spec_decode_tokens)
num_accepted_tokens = attn_metadata.num_accepted_tokens
mixed_qkv = mixed_qkv[:num_actual_tokens]
b = b[:num_actual_tokens]
a = a[:num_actual_tokens]
# 1. Convolution sequence transformation
conv_weights = self.conv1d.weight.view(self.conv1d.weight.size(0),
self.conv1d.weight.size(2))
if spec_sequence_masks is not None:
if (attn_metadata.num_prefills == 0
and attn_metadata.num_decodes == 0):
mixed_qkv_spec = mixed_qkv
mixed_qkv_non_spec = None
else:
mixed_qkv_spec = mixed_qkv.index_select(0, spec_token_indx)
mixed_qkv_non_spec = mixed_qkv.index_select(
0, non_spec_token_indx)
else:
mixed_qkv_spec = None
mixed_qkv_non_spec = mixed_qkv
# 1.1: Process the multi-query part
if spec_sequence_masks is not None:
mixed_qkv_spec = mixed_qkv_spec.view(
attn_metadata.num_spec_decodes, -1, mixed_qkv_spec.size(-1))
mixed_qkv_spec = rearrange(mixed_qkv_spec, 'b l d -> b d l')
mixed_qkv_spec = causal_conv1d.causal_conv1d_update(
mixed_qkv_spec,
conv_state,
conv_weights,
self.conv1d.bias,
self.activation,
conv_state_indices=spec_state_indices_tensor[:, 0]
[:attn_metadata.num_spec_decodes],
num_accepted_tokens=num_accepted_tokens,
validate_data=False,
)
mixed_qkv_spec = rearrange(mixed_qkv_spec, 'b d l -> (b l) d')
# 1.2: Process the remaining part
if attn_metadata.num_prefills > 0:
# - "cache_indices" updates the conv_state cache in positions
# pointed to by "mamba_cache_params.state_indices_tensor"
mixed_qkv_non_spec = causal_conv1d.causal_conv1d_fn(
mixed_qkv_non_spec.transpose(0, 1),
conv_weights,
self.conv1d.bias,
activation=self.activation,
conv_states=conv_state,
has_initial_state=has_initial_state,
cache_indices=non_spec_state_indices_tensor,
query_start_loc=non_spec_query_start_loc,
).transpose(0, 1)
elif attn_metadata.num_decodes > 0:
mixed_qkv_non_spec = causal_conv1d.causal_conv1d_update(
mixed_qkv_non_spec,
conv_state,
conv_weights,
self.conv1d.bias,
self.activation,
conv_state_indices=non_spec_state_indices_tensor[:attn_metadata
.num_decodes],
# validate_data=True,
)
else:
mixed_qkv_non_spec = None
query_spec, key_spec, value_spec = self.rearrange_mixed_qkv(
mixed_qkv_spec)
query_non_spec, key_non_spec, value_non_spec = self.rearrange_mixed_qkv(
mixed_qkv_non_spec)
beta = b.sigmoid()
g, beta = fused_gdn_gating(self.A_log, a, b, self.dt_bias)
if spec_sequence_masks is not None:
if (attn_metadata.num_prefills == 0
and attn_metadata.num_decodes == 0):
g_spec = g
beta_spec = beta
g_non_spec = None
beta_non_spec = None
else:
g_spec = g.index_select(1, spec_token_indx)
beta_spec = beta.index_select(1, spec_token_indx)
g_non_spec = g.index_select(1, non_spec_token_indx)
beta_non_spec = beta.index_select(1, non_spec_token_indx)
else:
g_spec = None
beta_spec = None
g_non_spec = g
beta_non_spec = beta
# 2. Recurrent attention
# 2.1: Process the multi-query part
if spec_sequence_masks is not None:
core_attn_out_spec, last_recurrent_state = (
fused_recurrent_gated_delta_rule(
q=query_spec,
k=key_spec,
v=value_spec,
g=g_spec,
beta=beta_spec,
initial_state=ssm_state,
inplace_final_state=True,
cu_seqlens=spec_query_start_loc[:attn_metadata.
num_spec_decodes + 1],
ssm_state_indices=spec_state_indices_tensor,
num_accepted_tokens=num_accepted_tokens,
use_qk_l2norm_in_kernel=True,
))
else:
core_attn_out_spec, last_recurrent_state = None, None
# 3.2: process the remaining part
if attn_metadata.num_prefills > 0:
initial_state = ssm_state[
non_spec_state_indices_tensor].contiguous()
initial_state[~has_initial_state, ...] = 0
batch_size = initial_state.shape[0]
core_attn_out = []
last_recurrent_state = []
for b_idx in range(batch_size):
start, end = non_spec_query_start_loc[
b_idx], non_spec_query_start_loc[b_idx + 1]
cur_q = query_non_spec[:, start:end, ...]
cur_k = key_non_spec[:, start:end, ...]
cur_v = value_non_spec[:, start:end, ...]
cur_g = g_non_spec[:, start:end, ...]
cur_b = beta_non_spec[:, start:end, ...]
cur_state = initial_state[b_idx].unsqueeze(0)
(
cur_core_attn_out_non_spec,
cur_last_recurrent_state,
) = chunk.chunk_gated_delta_rule(
query=cur_q,
key=cur_k,
value=cur_v,
g=cur_g,
beta=cur_b,
initial_state=cur_state,
output_final_state=True,
use_qk_l2norm_in_kernel=True,
)
core_attn_out.append(cur_core_attn_out_non_spec)
last_recurrent_state.append(cur_last_recurrent_state)
tar_dtype = core_attn_out[0].dtype
tar_device = core_attn_out[0].device
tar_shape = list(core_attn_out[0].shape)
tar_shape[1] = non_spec_query_start_loc[-1]
core_attn_out_non_spec = torch.empty(tar_shape,
dtype=tar_dtype,
device=tar_device)
for b_idx in range(batch_size):
cur_core_attn_out = core_attn_out[b_idx]
start, end = non_spec_query_start_loc[
b_idx], non_spec_query_start_loc[b_idx + 1]
core_attn_out_non_spec[:, start:end, ...] = cur_core_attn_out
last_recurrent_state = torch.cat(last_recurrent_state, dim=0)
# Init cache
ssm_state[non_spec_state_indices_tensor] = last_recurrent_state.to(
ssm_state.dtype)
elif attn_metadata.num_decodes > 0:
core_attn_out_non_spec, last_recurrent_state = (
fused_recurrent_gated_delta_rule(
q=query_non_spec,
k=key_non_spec,
v=value_non_spec,
g=g_non_spec,
beta=beta_non_spec,
initial_state=ssm_state,
inplace_final_state=True,
cu_seqlens=non_spec_query_start_loc[:attn_metadata.
num_decodes + 1],
ssm_state_indices=non_spec_state_indices_tensor,
use_qk_l2norm_in_kernel=True,
))
else:
core_attn_out_non_spec, last_recurrent_state = None, None
# 3. Merge core attention output
if spec_sequence_masks is not None and core_attn_out_non_spec is not None:
merged_out = torch.empty(
(1, num_actual_tokens, *core_attn_out_spec.shape[2:]),
dtype=core_attn_out_non_spec.dtype,
device=core_attn_out_non_spec.device,
)
merged_out.index_copy_(1, spec_token_indx, core_attn_out_spec)
merged_out.index_copy_(1, non_spec_token_indx,
core_attn_out_non_spec)
core_attn_out[:num_actual_tokens] = merged_out.squeeze(0)
elif spec_sequence_masks is not None:
core_attn_out[:num_actual_tokens] = core_attn_out_spec.squeeze(0)
else:
core_attn_out[:num_actual_tokens] = core_attn_out_non_spec.squeeze(
0)
class CustomQwen3NextDecoderLayer(Qwen3NextDecoderLayer):

7
vllm_ascend/ops/fused_moe/fused_moe.py
View File

@@ -444,6 +444,13 @@ class AscendSharedFusedMoE(SharedFusedMoE, AscendFusedMoE):
def is_internal_router(self) -> bool:
return False
@property
def use_dp_chunking(self) -> bool:
"""This func routes to the chunked forward path using the FlashInfer Cutlass kernel
only when data parallelism (DP) is enabled. Thus just returning False in vllm-ascend
"""
return False
def forward(
self,
hidden_states: torch.Tensor,

129
vllm_ascend/ops/triton/fla/sigmoid_gating.py
View File

@@ -169,132 +169,3 @@ def fused_recurrent_gated_delta_rule_fwd_kernel(
p_ht = ht + (bos + i_t) * stride_final_state_token
p_ht = p_ht + i_hv * K * V + o_k[:, None] * V + o_v[None, :]
tl.store(p_ht, b_h.to(p_ht.dtype.element_ty), mask=mask_h)
@triton.heuristics({
'USE_INITIAL_STATE':
lambda args: args['h0'] is not None,
'IS_VARLEN':
lambda args: args['cu_seqlens'] is not None,
"IS_CONTINUOUS_BATCHING":
lambda args: args['ssm_state_indices'] is not None,
"IS_SPEC_DECODING":
lambda args: args['num_accepted_tokens'] is not None,
})
@triton.jit(do_not_specialize=['N', 'T'])
def fused_recurrent_gated_delta_rule_fwd_kernel_0_11_0(
q,
k,
v,
g,
beta,
o,
h0,
ht,
cu_seqlens,
ssm_state_indices,
num_accepted_tokens,
scale,
N: tl.constexpr, # num of sequences
T: tl.constexpr, # num of tokens
B: tl.constexpr,
H: tl.constexpr,
HV: tl.constexpr,
K: tl.constexpr,
V: tl.constexpr,
BK: tl.constexpr,
BV: tl.constexpr,
stride_init_state_token: tl.constexpr,
stride_final_state_token: tl.constexpr,
stride_indices_seq: tl.constexpr,
stride_indices_tok: tl.constexpr,
USE_INITIAL_STATE: tl.constexpr, # whether to use initial state
INPLACE_FINAL_STATE: tl.constexpr, # whether to store final state inplace
IS_BETA_HEADWISE: tl.
constexpr, # whether beta is headwise vector or scalar,
USE_QK_L2NORM_IN_KERNEL: tl.constexpr,
IS_VARLEN: tl.constexpr,
IS_CONTINUOUS_BATCHING: tl.constexpr,
IS_SPEC_DECODING: tl.constexpr,
):
i_k, i_v, i_nh = tl.program_id(0), tl.program_id(1), tl.program_id(2)
i_n, i_hv = i_nh // HV, i_nh % HV
i_h = i_hv // (HV // H)
if IS_VARLEN:
bos, eos = tl.load(cu_seqlens + i_n).to(
tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64)
all = T
T = eos - bos
else:
bos, eos = i_n * T, i_n * T + T
all = B * T
if T == 0:
# no tokens to process for this sequence
return
o_k = i_k * BK + tl.arange(0, BK)
o_v = i_v * BV + tl.arange(0, BV)
mask_k = o_k < K
mask_v = o_v < V
mask_h = mask_k[:, None] & mask_v[None, :]
b_h = tl.zeros([BK, BV], dtype=tl.float32)
if USE_INITIAL_STATE:
if IS_CONTINUOUS_BATCHING:
if IS_SPEC_DECODING:
i_t = tl.load(num_accepted_tokens + i_n).to(tl.int64) - 1
else:
i_t = 0
p_h0 = h0 + tl.load(ssm_state_indices + i_n * stride_indices_seq +
i_t).to(tl.int64) * stride_init_state_token
else:
p_h0 = h0 + bos * HV * K * V
p_h0 = p_h0 + i_hv * K * V + o_k[:, None] * V + o_v[None, :]
b_h += tl.load(p_h0, mask=mask_h, other=0).to(tl.float32)
for i_t in range(0, T):
p_q = q + (bos * H + i_h) * K + o_k + H * K * i_t
p_k = k + (bos * H + i_h) * K + o_k + H * K * i_t
p_v = v + (bos * HV + i_hv) * V + o_v + HV * V * i_t
if IS_BETA_HEADWISE:
p_beta = beta + (bos * HV + i_hv) * V + o_v + HV * V * i_t
else:
p_beta = beta + bos * HV + i_hv + HV * i_t
p_g = g + bos * HV + i_hv + HV * i_t
p_o = o + ((i_k * all + bos) * HV + i_hv) * V + o_v + HV * V * i_t
b_q = tl.load(p_q, mask=mask_k, other=0).to(tl.float32)
b_k = tl.load(p_k, mask=mask_k, other=0).to(tl.float32)
b_v = tl.load(p_v, mask=mask_v, other=0).to(tl.float32)
b_g = tl.load(p_g).to(tl.float32)
if USE_QK_L2NORM_IN_KERNEL:
b_q = b_q / tl.sqrt(tl.sum(b_q * b_q) + 1e-6)
b_k = b_k / tl.sqrt(tl.sum(b_k * b_k) + 1e-6)
b_q = b_q * scale
# [BK, BV]
# b_h *= tl.exp(b_g)
b_h *= exp(b_g)
# [BV]
b_v -= tl.sum(b_h * b_k[:, None], 0)
if IS_BETA_HEADWISE:
b_beta = tl.load(p_beta, mask=mask_v, other=0).to(tl.float32)
else:
b_beta = tl.load(p_beta).to(tl.float32)
b_v *= b_beta
# [BK, BV]
b_h += b_k[:, None] * b_v[None, :]
# [BV]
b_o = tl.sum(b_h * b_q[:, None], 0)
tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=mask_v)
# keep the states for multi-query tokens
if INPLACE_FINAL_STATE:
p_ht = ht + tl.load(ssm_state_indices + i_n * stride_indices_seq +
i_t).to(tl.int64) * stride_final_state_token
else:
p_ht = ht + (bos + i_t) * stride_final_state_token
p_ht = p_ht + i_hv * K * V + o_k[:, None] * V + o_v[None, :]
tl.store(p_ht, b_h.to(p_ht.dtype.element_ty), mask=mask_h)

1
vllm_ascend/patch/platform/__init__.py
View File

@@ -18,6 +18,7 @@ import os
import vllm_ascend.patch.platform.patch_config # noqa
import vllm_ascend.patch.platform.patch_distributed # noqa
import vllm_ascend.patch.platform.patch_dynamo_vllm_backend # noqa
import vllm_ascend.patch.platform.patch_mamba_config # noqa
import vllm_ascend.patch.platform.patch_sched_yield # noqa

16
vllm_ascend/patch/platform/patch_dynamo_vllm_backend.py Normal file
View File

@@ -0,0 +1,16 @@
# mypy: ignore-errors
from typing import Any, Dict
import torch.fx as fx
from vllm.compilation.backends import VllmBackend
from vllm.compilation.caching import VllmSerializableFunction
_original_vllmbackend_call = VllmBackend.__call__
def __patch_call__(self, graph: fx.GraphModule, example_inputs,
options: Dict[str, Any]) -> VllmSerializableFunction:
return _original_vllmbackend_call(self, graph, example_inputs)
VllmBackend.__call__ = __patch_call__

90
vllm_ascend/patch/platform/patch_multiproc_executor.py
View File

@@ -1,24 +1,24 @@
import threading
import weakref
from concurrent.futures import ThreadPoolExecutor
from collections import deque
from collections.abc import Callable
from multiprocessing.synchronize import Lock as LockType
from typing import Optional
import vllm.v1.executor.multiproc_executor
from vllm import envs
from vllm.config import VllmConfig
from vllm.distributed.device_communicators.shm_broadcast import MessageQueue
from vllm.distributed.device_communicators.shm_broadcast import (Handle,
MessageQueue)
from vllm.utils.network_utils import (get_distributed_init_method,
get_loopback_ip, get_open_port)
from vllm.utils.system_utils import get_mp_context
from vllm.v1.executor.abstract import FailureCallback
from vllm.v1.executor.multiproc_executor import (
MultiprocExecutor, UnreadyWorkerProcHandle, WorkerProc,
FutureWrapper, MultiprocExecutor, UnreadyWorkerProcHandle, WorkerProc,
set_multiprocessing_worker_envs)
class AscendMultiprocExecutor(MultiprocExecutor):
supports_pp: bool = True
def _init_executor(self) -> None:
# Call self.shutdown at exit to clean up
@@ -26,10 +26,14 @@ class AscendMultiprocExecutor(MultiprocExecutor):
self._finalizer = weakref.finalize(self, self.shutdown)
self.is_failed = False
self.shutdown_event = threading.Event()
self.failure_callback: Optional[FailureCallback] = None
self.io_thread_pool: Optional[ThreadPoolExecutor] = None
self.failure_callback: FailureCallback | None = None
self.world_size = self.parallel_config.world_size
assert self.world_size % self.parallel_config.nnodes_within_dp == 0, (
f"global world_size ({self.parallel_config.world_size}) must be "
f"divisible by nnodes_within_dp "
f"({self.parallel_config.nnodes_within_dp}). ")
self.local_world_size = self.parallel_config.local_world_size
tensor_parallel_size = self.parallel_config.tensor_parallel_size
pp_parallel_size = self.parallel_config.pipeline_parallel_size
assert self.world_size == tensor_parallel_size * pp_parallel_size, (
@@ -45,27 +49,36 @@ class AscendMultiprocExecutor(MultiprocExecutor):
# get_loopback_ip() for communication.
distributed_init_method = get_distributed_init_method(
get_loopback_ip(), get_open_port())
self.rpc_broadcast_mq: MessageQueue | None = None
scheduler_output_handle: Handle | None = None
# Initialize worker and set up message queues for SchedulerOutputs
# and ModelRunnerOutputs
max_chunk_bytes = envs.VLLM_MQ_MAX_CHUNK_BYTES_MB * 1024 * 1024
self.rpc_broadcast_mq = MessageQueue(self.world_size,
self.world_size,
max_chunk_bytes=max_chunk_bytes)
scheduler_output_handle = self.rpc_broadcast_mq.export_handle()
if self.parallel_config.node_rank_within_dp == 0:
# For leader node within each dp rank,
# each dp will have its own leader multiproc executor.
max_chunk_bytes = envs.VLLM_MQ_MAX_CHUNK_BYTES_MB * 1024 * 1024
self.rpc_broadcast_mq = MessageQueue(
self.world_size,
self.local_world_size,
max_chunk_bytes=max_chunk_bytes,
connect_ip=self.parallel_config.master_addr,
)
scheduler_output_handle = self.rpc_broadcast_mq.export_handle()
# Create workers
context = get_mp_context()
shared_worker_lock = context.Lock()
unready_workers: list[UnreadyWorkerProcHandle] = []
success = False
try:
for rank in range(self.world_size):
global_start_rank = (self.local_world_size *
self.parallel_config.node_rank_within_dp)
for local_rank in range(self.local_world_size):
global_rank = global_start_rank + local_rank
unready_workers.append(
AscendWorkerProc.make_worker_process(
vllm_config=self.vllm_config,
local_rank=rank,
rank=rank,
local_rank=local_rank,
rank=global_rank,
distributed_init_method=distributed_init_method,
input_shm_handle=scheduler_output_handle,
shared_worker_lock=shared_worker_lock,
@@ -73,15 +86,38 @@ class AscendMultiprocExecutor(MultiprocExecutor):
# Workers must be created before wait_for_ready to avoid
# deadlock, since worker.init_device() does a device sync.
self.workers = WorkerProc.wait_for_ready(unready_workers)
# Wait for all local workers to be ready.
self.workers = AscendWorkerProc.wait_for_ready(unready_workers)
# Start background thread to monitor worker health if not in headless mode.
if self.monitor_workers:
self.start_worker_monitor()
self.response_mqs = []
# Only leader node have remote response mqs
if self.parallel_config.node_rank_within_dp == 0:
for rank in range(self.world_size):
if rank < self.local_world_size:
local_message_queue = self.workers[
rank].worker_response_mq
assert local_message_queue is not None
self.response_mqs.append(local_message_queue)
else:
remote_message_queue = self.workers[
0].peer_worker_response_mqs[rank]
assert remote_message_queue is not None
self.response_mqs.append(remote_message_queue)
# Ensure message queues are ready. Will deadlock if re-ordered
# Must be kept consistent with the WorkerProc.
self.rpc_broadcast_mq.wait_until_ready()
for w in self.workers:
w.worker_response_mq.wait_until_ready()
self.start_worker_monitor()
# Wait for all input mqs to be ready.
if self.rpc_broadcast_mq is not None:
self.rpc_broadcast_mq.wait_until_ready()
# Wait for all remote response mqs to be ready.
for response_mq in self.response_mqs:
response_mq.wait_until_ready()
success = True
finally:
if not success:
@@ -93,17 +129,9 @@ class AscendMultiprocExecutor(MultiprocExecutor):
self._ensure_worker_termination(
[uw.proc for uw in unready_workers])
# For pipeline parallel, we use a thread pool for asynchronous
# execute_model.
if self.max_concurrent_batches > 1:
# Note: must use only 1 IO thread to keep dequeue sequence
# from the response queue
# _async_aggregate_workers_output also assumes a single IO thread
self.io_thread_pool = ThreadPoolExecutor(
max_workers=1, thread_name_prefix="mp_exec_io")
self.futures_queue = deque[tuple[FutureWrapper, Callable]]()
self.output_rank = self._get_output_rank()
self.has_connector = self.vllm_config.kv_transfer_config is not None
class AscendWorkerProc(WorkerProc):

9
vllm_ascend/platform.py
View File

@@ -24,7 +24,7 @@ from vllm.logger import logger
from vllm.platforms import Platform, PlatformEnum
# todo: please remove it when solve cuda hard code in vllm
os.environ["VLLM_DISABLE_SHARED_EXPERTS_STREAM"] = "True"
os.environ["VLLM_DISABLE_SHARED_EXPERTS_STREAM"] = "1"
from vllm_ascend.ascend_config import (check_ascend_config, get_ascend_config,
init_ascend_config)
@@ -147,6 +147,8 @@ class NPUPlatform(Platform):
if enforce_eager:
logger.info("Compilation disabled, using eager mode by default")
compilation_config.mode = CompilationMode.NONE
if compilation_config.splitting_ops is None:
compilation_config.splitting_ops = []
compilation_config.cudagraph_num_of_warmups = 1
@@ -342,14 +344,11 @@ class NPUPlatform(Platform):
dtype,
kv_cache_dtype,
block_size,
use_v1,
use_mla,
has_sink=False,
use_sparse=False,
attn_type: str | None = None,
):
if not use_v1:
raise ValueError("vLLM Ascend does not support V0 engine.")
ascend_config = get_ascend_config()
if use_mla and ascend_config.enable_shared_expert_dp:

8
vllm_ascend/spec_decode/eagle_proposer.py
View File

@@ -136,7 +136,7 @@ class EagleProposer(Proposer):
)
def generate_token_ids(self,
valid_sampled_token_ids: list[list[int]],
valid_sampled_token_ids: list[np.ndarray],
sampling_metadata: SamplingMetadata = None,
scheduler_output: SchedulerOutput = None,
spec_decode_metadata: SpecDecodeMetadata = None,
@@ -149,7 +149,7 @@ class EagleProposer(Proposer):
attn_metadata = self._get_eagle_atten_dict(scheduler_output)
next_token_ids: list[int] = []
for i, token_ids in enumerate(valid_sampled_token_ids):
if token_ids:
if token_ids.shape[0] > 0:
# Common case.
next_token_id = token_ids[-1]
else:
@@ -161,7 +161,7 @@ class EagleProposer(Proposer):
scheduler_output.num_scheduled_tokens[req_id])
next_token_id = req_state.get_token_id(seq_len)
next_token_ids.append(next_token_id)
next_token_ids.append(next_token_id.item())
next_token_ids = torch.tensor(next_token_ids,
dtype=torch.int32,
device=self.device)
@@ -181,7 +181,7 @@ class EagleProposer(Proposer):
else:
num_draft_tokens = spec_decode_metadata.num_draft_tokens
num_rejected_tokens = [
n + 1 - len(valid_sampled_token_ids[i]) if n > 0 else 0
n + 1 - valid_sampled_token_ids[i].shape[0] if n > 0 else 0
for i, n in enumerate(num_draft_tokens)
]
num_rejected_tokens = torch.tensor(

3
vllm_ascend/spec_decode/interface.py
View File

@@ -1,6 +1,7 @@
import enum
from typing import Optional
import numpy as np
import torch
from vllm.config import CUDAGraphMode, VllmConfig
from vllm.v1.core.sched.output import SchedulerOutput
@@ -40,7 +41,7 @@ class Proposer:
raise NotImplementedError
def generate_token_ids(self,
valid_sampled_token_ids: list[list[int]],
valid_sampled_token_ids: list[np.ndarray],
sampling_metadata: SamplingMetadata = None,
scheduler_output: SchedulerOutput = None,
spec_decode_metadata: SpecDecodeMetadata = None,

12
vllm_ascend/spec_decode/mtp_proposer.py
View File

@@ -302,7 +302,8 @@ class MtpProposer(Proposer):
break
def generate_token_ids(self,
sampled_token_ids: list[list[int]],
sampled_token_ids: Union[torch.Tensor,
list[np.ndarray]],
sampling_metadata: SamplingMetadata = None,
scheduler_output: SchedulerOutput = None,
spec_decode_metadata: SpecDecodeMetadata = None,
@@ -379,6 +380,7 @@ class MtpProposer(Proposer):
common_attn_metadata.query_start_loc = \
query_start_loc_pcp_full[:num_reqs + 1]
if self.speculative_config.disable_padded_drafter_batch:
assert isinstance(sampled_token_ids, list)
# NOTE: Currently, MTP-fullgraph is incompatibility with pcp
token_indices_to_sample = None
common_attn_metadata, token_indices =\
@@ -437,7 +439,7 @@ class MtpProposer(Proposer):
def _prepare_inputs(
self,
common_attn_metadata: CommonAttentionMetadata,
sampled_token_ids: list[list[int]],
sampled_token_ids: list[np.ndarray],
num_draft_tokens: list[int],
) -> tuple[CommonAttentionMetadata, torch.Tensor]:
"""
@@ -895,7 +897,7 @@ class MtpProposer(Proposer):
def prepare_next_token_ids_cpu(
self,
sampled_token_ids: list[list[int]],
sampled_token_ids: list[np.ndarray],
requests: dict[str, CachedRequestState],
gpu_input_batch: InputBatch,
num_scheduled_tokens: dict[str, int],
@@ -910,7 +912,7 @@ class MtpProposer(Proposer):
req_ids = gpu_input_batch.req_ids
next_token_ids: list[int] = []
for i, token_ids in enumerate(sampled_token_ids):
if token_ids:
if token_ids.shape[0] > 0:
# Common case.
next_token_id = token_ids[-1]
else:
@@ -921,7 +923,7 @@ class MtpProposer(Proposer):
seq_len = req_state.num_computed_tokens + num_scheduled_tokens[
req_id]
next_token_id = req_state.get_token_id(seq_len)
next_token_ids.append(next_token_id)
next_token_ids.append(next_token_id.item())
next_token_ids = torch.tensor(next_token_ids,
dtype=torch.int32,
device=self.input_ids.device)

5
vllm_ascend/spec_decode/ngram_proposer.py
View File

@@ -1,3 +1,4 @@
import numpy as np
import torch
from vllm.config import CUDAGraphMode
from vllm.v1.spec_decode.ngram_proposer import \
@@ -30,7 +31,7 @@ class NgramProposer(VllmNgramProposer, Proposer):
pass
def generate_token_ids(self,
valid_sampled_token_ids,
valid_sampled_token_ids: list[np.ndarray],
sampling_metadata=None,
scheduler_output=None,
spec_decode_metadata=None,
@@ -41,7 +42,7 @@ class NgramProposer(VllmNgramProposer, Proposer):
aux_hidden_states=None) -> list[list[int]]:
valid_ngram_requests = []
for i, sampled_ids in enumerate(valid_sampled_token_ids):
num_sampled_ids = len(sampled_ids)
num_sampled_ids = sampled_ids.shape[0]
if not num_sampled_ids:
continue

6
vllm_ascend/torchair/models/qwen2.py
View File

@@ -248,7 +248,7 @@ class CustomQwen2Model(Qwen2Model):
if inputs_embeds is not None:
hidden_states = inputs_embeds
else:
hidden_states = self.get_input_embeddings(input_ids)
hidden_states = self.embed_input_ids(input_ids)
residual = None
else:
assert intermediate_tensors is not None
@@ -319,8 +319,8 @@ class CustomQwen2ForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
self.make_empty_intermediate_tensors = (
self.model.make_empty_intermediate_tensors)
def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
return self.model.get_input_embeddings(input_ids)
def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
return self.model.embed_input_ids(input_ids)
def forward(
self,

2
vllm_ascend/torchair/models/qwen3_moe.py
View File

@@ -426,7 +426,7 @@ class CustomQwen3MoeModel(Qwen3MoeModel):
if inputs_embeds is not None:
hidden_states = inputs_embeds
else:
hidden_states = self.get_input_embeddings(input_ids)
hidden_states = self.embed_input_ids(input_ids)
residual = None
else:
assert intermediate_tensors is not None

4
vllm_ascend/torchair/models/torchair_deepseek_v2.py
View File

@@ -1159,7 +1159,7 @@ class TorchairDeepseekV2Model(nn.Module):
make_empty_intermediate_tensors_factory(
["hidden_states", "residual"], config.hidden_size))
def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
return self.embed_tokens(input_ids)
def forward(
@@ -1175,7 +1175,7 @@ class TorchairDeepseekV2Model(nn.Module):
if inputs_embeds is not None:
hidden_states = inputs_embeds
else:
hidden_states = self.get_input_embeddings(input_ids)
hidden_states = self.embed_input_ids(input_ids)
residual = None
else:
assert intermediate_tensors is not None

8
vllm_ascend/torchair/models/torchair_pangu_moe.py
View File

@@ -808,7 +808,7 @@ class PanguProMoEModel(nn.Module):
make_empty_intermediate_tensors_factory(
["hidden_states", "residual"], config.hidden_size))
def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
return self.embed_tokens(input_ids)
def forward(
@@ -824,7 +824,7 @@ class PanguProMoEModel(nn.Module):
if inputs_embeds is not None:
hidden_states = inputs_embeds
else:
hidden_states = self.get_input_embeddings(input_ids)
hidden_states = self.embed_input_ids(input_ids)
residual = None
else:
assert intermediate_tensors is not None
@@ -916,8 +916,8 @@ class PanguProMoEForCausalLM(nn.Module, SupportsPP):
self.make_empty_intermediate_tensors = (
self.model.make_empty_intermediate_tensors)
def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
return self.model.get_input_embeddings(input_ids)
def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
return self.model.embed_input_ids(input_ids)
def forward(
self,

5
vllm_ascend/torchair/torchair_mla.py
View File

@@ -490,6 +490,11 @@ class AscendMLATorchairMetadataBuilder:
num_reqs_pad_size = (
graph_pad_size //
common_attn_metadata.decode_token_per_req - num_reqs)
# For the case when some request reach the max-tokens limit in this forward processing,
# so in this forward new_tokens scheduled is less than decode_token_per_req(1 + spec_token_num).
# Details can see PR:https://github.com/vllm-project/vllm/pull/27922
num_reqs_pad_size = max(0, num_reqs_pad_size)
padded_seq_lens = seq_lens.tolist(
) + [pad_value] * num_reqs_pad_size
else:

9
vllm_ascend/torchair/torchair_mtp_proposer.py
View File

@@ -1,5 +1,6 @@
import types
import numpy as np
import torch
import torch.nn as nn
import torchair
@@ -146,7 +147,7 @@ class TorchairMtpProposer(MtpProposer):
break
def generate_token_ids(self,
valid_sampled_token_ids: list[list[int]],
valid_sampled_token_ids: list[np.ndarray],
sampling_metadata: SamplingMetadata = None,
scheduler_output: SchedulerOutput = None,
spec_decode_metadata: SpecDecodeMetadata = None,
@@ -159,7 +160,7 @@ class TorchairMtpProposer(MtpProposer):
attn_metadata = attn_metadata['model.layers.0.self_attn.attn']
next_token_ids: list[int] = []
for i, token_ids in enumerate(valid_sampled_token_ids):
if token_ids:
if token_ids.shape[0] > 0:
# Common case.
next_token_id = token_ids[-1]
else:
@@ -170,7 +171,7 @@ class TorchairMtpProposer(MtpProposer):
seq_len = (req_state.num_computed_tokens +
scheduler_output.num_scheduled_tokens[req_id])
next_token_id = req_state.get_token_id(seq_len)
next_token_ids.append(next_token_id)
next_token_ids.append(next_token_id.item())
next_token_ids = torch.tensor(next_token_ids,
dtype=torch.int32,
device=self.device)
@@ -186,7 +187,7 @@ class TorchairMtpProposer(MtpProposer):
# TODO(woosuk): Refactor this.
num_draft_tokens = spec_decode_metadata.num_draft_tokens
num_rejected_tokens = [
n + 1 - len(valid_sampled_token_ids[i]) if n > 0 else 0
n + 1 - valid_sampled_token_ids[i].shape[0] if n > 0 else 0
for i, n in enumerate(num_draft_tokens)
]
num_rejected_tokens = torch.tensor(

192
vllm_ascend/worker/model_runner_v1.py
View File

@@ -152,7 +152,7 @@ if prefill_context_parallel_enable():
if TYPE_CHECKING:
import xgrammar as xgr # type: ignore[import-untyped]
from vllm.v1.core.sched.output import SchedulerOutput
from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput
else:
xgr = LazyLoader("xgr", globals(), "xgrammar")
@@ -243,15 +243,32 @@ class AsyncNPUModelRunnerOutput(AsyncModelRunnerOutput):
# Release the device tensor once the copy has completed
del self._sampled_token_ids
valid_sampled_token_ids = self._sampled_token_ids_cpu.tolist()
valid_sampled_token_ids: list[np.ndarray] = [
row for row in self._sampled_token_ids_cpu.numpy()
]
for i in self._invalid_req_indices:
valid_sampled_token_ids[i].clear()
valid_sampled_token_ids[i] = np.array([])
output = self._model_runner_output
output.sampled_token_ids = valid_sampled_token_ids
return output
class ExecuteModelState(NamedTuple):
"""Ephemeral cached state transferred between execute_model() and
sample_tokens(), after execute_model() returns None."""
scheduler_output: "SchedulerOutput"
logits: torch.Tensor
spec_decode_metadata: SpecDecodeMetadata | None
hidden_states: torch.Tensor
sample_hidden_states: torch.Tensor
aux_hidden_states: list[torch.Tensor] | None
kv_connector_output: KVConnectorOutput | None
attn_metadata: dict[str, Any]
positions: torch.Tensor
class NPUModelRunner(LoRAModelRunnerMixin):
def __init__(self, vllm_config: VllmConfig, device: torch.device):
@@ -604,6 +621,11 @@ class NPUModelRunner(LoRAModelRunnerMixin):
# TODO: EVS Support (Video tokens pruning) (see vllm#22980)
self.is_multimodal_pruning_enabled = False
# Ephemeral state transferred between execute_model() and sample_tokens().
self.execute_model_state: ExecuteModelState | None = None
self.transfer_event = torch.npu.Event()
def _set_up_drafter(self):
# Set up speculative decoding.
self.spec_attn_mask = None
@@ -865,39 +887,12 @@ class NPUModelRunner(LoRAModelRunnerMixin):
self.input_batch.refresh_metadata()
def _init_mrope_positions(self, req_state: CachedRequestState):
image_grid_thw = []
video_grid_thw = []
second_per_grid_ts = []
audio_feature_lengths = []
use_audio_in_video = False
assert req_state.mm_features is not None
for mm_feature in req_state.mm_features:
mm_item = mm_feature.data
if mm_item is None:
continue
mm_input = mm_item.get_data()
if (t := mm_input.get("image_grid_thw")) is not None:
image_grid_thw.append(t.tolist())
if (t := mm_input.get("video_grid_thw")) is not None:
video_grid_thw.append(t.tolist())
if (t := mm_input.get("second_per_grid_ts")) is not None:
second_per_grid_ts.append(t)
if (t := mm_input.get("audio_feature_lengths")) is not None:
audio_feature_lengths.append(t)
if mm_input.get("use_audio_in_video") is True:
use_audio_in_video = True
if supports_mrope(self.model):
req_state.mrope_positions, req_state.mrope_position_delta = \
self.model.get_mrope_input_positions(
req_state.prompt_token_ids,
hf_config=self.model_config.hf_config,
image_grid_thw=image_grid_thw,
video_grid_thw=video_grid_thw,
second_per_grid_ts=second_per_grid_ts,
audio_feature_lengths=audio_feature_lengths,
use_audio_in_video=use_audio_in_video,
)
assert supports_mrope(self.model), "MROPE is not supported"
req_state.mrope_positions, req_state.mrope_position_delta = \
self.model.get_mrope_input_positions(
req_state.prompt_token_ids,
req_state.mm_features,
)
def _sync_metadata_across_dp(
self, num_tokens: int,
@@ -1084,8 +1079,7 @@ class NPUModelRunner(LoRAModelRunnerMixin):
# 2. A list or tuple (length: num_items) of tensors, each of shape
# (feature_size, hidden_size) in case the feature size is dynamic
# depending on the input multimodal items.
curr_group_outputs = self.model.get_multimodal_embeddings(
**mm_kwargs_group)
curr_group_outputs = self.model.embed_multimodal(**mm_kwargs_group)
sanity_check_mm_encoder_outputs(
curr_group_outputs,
@@ -1636,7 +1630,7 @@ class NPUModelRunner(LoRAModelRunnerMixin):
mm_embeds, is_mm_embed = self._gather_mm_embeddings(
scheduler_output)
inputs_embeds = self.model.get_input_embeddings(
inputs_embeds = self.model.embed_input_ids(
input_ids,
multimodal_embeddings=mm_embeds,
is_multimodal=is_mm_embed,
@@ -1666,7 +1660,7 @@ class NPUModelRunner(LoRAModelRunnerMixin):
# Some tokens ids may need to become embeds
if token_ids_idx.numel() > 0:
token_ids = self.input_ids[token_ids_idx]
tokens_to_embeds = self.model.get_input_embeddings(
tokens_to_embeds = self.model.embed_input_ids(
input_ids=token_ids)
self.inputs_embeds.gpu[token_ids_idx] = tokens_to_embeds
@@ -2075,9 +2069,10 @@ class NPUModelRunner(LoRAModelRunnerMixin):
def apply_grammar_bitmask(
self,
scheduler_output: "SchedulerOutput",
grammar_output: "GrammarOutput",
logits: torch.Tensor,
) -> torch.Tensor:
grammar_bitmask = scheduler_output.grammar_bitmask
grammar_bitmask = grammar_output.grammar_bitmask
# We receive the structured output bitmask from the scheduler,
# compacted to contain bitmasks only for structured output requests.
@@ -2096,7 +2091,7 @@ class NPUModelRunner(LoRAModelRunnerMixin):
logit_index = batch_index + cumulative_offset
cumulative_offset += len(
scheduler_output.scheduled_spec_decode_tokens.get(req_id, []))
if req_id in scheduler_output.structured_output_request_ids:
if req_id in grammar_output.structured_output_request_ids:
struct_out_req_batch_indices[req_id] = logit_index
out_indices = []
@@ -2106,7 +2101,7 @@ class NPUModelRunner(LoRAModelRunnerMixin):
shape=(logits.shape[0],
grammar_bitmask.shape[1]))
cumulative_index = 0
for req_id in scheduler_output.structured_output_request_ids:
for req_id in grammar_output.structured_output_request_ids:
num_spec_tokens = len(
scheduler_output.scheduled_spec_decode_tokens.get(req_id, []))
if req_id in struct_out_req_batch_indices:
@@ -2137,7 +2132,7 @@ class NPUModelRunner(LoRAModelRunnerMixin):
def propose_draft_token_ids(
self,
valid_sampled_token_ids: Union[torch.Tensor, list[list[int]]],
valid_sampled_token_ids: Union[torch.Tensor, list[np.ndarray]],
sampling_metadata: SamplingMetadata,
scheduler_output: "SchedulerOutput",
spec_decode_metadata: SpecDecodeMetadata,
@@ -2270,7 +2265,11 @@ class NPUModelRunner(LoRAModelRunnerMixin):
self,
scheduler_output: "SchedulerOutput",
intermediate_tensors: Optional[IntermediateTensors] = None,
) -> Union[ModelRunnerOutput, AsyncModelRunnerOutput, IntermediateTensors]:
) -> Union[ModelRunnerOutput, IntermediateTensors] | None:
if self.execute_model_state is not None:
raise RuntimeError("State error: sample_tokens() must be called "
"after execute_model() returns None.")
with ProfileExecuteDuration().capture_async("prepare input"):
self._update_states(scheduler_output)
if not scheduler_output.total_num_scheduled_tokens:
@@ -2399,8 +2398,46 @@ class NPUModelRunner(LoRAModelRunnerMixin):
logits = model_output_broadcast_data["logits"]
# Apply structured output bitmasks if present
if scheduler_output.structured_output_request_ids:
logits = self.apply_grammar_bitmask(scheduler_output, logits)
self.execute_model_state = ExecuteModelState(
scheduler_output,
logits,
spec_decode_metadata,
hidden_states,
sample_hidden_states,
aux_hidden_states,
kv_connector_output,
attn_metadata,
positions,
)
return None
@torch.inference_mode
def sample_tokens(
self, grammar_output: "GrammarOutput | None"
) -> ModelRunnerOutput | AsyncModelRunnerOutput | IntermediateTensors:
if self.execute_model_state is None:
# Nothing to do (PP non-final rank case), output isn't used.
return None # noqa
need_dump = self.dump_enable and self.debugger is not None
# Unpack ephemeral state.
(
scheduler_output,
logits,
spec_decode_metadata,
hidden_states,
sample_hidden_states,
aux_hidden_states,
kv_connector_output,
attn_metadata,
positions,
) = self.execute_model_state
# Clear ephemeral state.
self.execute_model_state = None
# Apply structured output bitmasks if present.
if grammar_output is not None:
logits = self.apply_grammar_bitmask(scheduler_output,
grammar_output, logits)
with ProfileExecuteDuration().capture_async("Sample"):
# Sample the next token and get logprobs if needed.
@@ -2475,17 +2512,19 @@ class NPUModelRunner(LoRAModelRunnerMixin):
# Get the valid generated tokens.
max_gen_len = sampled_token_ids.shape[-1]
if max_gen_len == 1:
# No spec decode tokens.
valid_sampled_token_ids = sampled_token_ids.tolist()
# No spec decode tokens. It's a tensor.
valid_sampled_token_ids: list[np.ndarray] = [
row for row in sampled_token_ids.cpu().numpy()
]
else:
# Includes spec decode tokens.
# Includes spec decode tokens. It's a numpy array
valid_sampled_token_ids = self.rejection_sampler.parse_output(
sampled_token_ids,
self.input_batch.vocab_size,
)
# Mask out the sampled tokens that should not be sampled.
for i in discard_sampled_tokens_req_indices:
valid_sampled_token_ids[int(i)].clear()
valid_sampled_token_ids[int(i)] = np.array([])
else:
valid_sampled_token_ids = []
invalid_req_indices = discard_sampled_tokens_req_indices.tolist(
@@ -2511,16 +2550,17 @@ class NPUModelRunner(LoRAModelRunnerMixin):
# the sampled tokens back, because there's no direct communication
# between the first-stage worker and the last-stage worker.
for req_idx in range(num_sampled_tokens):
sampled_ids: np.ndarray | None
if self.use_async_scheduling:
sampled_ids = [-1] * 1 if \
req_idx not in invalid_req_indices_set else None
sampled_ids = (np.array([-1]) if req_idx
not in invalid_req_indices_set else None)
else:
sampled_ids = valid_sampled_token_ids[req_idx]
if not sampled_ids:
if sampled_ids is None or sampled_ids.shape[0] == 0:
continue
start_idx = self.input_batch.num_tokens_no_spec[req_idx]
end_idx = start_idx + len(sampled_ids)
end_idx = start_idx + sampled_ids.shape[0]
assert end_idx <= self.model_config.max_model_len, (
"Sampled token IDs exceed the max model length. "
f"Total number of tokens: {end_idx} > max_model_len: "
@@ -2534,7 +2574,7 @@ class NPUModelRunner(LoRAModelRunnerMixin):
self.input_batch.num_tokens[req_idx] = end_idx
req_id = self.input_batch.req_ids[req_idx]
req_state = self.requests[req_id]
req_state.output_token_ids.extend(sampled_ids)
req_state.output_token_ids.extend(sampled_ids.tolist())
def propose_draft_token_ids(sampled_token_ids):
assert self.spec_decode_common_attn_metadata is not None
@@ -2898,12 +2938,14 @@ class NPUModelRunner(LoRAModelRunnerMixin):
assert len(num_scheduled_tokens_list) == num_reqs
num_scheduled_tokens = np.array(num_scheduled_tokens_list,
dtype=np.int32)
num_sampled_tokens = np.ones(num_reqs, dtype=np.int32)
if not self.in_profile_run and self.dynamic_eplb:
self.eplb_updator.forward_before()
with self.maybe_dummy_run_with_lora(self.lora_config,
num_scheduled_tokens):
num_scheduled_tokens,
num_sampled_tokens):
if self.is_multimodal_model:
input_ids = None
inputs_embeds = self.inputs_embeds.gpu[:num_tokens]
@@ -3658,9 +3700,9 @@ class NPUModelRunner(LoRAModelRunnerMixin):
for k, v in attn_backend_layers.items()
}
def create_attn_groups(
attn_backends_map: dict[AttentionBackend, list[str]],
) -> list[AttentionGroup]:
def create_attn_groups(attn_backends_map: dict[AttentionBackend,
list[str]],
kv_cache_group_id: int) -> list[AttentionGroup]:
attn_groups: list[AttentionGroup] = []
for (attn_backend,
kv_cache_spec), layer_names in attn_backends_map.items():
@@ -3671,16 +3713,17 @@ class NPUModelRunner(LoRAModelRunnerMixin):
self.vllm_config,
self.device,
))
attn_group = AttentionGroup(attn_backend,
attn_metadata_builders,
layer_names, kv_cache_spec)
attn_group = AttentionGroup(attn_backend, layer_names,
kv_cache_spec, kv_cache_group_id,
attn_metadata_builders)
attn_groups.append(attn_group)
return attn_groups
for kv_cache_group_spec in kv_cache_config.kv_cache_groups:
for i, kv_cache_group_spec in enumerate(
kv_cache_config.kv_cache_groups):
attn_backends = get_attn_backends_for_group( # type: ignore
kv_cache_group_spec)
self.attn_groups.append(create_attn_groups(attn_backends))
self.attn_groups.append(create_attn_groups(attn_backends, i))
# Calculate reorder batch threshold (if needed)
self.calculate_reorder_batch_threshold()
@@ -3823,8 +3866,8 @@ class NPUModelRunner(LoRAModelRunnerMixin):
graph_support = builder.aclgraph_support.value
builder_aclgraph = builder.aclgraph_support
else:
graph_support = builder.cudagraph_support.value
builder_aclgraph = builder.cudagraph_support
graph_support = builder._cudagraph_support.value
builder_aclgraph = builder._cudagraph_support
if graph_support < min_ag_support.value:
min_ag_support = builder_aclgraph
min_ag_builder_name = builder.__class__.__name__
@@ -4422,3 +4465,18 @@ class NPUModelRunner(LoRAModelRunnerMixin):
self.input_ids_pcp_full_cpu[:total_num_scheduled_tokens_pcp_full],
non_blocking=True,
)
def _to_list(self, sampled_token_ids: torch.Tensor) -> list[np.ndarray]:
# This is a short term mitigation for issue mentioned in
# https://github.com/vllm-project/vllm/issues/22754.
# `tolist` would trigger a cuda wise stream sync, which
# would block other copy ops from other cuda streams.
# A cuda event sync would avoid such a situation. Since
# this is in the critical path of every single model
# forward loop, this has caused perf issue for a disagg
# setup.
pinned = self.sampled_token_ids_pinned_cpu[:sampled_token_ids.shape[0]]
pinned.copy_(sampled_token_ids, non_blocking=True)
self.transfer_event.record()
self.transfer_event.synchronize()
return [row for row in pinned.numpy()]

2
vllm_ascend/worker/npu_input_batch.py
View File

@@ -829,7 +829,7 @@ class InputBatch:
non_blocking=True)
def make_lora_inputs(
self, num_scheduled_tokens: np.ndarray
self, num_scheduled_tokens: np.ndarray, num_sampled_tokens: np.ndarray
) -> tuple[tuple[int, ...], tuple[int, ...], set[LoRARequest]]:
"""
Given the num_scheduled_tokens for each request in the batch, return

23
vllm_ascend/worker/worker_v1.py
View File

@@ -18,7 +18,8 @@
#
import copy
from typing import Optional, Union
from types import NoneType
from typing import Optional
import torch
import torch.nn as nn
@@ -37,7 +38,7 @@ from vllm.sequence import IntermediateTensors
from vllm.tasks import SupportedTask
from vllm.utils.mem_constants import GiB_bytes
from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE
from vllm.v1.core.sched.output import SchedulerOutput
from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput
from vllm.v1.kv_cache_interface import KVCacheConfig, KVCacheSpec
from vllm.v1.outputs import (EMPTY_MODEL_RUNNER_OUTPUT, AsyncModelRunnerOutput,
DraftTokenIds, ModelRunnerOutput)
@@ -206,6 +207,14 @@ class NPUWorker(WorkerBase):
device = torch.device(f"npu:{self.local_rank}")
NPUPlatform.set_device(device)
NPUPlatform.empty_cache()
visible_device_count = (torch.npu.device_count()
if torch.npu.is_available() else 0)
assert self.parallel_config.local_world_size <= visible_device_count, (
f"local_world_size ({self.parallel_config.local_world_size}) must be "
f"less than or equal to the number of visible devices "
f"({visible_device_count}).")
self.init_npu_memory = NPUPlatform.mem_get_info()[0]
# Initialize the distributed environment.
self._init_worker_distributed_environment()
@@ -266,7 +275,7 @@ class NPUWorker(WorkerBase):
def execute_model(
self,
scheduler_output: "SchedulerOutput",
) -> Optional[Union[ModelRunnerOutput, AsyncModelRunnerOutput]]:
) -> ModelRunnerOutput | None:
# enable msMonitor to monitor the performance of vllm-ascend
if envs_ascend.MSMONITOR_USE_DAEMON:
dp.step()
@@ -280,7 +289,7 @@ class NPUWorker(WorkerBase):
output = self.model_runner.execute_model(scheduler_output,
intermediate_tensors)
if isinstance(output, (ModelRunnerOutput, AsyncModelRunnerOutput)):
if isinstance(output, (ModelRunnerOutput, NoneType)):
return output
assert isinstance(output, IntermediateTensors)
@@ -304,6 +313,12 @@ class NPUWorker(WorkerBase):
output.kv_connector_output = kv_connector_output
return output
@torch.inference_mode()
def sample_tokens(
self, grammar_output: "GrammarOutput"
) -> ModelRunnerOutput | AsyncModelRunnerOutput:
return self.model_runner.sample_tokens(grammar_output)
def load_model(self) -> None:
if self.vllm_config.model_config.enable_sleep_mode:
allocator = CaMemAllocator.get_instance()