gridwise_sparse_embeddings_forward_layernorm.hpp File Reference#
gridwise_sparse_embeddings_forward_layernorm.hpp File Reference
#include "ck/utility/common_header.hpp"#include "ck/tensor_operation/gpu/thread/threadwise_tensor_slice_transfer.hpp"#include "ck/tensor_operation/gpu/thread/threadwise_welford.hpp"#include "ck/tensor_operation/gpu/block/blockwise_welford.hpp"Go to the source code of this file.
Namespaces | |
| namespace | ck |
Functions | |
| template<typename GridwiseSparseEmbedding, typename EmbType, typename IndexType, typename GammaDataType, typename BetaDataType, typename AccDataType, typename OutType, typename OutGridDesc, typename EmbElementwiseOperation, ck::index_t NumEmbeddings> | |
| __global__ void | ck::kernel_sparse_embeddings_forward_layernorm (OutType *p_out, const ck::Array< EmbType *, NumEmbeddings > p_embs, const ck::Array< IndexType *, NumEmbeddings > p_indexes, const GammaDataType *p_gamma, const BetaDataType *p_beta, const OutGridDesc out_grid_desc, const AccDataType epsilon, const EmbElementwiseOperation emb_elementwise_op) |