EngineX-Ascend/enginex-ascend-910-vllm
11
0
Fork 1
Code Issues Pull Requests Actions Projects Releases Wiki Activity

init v0.23.0

Browse Source 
Signed-off-by: Sun Ruoxi <sunruoxi@4paradigm.com>
This commit is contained in:
Sun Ruoxi
2026-08-27 15:11:51 +08:00
parent b582a8e7d1
commit 7f8a1b1f7a
2849 changed files with 712887 additions and 22001 deletions
Download Patch File Download Diff File Expand all files Collapse all files

0
vllm_ascend/xlite/__init__.py Normal file
View File

395
vllm_ascend/xlite/utils.py Normal file
View File

@@ -0,0 +1,395 @@
#
# 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.
#
"""Utility functions for xlite."""
import threading
from collections.abc import Callable, Generator, Sequence
from contextlib import contextmanager
from dataclasses import dataclass
from logging import Logger
from typing import Any, Literal, TypedDict
import torch
import torch.nn as nn
from vllm.logger import logger
from xlite._C import Model, ModelConfig
from vllm_ascend.attention.attention_v1 import AscendMetadata
from vllm_ascend.attention.mla_v1 import AscendMLAMetadata
from vllm_ascend.attention.sfa_v1 import AscendSFAMetadata
_MISSING = object()
"""Unique sentinel for missing attributes in this module."""
class AttributeSetterMixin:
"""A mixin that allows setting attributes safely without raising AttributeError for missing attributes. This is
useful for handling C++ extension objects that may not have all attributes defined in all versions. The class will
simply ignore attempts to set attributes that do not exist, while allowing setting existing attributes as usual.
Additionally, a context manager interface is provided for checking the incoming value before setting the attribute.
The value is only set if the attribute exists and the optional `match_condition` is satisfied.
Good for backwards compatibility. For subclasses, :mod:`AttributeSetterMixin` must be the first parent class in the
inheritance chain to work properly (i.e., the second object in the method resolution order (MRO)).
Example usage::
class Model:
def __init__(self):
self.some_existing_attr = 0
class MyModel(AttributeSetterMixin, Model):
_on_missing_attr = "ignore" # silently ignore missing attributes
model = MyModel(...)
model.some_existing_attr = 42 # sets the attribute as usual
model.some_missing_attr = "hello" # does nothing, no error raised
with model.condition(lambda v: isinstance(v, int) and v > 0):
model.some_existing_attr = -1 # does not set because condition is not met
model.some_existing_attr = 100 # sets because condition is met
"""
_on_missing_attr: Literal["raise", "warn", "ignore"] = "warn"
"""Behavior when attempting to set a missing attribute. If `warn`, a logger must be provided to log a warning."""
_logger: Logger | None = None
"""Optional logger for warning about missing attributes. If None, no warnings will be logged."""
def __init_subclass__(cls) -> None:
if cls.__mro__[1] is not AttributeSetterMixin:
raise TypeError(
f"{cls.__name__} inherits from AttributeSetterMixin but does not have AttributeSetterMixin as the first"
f" parent class. Use `class {cls.__name__}(AttributeSetterMixin, ...)` to define the subclass, instead."
)
def _get_thread_local(self) -> threading.local:
"""Lazily initialize a per-instance, per-thread local storage without going through __setattr__."""
try:
return object.__getattribute__(self, "_thread_local")
except AttributeError:
local = threading.local()
object.__setattr__(self, "_thread_local", local)
return local
def __setattr__(self, name: str, value: Any) -> None:
if not (hasattr(type(self), name) or name in self.__dict__):
if self._on_missing_attr == "raise":
raise AttributeError(f"{type(self).__name__} has no attribute {name}.")
elif self._on_missing_attr == "warn" and self._logger:
self._logger.warning(
"%s has no attribute %s. Your `xlite` version might be incompatible.", type(self).__name__, name
)
return
match_condition = getattr(self._get_thread_local(), "match_condition", None)
if match_condition is not None and not match_condition(value):
return
super().__setattr__(name, value)
@contextmanager
def condition(self, match_condition: Callable[..., bool]) -> Generator["AttributeSetterMixin", None, None]:
"""Context manager that gates attribute setting on `match_condition`.
Usage::
with obj.condition(lambda v: v > 0):
obj.some_attr = 42 # only set if 42 > 0
"""
local = self._get_thread_local()
previous = getattr(local, "match_condition", None) # save for nesting
local.match_condition = match_condition
try:
yield self
finally:
local.match_condition = previous # always restore, even on exception
class XModel(AttributeSetterMixin, Model):
""":mod:`xlite._C.Model` subclass with safe attribute setting for better backwards compatibility."""
if torch.distributed.get_rank() == 0:
_logger = logger
class XModelConfig(AttributeSetterMixin, ModelConfig):
""":mod:`xlite._C.ModelConfig` subclass with safe attribute setting for better backwards compatibility."""
if torch.distributed.get_rank() == 0:
_logger = logger
@dataclass
class AttnMetadataRouter:
"""A router for attention metadata objects of different types. This is used to handle the differences in attention
metadata across different model architectures and vLLM/vLLM-ascend versions in a more robust way.
The router provides unified access to commonly used attention metadata attributes (e.g., actual sequence lengths for
query and block tables) via properties.
Currently included metadata types:
- `AscendMetadata`
- `AscendMLAMetadata`
- `AscendSFAMetadata`
Typically, the attention metadata has the following notations::
|---------- N-1 iteration --------|
|---------------- N iteration ---------------------|
|- tokenA -|......................|-- newTokens ---|
|---------- context_len ----------|
|-------------------- seq_len ---------------------|
|-- query_len ---|
"""
attn_metadata: Any
"""The attention metadata object to route, e.g., an instance of `AscendMetadata` or `AscendSFAMetadata`."""
device: str | torch.device | int | None = "cpu"
"""Device specification for the returned tensors. If None, the tensors will be on the same device as the original
metadata tensors. The current implementation assumes `cpu` device for minimal data transfer."""
@contextmanager
def on_device(self, device: str | torch.device | int | None) -> Generator["AttnMetadataRouter", None, None]:
"""Context manager to temporarily set the device for the router. This is useful for cases where we want to
access multiple properties on the same device without repeatedly specifying the device.
Usage::
with router.on_device("cpu"):
query_lens = router.cu_query_lens # on cpu
block_tables = router.block_tables # also on cpu
"""
original_device = self.device
self.device = device
try:
yield self
finally:
self.device = original_device
def __getattr__(self, name: str) -> Any:
"""Route attribute access to the appropriate handler method based on the attribute name."""
if (value := getattr(self.attn_metadata, name, _MISSING)) is not _MISSING:
return value
raise AttributeError(f"{type(self.attn_metadata).__name__} has no attribute {name}.")
@property
def cu_query_lens(self) -> torch.Tensor:
"""Get the cumulative query lengths from the attention metadata, if available."""
if isinstance(self.attn_metadata, (AscendMetadata, AscendMLAMetadata)):
return torch.as_tensor(self.attn_metadata.query_start_loc, device=self.device)
if isinstance(self.attn_metadata, AscendSFAMetadata):
return torch.as_tensor(self.attn_metadata.cum_query_lens, device=self.device)
for candidate in ["query_start_loc", "cum_query_lens", "actual_seq_lengths_q"]:
if (lengths := getattr(self.attn_metadata, candidate, None)) is not None:
return torch.as_tensor(lengths, device=self.device)
raise ValueError(
f"Cannot find actual sequence lengths for query in attention metadata of type {type(self.attn_metadata)}."
)
@property
def block_tables(self) -> torch.Tensor:
"""Get the block tables from the attention metadata, if available."""
if isinstance(self.attn_metadata, AscendMetadata):
return torch.as_tensor(self.attn_metadata.block_tables, device=self.device)
if isinstance(self.attn_metadata, AscendSFAMetadata):
return torch.as_tensor(self.attn_metadata.block_table, device=self.device)
if isinstance(self.attn_metadata, AscendMLAMetadata):
# AscendMLAMetadataBuilder.build_decode_metadata breaks `AscendMLAMetadata.block_tables`
# thus we may need to patch together block tables from prefill and decode metadata if available
block_tables = []
if self.attn_metadata.decode is not None:
block_tables.append(torch.as_tensor(self.attn_metadata.decode.block_table, device=self.device))
if self.attn_metadata.prefill is not None:
block_tables.append(torch.as_tensor(self.attn_metadata.prefill.block_table, device=self.device))
if block_tables:
return torch.concat(block_tables, dim=0)
return torch.as_tensor(self.attn_metadata.block_tables, device=self.device)
for candidate in ["block_tables", "block_table"]:
if (tables := getattr(self.attn_metadata, candidate)) is not None:
return torch.as_tensor(tables, device=self.device)
raise ValueError(f"Cannot find block tables in attention metadata of type {type(self.attn_metadata)}.")
@property
def seq_lens(self) -> torch.Tensor:
"""Return the per-sequence `seq_lens` tensor in a device-safe torch.Tensor form."""
if isinstance(self.attn_metadata, (AscendMetadata, AscendSFAMetadata)):
return torch.as_tensor(self.attn_metadata.seq_lens_cpu, device=self.device)
if isinstance(self.attn_metadata, AscendMLAMetadata):
# AscendMLAMetadataBuilder.build_decode_metadata breaks `AscendMLAMetadata.seq_lens`
# thus prefill metadata's seq_lens is preferentially used if available
if self.attn_metadata.prefill is not None:
return torch.as_tensor(self.attn_metadata.prefill.seq_lens, device=self.device)
return torch.as_tensor(self.attn_metadata.seq_lens_cpu, device=self.device)
for candidate in ["seq_lens_cpu", "seq_lens"]:
if (s := getattr(self.attn_metadata, candidate)) is not None:
return torch.as_tensor(s, device=self.device)
raise ValueError(f"Cannot find seq_lens in attention metadata of type {type(self.attn_metadata)}.")
@property
def num_prefills(self) -> int:
for candidate in ["num_prefills"]:
if (num_prefills := getattr(self.attn_metadata, candidate)) is not None:
return int(num_prefills)
return 0
@property
def num_decodes(self) -> int:
for candidate in ["num_decodes"]:
if (num_decodes := getattr(self.attn_metadata, candidate)) is not None:
return int(num_decodes)
return 0
@property
def num_decode_tokens(self) -> int:
for candidate in ["num_decode_tokens"]:
if (num_decode_tokens := getattr(self.attn_metadata, candidate, None)) is not None:
return int(num_decode_tokens)
return 0
@property
def num_actual_tokens(self) -> int:
"""Return the number of actual tokens (excluding padding)."""
for candidate in ["num_actual_tokens"]:
if (num_actual_tokens := getattr(self.attn_metadata, candidate)) is not None:
return int(num_actual_tokens)
raise ValueError(f"Cannot find num_actual_tokens in attention metadata of type {type(self.attn_metadata)}.")
def get_nested_attr(obj: Any, /, *attrs: str, default: Any = None, raises: bool = False) -> Any:
"""Get/collect a nested attribute from an object.
The attribute path is specified as a sequence of attribute names. If any attribute in the path is missing, the
function returns the specified default value (which is None by default).
Args:
obj (Any): Root object.
*attrs (str): Sequence of attribute names to traverse.
default (Any, keyword-only, default=None): Default value to return if any attribute is missing.
raises (bool, keyword-only, default=False): Whether to raise an error if any attribute is missing.
Returns:
Any: The resolved nested attribute.
"""
current = obj
for attr in attrs:
if (current := getattr(current, attr, _MISSING)) is _MISSING:
if raises:
raise AttributeError(f"{type(obj).__name__} has no attribute {'.'.join(attrs)} (failed at {attr}).")
return default
return current
def get_dotted_attr(obj: Any, dotted_attr: str, /, *, default: Any = None, raises: bool = False) -> Any:
"""Get a nested attribute from an object using a dotted attribute string.
This is a convenience wrapper around :meth:`_get_nested_attr` that allows specifying the attribute path as a single
dotted string.
Args:
obj (Any): Root object.
dotted_attr (str): Dotted attribute string, e.g., "foo.bar.baz" to access `obj.foo.bar.baz`.
default (Any, keyword-only, default=None): Default value to return if any attribute is missing.
raises (bool, keyword-only, default=False): Whether to raise an error if any attribute is missing.
Returns:
Any: The resolved nested attribute.
"""
return get_nested_attr(obj, *dotted_attr.split("."), default=default, raises=raises)
class WeightGetterConfig(TypedDict):
"""Configuration dictionary for layer weight extraction in `get_layer_weights`.
This class is written as a TypedDict for better type checking with `mypy` in the `xlite` module.
"""
secondary_flattening: str | slice | None
post_processor: Callable[[torch.Tensor], torch.Tensor] | None
def get_layer_weights(
layers: Sequence[nn.Module],
layer_attr: str,
/,
*,
secondary_flattening: str | slice | None = None,
post_processor: Callable[[torch.Tensor], torch.Tensor] | None = None,
**kwargs: Any,
) -> list[torch.Tensor]:
"""Extract specified weights from a sequence of layers with optional secondary flattening and post-processing.
This function retrieves the specified attribute (e.g., "self_attn.q_proj.weight") from each layer in the provided
sequence. If `secondary_flattening` is specified, it will further expand the retrieved attribute as a list and
collect all items from these lists across layers. An optional `post_processor` can be applied to each retrieved
tensor before returning the final list of weights.
Args:
layers (Sequence[nn.Module]): Sequence of layers to retrieve weights from.
layer_attr (str): Dotted attribute string specifying the layer attribute to retrieve (`layers.[i].[layer_attr]`)
, e.g., "self_attn.q_norm.weight".
secondary_flattening (str | slice | None, optional): If specified, indicates that the retrieved layer attribute
is a list of tensors and we need to further flatten it. The expansion can be specified as:
- `str`: A dotted attribute string such that `layers.[i].[secondary_flattening]` gives the number of items
to flatten for that layer.
- `slice`: A slice specifying how to slice `layers.[i].[layer_attr]` and then flatten the sliced part.
- `None`: No secondary flattening; `layers.[i].[layer_attr]` is directly collected.
post_processor (Callable[[torch.Tensor], torch.Tensor] | None, optional): An optional function to apply to
each retrieved tensor before returning the final list of weights.
**kwargs: Additional keyword arguments for future extensions.
Returns:
list[torch.Tensor]: List of retrieved weights.
"""
if not secondary_flattening:
weights = [
weight for layer in layers if (weight := get_dotted_attr(layer, layer_attr, default=None)) is not None
]
elif isinstance(secondary_flattening, str):
weights = [
weight
for layer in layers
if (weight_lst := get_dotted_attr(layer, layer_attr, default=[])) is not None
for weight in weight_lst[: get_dotted_attr(layer, secondary_flattening, default=0)]
]
elif isinstance(secondary_flattening, slice):
weights = [
weight
for layer in layers
if (weight_lst := get_dotted_attr(layer, layer_attr, default=[])) is not None
for weight in weight_lst[secondary_flattening]
]
else:
raise ValueError(
f"Invalid type for secondary_flattening: {type(secondary_flattening)}. Expected str, slice, or None."
)
if not post_processor:
return weights
return [post_processor(weight) for weight in weights]

