add dispath_ffn_combine_bf16 (#5866)
### What this PR does / why we need it?
add dispath_ffn_combine_bf16
- vLLM version: v0.13.0
- vLLM main:
bde38c11df
---------
Signed-off-by: guanguan0308 <1546542263@qq.com>
This commit is contained in:
20
csrc/dispatch_ffn_combine_bf16/op_kernel/utils/layout3d.hpp
Normal file
20
csrc/dispatch_ffn_combine_bf16/op_kernel/utils/layout3d.hpp
Normal file
@@ -0,0 +1,20 @@
|
||||
#ifndef LAYOUT_3D_HPP
|
||||
#define LAYOUT_3D_HPP
|
||||
#include "kernel_operator.h"
|
||||
#include "catlass/catlass.hpp"
|
||||
class Layout3D {
|
||||
int64_t strides[2];
|
||||
public:
|
||||
CATLASS_DEVICE
|
||||
Layout3D() {}
|
||||
CATLASS_DEVICE
|
||||
Layout3D(int64_t stride0, int64_t stride1) {
|
||||
strides[0] = stride0;
|
||||
strides[1] = stride1;
|
||||
}
|
||||
CATLASS_DEVICE
|
||||
int64_t operator() (int64_t dim0, int64_t dim1, int64_t dim2) {
|
||||
return dim0 * strides[0] + dim1 * strides[1] + dim2;
|
||||
}
|
||||
};
|
||||
#endif // LAYOUT_3D_HPP
|
||||
Reference in New Issue
Block a user