File tree Expand file tree Collapse file tree 1 file changed +9
-11
lines changed Expand file tree Collapse file tree 1 file changed +9
-11
lines changed Original file line number Diff line number Diff line change 39
39
USING_LCCL_COM : bool = envs_ascend .USING_LCCL_COM
40
40
41
41
42
- def fused_experts_with_mc2 (
43
- hidden_states : torch .Tensor ,
44
- w1 : torch .Tensor ,
45
- w2 : torch .Tensor ,
46
- topk_weights : torch .Tensor ,
47
- topk_ids : torch .Tensor ,
48
- top_k : int ,
49
- expert_map : torch .Tensor = None ,
50
- moe_all_to_all_group_name : Optional [str ] = None ,
51
- ** kwargs
52
- ) -> torch .Tensor :
42
+ def fused_experts_with_mc2 (hidden_states : torch .Tensor ,
43
+ w1 : torch .Tensor ,
44
+ w2 : torch .Tensor ,
45
+ topk_weights : torch .Tensor ,
46
+ topk_ids : torch .Tensor ,
47
+ top_k : int ,
48
+ expert_map : torch .Tensor = None ,
49
+ moe_all_to_all_group_name : Optional [str ] = None ,
50
+ ** kwargs ) -> torch .Tensor :
53
51
global_bs = 0
54
52
moe_expert_num = len (expert_map )
55
53
kwargs_mc2 = {
You can’t perform that action at this time.
0 commit comments