InfiniTensor/include/cuda/softmax.h

7 lines
159 B
C++

#pragma once
namespace infini {
void softmax_kernel(int max_threadblock_size, int batch_size, float *x,
float *y, int dim, int stride);
}