github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-08-09 03:47:57 +00:00
Code Issues Packages Projects Releases Wiki Activity

[NFC] polish colossalai/kernel/cuda_native/csrc/scaled_upper_triang_masked_softmax.cpp code style (#959)

Browse Source
This commit is contained in:
superhao1995 2022-05-15 09:01:34 +08:00 committed by binmakeswell
parent 442a2975ab
commit 48c4a180c7
1 changed files with 20 additions and 25 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

15
colossalai/kernel/cuda_native/csrc/scaled_upper_triang_masked_softmax.cpp
View File

@ -3,18 +3,16 @@
#include <cuda_fp16.h> #include <cuda_fp16.h>
#include <torch/extension.h> #include <torch/extension.h>
#include <vector> #include <vector>
namespace multihead_attn { namespace multihead_attn {
namespace fused_softmax { namespace fused_softmax {
namespace scaled_upper_triang_masked_softmax { namespace scaled_upper_triang_masked_softmax {
torch::Tensor fwd_cuda( torch::Tensor fwd_cuda(torch::Tensor const& input, float scale_factor);
torch::Tensor const& input,
float scale_factor);
torch::Tensor bwd_cuda( torch::Tensor bwd_cuda(torch::Tensor const& output_grads,
torch::Tensor const& output_grads,
torch::Tensor const& softmax_results, torch::Tensor const& softmax_results,
float scale_factor); float scale_factor);
@ -27,11 +25,8 @@ torch::Tensor fwd(torch::Tensor const& input, float scale_factor) {
return fwd_cuda(input, scale_factor); return fwd_cuda(input, scale_factor);
} }
torch::Tensor bwd( torch::Tensor bwd(torch::Tensor const& output_grads,
torch::Tensor const& output_grads, torch::Tensor const& softmax_results, float scale_factor) {
torch::Tensor const& softmax_results,
float scale_factor) {
AT_ASSERTM(output_grads.dim() == 3, "expected 3D tensor"); AT_ASSERTM(output_grads.dim() == 3, "expected 3D tensor");
AT_ASSERTM(softmax_results.dim() == 3, "expected 3D tensor"); AT_ASSERTM(softmax_results.dim() == 3, "expected 3D tensor");