747
vllm_ascend/xlite/xlite.py Normal file
View File

@@ -0,0 +1,747 @@
#
# Copyright (c) 2025 Huawei Technologies Co., Ltd. All Rights Reserved.
# Copyright 2023 The vLLM team.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
"""Xlite integration module for vLLM-Ascend."""
from __future__ import annotations
from abc import ABC, abstractmethod
from collections.abc import Callable, Sequence
from typing import Any, TypeAlias, cast
import torch
import torch.nn as nn
import torch_npu
from transformers import PretrainedConfig
from vllm.config import VllmConfig
from vllm.distributed import get_ep_group, get_tensor_model_parallel_world_size, get_world_group
from vllm.forward_context import get_forward_context
from vllm.logger import logger
from vllm.sequence import IntermediateTensors
from xlite._C import AttnMeta, AttnMHA, Runtime, ScoringFuncSigmoid, ScoringFuncSoftmax
from vllm_ascend.ascend_config import get_ascend_config
from vllm_ascend.attention.attention_v1 import AscendAttentionState, AscendMetadata
from vllm_ascend.compilation.acl_graph import ACLGraphWrapper
from vllm_ascend.xlite.utils import (
AttnMetadataRouter,
WeightGetterConfig,
XModel,
XModelConfig,
get_dotted_attr,
get_layer_weights,
)
XliteInitResult: TypeAlias = tuple[XModel, torch.Tensor, int, torch.dtype]
XliteForwardResult: TypeAlias = torch.Tensor | IntermediateTensors | tuple[torch.Tensor, list[torch.Tensor]]
_architecture_strategy_map: dict[str, type[XliteModel]] = {}
"""Mapping from model architecture names in `config.json` to their corresponding xlite adapter classes."""
class XliteModel(ABC):
"""Base adapter for converting vLLM models into xlite runtime models.
Subclasses are responsible for mapping architecture-specific configuration and weights into the `xlite._C.Model`
interface.
Attributes:
runnable (nn.Module): The original runnable model used by vLLM. Used as the source of truth for weight
extraction for xlite model construction.
vllm_config (VllmConfig): The configuration object provided by vLLM. Used to build xlite configuration at
runtime.
xlite_config (ModelConfig): Native xlite configuration object populated by subclasses.
xlite_model (Model): Native xlite model container populated by subclasses.
"""
_attn_metadata_type: type | tuple[type, ...]
"""The expected type of attention metadata in the forward context for this architecture. Used for runtime checks
before forwarding. See :meth:`XliteWrapper.__call__` for usage."""
_supported_architectures: Sequence[str] | str
"""The list of model architecture names (from HuggingFace `config.json` "architectures" field) supported by this
adapter. Used for automatic adapter selection and registration."""
def __init_subclass__(cls, **kwargs: Any) -> None:
"""Automatically register subclasses in the architecture strategy map and metadata type set."""
ts = getattr(cls, "_attn_metadata_type", None)
if ts is None or (not isinstance(ts, type) and not all(isinstance(t, type) for t in ts)):
raise ValueError(
f"XliteModel subclass {cls.__name__} must define _attn_metadata_type as a type or a tuple of types."
)
arcs = getattr(cls, "_supported_architectures", None)
if arcs is None:
raise ValueError(f"XliteModel subclass {cls.__name__} must define _supported_architectures attribute.")
if isinstance(arcs, str):
arcs = [arcs]
for arc in arcs:
if arc in _architecture_strategy_map:
raise ValueError(f"Duplicate xlite adapter for architecture {arc}: {_architecture_strategy_map[arc]}")
_architecture_strategy_map[arc] = cls
super().__init_subclass__(**kwargs)
def __init__(self, runnable: nn.Module, vllm_config: VllmConfig) -> None:
"""Initialize the xlite model adapter.
Args:
runnable (nn.Module): The original runnable model used by vLLM.
vllm_config (VllmConfig): Runtime configuration used for model setup.
Notes:
The constructor stores the runnable model and vLLM config, and prepares empty xlite configuration and model
containers for subclass-specific population.
"""
self.runnable = runnable
self.vllm_config = vllm_config
self.xlite_config = XModelConfig()
self.xlite_model = XModel()
def initialize(self) -> XliteInitResult:
"""Initialize an xlite model and precomputed RoPE cache.
Returns:
XliteInitResult: A tuple of `(xlite_model, freq_cis, hidden_size, dtype)` required by `XliteWrapper`.
"""
self._build_model_config()
self._build_model()
rank = torch.distributed.get_rank()
self.xlite_model.init(self.xlite_config, rank)
freq_cis = self._precompute_freqs_cis()
return (self.xlite_model, freq_cis, self.xlite_config.hidden_size, self.vllm_config.model_config.dtype)
@abstractmethod
def _build_model_config(self) -> None:
"""Build architecture-specific xlite model configuration.
This method extracts necessary configuration attributes from the vLLM config (e.g., HuggingFace metadata) and
populates an xlite :class:`ModelConfig` object.
Returns:
None: `self` attribute :attr:`xlite_config` is updated in-place.
"""
@abstractmethod
def _build_model(self) -> None:
"""Build architecture-specific xlite model weights.
This method traverses the runnable model's parameters and maps them into the xlite :class:`Model` interface
according to the architecture's specific structure.
Returns:
None: `self` attribute :attr:`xlite_model` is updated in-place.
Notes:
:meth:`_build_model_config` should be called prior to this method to ensure the xlite configuration is
populated before weight mapping.
"""
def _get_layers_and_model_prefix(self) -> tuple[Sequence[nn.Module], str]:
"""Extract transformer layers and parameter prefix from runnable.
Returns:
tuple[Sequence[nn.Module], str]: A pair of `(layers, model_prefix)` for model traversal.
"""
if hasattr(self.runnable, "language_model"):
layers = cast(
Sequence[nn.Module], get_dotted_attr(self.runnable.language_model, "model.layers", default=[])
)
prefix = "language_model."
else:
layers = cast(Sequence[nn.Module], get_dotted_attr(self.runnable, "model.layers", default=[]))
prefix = ""
return layers, prefix
@abstractmethod
def _precompute_freqs_cis(self) -> torch.Tensor:
"""Precomputes frequency-based complex exponential values for rotary positional embeddings (RoPE).
This method generates the RoPE frequency cache (cosine and sine values) required by the xlite attention
implementation. The cache should be precomputed on the NPU device to avoid unnecessary host-device transfers
during inference.
Returns:
torch.Tensor: The precomputed RoPE frequency cache tensor ready for use in xlite attention computations.
Notes:
:meth:`_build_model_config` should be called prior to this method.
"""
@staticmethod
def is_tensor_nz(t: torch.Tensor) -> bool:
"""Check if a tensor is in NZ format.
Args:
t (torch.Tensor): The tensor to check.
Returns:
bool: True if the tensor is in NZ format, False otherwise.
"""
format = torch_npu.get_npu_format(t)
return format == torch_npu.Format.FRACTAL_NZ
@staticmethod
def all_tensors_zero(tensors: torch.Tensor | list[torch.Tensor] | tuple[torch.Tensor] | None) -> bool:
"""Check if all tensors in the list/tuple are zero tensors.
Args:
tensors (torch.Tensor | list[torch.Tensor] | tuple[torch.Tensor] | None): The tensors to check.
Returns:
bool: True if all tensors are zero tensors (or empty), False otherwise.
"""
if tensors is None:
return True
if not isinstance(tensors, (list, tuple)):
tensors = [tensors]
if len(tensors) == 0:
return True
return all(torch.allclose(t, t.new_zeros(1)) for t in tensors)
@staticmethod
def _transform_deq_scale(deq_scale: torch.Tensor) -> torch.Tensor:
"""
The data format required by the fixpipe hardware is as follows:
Data is stored in uint64_t, with the upper 32 bits being 0 and the lower 32 bits storing the FP32 format. The
lower 10 bits of the FP32 format are not involved in computation, and the actual data format is TF32.
"""
deq_scale_fp32 = deq_scale.to(torch.float32)
scale = deq_scale_fp32.new_zeros(deq_scale.shape[0] * 2)
scale[0::2] = deq_scale_fp32[0::1]
return scale
@property
def hf_text_config(self) -> PretrainedConfig:
"""Convenience property to access HuggingFace text configuration from vLLM config.
Returns:
PretrainedConfig: The HuggingFace text configuration object extracted from vLLM config.
"""
hf_config = self.vllm_config.model_config.hf_text_config
return cast(PretrainedConfig, getattr(hf_config, "text_config", hf_config))
@property
def hf_vision_config(self) -> PretrainedConfig | None:
"""Convenience property to access HuggingFace vision configuration from vLLM config, if exists.
Returns:
PretrainedConfig | None: The HuggingFace vision configuration object extracted from vLLM config, or None if
not present.
"""
return getattr(self.vllm_config.model_config.hf_config, "vision_config", None)
class LlamaXliteModel(XliteModel):
"""xlite adapter base for Llama-like architectures.
This is the *de facto* base adapter for all xlite-supported architectures and may contain configurations beyond
Llama-like dense models. `XliteModel` subclasses should inherit from this class unless there is a major divergence.
"""
_attn_metadata_type = AscendMetadata
_supported_architectures = [
"LlamaForCausalLM",
"Qwen2ForCausalLM",
"Qwen3ForCausalLM",
"Qwen3VLForConditionalGeneration",
]
def _build_model_config(self) -> None:
xlite_config, vllm_config, hf_config = self.xlite_config, self.vllm_config, self.hf_text_config
xlite_config.vocab_size = hf_config.vocab_size
xlite_config.hidden_size = hf_config.hidden_size
xlite_config.n_layers = hf_config.num_hidden_layers
xlite_config.n_heads = hf_config.num_attention_heads
xlite_config.n_kv_heads = hf_config.num_key_value_heads
if hasattr(hf_config, "head_dim"):
xlite_config.head_dim = hf_config.head_dim
else:
xlite_config.head_dim = hf_config.hidden_size // hf_config.num_attention_heads
xlite_config.rope_head_dim = xlite_config.head_dim
xlite_config.norm_eps = hf_config.rms_norm_eps
if hasattr(hf_config, "rope_theta"):
xlite_config.rope_theta = hf_config.rope_theta
else:
xlite_config.rope_theta = getattr(hf_config, "rope_parameters", {}).get("rope_theta", 10000.0)
xlite_config.softmax_scale = xlite_config.head_dim**-0.5
xlite_config.n_dense_layers = hf_config.num_hidden_layers
xlite_config.intermediate_size = hf_config.intermediate_size
xlite_config.def_tp_size = get_tensor_model_parallel_world_size()
xlite_config.def_dp_size = vllm_config.parallel_config.data_parallel_size
try:
ep_word_size = get_ep_group().world_size
xlite_config.moe_ep_size = ep_word_size if vllm_config.parallel_config.enable_expert_parallel else 1
xlite_config.moe_tp_size = 1 if vllm_config.parallel_config.enable_expert_parallel else ep_word_size
except AssertionError:
xlite_config.moe_ep_size, xlite_config.moe_tp_size = 1, 1
xlite_config.experts_weight_transpose = True
xlite_config.attn_type = AttnMHA
xlite_config.scoring_func = ScoringFuncSoftmax
xlite_config.weight_nz = get_ascend_config().weight_nz_mode == 2
xlite_config.max_m = (
vllm_config.scheduler_config.max_num_batched_tokens
if get_ascend_config().xlite_graph_config.full_mode
else vllm_config.scheduler_config.max_num_seqs
)
xlite_config.max_batch_size = vllm_config.scheduler_config.max_num_seqs
xlite_config.max_seq_len = vllm_config.model_config.max_model_len
xlite_config.block_size = vllm_config.cache_config.block_size
rope_parameters = getattr(hf_config, "rope_parameters", {})
xlite_config.deepstack_num_level = len(getattr(self.hf_vision_config, "deepstack_visual_indexes", []))
xlite_config.mrope_section = rope_parameters.get("mrope_section", [])
xlite_config.mrope_interleaved = rope_parameters.get("mrope_interleaved", False)
self.quantization = vllm_config.quant_config is not None
def _build_model(self) -> None:
xlite_model, xlite_config, hf_config = self.xlite_model, self.xlite_config, self.hf_text_config
layers, model_prefix = self._get_layers_and_model_prefix()
xlite_model.embed = get_dotted_attr(self.runnable, f"{model_prefix}model.embed_tokens.weight", raises=True)
xlite_model.norm = get_dotted_attr(self.runnable, f"{model_prefix}model.norm.weight", raises=True)
if hf_config.tie_word_embeddings:
xlite_model.head = xlite_model.embed
else:
xlite_model.head = get_dotted_attr(self.runnable, f"{model_prefix}lm_head.weight", raises=True)
xlite_model.attn_norm = get_layer_weights(layers, "input_layernorm.weight")
self.init_matmul_weights(layers, "mha_qkv", "self_attn.qkv_proj")
self.init_matmul_weights(layers, "attn_out", "self_attn.o_proj")
mha_qkv_bias = get_layer_weights(layers, "self_attn.qkv_proj.bias")
xlite_config.qkv_bias = len(mha_qkv_bias) == xlite_config.n_layers
xlite_model.mha_qkv_bias = mha_qkv_bias if xlite_config.qkv_bias else []
q_norm = get_layer_weights(layers, "self_attn.q_norm.weight")
k_norm = get_layer_weights(layers, "self_attn.k_norm.weight")
xlite_config.qk_norm = len(q_norm) == len(k_norm) == xlite_config.n_layers
xlite_model.mha_q_norm = q_norm if xlite_config.qk_norm else []
xlite_model.mha_k_norm = k_norm if xlite_config.qk_norm else []
self.init_matmul_weights(layers, "mlp_up_gate", "mlp.gate_up_proj")
self.init_matmul_weights(layers, "mlp_down", "mlp.down_proj")
xlite_model.mlp_norm = get_layer_weights(layers, "post_attention_layernorm.weight")
if not self.quantization:
return
if xlite_model.mha_qkv:
xlite_config.quant_attn_weight_nz = self.is_tensor_nz(xlite_model.mha_qkv[0])
xlite_config.quant_attn_weight_transpose = True
with xlite_model.condition(lambda tensors: not self.all_tensors_zero(tensors)):
xlite_model.norm_bias = get_dotted_attr(self.runnable, f"{model_prefix}model.norm.bias", raises=True)
xlite_model.attn_norm_bias = get_layer_weights(layers, "input_layernorm.bias")
xlite_model.mlp_norm_bias = get_layer_weights(layers, "post_attention_layernorm.bias")
if xlite_config.qk_norm:
xlite_model.mha_q_norm_bias = get_layer_weights(layers, "self_attn.q_norm.bias")
xlite_model.mha_k_norm_bias = get_layer_weights(layers, "self_attn.k_norm.bias")
def _precompute_freqs_cis(self) -> torch.Tensor:
"""Precompute rotary cosine/sine cache on NPU.
Returns:
torch.Tensor: Concatenated cosine/sine RoPE cache on NPU.
Raises:
ValueError: If rope dimensions, sequence length, or theta are invalid.
"""
base = self.xlite_config.rope_theta
rotary_dim = self.xlite_config.rope_head_dim
max_position_embeddings = self.xlite_config.max_seq_len
dtype = self.vllm_config.model_config.dtype
if rotary_dim <= 0 or max_position_embeddings <= 0 or base <= 0:
raise ValueError(
f"Invalid RoPE configuration: head_dim={rotary_dim}, max_seq_len={max_position_embeddings}, "
f"rope_theta={base}"
)
# Keep cache construction on CPU, then transfer once to NPU.
inv_freq = 1.0 / (base ** (torch.arange(0, rotary_dim, 2, dtype=torch.float32, device="cpu") / rotary_dim))
t = torch.arange(max_position_embeddings, dtype=torch.float32, device=inv_freq.device)
freqs = torch.outer(t, inv_freq).float()
cos_cache = freqs.cos().to(dtype)
sin_cache = freqs.sin().to(dtype)
freq_cis = torch.cat((cos_cache, sin_cache), dim=-1)
return freq_cis.to(device="npu")
def init_matmul_weights(self, layers: Sequence[torch.nn.Module], xlite_prefix: str, model_prefix: str) -> None:
"""
Initialize MatMul-related weights with quantization support.
Args:
layers (Sequence[torch.nn.Module]): The transformer layers to extract weights from.
xlite_prefix (str): The prefix for the xlite model attributes to set.
model_prefix (str): The prefix for the model attributes to look up in each layer.
"""
xlite_model = self.xlite_model
setattr(xlite_model, xlite_prefix, get_layer_weights(layers, f"{model_prefix}.weight"))
if not self.quantization:
return
def set_xlite_attr(xlite_attr: str, layer_attr: str):
setattr(xlite_model, xlite_attr, get_layer_weights(layers, layer_attr))
deq_scale = get_layer_weights(layers, f"{model_prefix}.deq_scale", post_processor=self._transform_deq_scale)
if len(deq_scale) > 0: # static quant
setattr(xlite_model, f"{xlite_prefix}_deq_scale", deq_scale)
set_xlite_attr(f"{xlite_prefix}_input_scale", f"{model_prefix}.aclnn_input_scale_reciprocal")
set_xlite_attr(f"{xlite_prefix}_input_offset", f"{model_prefix}.aclnn_input_offset")
set_xlite_attr(f"{xlite_prefix}_quant_bias", f"{model_prefix}.quant_bias")
else:
weight_scale = get_layer_weights(
layers, f"{model_prefix}.weight_scale", post_processor=self._transform_deq_scale
)
setattr(xlite_model, f"{xlite_prefix}_deq_scale", weight_scale)
class QwenMoeXliteModel(LlamaXliteModel):
"""xlite adapter for Qwen MoE architectures."""
_attn_metadata_type = AscendMetadata
_supported_architectures = ["Qwen3MoeForCausalLM", "Qwen3VLMoeForConditionalGeneration"]
def _build_model_config(self) -> None:
super()._build_model_config()
xlite_config, hf_config = self.xlite_config, self.hf_text_config
xlite_config.n_dense_layers = 0
xlite_config.n_routed_experts = hf_config.num_experts
xlite_config.n_shared_experts = 0
xlite_config.n_act_experts = hf_config.num_experts_per_tok
xlite_config.moe_intermediate_size = hf_config.moe_intermediate_size
xlite_config.norm_topk_prob = hf_config.norm_topk_prob
def _build_model(self) -> None:
super()._build_model()
xlite_model, xlite_config = self.xlite_model, self.xlite_config
layers, _ = self._get_layers_and_model_prefix()
xlite_model.gate = get_layer_weights(layers, "mlp.gate.weight")
prefix = "mlp.experts."
kwargs: WeightGetterConfig = {"secondary_flattening": f"{prefix}local_num_experts", "post_processor": None}
xlite_model.re_up_gate = get_layer_weights(layers, f"{prefix}w13_weight", **kwargs)
xlite_model.re_down = get_layer_weights(layers, f"{prefix}w2_weight", **kwargs)
xlite_config.experts_weight_nz = self.is_tensor_nz(xlite_model.re_up_gate[0])
if self.quantization:
kwargs["post_processor"] = self._transform_deq_scale
xlite_model.re_up_gate_scale = get_layer_weights(layers, f"{prefix}w13_weight_scale_fp32", **kwargs)
xlite_model.re_down_scale = get_layer_weights(layers, f"{prefix}w2_weight_scale", **kwargs)
class Glm4MoeXliteModel(LlamaXliteModel):
"""xlite adapter for GLM4 MoE architectures."""
_attn_metadata_type = AscendMetadata
_supported_architectures = ["Glm4MoeForCausalLM"]
def _build_model_config(self) -> None:
super()._build_model_config()
xlite_config, hf_config = self.xlite_config, self.hf_text_config
if hasattr(hf_config, "partial_rotary_factor"):
partial_rotary_factor = hf_config.partial_rotary_factor
else:
partial_rotary_factor = getattr(hf_config, "rope_parameters", {}).get("partial_rotary_factor", 1.0)
xlite_config.rope_head_dim = int(xlite_config.head_dim * partial_rotary_factor)
xlite_config.n_dense_layers = getattr(hf_config, "first_k_dense_replace", 0)
xlite_config.n_routed_experts = hf_config.n_routed_experts
xlite_config.n_shared_experts = hf_config.n_shared_experts
xlite_config.n_act_experts = hf_config.num_experts_per_tok
xlite_config.moe_intermediate_size = hf_config.moe_intermediate_size
xlite_config.norm_topk_prob = hf_config.norm_topk_prob
xlite_config.scoring_func = ScoringFuncSigmoid
xlite_config.route_scale = hf_config.routed_scaling_factor
xlite_config.gate_captured = False
def _build_model(self) -> None:
super()._build_model()
xlite_model, xlite_config = self.xlite_model, self.xlite_config
layers, _ = self._get_layers_and_model_prefix()
xlite_model.gate = get_layer_weights(layers, "mlp.gate.weight")
# NOTE: type conversion for numerical stability in xlite's implementation
xlite_model.gate_bias = get_layer_weights(
layers, "mlp.gate.e_score_correction_bias", post_processor=lambda b: b.to(torch.float32)
)
self.init_matmul_weights(layers, "se_up_gate", "mlp.shared_experts.gate_up_proj")
self.init_matmul_weights(layers, "se_down", "mlp.shared_experts.down_proj")
prefix = "mlp.experts."
kwargs: WeightGetterConfig = {"secondary_flattening": f"{prefix}local_num_experts", "post_processor": None}
xlite_model.re_up_gate = get_layer_weights(layers, f"{prefix}w13_weight", **kwargs)
xlite_model.re_down = get_layer_weights(layers, f"{prefix}w2_weight", **kwargs)
if xlite_model.re_up_gate:
xlite_config.experts_weight_nz = self.is_tensor_nz(xlite_model.re_up_gate[0])
if self.quantization:
kwargs["post_processor"] = self._transform_deq_scale
xlite_model.re_up_gate_scale = get_layer_weights(layers, f"{prefix}w13_weight_scale_fp32", **kwargs)
xlite_model.re_down_scale = get_layer_weights(layers, f"{prefix}w2_weight_scale", **kwargs)
class MiniMaxM2XliteModel(LlamaXliteModel):
"""xlite adapter for MiniMax M2 architectures."""
_attn_metadata_type = AscendMetadata
_supported_architectures = ["MiniMaxM2ForCausalLM"]
def _build_model_config(self) -> None:
super()._build_model_config()
xlite_config, hf_config = self.xlite_config, self.hf_text_config
xlite_config.rope_head_dim = hf_config.rotary_dim
xlite_config.n_dense_layers = 0
xlite_config.n_routed_experts = hf_config.num_local_experts
xlite_config.n_shared_experts = 0
xlite_config.n_act_experts = hf_config.num_experts_per_tok
xlite_config.moe_intermediate_size = hf_config.intermediate_size
xlite_config.norm_topk_prob = True
xlite_config.qk_norm_full = True
xlite_config.scoring_func = ScoringFuncSigmoid
def _build_model(self) -> None:
super()._build_model()
xlite_model, xlite_config = self.xlite_model, self.xlite_config
layers, _ = self._get_layers_and_model_prefix()
xlite_model.gate = get_layer_weights(layers, "block_sparse_moe.gate.weight")
# NOTE: type conversion for numerical stability in xlite's implementation
xlite_model.gate_bias = get_layer_weights(
layers, "block_sparse_moe.e_score_correction_bias", post_processor=lambda b: b.to(torch.float32)
)
prefix = "block_sparse_moe.experts."
kwargs: WeightGetterConfig = {"secondary_flattening": f"{prefix}local_num_experts", "post_processor": None}
xlite_model.re_up_gate = get_layer_weights(layers, f"{prefix}w13_weight", **kwargs)
xlite_model.re_down = get_layer_weights(layers, f"{prefix}w2_weight", **kwargs)
if xlite_model.re_up_gate:
xlite_config.experts_weight_nz = self.is_tensor_nz(xlite_model.re_up_gate[0])
if self.quantization:
kwargs["post_processor"] = self._transform_deq_scale
xlite_model.re_up_gate_scale = get_layer_weights(layers, f"{prefix}w13_weight_scale_fp32", **kwargs)
xlite_model.re_down_scale = get_layer_weights(layers, f"{prefix}w2_weight_scale", **kwargs)
def get_adapter_xlite_model(runnable: nn.Module, vllm_config: VllmConfig) -> XliteModel:
"""Look up and initialize the appropriate xlite model adapter based on the architecture specified in vLLM config and
the runnable model.
Args:
runnable (nn.Module): The runnable model instance.
vllm_config (VllmConfig): Runtime configuration for model execution.
Raises:
ValueError: If the model architecture is not supported by xlite.
Returns:
XliteModel: An initialized xlite model adapter ready for inference.
"""
architecture = vllm_config.model_config.architectures[0]
if not (strategy_class := _architecture_strategy_map.get(architecture)):
raise ValueError(f"{architecture} not supported!")
return strategy_class(runnable, vllm_config)
class XliteWrapper:
"""A graph-based wrapper that dispatches between xlite and runnable paths."""
def __init__(self, runnable: nn.Module, vllm_config: VllmConfig, device: torch.device) -> None:
"""Initialize xlite runtime, model tensors, and hidden-state workspace.
Args:
runnable (nn.Module): The runnable model implementation.
vllm_config (VllmConfig): Runtime configuration for execution.
device (torch.device): The device to initialize the xlite model on.
Raises:
ValueError: If xlite runtime tensor-pool initialization fails.
"""
self.runnable = runnable
self.device = device
self.full_mode = get_ascend_config().xlite_graph_config.full_mode
rank = torch.distributed.get_rank()
local_rank = get_world_group().local_rank
self.data_parallel_size = vllm_config.parallel_config.data_parallel_size
self.adapter_xlite_model = get_adapter_xlite_model(runnable, vllm_config)
(self.xlite_model, self.freq_cis, hidden_size, dtype) = self.adapter_xlite_model.initialize()
xlite_config = self.adapter_xlite_model.xlite_config
self.xlite_rt = Runtime(
devid=local_rank,
size=0,
rank=rank,
tp_size=xlite_config.def_tp_size,
dp_size=xlite_config.def_dp_size,
moe_tp_size=xlite_config.moe_tp_size,
moe_ep_size=xlite_config.moe_ep_size,
)
rt_pool_size = self.xlite_model.get_tensor_pool_size()
if rank == 0:
logger.info("xlite runtime pool size: %s MB", rt_pool_size)
if self.xlite_rt.init_tensor_pool(rt_pool_size) != 0:
raise ValueError(f"xlite wrapper init failed! runtime pool size: {rt_pool_size} MB")
max_num_tokens = vllm_config.scheduler_config.max_num_batched_tokens
self.hidden_states = torch.empty(max_num_tokens, hidden_size, device=self.device, dtype=dtype)
def __getattr__(self, key: str) -> Any:
"""Proxy unknown attributes to the wrapped runnable model.
Args:
key (str): The attribute name requested by the caller.
Raises:
AttributeError: If neither wrapper nor runnable has the attribute.
Returns:
Any: Attribute value resolved from the runnable.
"""
try:
return getattr(self.runnable, key)
except Exception: # runnable may raise various exceptions
raise AttributeError(f"{self.__class__.__name__} object has no attribute {key}") from None
def unwrap(self) -> Callable:
"""Return the original runnable callable. See :meth:`ACLGraphWrapper.unwrap` for details.
Returns:
Callable: Original model runnable.
"""
# in case we need to access the original runnable.
if isinstance(runnable := self.runnable, ACLGraphWrapper):
return runnable.unwrap()
return runnable
def register_kv_caches(self, kv_caches: Any) -> None:
"""Register KV cache references used by xlite runtime.
Args:
kv_caches (Any): Runtime KV cache handles or tensors.
"""
self.kv_caches = kv_caches
def __call__(
self,
input_ids: torch.Tensor,
positions: torch.Tensor,
intermediate_tensors: IntermediateTensors | None = None,
inputs_embeds: torch.Tensor | None = None,
**model_kwargs: Any,
) -> XliteForwardResult:
"""Run one forward step through xlite graph or fallback runnable path.
Args:
input_ids (torch.Tensor): Token IDs for current step.
positions (torch.Tensor): Position IDs used by attention.
intermediate_tensors (Optional[IntermediateTensors]): Optional intermediate tensors from pipeline stages.
inputs_embeds (Optional[torch.Tensor]): Optional external input embeddings (e.g. multimodal/deepstack
scenarios).
**model_kwargs (Any): Additional keyword arguments for the runnable.
Returns:
XliteForwardResult: Forward outputs from xlite graph or the original runnable implementation.
"""
forward_context = get_forward_context()
if getattr(forward_context, "in_profile_run", False):
if self.full_mode:
# In full mode, xlite handles both prefill and decode, and aclgraph runnable should not reserve memory.
# This is to avoid redundant memory allocation that reduces KV cache capacity and regresses performance.
# NOTE: returning a single hidden state tensor may break the vLLM pipeline if the runnable expects a
# tuple of outputs, e.g., (hidden_states, aux_hidden_states) under certain speculative scenarios
return self.hidden_states
return self.runnable(input_ids, positions, intermediate_tensors, inputs_embeds, **model_kwargs)
attn_metadata: Any = forward_context.attn_metadata
if attn_metadata is None:
return self.runnable(input_ids, positions, intermediate_tensors, inputs_embeds, **model_kwargs)
attn_metadata = attn_metadata[0] if isinstance(attn_metadata, list) else attn_metadata
attn_metadata = next(iter(attn_metadata.values()), None)
if not isinstance(attn_metadata, self.adapter_xlite_model._attn_metadata_type):
return self.runnable(input_ids, positions, intermediate_tensors, inputs_embeds, **model_kwargs)
with_prefill = attn_metadata.attn_state not in (
AscendAttentionState.DecodeOnly,
AscendAttentionState.SpecDecoding,
)
# Full: graph for prefill and decode
# Decode-Only: runnable for prefill, graph for decode
if not self.full_mode and self.data_parallel_size > 1:
num_tokens = forward_context.batch_descriptor.num_tokens
num_reqs = forward_context.batch_descriptor.num_reqs
use_xlite_graph = num_reqs is not None and num_tokens <= num_reqs
else:
use_xlite_graph = not with_prefill or self.full_mode
if not use_xlite_graph:
# fall back to runnable for prefill in decode-only mode
# or when the number of tokens exceeds the graph capacity in non-full mode
return self.runnable(input_ids, positions, intermediate_tensors, inputs_embeds, **model_kwargs)
attn_metadata_router = AttnMetadataRouter(attn_metadata=attn_metadata, device="cpu")
seq_lens = attn_metadata_router.seq_lens
cum_query_lens = attn_metadata_router.cu_query_lens[-seq_lens.size(0) :].to(device=seq_lens.device)
query_lens = torch.diff(cum_query_lens, prepend=seq_lens.new_zeros(1))
cached_lens = torch.clamp(seq_lens - query_lens, min=0)
num_tokens = forward_context.batch_descriptor.num_tokens
num_actual_tokens = attn_metadata.num_actual_tokens
xlite_attn_metadata = AttnMeta()
xlite_attn_metadata.lens = query_lens.tolist()
xlite_attn_metadata.cached_lens = cached_lens.tolist()
xlite_attn_metadata.block_tables_cpu = attn_metadata_router.block_tables.tolist()
if positions.ndim == 2:
xlite_attn_metadata.positions = positions[:, :num_actual_tokens].contiguous()
positions = positions[0]
else:
xlite_attn_metadata.positions = positions
# Compatibility between DP and Non-DP scenarios
h = self.hidden_states[:num_tokens]
stream = torch.npu.current_stream().npu_stream
if inputs_embeds is None:
self.xlite_model.forward(
self.xlite_rt, input_ids, xlite_attn_metadata, self.kv_caches, self.freq_cis, h, stream
)
else:
deepstack_input_embeds = getattr(self.runnable, "deepstack_input_embeds", [])
xlite_deepstack_input_embeds = [
deepstack_input[: inputs_embeds.size(0)] for deepstack_input in deepstack_input_embeds
]
self.xlite_model.forward_with_inputs_embeds(
self.xlite_rt,
inputs_embeds,
xlite_attn_metadata,
self.kv_caches,
self.freq_cis,
h,
stream,
xlite_deepstack_input_embeds,
)
if xlite_deepstack_input_embeds and hasattr(self.runnable, "_clear_deepstack_input_embeds"):
self.runnable._clear_deepstack_input_embeds(inputs_embeds.size(0))
return h[:num_actual_tokens]

