Files
enginex-ascend-910-vllm/vllm_ascend/ascend_config.py
Sun Ruoxi 7f8a1b1f7a init v0.23.0
Signed-off-by: Sun Ruoxi <sunruoxi@4paradigm.com>
2026-08-27 15:11:51 +08:00

945 lines
46 KiB
Python

#
# Copyright (c) 2025 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 json
import os
from typing import TYPE_CHECKING, Any
from vllm.logger import logger
from vllm.utils.math_utils import cdiv
if TYPE_CHECKING:
from vllm.config import VllmConfig
class AscendConfig:
"""
Configuration Object for additional_config from vllm.configs.
"""
def __init__(self, vllm_config: "VllmConfig"):
self.vllm_config = vllm_config
additional_config = vllm_config.additional_config if vllm_config.additional_config is not None else {}
self._check_mooncake_c8_kv_cache_quant(vllm_config)
xlite_graph_config = additional_config.get("xlite_graph_config", {})
self.xlite_graph_config = XliteGraphConfig(xlite_graph_config, vllm_config)
ascend_compilation_config = additional_config.get("ascend_compilation_config", {})
self.ascend_compilation_config = AscendCompilationConfig(**ascend_compilation_config)
ascend_fusion_config = additional_config.get("ascend_fusion_config", {})
self.ascend_fusion_config = AscendFusionConfig(**ascend_fusion_config)
finegrained_tp_config = additional_config.get("finegrained_tp_config", {})
self.finegrained_tp_config = FinegrainedTPConfig(finegrained_tp_config, vllm_config)
eplb_config = additional_config.get("eplb_config", {})
self.eplb_config = EplbConfig(eplb_config)
weight_prefetch_config = additional_config.get("weight_prefetch_config", {})
self.weight_prefetch_config = WeightPrefetchConfig(weight_prefetch_config)
profiling_chunk_config = additional_config.get("profiling_chunk_config", {})
self.profiling_chunk_config = ProfilingChunkConfig(profiling_chunk_config)
if self.profiling_chunk_config.enabled:
max_batched = vllm_config.scheduler_config.max_num_batched_tokens
if max_batched < self.profiling_chunk_config.min_chunk:
logger.warning(
"max_num_batched_tokens is smaller than profiling_chunk_config.min_chunk. "
"max_num_batched_tokens=%d, min_chunk=%d. "
"Clamping min_chunk to %d to avoid it being silently ignored.",
max_batched,
self.profiling_chunk_config.min_chunk,
max_batched,
)
self.profiling_chunk_config.min_chunk = max_batched
if self.profiling_chunk_config.enabled and vllm_config.parallel_config.pipeline_parallel_size <= 1:
raise ValueError(
"profiling_chunk_config requires pipeline parallelism (pp > 1). "
"Please set --pipeline-parallel-size to a value greater than 1, "
"or disable profiling_chunk_config."
)
from vllm_ascend import envs as ascend_envs
self.enable_balance_scheduling = self._get_config_value(
additional_config,
"enable_balance_scheduling",
"VLLM_ASCEND_BALANCE_SCHEDULING",
ascend_envs.VLLM_ASCEND_BALANCE_SCHEDULING,
)
self.enable_flashcomm1 = self._get_config_value(
additional_config,
"enable_flashcomm1",
"VLLM_ASCEND_ENABLE_FLASHCOMM1",
ascend_envs.VLLM_ASCEND_ENABLE_FLASHCOMM1,
)
if self.profiling_chunk_config.enabled and self.enable_balance_scheduling:
raise ValueError(
"profiling_chunk_config and balance scheduling (enable_balance_scheduling) "
"cannot be enabled at the same time. Please disable one of them."
)
# Dump / PrecisionDebugger configuration
self.dump_config_path = self._resolve_dump_config_path(additional_config)
# Log configuration
self.ascend_log_path = additional_config.get(
"ascend_log_path",
os.path.join(os.path.expanduser("~"), "ascend", "log", "vllm_ascend"),
)
self.layer_sharding = additional_config.get("layer_sharding", None)
if self.layer_sharding:
logger.info_once(
"Linear layer sharding enabled with config: %s. "
"Note: This feature works optimally with FLASHCOMM2 and DSA-CP enabled; "
"using it without these features may result in significant performance degradation.",
str(self.layer_sharding),
)
self.enable_shared_expert_dp = (
additional_config.get("enable_shared_expert_dp", False)
and vllm_config.parallel_config.enable_expert_parallel
and vllm_config.parallel_config.tensor_parallel_size > 1
)
from vllm_ascend.utils import enable_sp
if self.enable_shared_expert_dp:
assert enable_sp(vllm_config=vllm_config, enable_shared_expert_dp=True)
if vllm_config.parallel_config.prefill_context_parallel_size > 1 and enable_sp(vllm_config=vllm_config):
tp_pcp_size = (
vllm_config.parallel_config.tensor_parallel_size
* vllm_config.parallel_config.prefill_context_parallel_size
)
if vllm_config.scheduler_config.max_num_batched_tokens % tp_pcp_size != 0:
vllm_config.scheduler_config.max_num_batched_tokens = (
cdiv(vllm_config.scheduler_config.max_num_batched_tokens, tp_pcp_size) * tp_pcp_size
)
logger.warning_once(
"When using FLASHCOMM1, the max_num_batched_tokens should be divisible "
"by tp_size * pcp_size (%s). It has been adjusted to %s.",
str(tp_pcp_size),
str(vllm_config.scheduler_config.max_num_batched_tokens),
)
self.multistream_overlap_shared_expert = additional_config.get("multistream_overlap_shared_expert", False)
self.multistream_overlap_gate = additional_config.get("multistream_overlap_gate", False)
# PD-disaggregated D node only (kv_consumer); invalid on P nodes and in PD-mixed mode.
self.recompute_scheduler_enable = additional_config.get("recompute_scheduler_enable", False)
# DSV4 oproj / embedding fine-grained TP (oproj_tensor_parallel_size /
# embedding_tensor_parallel_size) use static, graph-stable exchange
# buffers and run cross-DP HCCL collectives (all_to_all / all_gather /
# reduce_scatter) that require uniform num_tokens across all DP ranks.
# Only the recompute scheduler balances num_tokens across DP ranks; the
# MC2 uneven-token skip path leaves each rank at its own num_tokens and
# would deadlock these collectives on a shape mismatch. So bind both
# features to recompute_scheduler_enable: they are only supported when
# recompute is on.
if (
self.finegrained_tp_config.oproj_tensor_parallel_size > 0
or self.finegrained_tp_config.embedding_tensor_parallel_size > 0
) and not self.recompute_scheduler_enable:
raise AssertionError(
"oproj_tensor_parallel_size / embedding_tensor_parallel_size "
"require recompute_scheduler_enable=true: their cross-DP HCCL "
"collectives need uniform num_tokens across DP ranks, which is "
"only guaranteed when the recompute scheduler is enabled."
)
self.enable_cpu_binding = additional_config.get("enable_cpu_binding", True)
self.enable_sleep_mode_extra_cleanup = additional_config.get("enable_sleep_mode_extra_cleanup", False)
self.multistream_dsv4_dsa_overlap = additional_config.get("multistream_dsv4_dsa_overlap", True)
self.enable_prefill_mc2 = bool(additional_config.get("enable_prefill_mc2", False))
self.enable_matmul_allreduce = self._get_config_value(
additional_config,
"enable_matmul_allreduce",
"VLLM_ASCEND_ENABLE_MATMUL_ALLREDUCE",
ascend_envs.VLLM_ASCEND_ENABLE_MATMUL_ALLREDUCE,
)
self.enable_fused_mc2 = self._get_config_value(
additional_config,
"enable_fused_mc2",
"VLLM_ASCEND_ENABLE_FUSED_MC2",
ascend_envs.VLLM_ASCEND_ENABLE_FUSED_MC2,
)
if self.enable_fused_mc2 == 1 and self.multistream_overlap_shared_expert:
self.multistream_overlap_shared_expert = False
logger.warning_once(
"VLLM_ASCEND_ENABLE_FUSED_MC2 (fused mc2) and multistream_overlap_shared_expert "
"cannot be enabled at the same time. Setting multistream_overlap_shared_expert to False."
)
self.enable_mlapo = self._get_config_value(
additional_config,
"enable_mlapo",
"VLLM_ASCEND_ENABLE_MLAPO",
ascend_envs.VLLM_ASCEND_ENABLE_MLAPO,
)
self.enable_flashcomm2_parallel_size = self._get_config_value(
additional_config,
"enable_flashcomm2_parallel_size",
"VLLM_ASCEND_FLASHCOMM2_PARALLEL_SIZE",
ascend_envs.VLLM_ASCEND_FLASHCOMM2_PARALLEL_SIZE,
)
self.msmonitor_use_daemon = self._get_config_value(
additional_config,
"msmonitor_use_daemon",
"MSMONITOR_USE_DAEMON",
ascend_envs.MSMONITOR_USE_DAEMON,
)
self.enable_transpose_kv_cache_by_block = self._get_config_value(
additional_config,
"enable_transpose_kv_cache_by_block",
"VLLM_ASCEND_FUSION_OP_TRANSPOSE_KV_CACHE_BY_BLOCK",
ascend_envs.VLLM_ASCEND_FUSION_OP_TRANSPOSE_KV_CACHE_BY_BLOCK,
)
self.pd_tp_ratio = 1
self.pd_head_ratio = 1
self.num_head_replica = 1
if (
vllm_config.kv_transfer_config is not None
and vllm_config.model_config is not None
and not vllm_config.model_config.is_deepseek_mla
):
prefill_tp_size = vllm_config.kv_transfer_config.get_from_extra_config("prefill", {"tp_size": 1})["tp_size"]
decode_tp_size = vllm_config.kv_transfer_config.get_from_extra_config("decode", {"tp_size": 1})["tp_size"]
assert prefill_tp_size % decode_tp_size == 0, "Prefill TP size must be divisible by Decode TP size."
self.pd_tp_ratio = prefill_tp_size // decode_tp_size
if self.pd_tp_ratio > 1:
# Total KV heads from vLLM's resolved architecture (ModelArchConfigConvertor).
num_kv_head = vllm_config.model_config.get_total_num_kv_heads()
if not num_kv_head or num_kv_head < 1:
raise ValueError(
"Could not determine a positive total KV head count for PD "
"disaggregation (pd_tp_ratio > 1). Check that the model config "
"is compatible with vLLM."
)
self.num_head_replica = prefill_tp_size // num_kv_head if prefill_tp_size >= num_kv_head else 1
prefill_tp_size = min(prefill_tp_size, num_kv_head)
decode_tp_size = min(decode_tp_size, num_kv_head)
self.pd_head_ratio = prefill_tp_size // decode_tp_size
if self.pd_tp_ratio == 0:
raise AssertionError("Only support P node tp size lagger then D node tp size")
self.SLO_limits_for_dynamic_batch = additional_config.get("SLO_limits_for_dynamic_batch", -1)
from vllm_ascend.utils import get_flashcomm2_config_and_validate
self.flashcomm2_oproj_tensor_parallel_size = get_flashcomm2_config_and_validate(self, vllm_config)
# We find that _npu_paged_attention still performs better than
# npu_fused_infer_attention_score in some cases. We allow to execute
# _npu_paged_attention in this cases. This should be removed once
# npu_fused_infer_attention_score performs better on all scenarios.
self.pa_shape_list = additional_config.get("pa_shape_list", [])
# Weight NZ mode configuration.
# 0: disabled, 1: only quant case enable nz (default), 2: BF16/FP16 also enable nz
self.weight_nz_mode = self._get_config_value(
additional_config,
"weight_nz_mode",
"VLLM_ASCEND_ENABLE_NZ",
ascend_envs.VLLM_ASCEND_ENABLE_NZ,
)
# when enable_async_exponential is True, AscendSampler will be different from vllm Sampler,
# which make batch_invariant mode not working.
# so we disable async exponential when batch_invariant mode is enabled.
import vllm.envs as envs
self.enable_async_exponential = (
bool(additional_config.get("enable_async_exponential", False)) and not envs.VLLM_BATCH_INVARIANT
)
from vllm_ascend.utils import model_uses_sfa_sparse
use_sparse = model_uses_sfa_sparse(vllm_config.model_config)
self.enable_kv_nz = additional_config.get("enable_kv_nz", False)
if self.enable_kv_nz:
if vllm_config.model_config is None:
raise RuntimeError("enable_kv_nz requires a valid model_config.")
if not vllm_config.model_config.is_deepseek_mla or use_sparse:
raise RuntimeError("enable_kv_nz is only supported for mla currently.")
if vllm_config.kv_transfer_config is None or not vllm_config.kv_transfer_config.is_kv_consumer:
raise NotImplementedError(
"enable_kv_nz is only supported in pd scenario and can only be used in D node."
)
self.enable_sparse_sfa_c8 = additional_config.get("enable_sparse_sfa_c8", False) and use_sparse
self.enable_sparse_li_c8 = additional_config.get("enable_sparse_li_c8", False) and use_sparse
self.c8_enable_reshape_optim = self.enable_sparse_li_c8 and additional_config.get(
"c8_enable_reshape_optim", False
)
quant_config = getattr(vllm_config, "quant_config", None)
(
self._sparse_li_c8_layer_ids,
self._sparse_li_c8_layer_names,
) = self._parse_sparse_li_c8_layers_from_quant_config(quant_config)
self._sparse_li_c8_layer_filter_enabled = self._has_sparse_li_c8_layer_config(quant_config)
self.enable_sp_by_pass = (
vllm_config.model_config is not None
and not vllm_config.model_config.enforce_eager
and vllm_config.compilation_config.pass_config.enable_sp
)
# Enable dispatch/combine op inter-node communication by ROCE
self.enable_mc2_hierarchy_comm = additional_config.get("enable_mc2_hierarchy_comm", False)
# Per-rank token capacity after dispatch in the mega moe (dispatch_ffn_combine) fused operator.
# When load imbalance causes a rank to receive more tokens than this limit, the excess tokens
# are dropped and skipped from computation, degrading accuracy.
# Do not set this too large: workspace memory scales linearly with this value, which matters
# especially under long-context scenarios where the operator should not hold too much memory.
# Default 65536.
self.mega_moe_max_tokens = additional_config.get("mega_moe_max_tokens", 65536)
if not isinstance(self.mega_moe_max_tokens, int):
raise ValueError(
f"mega_moe_max_tokens must be an integer, got {type(self.mega_moe_max_tokens).__name__}: "
f"{self.mega_moe_max_tokens}"
)
if self.mega_moe_max_tokens <= 0:
raise ValueError(f"mega_moe_max_tokens must be a positive integer, got {self.mega_moe_max_tokens}")
# Whether to use NPU device group for DP metadata all_reduce.
# "True": use NPU device group, "False" (default): use CPU group.
self.dp_allreduce_on_npu = additional_config.get("dp_allreduce_on_npu", False)
# Enable optimized reduce sampling scheme
# NOTE: reduce sample is an experimental feature. It is incompatible with
# lmhead TP and PD-disaggregated P nodes (kv_role='kv_producer'); raising
# ValueError on those to avoid silent correctness issues. PD-disaggregated
# D nodes (kv_role='kv_consumer') are allowed for backward compatibility.
self.enable_reduce_sample = additional_config.get("enable_reduce_sample", False)
if self.enable_reduce_sample:
logger.warning_once("enable_reduce_sample is an experimental feature. Use with caution.")
if self.finegrained_tp_config.lmhead_tensor_parallel_size > 0:
raise ValueError(
"enable_reduce_sample is incompatible with "
"finegrained_tp_config.lmhead_tensor_parallel_size. "
"Please disable one of them."
)
kv_transfer_config = getattr(vllm_config, "kv_transfer_config", None)
kv_role = getattr(kv_transfer_config, "kv_role", None)
if kv_role == "kv_producer":
raise ValueError(
"enable_reduce_sample is not supported on PD-disaggregated "
"scenarios. Please disable enable_reduce_sample."
)
self.mix_placement = additional_config.get("mix_placement", False)
self._check_mix_placement()
self.hamming_sparse = additional_config.get("hamming_sparse", {"enabled": False, "sparse_json_location": ""})
self.enable_hamming_sparse = self.hamming_sparse["enabled"]
self.sparse_json = self.hamming_sparse["sparse_json_location"]
self._check_enable_hamming_sparse()
# Enable Block Verify and Entropy Verify in Rejection Sampler
rejection_sampler_config = additional_config.get("rejection_sampler_config", {})
self.rejection_sampler_config = RejectionSamplerConfig(rejection_sampler_config)
@staticmethod
def _get_config_value(additional_config: dict[str, Any], config_key: str, env_key: str, env_value: Any) -> Any:
if config_key in additional_config:
value = additional_config[config_key]
logger.info_once(f"AscendConfig.{config_key} is set from additional_config with value {value}.")
return value
if env_key in os.environ:
logger.info_once(
f"AscendConfig.{config_key} falls back to environment variable {env_key} with value {env_value}. "
f"Please use additional_config.{config_key} instead, because {env_key} will be removed in the "
"next release."
)
return env_value
@classmethod
def _check_mooncake_c8_kv_cache_quant(cls, vllm_config: "VllmConfig") -> None:
kv_transfer_config = getattr(vllm_config, "kv_transfer_config", None)
if kv_transfer_config is None:
return
quant_config = getattr(vllm_config, "quant_config", None)
enable_c8_quant = getattr(quant_config, "enable_c8_quant", False)
if enable_c8_quant is not True:
return
from vllm_ascend.utils import is_gqa_backend, uses_mooncake_connector
if not is_gqa_backend(vllm_config):
return
if not uses_mooncake_connector(kv_transfer_config):
return
raise ValueError(
"MooncakeConnector does not support C8 KV cache quantization on GQA models. "
"The producer keeps KV cache in bf16 while the consumer allocates int8 KV cache, so raw "
"Mooncake transfer would reinterpret bf16 bytes as int8. Please disable C8 KV cache quantization "
"or use MooncakeLayerwiseConnector, which quantizes KV cache before transfer."
)
def _check_mix_placement(self):
if self.mix_placement:
if self.enable_shared_expert_dp or self.multistream_overlap_shared_expert:
raise ValueError("Mix placement is not supported with shared expert DP or multistream overlap.")
def _check_enable_hamming_sparse(self):
if self.enable_hamming_sparse:
if isinstance(self.sparse_json, str) and not os.path.isfile(self.sparse_json):
raise ValueError("Hamming sparse config json file doesn't exist.")
@staticmethod
def _materialize_dump_config_to_file(dump_config: dict[str, Any]) -> str:
dump_config_dir = os.path.join(os.getcwd(), ".vllm_ascend", "msprobe")
os.makedirs(dump_config_dir, exist_ok=True)
dump_config_file_path = os.path.join(dump_config_dir, "msprobe_dump_config.json")
with open(dump_config_file_path, "w", encoding="utf-8") as file:
json.dump(dump_config, file, ensure_ascii=False, indent=2)
logger.info("Materialized additional_config.dump_config to file: %s", dump_config_file_path)
return dump_config_file_path
@classmethod
def _resolve_dump_config_path(cls, additional_config: dict[str, Any]) -> str | None:
dump_config_path = additional_config.get("dump_config_path")
dump_config = additional_config.get("dump_config")
if dump_config_path is not None and dump_config is not None:
raise ValueError(
"Only one of additional_config.dump_config_path or additional_config.dump_config can be set."
)
if dump_config is not None:
if not isinstance(dump_config, dict):
raise ValueError(f"additional_config.dump_config must be a dict, got {type(dump_config).__name__}.")
return cls._materialize_dump_config_to_file(dump_config)
if dump_config_path is not None and not isinstance(dump_config_path, str):
raise ValueError(
f"additional_config.dump_config_path must be a string, got {type(dump_config_path).__name__}."
)
return dump_config_path
@staticmethod
def _has_sparse_li_c8_layer_config(quant_config: Any) -> bool:
quant_description = getattr(quant_config, "quant_description", None)
if not isinstance(quant_description, dict):
return False
return any(isinstance(key, str) and key.endswith(".indexer.quant_type") for key in quant_description)
@classmethod
def _parse_sparse_li_c8_layers_from_quant_config(cls, quant_config: Any) -> tuple[set[int], set[str]]:
quant_description = getattr(quant_config, "quant_description", None)
if not isinstance(quant_description, dict):
return set(), set()
QUANT_SUFFIXES = (".indexer.quant_type", ".indexer.wq_b_weight")
VALID_QUANT_TYPES = ("INT8_DYNAMIC", "W8A8_MXFP8")
layer_ids: set[int] = set()
layer_names: set[str] = set()
from vllm.model_executor.models.utils import extract_layer_index
for key, value in quant_description.items():
if not isinstance(key, str):
continue
matched_suffix = next((s for s in QUANT_SUFFIXES if key.endswith(s)), None)
if matched_suffix is None or value not in VALID_QUANT_TYPES:
continue
layer_name = key[: -len(matched_suffix)].rstrip(".")
if not layer_name:
continue
layer_names.add(layer_name)
layer_ids.add(extract_layer_index(layer_name))
return layer_ids, layer_names
def is_sparse_li_c8_layer(self, layer_name: str | None) -> bool:
if not self.enable_sparse_li_c8:
return False
if not self._sparse_li_c8_layer_filter_enabled:
return True
if layer_name is None:
return False
normalized_layer_name = layer_name.rstrip(".")
if any(
normalized_layer_name == candidate or normalized_layer_name.startswith(f"{candidate}.")
for candidate in self._sparse_li_c8_layer_names
):
return True
from vllm.model_executor.models.utils import extract_layer_index
layer_ids = {extract_layer_index(normalized_layer_name)}
return any(layer_id in self._sparse_li_c8_layer_ids for layer_id in layer_ids)
@staticmethod
def _get_compile_ranges(compilation_config):
return compilation_config.compile_ranges_endpoints or []
@staticmethod
def _set_compile_ranges(compilation_config, value):
compilation_config.compile_ranges_endpoints = value
def update_compile_ranges_split_points(self):
vllm_config = self.vllm_config
if self.ascend_compilation_config.enable_npugraph_ex:
if self.ascend_compilation_config.fuse_allreduce_rms:
from vllm_ascend.compilation.passes.allreduce_rmsnorm_fusion_pass import ALLREDUCE_NORM_FUSE_THRESHOLD
new_compile_ranges_split_points = self._get_compile_ranges(vllm_config.compilation_config)
new_compile_ranges_split_points.append(ALLREDUCE_NORM_FUSE_THRESHOLD)
new_compile_ranges_split_points = sorted(new_compile_ranges_split_points)
self._set_compile_ranges(vllm_config.compilation_config, new_compile_ranges_split_points)
logger.debug(
"Set compile_ranges_split_points to %s for matmul and allreduce fusion",
new_compile_ranges_split_points,
)
else:
new_compile_ranges_split_points = self._get_compile_ranges(vllm_config.compilation_config)
if vllm_config.additional_config.get("ascend_compilation_config", {}).get("fuse_allreduce_rms", True):
from vllm_ascend.compilation.passes.allreduce_rmsnorm_fusion_pass import ALLREDUCE_NORM_FUSE_THRESHOLD
new_compile_ranges_split_points.append(ALLREDUCE_NORM_FUSE_THRESHOLD)
new_compile_ranges_split_points = sorted(new_compile_ranges_split_points)
self._set_compile_ranges(vllm_config.compilation_config, new_compile_ranges_split_points)
logger.debug(
"Set compile_ranges_split_points to %s for matmul and allreduce fusion",
new_compile_ranges_split_points,
)
if len(new_compile_ranges_split_points) > len(self._get_compile_ranges(vllm_config.compilation_config)):
new_compile_ranges_split_points = sorted(new_compile_ranges_split_points)
self._set_compile_ranges(vllm_config.compilation_config, new_compile_ranges_split_points)
class FinegrainedTPConfig:
"""
Configuration Object for finegrained_tp_config from additional_config
"""
def __init__(self, finegrained_tp_config: dict, vllm_config):
self.oproj_tensor_parallel_size = finegrained_tp_config.get("oproj_tensor_parallel_size", 0)
self.lmhead_tensor_parallel_size = finegrained_tp_config.get("lmhead_tensor_parallel_size", 0)
self.embedding_tensor_parallel_size = finegrained_tp_config.get("embedding_tensor_parallel_size", 0)
self.mlp_tensor_parallel_size = finegrained_tp_config.get("mlp_tensor_parallel_size", 0)
self.olora_tensor_parallel_size = finegrained_tp_config.get("olora_tensor_parallel_size", 0)
enabled_configs = []
if self.oproj_tensor_parallel_size > 0:
enabled_configs.append(f"oproj_tensor_parallel_size={self.oproj_tensor_parallel_size}")
# wo_a/wo_b are sharded solely by the OTP group (which splits DP,
# orthogonal to the standard TP group), but _forward_o_proj reshapes
# the attention output with n_local_groups = n_groups // tp_size
# (standard TP). When tp_size > 1 the weight-shard and input-shard
# operate on different axes of the rank grid and no longer align,
# so oproj TP currently requires standard tp_size == 1.
if vllm_config.parallel_config.tensor_parallel_size > 1:
raise AssertionError(
"oproj_tensor_parallel_size currently requires "
"tensor_parallel_size == 1, got "
f"{vllm_config.parallel_config.tensor_parallel_size}."
)
# The static all_to_all / reduce_scatter exchange buffers used by
# _forward_o_proj are sized for graph replay and require ACL graph
# capture; dummy_run does not run the entire attention module in
# eager mode, so o_proj tp split can only be used in graph mode.
if vllm_config.model_config and vllm_config.model_config.enforce_eager:
raise AssertionError("oproj_tensor_parallel_size is only supported in graph mode")
if vllm_config.kv_transfer_config is None or not vllm_config.kv_transfer_config.is_kv_consumer:
raise AssertionError(
"oproj_tensor_parallel_size is only supported in pd scenario and can only be used in D node."
)
if self.olora_tensor_parallel_size > 0:
enabled_configs.append(f"olora_tensor_parallel_size={self.olora_tensor_parallel_size}")
# dummy_run does not run the entire attention module in eager mode,
# so the o_lora tp split can only be used in graph mode.
if vllm_config.model_config and vllm_config.model_config.enforce_eager:
raise AssertionError("olora_tensor_parallel_size is only supported in graph mode")
if vllm_config.kv_transfer_config is None or not vllm_config.kv_transfer_config.is_kv_consumer:
raise AssertionError(
"olora_tensor_parallel_size is only supported in pd scenario and can only be used in D node."
)
if self.lmhead_tensor_parallel_size > 0:
enabled_configs.append(f"lmhead_tensor_parallel_size={self.lmhead_tensor_parallel_size}")
if self.embedding_tensor_parallel_size > 0:
enabled_configs.append(f"embedding_tensor_parallel_size={self.embedding_tensor_parallel_size}")
if self.mlp_tensor_parallel_size > 0:
enabled_configs.append(f"mlp_tensor_parallel_size={self.mlp_tensor_parallel_size}")
module_tp_sizes = [
self.oproj_tensor_parallel_size,
self.lmhead_tensor_parallel_size,
self.embedding_tensor_parallel_size,
self.mlp_tensor_parallel_size,
self.olora_tensor_parallel_size,
]
for module_tp_size in module_tp_sizes:
# If it is a dense model, then expert parallel is not needed,
# and data parallel is also not needed. If the data parallel size is set
# to greater than 1 in the model launch configuration, its value will be changed to 1 later.
# This will cause an issue when finegrained tp is enabled, as it
# cannot be split into the data parallel communication group, leading to an error.
if module_tp_size > 0 and not vllm_config.model_config.is_moe:
raise AssertionError("The finegrained tp sizes can be enabled only for MOE models.")
if module_tp_size > 0 and vllm_config.parallel_config.data_parallel_size % module_tp_size != 0:
raise AssertionError("finegrained tp sizes must divide by data_parallel_size.")
if any(size > 0 for size in module_tp_sizes) and enabled_configs:
logger.info("finegrained_tp_config enabled: %s", ", ".join(enabled_configs))
class AscendCompilationConfig:
"""
Configuration for controlling the behavior of Ascend graph optimization.
This class provides a way to configure graph fusion optimizations.
These configurations directly impact the performance and behavior of models
deployed on Ascend platforms.
"""
def __init__(
self,
enable_npugraph_ex: bool = True,
enable_static_kernel: bool = False,
fuse_norm_quant: bool = True,
fuse_qknorm_rope: bool = True,
fuse_allreduce_rms: bool = False,
**kwargs,
):
"""
Initialize the configuration.
Args:
enable_npugraph_ex (bool): Whether to enable npugraph_ex backend.
When set to True, the Fx graph generated by Dymano will be
optimized and compiled by the npugraph_ex backend.
Default: True
enable_static_kernel (bool): Whether to enable static kernel.
Static kernel is suitable for scenarios with purely static shapes
or minimal shape changes, and can improve network performance.
When set to True, when during graph capture, it will compile operator
binary files with the corresponding shapes based on the current batch_size,
which usually takes some time.
Default: False
fuse_norm_quant (bool): Whether to enable norm and quant fusion optimization.
When set to True, the system will optimize norm and quant operations.
Default: True
fuse_qknorm_rope (bool): Whether to enable qknorm and rope fusion optimization.
Default: True
fuse_allreduce_rms (bool): Whether to enable allreduce and addrmsnorm fusion optimization.
Default: False
**kwargs: Additional optional parameters for forward compatibility and configuration extension.
"""
from vllm_ascend.utils import is_310p
if is_310p():
if enable_npugraph_ex:
logger.warning("npugraph_ex is not supported on Ascend 310P. Disabling it.")
if enable_static_kernel:
logger.warning(
"static kernel requires npugraph_ex, which is not supported on Ascend 310P. Disabling it."
)
enable_npugraph_ex = False
enable_static_kernel = False
self.fuse_norm_quant = fuse_norm_quant
self.fuse_qknorm_rope = fuse_qknorm_rope
self.fuse_allreduce_rms = fuse_allreduce_rms
self.enable_npugraph_ex = enable_npugraph_ex
self.enable_static_kernel = enable_static_kernel
self.fuse_muls_add = kwargs.get("fuse_muls_add", True)
if self.enable_static_kernel:
assert self.enable_npugraph_ex, "Static kernel generation requires npugraph_ex to be enabled."
class AscendFusionConfig:
"""
Configuration for controlling whether to use a fused operator gmmswigluquant.
"""
def __init__(self, fusion_ops_gmmswigluquant: bool = True, **kwargs):
"""
Initialize the configuration.
Args:
fusion_ops_gmmswigluquant (bool): Whether to use a fused operator gmmswigluquant.
When set to True, the system will use a fused operator gmmswigluquant.
Default: True
**kwargs: Additional optional parameters for forward compatibility and configuration extension.
"""
self.fusion_ops_gmmswigluquant = fusion_ops_gmmswigluquant
class XliteGraphConfig:
"""
Configuration Object for xlite_graph_config from additional_config
"""
def __init__(self, xlite_graph_config, vllm_config):
self.enabled = xlite_graph_config.get("enabled", False)
self.full_mode = xlite_graph_config.get("full_mode", False)
if self.enabled:
if bool(vllm_config.speculative_config) and vllm_config.speculative_config.num_speculative_tokens != 1:
raise RuntimeError("Xlite graph mode only support speculative decoding with num_speculative_tokens=1.")
if vllm_config.parallel_config.pipeline_parallel_size > 1:
raise RuntimeError(
"Xlite graph mode is not compatible with pipeline parallelism. "
"Please set pipeline_parallel_size to 1."
)
if vllm_config.cache_config.block_size != 128:
logger.warning(
"Current cache block size may not be optimal for xlite graph mode. "
"current_block_size=%d, recommended_block_size=128.",
vllm_config.cache_config.block_size,
)
class WeightPrefetchConfig:
"""
Configuration Object for weight_prefetch_config from additional_config
"""
prefetch_ratio: dict = {
"attn": {
"qkv": 1.0,
"o": 1.0,
},
"moe": {"gate_up": 0.8},
"mlp": {"gate_up": 1.0, "down": 1.0},
}
def __init__(self, weight_prefetch_config: dict):
self.enabled = weight_prefetch_config.get("enabled", False)
self.prefetch_ratio = weight_prefetch_config.get("prefetch_ratio", self.prefetch_ratio)
class ProfilingChunkConfig:
"""Configuration for profiling-based dynamic chunk sizing.
When enabled, the scheduler profiles prefill latency during initialization
and uses a quadratic model to predict optimal chunk sizes at runtime.
Usage (online)::
vllm serve <model> --additional-config '{"profiling_chunk_config": {"enabled": true}}'
Usage (offline)::
llm = LLM(model, additional_config={"profiling_chunk_config": {"enabled": true}})
"""
def __init__(self, config: dict | None = None):
if config is None:
config = {}
self.enabled: bool = config.get("enabled", False)
self.smooth_factor: float = float(config.get("smooth_factor", 1.0))
self.min_chunk: int = int(config.get("min_chunk", 4096))
# Controls online history-aware calibration. When True, the model
# runner synchronizes the device each step to measure execution time
# and feeds it back for incremental refitting. Automatically set to
# False once calibration completes. Users can set it to False from
# the start to skip online calibration entirely and rely solely on
# the startup profiling model (avoids per-step sync overhead).
self.need_timing: bool = config.get("need_timing", self.enabled)
self.max_fit_chunk: int = int(config.get("max_fit_chunk", 30))
self._validate()
def _validate(self):
if not (0 < self.smooth_factor <= 1.0):
raise ValueError(f"profiling_chunk_config.smooth_factor must be in (0, 1], got {self.smooth_factor}")
if self.min_chunk <= 0:
raise ValueError(f"profiling_chunk_config.min_chunk must be positive, got {self.min_chunk}")
if self.max_fit_chunk <= 5:
raise ValueError(f"Recommend to use at least 30 data points for fitting, got {self.max_fit_chunk}")
class RejectionSamplerConfig:
"""Configuration for Block Verify and Entropy Verify in Rejection Sampler.
Block Verify improves acceptance rate by evaluating all draft tokens
as a block using cumulative probability products. Entropy Verify
adjusts the acceptance threshold based on the entropy of the target
distribution, allowing higher acceptance for high-entropy (uncertain)
tokens and stricter rejection for low-entropy (confident) tokens.
Usage (online)::
vllm serve <model> --additional-config \
'{"rejection_sampler_config": {"enable_block_verify": true, \
"enable_entropy_verify": true, "posterior_threshold": 0.95, \
"posterior_alpha": 0.4}}'
Usage (offline)::
llm = LLM(
model,
additional_config={
"rejection_sampler_config": {
"enable_block_verify": true,
"enable_entropy_verify": true,
"posterior_threshold": 0.95,
"posterior_alpha": 0.4,
}
},
)
"""
def __init__(self, config: dict | None = None):
if config is None:
config = {}
self.enable_block_verify: bool = config.get("enable_block_verify", False)
self.enable_entropy_verify: bool = config.get("enable_entropy_verify", False)
self.posterior_threshold: float = config.get("posterior_threshold", 0.95)
self.posterior_alpha: float = config.get("posterior_alpha", 0.4)
self._validate()
def _validate(self):
if not isinstance(self.enable_block_verify, bool):
raise ValueError(
f"rejection_sampler_config.enable_block_verify must be a bool, "
f"got {type(self.enable_block_verify).__name__}"
)
if not isinstance(self.enable_entropy_verify, bool):
raise ValueError(
f"rejection_sampler_config.enable_entropy_verify must be a bool, "
f"got {type(self.enable_entropy_verify).__name__}"
)
if not isinstance(self.posterior_threshold, (int, float)):
raise ValueError(
f"rejection_sampler_config.posterior_threshold must be a float, "
f"got {type(self.posterior_threshold).__name__}"
)
if not isinstance(self.posterior_alpha, (int, float)):
raise ValueError(
f"rejection_sampler_config.posterior_alpha must be a float, got {type(self.posterior_alpha).__name__}"
)
if not (0 < self.posterior_threshold <= 1):
raise ValueError(
f"rejection_sampler_config.posterior_threshold must be in (0, 1], got {self.posterior_threshold}"
)
if self.posterior_alpha < 0:
raise ValueError(f"rejection_sampler_config.posterior_alpha must be >= 0, got {self.posterior_alpha}")
class EplbConfig:
"""
Configuration Object for xlite_graph_config from additional_config
"""
_defaults = {
"dynamic_eplb": False,
"expert_map_path": None,
"expert_heat_collection_interval": 600,
"algorithm_execution_interval": 50,
"expert_map_record_path": None,
"num_redundant_experts": 0,
"eplb_policy_type": 2,
"eplb_heat_collection_stage": "all",
}
def __init__(self, user_config: dict | None = None):
if user_config is None:
user_config = {}
self.config = self._defaults.copy()
if user_config and isinstance(user_config, dict):
for key, value in user_config.items():
if key in self.config:
self.config[key] = value
else:
raise ValueError(f"Config has no attribute '{key}'")
self._validate_config()
def __getattr__(self, key):
if key in self.config:
return self.config[key]
raise AttributeError(f"Config has no attribute '{key}'")
def _validate_config(self):
if self.expert_map_path is not None:
logger.info("The expert_map is %s", self.expert_map_path)
if self.expert_map_path[-5:] != ".json":
raise TypeError("The expert_map is not json.")
if not (os.path.exists(self.expert_map_path) and os.access(self.expert_map_path, os.R_OK)):
raise ValueError("The expert_map is not exist.")
if self.expert_map_record_path is not None:
self.config["dynamic_eplb"] = True
if self.expert_map_record_path[-5:] != ".json":
raise TypeError("The expert_map_record_path is not json.")
dirname = os.path.dirname(self.expert_map_record_path)
os.makedirs(dirname, exist_ok=True)
for key in ["expert_heat_collection_interval", "algorithm_execution_interval", "num_redundant_experts"]:
if not isinstance(self.config[key], int):
raise TypeError(f"{key} must be an integer")
if self.config[key] < 0: # type: ignore
raise ValueError(f"{key} must greater than 0; got {self.config[key]} instead")
if self.eplb_policy_type not in [0, 1, 2, 3]:
raise ValueError("eplb_policy_type must in [0, 1, 2, 3]")
if self.config["dynamic_eplb"]:
assert (
os.getenv("DYNAMIC_EPLB", "false").lower() in ("true", "1")
or os.getenv("EXPERT_MAP_RECORD", "false") == "true"
), "The environment variable DYNAMIC_EPLB or EXPERT_MAP_RECORD of the EPLB must be set to true."
if self.eplb_heat_collection_stage not in ["all", "prefill", "decode"]:
raise ValueError('eplb_heat_collection_stage must be one of ["all", "prefill", "decode"]')
logger.info("Dynamic EPLB is %s", self.config["dynamic_eplb"])
logger.info("The number of redundant experts is %s", self.config["num_redundant_experts"])
_ASCEND_CONFIG: AscendConfig | None = None
def _is_ascend_config_initialized(config: AscendConfig | None) -> bool:
"""Check whether a config object has essential initialized fields.
Some unit tests monkeypatch ``AscendConfig.__init__`` to bypass heavy
initialization. In that case, the singleton cache can be polluted with a
partially initialized instance. This guard prevents reusing such instances
across tests.
"""
if config is None:
return False
return hasattr(config, "ascend_compilation_config") and hasattr(config, "eplb_config")
def init_ascend_config(vllm_config):
additional_config = vllm_config.additional_config if vllm_config.additional_config is not None else {}
refresh = additional_config.get("refresh", False) if additional_config else False
global _ASCEND_CONFIG
if (
_ASCEND_CONFIG is not None
and not refresh
and _is_ascend_config_initialized(_ASCEND_CONFIG)
and getattr(_ASCEND_CONFIG, "vllm_config", None) is vllm_config
):
return _ASCEND_CONFIG
new_config = AscendConfig(vllm_config)
if _is_ascend_config_initialized(new_config):
_ASCEND_CONFIG = new_config
else:
logger.warning("Ascend config instance is not fully initialized. action: skip singleton cache update. ")
return new_config
def clear_ascend_config():
global _ASCEND_CONFIG
_ASCEND_CONFIG = None
from vllm_ascend.utils import clear_enable_sp
clear_enable_sp()
def get_ascend_config():
global _ASCEND_CONFIG
if _ASCEND_CONFIG is None or not _is_ascend_config_initialized(_ASCEND_CONFIG):
raise RuntimeError("Ascend config is not initialized. Please call init_ascend_config first.")
return _ASCEND_CONFIG