From 394a1c2e39324b584eb7cecc7ed409066da89199 Mon Sep 17 00:00:00 2001 From: weijinqian_v1 Date: Thu, 8 Jan 2026 13:39:36 +0800 Subject: [PATCH 01/11] [Refactor] refactor mask in attention_v1 and mla_v1. Signed-off-by: weijinqian_v1 --- vllm_ascend/attention/attention_v1.py | 14 ++-------- vllm_ascend/device/__init__.py | 0 vllm_ascend/device/device_op.py | 37 +++++++++++++++++++++++++++ vllm_ascend/utils.py | 2 ++ 4 files changed, 41 insertions(+), 12 deletions(-) create mode 100644 vllm_ascend/device/__init__.py create mode 100644 vllm_ascend/device/device_op.py diff --git a/vllm_ascend/attention/attention_v1.py b/vllm_ascend/attention/attention_v1.py index d19d3369b99..a7da06e099f 100644 --- a/vllm_ascend/attention/attention_v1.py +++ b/vllm_ascend/attention/attention_v1.py @@ -43,6 +43,7 @@ from vllm_ascend.compilation.acl_graph import ( get_draft_graph_params, get_graph_params, update_draft_graph_params_workspaces, update_graph_params_workspaces) +from vllm_ascend.device.device_op import DeviceOperator from vllm_ascend.utils import (AscendDeviceType, get_ascend_device_type, weak_ref_tensors) @@ -669,18 +670,7 @@ def reshape_and_cache( if self.key_cache is None: self.key_cache, self.value_cache = kv_cache[0], kv_cache[1] slots = attn_metadata.slot_mapping - if get_ascend_device_type() == AscendDeviceType.A5: - # TODO: Once eagle running to here, it may has error because of the 0 dim of slot_mapping. - # Should check if the 0 dim of slot_mapping must equal to the 0 dim of key. - # If it's necessary, the slots should be sliced. - torch_npu.npu_scatter_pa_kv_cache( - key=key[:attn_metadata.num_actual_tokens], - value=value[:attn_metadata.num_actual_tokens].contiguous(), - key_cache=self.key_cache, - value_cache=self.value_cache, - slot_mapping=slots) - else: - torch_npu._npu_reshape_and_cache( + DeviceOperator.reshape_and_cache( key=key[:attn_metadata.num_actual_tokens], value=value[:attn_metadata.num_actual_tokens], key_cache=self.key_cache, diff --git a/vllm_ascend/device/__init__.py b/vllm_ascend/device/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/vllm_ascend/device/device_op.py b/vllm_ascend/device/device_op.py new file mode 100644 index 00000000000..5347499c97c --- /dev/null +++ b/vllm_ascend/device/device_op.py @@ -0,0 +1,37 @@ + +from typing import Optional + + +import torch_npu + +from vllm_ascend.utils import AscendDeviceType + + +class CommonDeviceOperator(object): + @classmethod + def reshape_and_cache(cls, key, value, key_cache, value_cache, slot_mapping): + torch_npu._npu_reshape_and_cache( + key=key, + value=value, + key_cache=key_cache, + value_cache=value_cache, + slot_indices=slot_mapping) + +class A5Operator(CommonDeviceOperator): + @classmethod + def reshape_and_cache(cls, key, value, key_cache, value_cache, slot_mapping): + torch_npu.npu_scatter_pa_kv_cache( + key=key, + value=value, + key_cache=key_cache, + value_cache=value_cache, + slot_mapping=slot_mapping) + + +DeviceOperator: Optional[CommonDeviceOperator.__class__] = None + +def set_device(ascend_device_type): + global DeviceOperator + if ascend_device_type == AscendDeviceType.A5: + DeviceOperator = A5Operator + DeviceOperator = CommonDeviceOperator \ No newline at end of file diff --git a/vllm_ascend/utils.py b/vllm_ascend/utils.py index 8056555461a..e44340ea1d7 100644 --- a/vllm_ascend/utils.py +++ b/vllm_ascend/utils.py @@ -35,6 +35,7 @@ import vllm_ascend.envs as envs_ascend from vllm_ascend.ascend_config import WeightPrefetchConfig, get_ascend_config +from vllm_ascend.device.device_op import set_device if TYPE_CHECKING: from vllm.config import VllmConfig @@ -737,6 +738,7 @@ def _init_ascend_device_type(): global _ascend_device_type from vllm_ascend import _build_info # type: ignore _ascend_device_type = AscendDeviceType[_build_info.__device_type__] + set_device(_ascend_device_type) def check_ascend_device_type(): From 1039ee8c8215914f246cea5f9330656a98e27796 Mon Sep 17 00:00:00 2001 From: weijinqian_v1 Date: Thu, 8 Jan 2026 13:40:18 +0800 Subject: [PATCH 02/11] [Refactor] refactor mask in attention_v1 and mla_v1. Signed-off-by: weijinqian_v1 --- vllm_ascend/device/device_op.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/vllm_ascend/device/device_op.py b/vllm_ascend/device/device_op.py index 5347499c97c..35f78748370 100644 --- a/vllm_ascend/device/device_op.py +++ b/vllm_ascend/device/device_op.py @@ -17,7 +17,7 @@ def reshape_and_cache(cls, key, value, key_cache, value_cache, slot_mapping): value_cache=value_cache, slot_indices=slot_mapping) -class A5Operator(CommonDeviceOperator): +class A5DeviceOperator(CommonDeviceOperator): @classmethod def reshape_and_cache(cls, key, value, key_cache, value_cache, slot_mapping): torch_npu.npu_scatter_pa_kv_cache( @@ -33,5 +33,5 @@ def reshape_and_cache(cls, key, value, key_cache, value_cache, slot_mapping): def set_device(ascend_device_type): global DeviceOperator if ascend_device_type == AscendDeviceType.A5: - DeviceOperator = A5Operator + DeviceOperator = A5DeviceOperator DeviceOperator = CommonDeviceOperator \ No newline at end of file From ff6d2e9bdf04f835f7e80ae9a4aca374e02f356a Mon Sep 17 00:00:00 2001 From: weijinqian_v1 Date: Thu, 8 Jan 2026 20:39:24 +0800 Subject: [PATCH 03/11] [Refactor] refactor mask in attention_v1 and mla_v1. Signed-off-by: weijinqian_v1 --- vllm_ascend/attention/attention_v1.py | 13 +++--- vllm_ascend/device/device_op.py | 67 ++++++++++++++++++--------- 2 files changed, 50 insertions(+), 30 deletions(-) diff --git a/vllm_ascend/attention/attention_v1.py b/vllm_ascend/attention/attention_v1.py index a7da06e099f..d8b60f1dc44 100644 --- a/vllm_ascend/attention/attention_v1.py +++ b/vllm_ascend/attention/attention_v1.py @@ -44,8 +44,7 @@ get_draft_graph_params, get_graph_params, update_draft_graph_params_workspaces, update_graph_params_workspaces) from vllm_ascend.device.device_op import DeviceOperator -from vllm_ascend.utils import (AscendDeviceType, get_ascend_device_type, - weak_ref_tensors) +from vllm_ascend.utils import weak_ref_tensors # default max value of sliding window size SWA_INT_MAX = 2147483647 @@ -671,11 +670,11 @@ def reshape_and_cache( self.key_cache, self.value_cache = kv_cache[0], kv_cache[1] slots = attn_metadata.slot_mapping DeviceOperator.reshape_and_cache( - key=key[:attn_metadata.num_actual_tokens], - value=value[:attn_metadata.num_actual_tokens], - key_cache=self.key_cache, - value_cache=self.value_cache, - slot_indices=slots[:attn_metadata.num_actual_tokens]) + key=key[:attn_metadata.num_actual_tokens], + value=value[:attn_metadata.num_actual_tokens], + key_cache=self.key_cache, + value_cache=self.value_cache, + slot_mapping=slots[:attn_metadata.num_actual_tokens]) if self.is_kv_producer: attn_metadata.reshape_cache_event.record() return key, value diff --git a/vllm_ascend/device/device_op.py b/vllm_ascend/device/device_op.py index 35f78748370..84f416518f8 100644 --- a/vllm_ascend/device/device_op.py +++ b/vllm_ascend/device/device_op.py @@ -1,37 +1,58 @@ - +# Adapt from https://github.com/vllm-project/vllm/blob/main/vllm/v1/worker/gpu/model_runner.py +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Copyright (c) 2025 Huawei Technologies Co., Ltd. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# This file is a part of the vllm-ascend project. +# from typing import Optional - import torch_npu -from vllm_ascend.utils import AscendDeviceType +from vllm_ascend.utils import AscendDeviceType, get_ascend_device_type class CommonDeviceOperator(object): + @classmethod - def reshape_and_cache(cls, key, value, key_cache, value_cache, slot_mapping): - torch_npu._npu_reshape_and_cache( - key=key, - value=value, - key_cache=key_cache, - value_cache=value_cache, - slot_indices=slot_mapping) + def reshape_and_cache(cls, key, value, key_cache, value_cache, + slot_mapping): + torch_npu._npu_reshape_and_cache(key=key, + value=value, + key_cache=key_cache, + value_cache=value_cache, + slot_indices=slot_mapping) + class A5DeviceOperator(CommonDeviceOperator): - @classmethod - def reshape_and_cache(cls, key, value, key_cache, value_cache, slot_mapping): - torch_npu.npu_scatter_pa_kv_cache( - key=key, - value=value, - key_cache=key_cache, - value_cache=value_cache, - slot_mapping=slot_mapping) + @classmethod + def reshape_and_cache(cls, key, value, key_cache, value_cache, + slot_mapping): + torch_npu.npu_scatter_pa_kv_cache(key=key, + value=value, + key_cache=key_cache, + value_cache=value_cache, + slot_mapping=slot_mapping) -DeviceOperator: Optional[CommonDeviceOperator.__class__] = None -def set_device(ascend_device_type): - global DeviceOperator +def get_device_operator(): + ascend_device_type = get_ascend_device_type() if ascend_device_type == AscendDeviceType.A5: - DeviceOperator = A5DeviceOperator - DeviceOperator = CommonDeviceOperator \ No newline at end of file + return A5DeviceOperator + return CommonDeviceOperator + + +DeviceOperator: Optional[ + CommonDeviceOperator.__class__] = get_device_operator() From 66ed782be0f369154189692dcca3d7c4e0dffff5 Mon Sep 17 00:00:00 2001 From: weijinqian_v1 Date: Thu, 8 Jan 2026 20:41:21 +0800 Subject: [PATCH 04/11] [Refactor] refactor mask in attention_v1 and mla_v1. Signed-off-by: weijinqian_v1 --- vllm_ascend/utils.py | 1 - 1 file changed, 1 deletion(-) diff --git a/vllm_ascend/utils.py b/vllm_ascend/utils.py index e44340ea1d7..ca4df6d545e 100644 --- a/vllm_ascend/utils.py +++ b/vllm_ascend/utils.py @@ -35,7 +35,6 @@ import vllm_ascend.envs as envs_ascend from vllm_ascend.ascend_config import WeightPrefetchConfig, get_ascend_config -from vllm_ascend.device.device_op import set_device if TYPE_CHECKING: from vllm.config import VllmConfig From 06f56ad400cecc7ef3165764c5d0aee5148d3d6c Mon Sep 17 00:00:00 2001 From: weijinqian_v1 Date: Thu, 8 Jan 2026 20:42:54 +0800 Subject: [PATCH 05/11] [Refactor] refactor mask in attention_v1 and mla_v1. Signed-off-by: weijinqian_v1 --- vllm_ascend/utils.py | 1 - 1 file changed, 1 deletion(-) diff --git a/vllm_ascend/utils.py b/vllm_ascend/utils.py index ca4df6d545e..8056555461a 100644 --- a/vllm_ascend/utils.py +++ b/vllm_ascend/utils.py @@ -737,7 +737,6 @@ def _init_ascend_device_type(): global _ascend_device_type from vllm_ascend import _build_info # type: ignore _ascend_device_type = AscendDeviceType[_build_info.__device_type__] - set_device(_ascend_device_type) def check_ascend_device_type(): From 8ff0923b356cfe8eb99a829863ceb4673d28e097 Mon Sep 17 00:00:00 2001 From: weijinqian_v1 Date: Thu, 8 Jan 2026 20:43:31 +0800 Subject: [PATCH 06/11] [Refactor] refactor mask in attention_v1 and mla_v1. Signed-off-by: weijinqian_v1 --- vllm_ascend/device/device_op.py | 1 - 1 file changed, 1 deletion(-) diff --git a/vllm_ascend/device/device_op.py b/vllm_ascend/device/device_op.py index 84f416518f8..06288acd5c8 100644 --- a/vllm_ascend/device/device_op.py +++ b/vllm_ascend/device/device_op.py @@ -1,4 +1,3 @@ -# Adapt from https://github.com/vllm-project/vllm/blob/main/vllm/v1/worker/gpu/model_runner.py # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project # Copyright (c) 2025 Huawei Technologies Co., Ltd. All Rights Reserved. From f0e49e5f8a0a67cc5d8d525863069013ac40bfae Mon Sep 17 00:00:00 2001 From: weijinqian_v1 Date: Fri, 9 Jan 2026 09:18:25 +0800 Subject: [PATCH 07/11] [Refactor] refactor mask in attention_v1 and mla_v1. Signed-off-by: weijinqian_v1 --- vllm_ascend/device/device_op.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/vllm_ascend/device/device_op.py b/vllm_ascend/device/device_op.py index 06288acd5c8..0fd9e33d306 100644 --- a/vllm_ascend/device/device_op.py +++ b/vllm_ascend/device/device_op.py @@ -15,7 +15,7 @@ # limitations under the License. # This file is a part of the vllm-ascend project. # -from typing import Optional +from typing import Optional, Type import torch_npu @@ -53,5 +53,4 @@ def get_device_operator(): return CommonDeviceOperator -DeviceOperator: Optional[ - CommonDeviceOperator.__class__] = get_device_operator() +DeviceOperator: Optional[Type['CommonDeviceOperator']] = get_device_operator() From 3b5299a0ffbf9568874e9102b494c69cd3f9f3b4 Mon Sep 17 00:00:00 2001 From: weijinqian_v1 Date: Sat, 10 Jan 2026 09:39:50 +0800 Subject: [PATCH 08/11] [Refactor] rename class name. Signed-off-by: weijinqian_v1 --- vllm_ascend/device/device_op.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/vllm_ascend/device/device_op.py b/vllm_ascend/device/device_op.py index 0fd9e33d306..d74bd847416 100644 --- a/vllm_ascend/device/device_op.py +++ b/vllm_ascend/device/device_op.py @@ -22,7 +22,7 @@ from vllm_ascend.utils import AscendDeviceType, get_ascend_device_type -class CommonDeviceOperator(object): +class BaseDeviceAdaptor(object): @classmethod def reshape_and_cache(cls, key, value, key_cache, value_cache, @@ -34,7 +34,7 @@ def reshape_and_cache(cls, key, value, key_cache, value_cache, slot_indices=slot_mapping) -class A5DeviceOperator(CommonDeviceOperator): +class A5DeviceAdaptor(BaseDeviceAdaptor): @classmethod def reshape_and_cache(cls, key, value, key_cache, value_cache, @@ -49,8 +49,8 @@ def reshape_and_cache(cls, key, value, key_cache, value_cache, def get_device_operator(): ascend_device_type = get_ascend_device_type() if ascend_device_type == AscendDeviceType.A5: - return A5DeviceOperator - return CommonDeviceOperator + return A5DeviceAdaptor + return BaseDeviceAdaptor -DeviceOperator: Optional[Type['CommonDeviceOperator']] = get_device_operator() +DeviceOperator: Optional[Type['BaseDeviceAdaptor']] = get_device_operator() From 11801205c6ac3ca60e9620665c48f8461e76ee7d Mon Sep 17 00:00:00 2001 From: weijinqian_v1 Date: Sat, 10 Jan 2026 09:43:59 +0800 Subject: [PATCH 09/11] [Refactor] rename class name. Signed-off-by: weijinqian_v1 --- vllm_ascend/device/device_op.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/vllm_ascend/device/device_op.py b/vllm_ascend/device/device_op.py index d74bd847416..022513c93f3 100644 --- a/vllm_ascend/device/device_op.py +++ b/vllm_ascend/device/device_op.py @@ -46,11 +46,11 @@ def reshape_and_cache(cls, key, value, key_cache, value_cache, slot_mapping=slot_mapping) -def get_device_operator(): +def get_device_adaptor(): ascend_device_type = get_ascend_device_type() if ascend_device_type == AscendDeviceType.A5: return A5DeviceAdaptor return BaseDeviceAdaptor -DeviceOperator: Optional[Type['BaseDeviceAdaptor']] = get_device_operator() +DeviceOperator: Optional[Type['BaseDeviceAdaptor']] = get_device_adaptor() From fcef38736acd513fb0f9221660e5ac2e0bfb8b27 Mon Sep 17 00:00:00 2001 From: weijinqian_v1 Date: Sat, 10 Jan 2026 19:26:12 +0800 Subject: [PATCH 10/11] [Refactor] rename class name. Signed-off-by: weijinqian_v1 --- vllm_ascend/device/device_op.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vllm_ascend/device/device_op.py b/vllm_ascend/device/device_op.py index 022513c93f3..ccd874dd618 100644 --- a/vllm_ascend/device/device_op.py +++ b/vllm_ascend/device/device_op.py @@ -40,7 +40,7 @@ class A5DeviceAdaptor(BaseDeviceAdaptor): def reshape_and_cache(cls, key, value, key_cache, value_cache, slot_mapping): torch_npu.npu_scatter_pa_kv_cache(key=key, - value=value, + value=value.contiguous(), key_cache=key_cache, value_cache=value_cache, slot_mapping=slot_mapping) From 6bed04f400db645e0337acb01f6db70c80359c61 Mon Sep 17 00:00:00 2001 From: weijinqian_v1 Date: Sun, 11 Jan 2026 21:37:35 +0800 Subject: [PATCH 11/11] [Refactor] rename class name. Signed-off-by: weijinqian_v1 --- vllm_ascend/attention/attention_v1.py | 9 ++++----- 1 file changed, 4 insertions(+), 5 deletions(-) diff --git a/vllm_ascend/attention/attention_v1.py b/vllm_ascend/attention/attention_v1.py index a2a8e24ea90..0479de6fd10 100644 --- a/vllm_ascend/attention/attention_v1.py +++ b/vllm_ascend/attention/attention_v1.py @@ -45,8 +45,7 @@ update_draft_graph_params_workspaces, update_graph_params_workspaces) from vllm_ascend.device.device_op import DeviceOperator from vllm_ascend.ops.flashcomm2_oshard_manager import flashcomm2_oshard_manager -from vllm_ascend.utils import (AscendDeviceType, get_ascend_device_type, - weak_ref_tensors) +from vllm_ascend.utils import weak_ref_tensors # default max value of sliding window size SWA_INT_MAX = 2147483647 @@ -685,13 +684,13 @@ def reshape_and_cache( encoder_decoder = (self.attn_type == AttentionType.ENCODER_DECODER) DeviceOperator.reshape_and_cache( key=key[:attn_metadata.num_actual_tokens] - if not encoder_decoder else key, + if not encoder_decoder else key, value=value[:attn_metadata.num_actual_tokens] - if not encoder_decoder else value, + if not encoder_decoder else value, key_cache=self.key_cache, value_cache=self.value_cache, slot_mapping=slots[:attn_metadata.num_actual_tokens] - if not encoder_decoder else slots) + if not encoder_decoder else slots) if self.is_kv_producer: attn_metadata.reshape_cache_event.record() return key, value