ViperEkura
20041d7aa9
- Multi-pass MMA: encode pass in grid blockIdx.x, compute q_head0/G in-kernel - Fixes crash for G>32 (previously block(32,G) exceeded 1024 threads) - Fixes alloc_split_partials using uninitialized num_splits (MAX_SPLITS=32) - __launch_bounds__ on all MMA and prefill kernels for better register allocation - 4x vectorized combine kernel (4 head_dim per thread) - uint4 vectorized K loads in scalar decode kernels - cp.async .L2::128B cache hint for K/V tile streaming - Extract warp_reduce_sum, bf16, MAX_SPLITS to attn_warp_utils.cuh
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;
|
|
}
|