# # 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. # """Ascend W4A16_MXFP4 quantization helpers and fused MoE method.""" from collections.abc import Callable from typing import Any import torch import torch_npu from vllm.config import CompilationMode, get_current_vllm_config from vllm.distributed import get_ep_group from vllm_ascend.ascend_config import get_ascend_config from vllm_ascend.ascend_forward_context import _EXTRA_CTX from vllm_ascend.device.mxfp_compat import ( FLOAT8_E8M0FNU_DTYPE, ensure_mxfp4_moe_available, ) from vllm_ascend.ops.fused_moe.experts_selector import select_experts from vllm_ascend.ops.fused_moe.moe_runtime_args import build_fused_experts_input from .base import AscendMoEScheme, QuantType, get_moe_num_logical_experts from .registry import register_scheme # Unpack the weights to FP4 and return them in float32 format def unpack_uint8_to_fp4_return_float32(packed: torch.Tensor) -> torch.Tensor: low = packed & 0x0F high = packed // 16 # The high 4 bits and low 4 bits are arranged alternately, with the low 4 bits in front. unpacked = torch.stack([low, high], dim=-1).reshape(*packed.shape[:-1], -1) # A 4-digit integer is mapped to mxfp4 based on its value. fp4_values = torch.tensor( [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, -0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0], dtype=torch.float32, device=packed.device, ) return fp4_values[unpacked.to(torch.long)] @register_scheme("W4A16_MXFP4", "moe") class AscendW4A16MXFP4FusedMoEMethod(AscendMoEScheme): """FusedMoE method for Ascend W4A16_MXFP4.""" quant_type: QuantType = QuantType.W4A16MXFP4 def __init__(self) -> None: ensure_mxfp4_moe_available("W4A16_MXFP4 MoE quantization") self.ep_group = get_ep_group() vllm_config = get_current_vllm_config() self.group_size = vllm_config.quant_config.quant_description.get("group_size", 32) ascend_config = get_ascend_config() self.use_aclgraph = ( vllm_config.compilation_config.mode == CompilationMode.VLLM_COMPILE and not vllm_config.model_config.enforce_eager ) self.dynamic_eplb = ascend_config.eplb_config.dynamic_eplb def get_weight( self, num_experts: int, intermediate_size_per_partition: int, hidden_sizes: int, params_dtype: torch.dtype, ) -> dict[str, Any]: param_dict = {} param_dict["w13_weight"] = torch.empty( num_experts, 2 * intermediate_size_per_partition, hidden_sizes // 2, dtype=torch.uint8, ) param_dict["w2_weight"] = torch.empty( num_experts, hidden_sizes, intermediate_size_per_partition // 2, dtype=torch.uint8, ) return param_dict def get_dynamic_quant_param( self, num_experts: int, intermediate_size_per_partition: int, hidden_sizes: int, params_dtype: torch.dtype, ) -> dict[str, Any]: param_dict = {} param_dict["w13_weight_scale"] = torch.empty( num_experts, 2 * intermediate_size_per_partition, hidden_sizes // self.group_size, dtype=torch.uint8 ) param_dict["w2_weight_scale"] = torch.empty( num_experts, hidden_sizes, intermediate_size_per_partition // self.group_size, dtype=torch.uint8 ) return param_dict def apply( self, layer: torch.nn.Module, x: torch.Tensor, router_logits: torch.Tensor, top_k: int, renormalize: bool, use_grouped_topk: bool = False, num_experts: int = -1, expert_map: torch.Tensor | None = None, topk_group: int | None = None, num_expert_group: int | None = None, custom_routing_function: Callable | None = None, scoring_func: str = "softmax", routed_scaling_factor: float = 1.0, e_score_correction_bias: torch.Tensor | None = None, is_prefill: bool = True, enable_force_load_balance: bool = True, log2phy: torch.Tensor | None = None, global_redundant_expert_num: int = 0, pertoken_scale: Any | None = None, activation: str = "silu", apply_router_weight_on_input: bool = False, mc2_mask: torch.Tensor | None = None, tid2eid: torch.Tensor | None = None, ) -> torch.Tensor: num_shared_experts = getattr(layer, "n_shared_experts", 0) if num_shared_experts is None: num_shared_experts = 0 num_logical_experts = get_moe_num_logical_experts( layer, num_experts, global_redundant_expert_num=global_redundant_expert_num, num_shared_experts=num_shared_experts, ) assert router_logits.shape[1] == num_logical_experts, ( "Number of global experts mismatch (excluding redundancy): " f"router_logits.shape[1]={router_logits.shape[1]}, num_logical_experts={num_logical_experts}" ) topk_weights, topk_ids = select_experts( hidden_states=x, router_logits=router_logits, top_k=top_k, use_grouped_topk=use_grouped_topk, renormalize=renormalize, topk_group=topk_group, num_expert_group=num_expert_group, custom_routing_function=custom_routing_function, scoring_func=scoring_func, e_score_correction_bias=e_score_correction_bias, routed_scaling_factor=routed_scaling_factor, num_experts=num_logical_experts, tid2eid=tid2eid, ) if enable_force_load_balance: random_matrix = torch.rand(topk_ids.size(0), num_logical_experts, device=topk_ids.device) topk_ids = torch.argsort(random_matrix, dim=1)[:, : topk_ids.size(1)].to(topk_ids.dtype) topk_weights = topk_weights.to(x.dtype) moe_comm_method = _EXTRA_CTX.moe_comm_method return moe_comm_method.fused_experts( fused_experts_input=build_fused_experts_input( hidden_states=x, topk_weights=topk_weights, topk_ids=topk_ids, w1=layer.w13_weight, w2=layer.w2_weight, quant_type=self.quant_type, dynamic_eplb=self.dynamic_eplb, expert_map=expert_map, global_redundant_expert_num=global_redundant_expert_num, mc2_mask=mc2_mask, apply_router_weight_on_input=apply_router_weight_on_input, log2phy=log2phy, pertoken_scale=pertoken_scale, activation=activation, mxfp_act_quant_type=None, mxfp_weight_quant_type=torch_npu.float4_e2m1fn_x2, mxfp_scale_dtype=FLOAT8_E8M0FNU_DTYPE, mxfp_per_token_scale_dtype=None, mxfp_use_bf16=(x.dtype == torch.bfloat16), w1_scale=layer.w13_weight_scale, w2_scale=layer.w2_weight_scale, swiglu_limit=layer.swiglu_limit, ) ) def process_weights_after_loading(self, layer): layer.w13_weight.data = unpack_uint8_to_fp4_return_float32(layer.w13_weight.data) layer.w13_weight.data = layer.w13_weight.data.transpose(1, 2) layer.w13_weight.data = torch_npu.npu_format_cast(layer.w13_weight.data, 29, customize_dtype=torch.bfloat16) layer.w13_weight.data = torch_npu.npu_convert_weight_to_int4pack(layer.w13_weight.data).contiguous() layer.w2_weight.data = unpack_uint8_to_fp4_return_float32(layer.w2_weight.data) layer.w2_weight.data = layer.w2_weight.data.transpose(1, 2) layer.w2_weight.data = torch_npu.npu_format_cast(layer.w2_weight.data, 29, customize_dtype=torch.bfloat16) layer.w2_weight.data = torch_npu.npu_convert_weight_to_int4pack(layer.w2_weight.data).contiguous() layer.w13_weight_scale.data = layer.w13_weight_scale.data.transpose(1, 2).contiguous() layer.w2_weight_scale.data = layer.w2_weight_scale.data.transpose(1, 2).contiguous()