0
vllm_ascend/patch/worker/patch_v2/__init__.py
Normal file
0
vllm_ascend/patch/worker/patch_v2/__init__.py
Normal file
11
vllm_ascend/patch/worker/patch_v2/patch_attn_utils.py
Normal file
11
vllm_ascend/patch/worker/patch_v2/patch_attn_utils.py
Normal file
@@ -0,0 +1,11 @@
|
||||
import vllm
|
||||
|
||||
from vllm_ascend.worker.v2.attn_utils import (
|
||||
_allocate_kv_cache,
|
||||
_reshape_kv_cache_v2,
|
||||
get_kv_cache_spec,
|
||||
)
|
||||
|
||||
vllm.v1.worker.gpu.attn_utils._allocate_kv_cache = _allocate_kv_cache
|
||||
vllm.v1.worker.gpu.attn_utils._reshape_kv_cache = _reshape_kv_cache_v2
|
||||
vllm.v1.worker.gpu.model_runner.get_kv_cache_spec = get_kv_cache_spec
|
||||
25
vllm_ascend/patch/worker/patch_v2/patch_block_table.py
Normal file
25
vllm_ascend/patch/worker/patch_v2/patch_block_table.py
Normal file
@@ -0,0 +1,25 @@
|
||||
# Adapt from https://github.com/vllm-project/vllm/blob/main/vllm/v1/worker/gpu/block_table.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 vllm.v1.worker.gpu import model_runner
|
||||
|
||||
from vllm_ascend.worker.v2.block_table import AscendBlockTables
|
||||
|
||||
# vllm-ascend need to initialize slot mapping as torch.int32 dtype,
|
||||
# but vllm default is torch.int64 dtype.
|
||||
model_runner.BlockTables = AscendBlockTables
|
||||
27
vllm_ascend/patch/worker/patch_v2/patch_input_batch.py
Normal file
27
vllm_ascend/patch/worker/patch_v2/patch_input_batch.py
Normal file
@@ -0,0 +1,27 @@
|
||||
# Adapt from https://github.com/vllm-project/vllm/blob/main/vllm/v1/worker/gpu/input_batch.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.
|
||||
#
|
||||
|
||||
|
||||
# 显式导入模块,确保模块被加载后再进行 patch
|
||||
from vllm.v1.worker.gpu import cudagraph_utils, model_runner
|
||||
|
||||
from vllm_ascend.worker.v2.input_batch import AscendInputBatch
|
||||
|
||||
cudagraph_utils.InputBatch = AscendInputBatch
|
||||
model_runner.InputBatch = AscendInputBatch
|
||||
26
vllm_ascend/patch/worker/patch_v2/patch_model_state.py
Normal file
26
vllm_ascend/patch/worker/patch_v2/patch_model_state.py
Normal file
@@ -0,0 +1,26 @@
|
||||
# Adapt from https://github.com/vllm-project/vllm/blob/main/vllm/v1/worker/gpu/model_states/default.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 vllm.v1.worker.gpu import model_runner
|
||||
|
||||
from vllm_ascend.worker.v2.model_states import init_asecnd_model_state
|
||||
|
||||
# prepare_attn in AscendModelState is different from vllm,
|
||||
# we need to override init_model_state.
|
||||
model_runner.init_model_state = init_asecnd_model_state
|
||||
34
vllm_ascend/patch/worker/patch_v2/patch_triton.py
Normal file
34
vllm_ascend/patch/worker/patch_v2/patch_triton.py
Normal file
@@ -0,0 +1,34 @@
|
||||
from vllm.v1.worker.gpu import input_batch, model_runner, structured_outputs
|
||||
from vllm.v1.worker.gpu.sample import bad_words, gumbel, logprob, penalties, prompt_logprob, sampler, states
|
||||
from vllm.v1.worker.gpu.spec_decode import rejection_sampler, rejection_sampler_utils
|
||||
from vllm.v1.worker.gpu.spec_decode.eagle import speculator
|
||||
|
||||
from vllm_ascend.worker.v2.input_batch import post_update
|
||||
from vllm_ascend.worker.v2.sample.bad_words import apply_bad_words
|
||||
from vllm_ascend.worker.v2.sample.gumbel import apply_temperature, gumbel_sample
|
||||
from vllm_ascend.worker.v2.sample.logprob import compute_token_logprobs, compute_topk_logprobs
|
||||
from vllm_ascend.worker.v2.sample.min_p import apply_min_p
|
||||
from vllm_ascend.worker.v2.sample.penalties import apply_penalties, bincount
|
||||
from vllm_ascend.worker.v2.spec_decode.rejection_sampler_utils import (
|
||||
rejection_sample as npu_rejection_sample,
|
||||
)
|
||||
from vllm_ascend.worker.v2.structured_outputs import _apply_grammar_bitmask_kernel
|
||||
|
||||
penalties.apply_penalties = apply_penalties
|
||||
# because sampler.py and speculator.py are imported before this patch, they must be overridden
|
||||
sampler.gumbel_sample = gumbel_sample
|
||||
input_batch.post_update = post_update
|
||||
prompt_logprob.compute_topk_logprobs = compute_topk_logprobs
|
||||
sampler.compute_topk_logprobs = compute_topk_logprobs
|
||||
rejection_sampler.compute_topk_logprobs = compute_topk_logprobs
|
||||
states.apply_min_p = apply_min_p
|
||||
penalties.bincount = bincount
|
||||
speculator.gumbel_sample = gumbel_sample
|
||||
model_runner.post_update = post_update
|
||||
bad_words.apply_bad_words = apply_bad_words
|
||||
gumbel.apply_temperature = apply_temperature
|
||||
states.apply_temperature = apply_temperature
|
||||
logprob.compute_token_logprobs = compute_token_logprobs
|
||||
structured_outputs._apply_grammar_bitmask_kernel = _apply_grammar_bitmask_kernel
|
||||
rejection_sampler_utils.rejection_sample = npu_rejection_sample
|
||||
rejection_sampler.rejection_sample = npu_rejection_sample
|
||||
@@ -0,0 +1,3 @@
|
||||
# Reuse the platform patch. EngineCore subprocesses only load global/platform
|
||||
# patches, while workers also import this compatibility module.
|
||||
import vllm_ascend.patch.platform.patch_use_v2_model_runner # noqa: F401
|
||||
161
vllm_ascend/patch/worker/patch_v2/patch_uva.py
Normal file
161
vllm_ascend/patch/worker/patch_v2/patch_uva.py
Normal file
@@ -0,0 +1,161 @@
|
||||
# Adapt from https://github.com/vllm-project/vllm/blob/main/vllm/v1/worker/gpu/block_table.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.
|
||||
#
|
||||
import os
|
||||
from collections.abc import Callable, Sequence
|
||||
from importlib.metadata import version
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
import vllm.v1.worker.gpu.buffer_utils
|
||||
from vllm.logger import logger
|
||||
|
||||
|
||||
def check_triton_ascend_version_valid() -> bool:
|
||||
"""
|
||||
Check triton-ascend version and warn about UVA feature disablement.
|
||||
If the installed version isn't affected by the UVA issue, return True.
|
||||
"""
|
||||
# Triton Ascend versions affected by the UVA pointer validation issue.
|
||||
UVA_INCOMPATIBLE_VERSIONS = ("3.2.1", "3.2.2")
|
||||
installed_version = version("triton-ascend")
|
||||
if installed_version in UVA_INCOMPATIBLE_VERSIONS:
|
||||
logger.warning(
|
||||
"triton-ascend %s disables the UVA feature.\n"
|
||||
"Related bug issue: https://github.com/triton-lang/triton-ascend/issues/783",
|
||||
installed_version,
|
||||
)
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
def is_uva_available() -> bool:
|
||||
"""check if uva feature is supported in this environment"""
|
||||
# FIXME(chenboxun): Some triton-ascend versions reject pinned CPU tensors.
|
||||
# Thus UVA is disabled for affected versions.
|
||||
# (Related bug issue link: https://github.com/triton-lang/triton-ascend/issues/783)
|
||||
return (
|
||||
"pinned_mem_register:True" in os.environ.get("PYTORCH_NPU_ALLOC_CONF", {})
|
||||
and check_triton_ascend_version_valid()
|
||||
)
|
||||
|
||||
|
||||
def get_row_indices_from_key(key: int | slice | tuple, dim_size: int) -> set[int]:
|
||||
"""get the set of row indices involved in the given key."""
|
||||
if isinstance(key, int):
|
||||
# parse index such as np[1]
|
||||
key = key if key >= 0 else dim_size + key
|
||||
# handle negative index
|
||||
if key < 0 or key >= dim_size:
|
||||
raise IndexError(f"row index {key} out of [0, {dim_size})")
|
||||
return {key}
|
||||
elif isinstance(key, slice):
|
||||
# parse slice such as np[1:3]
|
||||
start, stop, step = key.indices(dim_size)
|
||||
return set(range(start, stop, step))
|
||||
elif isinstance(key, tuple):
|
||||
# parse row slice such as np[1,:100]
|
||||
if len(key) == 0:
|
||||
return set(range(dim_size))
|
||||
return get_row_indices_from_key(key[0], dim_size)
|
||||
else:
|
||||
# for other types such as list/ndarray, we return all rows.
|
||||
return set(range(dim_size))
|
||||
|
||||
|
||||
class MonitoredNumPyArray:
|
||||
"""A wrapper around a NumPy array that monitors modifications."""
|
||||
|
||||
def __init__(self, array: np.ndarray, callback: Callable):
|
||||
self._array = array
|
||||
self._callback = callback
|
||||
|
||||
def __setitem__(self, key, value):
|
||||
self._array[key] = value
|
||||
dim_size = self._array.shape[0]
|
||||
row_indices = get_row_indices_from_key(key, dim_size)
|
||||
for row in row_indices:
|
||||
self._callback(row)
|
||||
|
||||
def __getitem__(self, key):
|
||||
return self._array[key]
|
||||
|
||||
def __getattr__(self, name):
|
||||
return getattr(self._array, name)
|
||||
|
||||
|
||||
class MonitoredTorchTensor:
|
||||
"""A wrapper around a torch tensor that monitors modifications."""
|
||||
|
||||
def __init__(self, tensor: torch.Tensor, callback: Callable):
|
||||
self._tensor = tensor
|
||||
self._callback = callback
|
||||
|
||||
def __setitem__(self, key, value):
|
||||
self._tensor[key] = value
|
||||
dim_size = self._tensor.size(0)
|
||||
row_indices = get_row_indices_from_key(key, dim_size)
|
||||
for row in row_indices:
|
||||
self._callback(row)
|
||||
|
||||
def __getitem__(self, key):
|
||||
return self._tensor[key]
|
||||
|
||||
def __getattr__(self, name):
|
||||
return getattr(self._tensor, name)
|
||||
|
||||
|
||||
class UvaBufferWrapper:
|
||||
"""
|
||||
Ascend NPU doesn't support UVA tensors directly.
|
||||
This is a wrapper class that provides CPU and NPU views of a UVA tensor.
|
||||
However if users add environment parameter below, UVA feature is Supported.
|
||||
os.environ['PYTORCH_NPU_ALLOC_CONF'] = 'pinned_mem_register:True'
|
||||
"""
|
||||
|
||||
def __init__(self, size: int | Sequence[int], dtype: torch.dtype):
|
||||
self._cpu: torch.Tensor = torch.zeros(size, dtype=dtype, device="cpu", pin_memory=True)
|
||||
self._np: np.ndarray = self._cpu.numpy()
|
||||
self._modified_indices: set[int] = set()
|
||||
self._uva: torch.Tensor = self._cpu if is_uva_available() else torch.zeros_like(self._cpu, device="npu")
|
||||
|
||||
def _mark_cpu_modified(self, key: int):
|
||||
self._modified_indices.add(key)
|
||||
|
||||
@property
|
||||
def cpu(self):
|
||||
return self._cpu if is_uva_available() else MonitoredTorchTensor(self._cpu, self._mark_cpu_modified)
|
||||
|
||||
@property
|
||||
def np(self):
|
||||
return self._np if is_uva_available() else MonitoredNumPyArray(self._np, self._mark_cpu_modified)
|
||||
|
||||
@property
|
||||
def uva(self):
|
||||
"""Get the device data of the buffer."""
|
||||
if not is_uva_available() and self._modified_indices:
|
||||
# Sort for better memory access locality
|
||||
dirty_rows = sorted(self._modified_indices)
|
||||
# can't use copy_ method, because copy_ for index tensor
|
||||
# will malloc new memory.
|
||||
self._uva[dirty_rows] = self._cpu[dirty_rows].to(device="npu", non_blocking=True)
|
||||
self._modified_indices.clear()
|
||||
return self._uva
|
||||
|
||||
|
||||
vllm.v1.worker.gpu.buffer_utils.UvaBuffer = UvaBufferWrapper
|
||||
Reference in New Issue
Block a user