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
f6970ef8
Commit
f6970ef8
authored
May 13, 2022
by
Luxios22
Committed by
binmakeswell
May 17, 2022
Browse files
[NFC] polish colossalai/kernel/cuda_native/csrc/kernels/softmax_kernels.cu code style (#954)
parent
0b86a634
Changes
1
Hide whitespace changes
Inline
Side-by-side
Showing
1 changed file
with
4 additions
and
6 deletions
+4
-6
colossalai/kernel/cuda_native/csrc/kernels/softmax_kernels.cu
...ssalai/kernel/cuda_native/csrc/kernels/softmax_kernels.cu
+4
-6
No files found.
colossalai/kernel/cuda_native/csrc/kernels/softmax_kernels.cu
View file @
f6970ef8
#include <cooperative_groups.h>
#include <math.h>
#include <math.h>
#include <cub/block/block_load.cuh>
#include <cub/block/block_load.cuh>
...
@@ -6,8 +7,6 @@
...
@@ -6,8 +7,6 @@
#include "block_reduce.h"
#include "block_reduce.h"
#include "kernels.h"
#include "kernels.h"
#include <cooperative_groups.h>
namespace
cg
=
cooperative_groups
;
namespace
cg
=
cooperative_groups
;
const
float
EPSILON
=
1e-8
f
;
const
float
EPSILON
=
1e-8
f
;
...
@@ -120,7 +119,7 @@ __global__ void ker_attn_softmax(T *inp, const T *attn_mask, int from_len,
...
@@ -120,7 +119,7 @@ __global__ void ker_attn_softmax(T *inp, const T *attn_mask, int from_len,
BlockStore
(
ts_store
).
Store
(
inp
+
(
token_id
+
i
)
*
to_len
,
inp_val
[
i
],
BlockStore
(
ts_store
).
Store
(
inp
+
(
token_id
+
i
)
*
to_len
,
inp_val
[
i
],
to_len
);
to_len
);
}
}
}
// blockIdx.x
}
// blockIdx.x
}
}
template
<
typename
T
,
int
block_dim
,
int
ele_per_thread
>
template
<
typename
T
,
int
block_dim
,
int
ele_per_thread
>
...
@@ -198,7 +197,7 @@ __global__ void ker_attn_softmax_lt32(T *inp, const T *attn_mask, int from_len,
...
@@ -198,7 +197,7 @@ __global__ void ker_attn_softmax_lt32(T *inp, const T *attn_mask, int from_len,
BlockStore
(
ts_store
).
Store
(
inp
+
(
token_id
+
i
)
*
to_len
,
inp_val
[
i
],
BlockStore
(
ts_store
).
Store
(
inp
+
(
token_id
+
i
)
*
to_len
,
inp_val
[
i
],
to_len
);
to_len
);
}
}
}
// blockIdx.x
}
// blockIdx.x
}
}
/*
/*
...
@@ -304,8 +303,7 @@ __global__ void ker_attn_softmax_bw(T *grad, const T *inp, int softmax_length) {
...
@@ -304,8 +303,7 @@ __global__ void ker_attn_softmax_bw(T *grad, const T *inp, int softmax_length) {
cg
::
thread_block
b
=
cg
::
this_thread_block
();
cg
::
thread_block
b
=
cg
::
this_thread_block
();
cg
::
thread_block_tile
<
WARP_SIZE
>
g
=
cg
::
tiled_partition
<
WARP_SIZE
>
(
b
);
cg
::
thread_block_tile
<
WARP_SIZE
>
g
=
cg
::
tiled_partition
<
WARP_SIZE
>
(
b
);
for
(
int
i
=
1
;
i
<
WARP_SIZE
;
i
<<=
1
)
for
(
int
i
=
1
;
i
<
WARP_SIZE
;
i
<<=
1
)
sum
+=
g
.
shfl_xor
(
sum
,
i
);
sum
+=
g
.
shfl_xor
(
sum
,
i
);
#pragma unroll
#pragma unroll
for
(
int
i
=
0
;
i
<
ITERATIONS
;
++
i
)
{
for
(
int
i
=
0
;
i
<
ITERATIONS
;
++
i
)
{
...
...
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