ViperEkura
16a55bb474
- move attention kernels to csrc/kernels/attention/ and rotary to rotary/ - add shared common/mma.cuh (mma_sync, ldmatrix) and device.cuh (sm checks) - split fp8_mm into three-layer fp8/common.h, gemm.cuh, mm.cu - fix fused FP8 GEMM ldmatrix lane indexing to fix OOB shared reads - update extension ops, loader, and kernel tests
14 lines
304 B
Plaintext
14 lines
304 B
Plaintext
#pragma once
|
|
#include <cuda_bf16.h>
|
|
|
|
using bf16 = __nv_bfloat16;
|
|
|
|
static constexpr int MAX_SPLITS = 32;
|
|
|
|
__device__ inline float warp_reduce_sum(float val) {
|
|
#pragma unroll
|
|
for (int offset = 16; offset > 0; offset >>= 1)
|
|
val += __shfl_xor_sync(0xFFFFFFFF, val, offset);
|
|
return val;
|
|
}
|