turboderp-org/exllamav2
1
0
Fork 0
mirror of https://github.com/turboderp-org/exllamav2.git synced 2026-04-20 14:29:28 +00:00
Code Issues Packages Projects Releases Wiki Activity

Fix layernorm kernels for wave64 GPUs

Browse Source
This commit is contained in:
turboderp
2024-06-24 02:09:46 +02:00
parent 05b1f2194e
commit 6feebfb56e
2 changed files with 15 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
exllamav2/exllamav2_ext/cuda/layer_norm.cu
View File

@@ -81,7 +81,11 @@ __global__ void layer_norm_kernel
// Load partial sums from across warps, shuffle again across lanes
sum = sums[lane_id];
#if defined(USE_ROCM)
sum = lane_id < NUM_WARPS ? sums[lane_id] : 0.0f;
#else
sum = sums[lane_id];
#endif
for(int offset = warpSize / 2; offset > 0; offset /= 2) sum += __shfl_xor_sync(0xffffffff, sum, offset);
// Compute mean
@@ -116,7 +120,11 @@ __global__ void layer_norm_kernel
// Load partial sums from across warps, shuffle again across lanes
sum = sums[lane_id];
#if defined(USE_ROCM)
sum = lane_id < NUM_WARPS ? sums[lane_id] : 0.0f;
#else
sum = sums[lane_id];
#endif
for(int offset = warpSize / 2; offset > 0; offset /= 2) sum += __shfl_xor_sync(0xffffffff, sum, offset);
// Get 1/sqrt(variance)

6
exllamav2/exllamav2_ext/cuda/rms_norm.cu
View File

@@ -77,7 +77,11 @@ __global__ void rms_norm_kernel
// Load partial sums from across warps, shuffle again across lanes
sum = sums[lane_id];
#if defined(USE_ROCM)
sum = lane_id < NUM_WARPS ? sums[lane_id] : 0.0f;
#else
sum = sums[lane_id];
#endif
for(int offset = warpSize / 2; offset > 0; offset /= 2) sum += __shfl_xor_sync(0xffffffff, sum, offset);
// Get norm