Skip to content

Commit cea4675

Browse files
committed
try hack in missing hmax2 functions (+1 squashed commits)
Squashed commits: [c98d0ab] try hack in missing hmax2 functions (+1 squashed commits) Squashed commits: [9ba8599] try hack in missing hmax2 functions (+2 squashed commit) Squashed commit: [be49749] try hack in missing hmax2 functions [159ee4c] bypass missing hmax functions on old cuda
1 parent b48ea96 commit cea4675

File tree

1 file changed

+26
-3
lines changed

1 file changed

+26
-3
lines changed

ggml-cuda/fattn.cu

+26-3
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,29 @@
1111
#define HALF_MAX_HALF __float2half(65504.0f/2) // Use neg. of this instead of -INFINITY to initialize KQ max vals to avoid NaN upon subtraction.
1212
#define SOFTMAX_FTZ_THRESHOLD -20.0f // Softmax exp. of values smaller than this are flushed to zero to avoid NaNs.
1313

14+
//hack: polyfill hmax and hmax2 for older cuda version
15+
#if CUDART_VERSION < CUDART_HMAX
16+
__device__ __inline__ __half hmax(const __half a, const __half b) {
17+
const float fa = __half2float(a);
18+
const float fb = __half2float(b);
19+
return __float2half(fa > fb ? fa : fb);
20+
}
21+
__device__ __inline__ __half2 hmax2(const __half2 a, const __half2 b) {
22+
__half2 result;
23+
result.x = hmax(a.x, b.x);
24+
result.y = hmax(a.y, b.y);
25+
return result;
26+
}
27+
#else
28+
__device__ __inline__ __half hmax(const __half a, const __half b) {
29+
return __hmax(a,b);
30+
}
31+
__device__ __inline__ __half2 hmax2(const __half2 a, const __half2 b) {
32+
return __hmax2(a,b);
33+
}
34+
#endif
35+
36+
1437
template<int D, int parallel_blocks> // D == head size
1538
__launch_bounds__(((D + WARP_SIZE - 1) / WARP_SIZE)*WARP_SIZE, 1)
1639
static __global__ void flash_attn_vec_ext_f16(
@@ -116,7 +139,7 @@ static __global__ void flash_attn_vec_ext_f16(
116139
sum2 = warp_reduce_sum(sum2);
117140
half sum = __low2half(sum2) + __high2half(sum2);
118141
sum += mask ? maskh[k_VKQ_0 + i_KQ] : __float2half(0.0f);
119-
kqmax_new = __hmax(kqmax_new, sum);
142+
kqmax_new = hmax(kqmax_new, sum);
120143
if (threadIdx.x == 0) {
121144
KQ[i_KQ] = sum;
122145
}
@@ -416,9 +439,9 @@ static __global__ void flash_attn_ext_f16(
416439
const int k = k0 + threadIdx.x;
417440

418441
KQ2_tmp[k0/WARP_SIZE] += mask ? mask2[(j*ne11 + k_VKQ_0)/2 + k] : make_half2(0.0f, 0.0f);
419-
KQ_max_new = __hmax2(KQ_max_new, KQ2_tmp[k0/WARP_SIZE]);
442+
KQ_max_new = hmax2(KQ_max_new, KQ2_tmp[k0/WARP_SIZE]);
420443
}
421-
KQ_max_new = __half2half2(warp_reduce_max(__hmax(__low2half(KQ_max_new), __high2half(KQ_max_new))));
444+
KQ_max_new = __half2half2(warp_reduce_max(hmax(__low2half(KQ_max_new), __high2half(KQ_max_new))));
422445
const half2 diff = KQ_max_h2[j0/nwarps] - KQ_max_new;
423446
KQ_max_scale_h2[j0/nwarps] = h2exp(diff);
424447
const uint32_t ftz_mask = __hgt2_mask(diff, make_half2(SOFTMAX_FTZ_THRESHOLD, SOFTMAX_FTZ_THRESHOLD));

0 commit comments

Comments
 (0)