Skip to content

GitLab

Commit 24c52aee authored by zhanghj2's avatar zhanghj2
Browse files

优化tp1 warp求和部分

parent 6a04965a
Hide whitespace changes
Inline Side-by-side
Showing with 12 additions and 5 deletions
csrc/extension/softmax.h
View file @ 24c52aee
...@@ -77,12 +77,19 @@ __device__ __forceinline__ void warp_allreduce_tp1(Tensor<Engine0, Layout0> &dst ...@@ -77,12 +77,19 @@ __device__ __forceinline__ void warp_allreduce_tp1(Tensor<Engine0, Layout0> &dst
smem_reduce[row * 2 + (warp_id / 4)] = dst[0]; smem_reduce[row * 2 + (warp_id / 4)] = dst[0];
} }
__syncthreads(); __syncthreads();
if (col == 0 && warp_id < 4) { // if (col == 0 && warp_id < 4) {
// printf("sum %d %d %d %.2f %.2f \n", row, tidx, warp_id, smem_reduce[row * 2], smem_reduce[row * 2 + warp_id / 4]); // // printf("sum %d %d %d %.2f %.2f \n", row, tidx, warp_id, smem_reduce[row * 2], smem_reduce[row * 2 + warp_id / 4]);
smem_reduce[128 + row] = op(smem_reduce[row * 2], smem_reduce[row * 2 + 1]); // smem_reduce[128 + row] = op(smem_reduce[row * 2], smem_reduce[row * 2 + 1]);
// }
// __syncthreads();
// dst(0) = smem_reduce(128 + row);
if (warp_id < 4) {
dst(0) = op(dst(0), smem_reduce[row * 2 + 1]);
} else {
dst(0) = op(dst(0), smem_reduce[row * 2]);
} }
__syncthreads(); // dst(0) = op(smem_reduce[row * 2], smem_reduce[row * 2 + 1]);
dst(0) = smem_reduce(128 + row);
} }
template<typename Engine0, typename Layout0, typename Engine1, typename Layout1, typename Operator> template<typename Engine0, typename Layout0, typename Engine1, typename Layout1, typename Operator>
__device__ __forceinline__ void warp_allreduce_tp4(Tensor<Engine0, Layout0> &dst, Tensor<Engine1, Layout1> &smem_reduce, Operator &op) { __device__ __forceinline__ void warp_allreduce_tp4(Tensor<Engine0, Layout0> &dst, Tensor<Engine1, Layout1> &smem_reduce, Operator &op) {
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment