InfiniTensor/include/cuda/gbmm_g2bmm.h

12 lines
322 B
C++

#pragma once
namespace infini {
void _sg2bmm(float *__restrict__ q, float *__restrict__ k,
float *__restrict__ y, int bs, int n, int m, int w, int d);
void _sgbmml(float *__restrict__ q, float *__restrict__ k,
float *__restrict__ y, int bs, int n, int m, int w, int d);
} // namespace infini