a0315f6697
### What this PR does / why we need it? This pull request enables the `npugraph_ex` backend by default to improve performance on Ascend NPUs, as proposed in the [RFC](https://github.com/vllm-project/vllm-ascend/issues/6214). ### Does this PR introduce _any_ user-facing change? Yes. `npugraph_ex` is now enabled by default. Users can disable it by setting `enable: false` in the `npugraph_ex_config` section of the `additional_config`. ### How was this patch tested? CI passed. The changes are covered by existing and new E2E tests (`test_aclgraph_accuracy.py`) and unit tests (`test_ascend_config.py`) that have been updated to reflect the new default behavior. The tests verify correctness and consistency with `npugraph_ex` enabled and disabled, as well as with the new static kernel option. Signed-off-by: huyuanquan1 <huyuanquan1@huawei.com> Co-authored-by: huyuanquan1 <huyuanquan1@huawei.com>
117 lines
4.0 KiB
Python
117 lines
4.0 KiB
Python
#
|
|
# Copyright (c) 2026 Huawei Technologies Co., Ltd. All Rights Reserved.
|
|
# This file is a part of the vllm-ascend project.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
#
|
|
|
|
import importlib
|
|
import sys
|
|
|
|
import torch
|
|
import torchair
|
|
from torch._subclasses.fake_tensor import FakeTensor
|
|
from torchair.core._concrete_graph import _is_symlist
|
|
from torchair.npu_fx_compiler import _unpack_meta_list
|
|
|
|
|
|
class ValuePack:
|
|
def __init__(self, meta, npu_meta=None) -> None:
|
|
self._meta = meta
|
|
self._npu_meta = meta if npu_meta is None else npu_meta
|
|
|
|
@property
|
|
def meta(self):
|
|
return self._meta
|
|
|
|
@property
|
|
def npu(self):
|
|
return self._npu_meta
|
|
|
|
def __getitem__(self, key):
|
|
if isinstance(self._meta, dict):
|
|
return self._meta.get(key)
|
|
raise ValueError(f"Unsupported meta type for ValuePack __getitem__, key:{key}, type: {type(self._meta)}")
|
|
|
|
def __repr__(self) -> str:
|
|
if isinstance(self._meta, FakeTensor):
|
|
meta_str = f"FakeTensor(dtype={self._meta.dtype}, size={list(self._meta.size())}"
|
|
elif isinstance(self._meta, torch.Tensor):
|
|
meta_str = f"torch.Tensor(dtype={self._meta.dtype}, size={list(self._meta.size())}"
|
|
elif isinstance(self._meta, torch.SymInt):
|
|
meta_str = f"torch.SymInt({self._meta})"
|
|
else:
|
|
try:
|
|
meta_str = f"{type(self._meta)}({self._meta})"
|
|
except Exception:
|
|
meta_str = f"{type(self._meta)}"
|
|
return f"Pack(meta:{meta_str} npu:{self._npu_meta})"
|
|
|
|
|
|
def _unpack_meta(args, kwargs):
|
|
unpacked_args = []
|
|
unpacked_kwargs = {}
|
|
|
|
def _get_meta_part(arg):
|
|
if isinstance(arg, (list, tuple)) and any(isinstance(v, ValuePack) for v in arg):
|
|
return _unpack_meta_list(arg)
|
|
elif isinstance(arg, dict):
|
|
return {k: v.meta if isinstance(v, ValuePack) else v for k, v in arg.items()}
|
|
elif isinstance(arg, ValuePack):
|
|
return arg.meta
|
|
else:
|
|
return arg
|
|
|
|
for arg in args:
|
|
unpacked_args.append(_get_meta_part(arg))
|
|
|
|
for key, value in kwargs.items():
|
|
unpacked_kwargs[key] = _get_meta_part(value)
|
|
|
|
return list(unpacked_args), unpacked_kwargs
|
|
|
|
|
|
def _unpack_npu(self, args, kwargs):
|
|
unpacked = []
|
|
unpacked_kwargs = {}
|
|
|
|
def _get_npu_part(arg):
|
|
if isinstance(arg, (list, tuple)) and len(arg):
|
|
if _is_symlist(arg):
|
|
arg = self._graph.parse_symlist(arg)
|
|
else:
|
|
arg = [(v.npu if isinstance(v, ValuePack) else v) for v in arg]
|
|
return arg
|
|
elif isinstance(arg, dict):
|
|
return {k: v.npu if isinstance(v, ValuePack) else v for k, v in arg.items()}
|
|
elif isinstance(arg, ValuePack):
|
|
return arg.npu
|
|
else:
|
|
return arg
|
|
|
|
for arg in args:
|
|
unpacked.append(_get_npu_part(arg))
|
|
|
|
for key, value in kwargs.items():
|
|
unpacked_kwargs[key] = _get_npu_part(value)
|
|
|
|
return unpacked, unpacked_kwargs
|
|
|
|
|
|
torchair.core._concrete_graph.ValuePack = ValuePack
|
|
# The ValuePack class is referenced in these two modules, and after the patch, these two modules need to be reloaded.
|
|
importlib.reload(sys.modules["torchair.fx_summary"])
|
|
importlib.reload(sys.modules["torchair.npu_fx_compiler"])
|
|
torchair.npu_fx_compiler._unpack_meta = _unpack_meta
|
|
torchair.npu_fx_compiler._NpuGraphConverter._unpack_npu = _unpack_npu
|