github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-08-10 20:32:40 +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

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

@ -3,57 +3,52 @@
#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);
torch::Tensor fwd(torch::Tensor const& input, float scale_factor) { torch::Tensor fwd(torch::Tensor const& input, float scale_factor) {
AT_ASSERTM(input.dim() == 3, "expected 3D tensor"); AT_ASSERTM(input.dim() == 3, "expected 3D tensor");
AT_ASSERTM((input.scalar_type() == at::ScalarType::Half) || AT_ASSERTM((input.scalar_type() == at::ScalarType::Half) ||
(input.scalar_type() == at::ScalarType::BFloat16), (input.scalar_type() == at::ScalarType::BFloat16),
"Only fp16 and bf16 are supported"); "Only fp16 and bf16 are supported");
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");
AT_ASSERTM((output_grads.scalar_type() == at::ScalarType::Half) || AT_ASSERTM((output_grads.scalar_type() == at::ScalarType::Half) ||
(output_grads.scalar_type() == at::ScalarType::BFloat16), (output_grads.scalar_type() == at::ScalarType::BFloat16),
"Only fp16 and bf16 are supported"); "Only fp16 and bf16 are supported");
AT_ASSERTM((softmax_results.scalar_type() == at::ScalarType::Half) || AT_ASSERTM((softmax_results.scalar_type() == at::ScalarType::Half) ||
(softmax_results.scalar_type() == at::ScalarType::BFloat16), (softmax_results.scalar_type() == at::ScalarType::BFloat16),
"Only fp16 and bf16 are supported"); "Only fp16 and bf16 are supported");
return bwd_cuda(output_grads, softmax_results, scale_factor); return bwd_cuda(output_grads, softmax_results, scale_factor);
} }
} // end namespace scaled_upper_triang_masked_softmax } // end namespace scaled_upper_triang_masked_softmax
} // end namespace fused_softmax } // end namespace fused_softmax
} // end namespace multihead_attn } // end namespace multihead_attn
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
m.def("forward", m.def("forward",
&multihead_attn::fused_softmax::scaled_upper_triang_masked_softmax::fwd, &multihead_attn::fused_softmax::scaled_upper_triang_masked_softmax::fwd,
"Self Multihead Attention scaled, time masked softmax -- Forward."); "Self Multihead Attention scaled, time masked softmax -- Forward.");
m.def("backward", m.def("backward",
&multihead_attn::fused_softmax::scaled_upper_triang_masked_softmax::bwd, &multihead_attn::fused_softmax::scaled_upper_triang_masked_softmax::bwd,
"Self Multihead Attention scaled, time masked softmax -- Backward."); "Self Multihead Attention scaled, time masked softmax -- Backward.");
} }