Skip to content
GitLab
Menu
Projects
Groups
Snippets
Loading...
Help
Help
Support
Community forum
Keyboard shortcuts
?
Submit feedback
Contribute to GitLab
Sign in / Register
Toggle navigation
Menu
Open sidebar
OpenDAS
ColossalAI
Commits
8ffdc383
Commit
8ffdc383
authored
May 13, 2022
by
wky
Committed by
binmakeswell
May 17, 2022
Browse files
[NFC] polish colossalai/kernel/cuda_native/csrc/moe_cuda.cpp code style (#942)
parent
c0f373db
Changes
1
Show whitespace changes
Inline
Side-by-side
Showing
1 changed file
with
14 additions
and
16 deletions
+14
-16
colossalai/kernel/cuda_native/csrc/moe_cuda.cpp
colossalai/kernel/cuda_native/csrc/moe_cuda.cpp
+14
-16
No files found.
colossalai/kernel/cuda_native/csrc/moe_cuda.cpp
View file @
8ffdc383
...
@@ -15,10 +15,10 @@ torch::Tensor moe_combine_cuda_forward(int s, int e, int c, int h,
...
@@ -15,10 +15,10 @@ torch::Tensor moe_combine_cuda_forward(int s, int e, int c, int h,
torch
::
Tensor
logits
,
torch
::
Tensor
mask
,
torch
::
Tensor
logits
,
torch
::
Tensor
mask
,
torch
::
Tensor
dest_idx
);
torch
::
Tensor
dest_idx
);
std
::
vector
<
torch
::
Tensor
>
std
::
vector
<
torch
::
Tensor
>
moe_combine_cuda_backward
(
moe_combine_cuda_backward
(
int
s
,
int
e
,
int
c
,
int
h
,
torch
::
Tensor
tokens_grad
,
int
s
,
int
e
,
int
c
,
int
h
,
torch
::
Tensor
tokens_grad
,
torch
::
Tensor
expert_tokens
,
torch
::
Tensor
logits
,
torch
::
Tensor
expert_tokens
,
torch
::
Tensor
logits
,
torch
::
Tensor
mask
,
torch
::
Tensor
mask
,
torch
::
Tensor
dest_idx
);
torch
::
Tensor
dest_idx
);
torch
::
Tensor
cumsum_sub_one_in_dim0
(
torch
::
Tensor
mask
);
torch
::
Tensor
cumsum_sub_one_in_dim0
(
torch
::
Tensor
mask
);
...
@@ -33,7 +33,6 @@ torch::Tensor cumsum_sub_one_in_dim0(torch::Tensor mask);
...
@@ -33,7 +33,6 @@ torch::Tensor cumsum_sub_one_in_dim0(torch::Tensor mask);
torch
::
Tensor
moe_dispatch_forward
(
int
s
,
int
ec
,
int
h
,
torch
::
Tensor
moe_dispatch_forward
(
int
s
,
int
ec
,
int
h
,
torch
::
Tensor
batch_tokens
,
torch
::
Tensor
batch_tokens
,
torch
::
Tensor
mask
,
torch
::
Tensor
dest_idx
)
{
torch
::
Tensor
mask
,
torch
::
Tensor
dest_idx
)
{
CHECK_INPUT
(
batch_tokens
);
CHECK_INPUT
(
batch_tokens
);
CHECK_CUDA
(
mask
);
CHECK_CUDA
(
mask
);
CHECK_CUDA
(
dest_idx
);
CHECK_CUDA
(
dest_idx
);
...
@@ -45,7 +44,6 @@ torch::Tensor moe_dispatch_backward(int s, int ec, int h,
...
@@ -45,7 +44,6 @@ torch::Tensor moe_dispatch_backward(int s, int ec, int h,
torch
::
Tensor
expert_grad
,
torch
::
Tensor
expert_grad
,
torch
::
Tensor
mask
,
torch
::
Tensor
mask
,
torch
::
Tensor
dest_idx
)
{
torch
::
Tensor
dest_idx
)
{
CHECK_INPUT
(
expert_grad
);
CHECK_INPUT
(
expert_grad
);
CHECK_CUDA
(
mask
);
CHECK_CUDA
(
mask
);
CHECK_CUDA
(
dest_idx
);
CHECK_CUDA
(
dest_idx
);
...
@@ -57,7 +55,6 @@ torch::Tensor moe_combine_forward(int s, int e, int c, int h,
...
@@ -57,7 +55,6 @@ torch::Tensor moe_combine_forward(int s, int e, int c, int h,
torch
::
Tensor
expert_tokens
,
torch
::
Tensor
expert_tokens
,
torch
::
Tensor
logits
,
torch
::
Tensor
mask
,
torch
::
Tensor
logits
,
torch
::
Tensor
mask
,
torch
::
Tensor
dest_idx
)
{
torch
::
Tensor
dest_idx
)
{
CHECK_INPUT
(
expert_tokens
);
CHECK_INPUT
(
expert_tokens
);
CHECK_INPUT
(
logits
);
CHECK_INPUT
(
logits
);
CHECK_CUDA
(
mask
);
CHECK_CUDA
(
mask
);
...
@@ -67,11 +64,12 @@ torch::Tensor moe_combine_forward(int s, int e, int c, int h,
...
@@ -67,11 +64,12 @@ torch::Tensor moe_combine_forward(int s, int e, int c, int h,
dest_idx
);
dest_idx
);
}
}
std
::
vector
<
torch
::
Tensor
>
std
::
vector
<
torch
::
Tensor
>
moe_combine_backward
(
int
s
,
int
e
,
int
c
,
int
h
,
moe_combine_backward
(
int
s
,
int
e
,
int
c
,
int
h
,
torch
::
Tensor
tokens_grad
,
torch
::
Tensor
tokens_grad
,
torch
::
Tensor
expert_tokens
,
torch
::
Tensor
logits
,
torch
::
Tensor
expert_tokens
,
torch
::
Tensor
mask
,
torch
::
Tensor
dest_idx
)
{
torch
::
Tensor
logits
,
torch
::
Tensor
mask
,
torch
::
Tensor
dest_idx
)
{
CHECK_INPUT
(
tokens_grad
);
CHECK_INPUT
(
tokens_grad
);
CHECK_INPUT
(
logits
);
CHECK_INPUT
(
logits
);
CHECK_CUDA
(
mask
);
CHECK_CUDA
(
mask
);
...
...
Write
Preview
Markdown
is supported
0%
Try again
or
attach a new file
.
Attach a file
Cancel
You are about to add
0
people
to the discussion. Proceed with caution.
Finish editing this message first!
Cancel
Please
register
or
sign in
to comment