|
11 | 11 | #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.
|
12 | 12 | #define SOFTMAX_FTZ_THRESHOLD -20.0f // Softmax exp. of values smaller than this are flushed to zero to avoid NaNs.
|
13 | 13 |
|
| 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 | + |
14 | 37 | template<int D, int parallel_blocks> // D == head size
|
15 | 38 | __launch_bounds__(((D + WARP_SIZE - 1) / WARP_SIZE)*WARP_SIZE, 1)
|
16 | 39 | static __global__ void flash_attn_vec_ext_f16(
|
@@ -116,7 +139,7 @@ static __global__ void flash_attn_vec_ext_f16(
|
116 | 139 | sum2 = warp_reduce_sum(sum2);
|
117 | 140 | half sum = __low2half(sum2) + __high2half(sum2);
|
118 | 141 | 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); |
120 | 143 | if (threadIdx.x == 0) {
|
121 | 144 | KQ[i_KQ] = sum;
|
122 | 145 | }
|
@@ -416,9 +439,9 @@ static __global__ void flash_attn_ext_f16(
|
416 | 439 | const int k = k0 + threadIdx.x;
|
417 | 440 |
|
418 | 441 | 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]); |
420 | 443 | }
|
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)))); |
422 | 445 | const half2 diff = KQ_max_h2[j0/nwarps] - KQ_max_new;
|
423 | 446 | KQ_max_scale_h2[j0/nwarps] = h2exp(diff);
|
424 | 447 | const uint32_t ftz_mask = __hgt2_mask(diff, make_half2(SOFTMAX_FTZ_THRESHOLD, SOFTMAX_FTZ_THRESHOLD));
|
|
0 commit comments