748 lines
34 KiB
Python
748 lines
34 KiB
Python
#
|
|
# 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]
|