File tree Expand file tree Collapse file tree 2 files changed +8
-1
lines changed
python/paddle/incubate/nn/functional Expand file tree Collapse file tree 2 files changed +8
-1
lines changed Original file line number Diff line number Diff line change 17
17
#include " paddle/common/exception.h"
18
18
#include " paddle/phi/kernels/moe_kernel_impl.h"
19
19
20
-
20
+ p/
21
21
template <typename T, int64_t vec_size>
22
22
__global__ void gather_with_mask_permute_kernel (const T* dy, // [s*k, d]
23
23
const int * scatter_index, // [s, k]
Original file line number Diff line number Diff line change 50
50
from .build_src_rank_and_local_expert_id import build_src_rank_and_local_expert_id
51
51
from .int_bincount import int_bincount
52
52
from .fused_rms_norm_ext import fused_rms_norm_ext
53
+ from .moe_gate_dispatch import moe_gate_dispatch
54
+ from .moe_gate_dispatch_permute import moe_gate_dispatch_permute
55
+ from .moe_ops_partial_nosoftmaxtopk import moe_gate_dispatch_partial_nosoftmaxtopk
53
56
54
57
__all__ = [
55
58
'fused_multi_head_attention' ,
75
78
"build_src_rank_and_local_expert_id"
76
79
"int_bincount" ,
77
80
"fused_rms_norm_ext" ,
81
+ "moe_gate_dispatch" ,
82
+ "moe_gate_dispatch_permute" ,
83
+ "moe_gate_dispatch_partial_nosoftmaxtopk" ,
84
+
78
85
]
You can’t perform that action at this time.
0 commit comments