diff --git a/vllm_ascend/eplb/core/eplb_device_transfer_loader.py b/vllm_ascend/eplb/core/eplb_device_transfer_loader.py index 9a8a323..67e4d56 100644 --- a/vllm_ascend/eplb/core/eplb_device_transfer_loader.py +++ b/vllm_ascend/eplb/core/eplb_device_transfer_loader.py @@ -50,10 +50,6 @@ class D2DExpertWeightLoader: ) return - # If neither send nor receive task is needed for this layer on this rank, return - if not (expert_send_info or expert_recv_info): - return - self.updated_expert_map = updated_expert_map self.layer_id = layer_id diff --git a/vllm_ascend/ops/expert_load_balancer.py b/vllm_ascend/ops/expert_load_balancer.py index 5b03ffa..08c0106 100644 --- a/vllm_ascend/ops/expert_load_balancer.py +++ b/vllm_ascend/ops/expert_load_balancer.py @@ -48,7 +48,7 @@ class ExpertLoadBalancer(object): def generate_expert_placement_map(self): expert_placement_map = torch.full( - (self.layers_num, self.ranks_num, self.global_expert_num), + (self.layers_num, self.ranks_num, self.num_experts), -1, dtype=torch.int32, ) @@ -71,7 +71,7 @@ class ExpertLoadBalancer(object): result_dict[key] = [] result_dict[key].append(idx) - log2phy_map = torch.full((self.ranks_num, self.global_expert_num), + log2phy_map = torch.full((self.ranks_num, self.num_experts), -1, dtype=torch.int32) for rank in range(self.ranks_num): diff --git a/vllm_ascend/ops/moe/moe_mlp.py b/vllm_ascend/ops/moe/moe_mlp.py index 5ee7d70..2c73421 100644 --- a/vllm_ascend/ops/moe/moe_mlp.py +++ b/vllm_ascend/ops/moe/moe_mlp.py @@ -105,6 +105,9 @@ def quant_apply_mlp(hidden_states: torch.Tensor, group_list=group_list, output_dtype=torch.int32)[0] # act_fn: swiglu + group_diff = torch.diff(group_list) + new_group = torch.cat([group_list[0].unsqueeze(0), group_diff], + dim=0) hidden_states, swiglu_out_scale = torch_npu.npu_dequant_swiglu_quant( x=hidden_states, weight_scale=w1_scale, @@ -112,7 +115,7 @@ def quant_apply_mlp(hidden_states: torch.Tensor, bias=None, quant_scale=None, quant_offset=None, - group_index=group_list, + group_index=new_group, activate_left=True, quant_mode=1, ) diff --git a/vllm_ascend/ops/moe/token_dispatcher.py b/vllm_ascend/ops/moe/token_dispatcher.py index da9b2be..a6e1f55 100644 --- a/vllm_ascend/ops/moe/token_dispatcher.py +++ b/vllm_ascend/ops/moe/token_dispatcher.py @@ -122,18 +122,14 @@ class TokenDispatcherWithMC2(MoETokenDispatcher): expert_map: torch.Tensor, global_redundant_expert_num: int = 0, ): - if self.with_quant: - quant_mode = 2 - moe_expert_num = len(expert_map) - else: - quant_mode = 0 - moe_expert_num = len(expert_map) + quant_mode = 2 if self.with_quant else 0 + self.moe_expert_num = len(expert_map) + global_redundant_expert_num kwargs_mc2 = { "x": hidden_states, "expert_ids": topk_ids, "expert_shard_type": 0, "shared_expert_rank_num": 0, - "moe_expert_num": moe_expert_num, + "moe_expert_num": self.moe_expert_num, "global_bs": 0, "expert_token_nums_type": 0, } @@ -229,7 +225,6 @@ class TokenDispatcherWithMC2(MoETokenDispatcher): assert self.topk_weights is not None assert self.topk_ids is not None assert self.output is not None - moe_expert_num = len(self.expert_map) # moeCombine kwargs_mc2 = { "expand_x": hidden_states, @@ -237,7 +232,7 @@ class TokenDispatcherWithMC2(MoETokenDispatcher): "expert_scales": self.topk_weights.to(torch.float32), "expert_shard_type": 0, "shared_expert_rank_num": 0, - "moe_expert_num": moe_expert_num, + "moe_expert_num": self.moe_expert_num, "global_bs": 0, } if self.with_quant: @@ -359,7 +354,7 @@ class TokenDispatcherWithAllGather(MoETokenDispatcher): hidden_states = hidden_states * \ topk_weights.to(hidden_states.dtype) if expert_map is not None: - global_num_experts = len(expert_map) + global_num_experts = len(expert_map) + global_redundant_expert_num mask = (expert_map[topk_ids] != -1) self.topk_weights = topk_weights * mask first_expert_idx = get_ep_group( diff --git a/vllm_ascend/quantization/w8a8_dynamic.py b/vllm_ascend/quantization/w8a8_dynamic.py index 59b0479..709f488 100644 --- a/vllm_ascend/quantization/w8a8_dynamic.py +++ b/vllm_ascend/quantization/w8a8_dynamic.py @@ -250,7 +250,7 @@ class AscendW8A8DynamicFusedMoEMethod: return moe_comm_method.fused_experts( hidden_states=x, w1=layer.w13_weight, - w1_scale=layer.w13_weight_scale_fp32, + w1_scale=layer.w13_weight_scale.to(torch.float32), w2=layer.w2_weight, w2_scale=layer.w2_weight_scale, topk_weights=topk_weights,