55
vllm_ascend/xlite/xlite_model_runner.py Normal file
View File

@@ -0,0 +1,55 @@
#
# Copyright (c) 2025 Huawei Technologies Co., Ltd. All Rights Reserved.
# Copyright 2023 The vLLM team.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# This file is a part of the vllm-ascend project.
# Adapted from vllm-project/vllm/vllm/worker/gpu_model_runner.py
# isort: skip_file
import torch.nn as nn
from vllm.config import CUDAGraphMode
from vllm.v1.kv_cache_interface import KVCacheConfig
from vllm_ascend.worker.model_runner_v1 import NPUModelRunner
class XliteModelRunner(NPUModelRunner):
def get_model(self) -> nn.Module:
"""See :meth:`NPUModelRunner.get_model` and :meth:`XliteWrapper.unwrap` for details."""
if not hasattr(self, "xlite_model"):
return super().get_model()
return self.xlite_model.unwrap()
def load_model(self) -> None:
from vllm_ascend.xlite.xlite import XliteWrapper
super().load_model()
self.model = self.xlite_model = XliteWrapper(self.model, self.vllm_config, device=self.device)
def initialize_kv_cache(self, kv_cache_config: KVCacheConfig) -> None:
super().initialize_kv_cache(kv_cache_config)
self.xlite_model.register_kv_caches(self.kv_caches)
def _should_build_dummy_attn_metadata(
self,
force_attention: bool = False,
is_profile: bool = False,
cudagraph_runtime_mode: CUDAGraphMode | None = None,
) -> bool:
"""
Override to build attention metadata during dummy_run when xlite is enable.
For xlite, we need to build metadata during DP dummy_run to ensure all ranks
have consistent metadata, even when some ranks have no requests.
"""
base_condition = super()._should_build_dummy_attn_metadata(force_attention, is_profile, cudagraph_runtime_mode)
xlite_condition = self.ascend_config.xlite_graph_config.enabled and not is_profile
return base_condition or xlite_condition

30
vllm_ascend/xlite/xlite_worker.py Normal file
View File

@@ -0,0 +1,30 @@
#
# Copyright (c) 2025 Huawei Technologies Co., Ltd. All Rights Reserved.
# Copyright 2023 The vLLM team.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from vllm.v1.worker.workspace import init_workspace_manager
from vllm_ascend.worker.worker import NPUWorker
from vllm_ascend.xlite.xlite_model_runner import XliteModelRunner
class XliteWorker(NPUWorker):
"""Xlite worker bases on NPUWorker. Only xlite specified code should be added in this class."""
def init_device(self):
"""Override init_device to init xlite model runner"""
self.device = self._init_device()
num_ubatches = 1
init_workspace_manager(self.device, num_ubatches)
self.model_runner = XliteModelRunner(self.vllm_config, self.device)