diff --git a/ggml/src/ggml-cuda/fattn-tile-f32.cu b/ggml/src/ggml-cuda/fattn-tile-f32.cu index 7a1d36e5..66582368 100644 --- a/ggml/src/ggml-cuda/fattn-tile-f32.cu +++ b/ggml/src/ggml-cuda/fattn-tile-f32.cu @@ -11,9 +11,35 @@ #define FATTN_KQ_STRIDE_TILE_F32 32 +// S3 exact-retile (sm_60/P100). Same fp32 arithmetic as the original all-fp32 tile kernel -- fp32 Q +// (scaled in fp32), exact fp16->fp32 K/V conversion, fp32 products, fp32 scores and online softmax, +// fp32 post-exp probability store, fp32 P.V accumulate -- restructured to run a leaner inner loop: +// * K/V tiles are staged in shared memory as raw half2 exactly as they arrive from the fp16 KV +// cache and converted to fp32 in registers at the point of use. fp16->fp32 conversion is exact, +// so every FFMA sees bit-identical inputs to the original fp32-staged tiles. This halves the K/V +// tile (16512 B -> 8320 B), roughly halving the QK loop's shared-memory loads. +// * Raw KQ scores never round-trip through shared memory: the (threadIdx.x, threadIdx.y) -> +// (i_KQ, j_KQ) mapping of the score pass is identical to the softmax pass, so the scores stay +// in registers and only the post-exp fp32 probabilities are stored (unchanged values, one +// tile-sized smem round trip AND one __syncthreads() fewer per KV tile). +// The measured +4-9% vs the un-retiled fp32 tile kernel on P100 is a COMPOUND of two effects the half2 +// staging produces together, and we did not isolate their individual shares: (1) an occupancy gain -- the +// smem drop 36992 -> 28800 B/block lets 2 blocks fit per SM where the un-retiled kernel fits only 1 +// (2 x 28800 = 57600 <= 65536; 80 regs would allow 3, so smem is the ceiling); and (2) a leaner inner +// loop (fewer smem loads, the dropped score round-trip, one fewer barrier). __launch_bounds__(..., 2) +// only PINS the target the smem reduction already reaches (a (..., 1) bound compiles bit-identically), so +// the directive is not itself a lever; and a 3rd block does not help (the branch cost-curve and an +// implemented Q_f-to-registers attempt that pushed regs 80 -> 124 both confirm smem caps occupancy at 2). +// Scope: this kernel is only the explicit GGML_PREC_F32 tile path. On P100 that is where the fp16 +// fast path is routed off; standalone it also affects sm_61 consumer-Pascal prefill, not benchmarked. +// The ONLY numerical difference vs the original kernel: the QK dot product accumulates its D terms +// in natural order 0,1,...,D-1, while the original accumulated them even/odd-deinterleaved within +// each 64-element chunk (an artifact of its fp32 staging layout). Both are all-fp32 sums of the same +// exact products; individual scores may differ by ~1 ulp. The P.V accumulation order is bit-identical. + template // D == head size #if !(defined(GGML_USE_HIPBLAS) && defined(__HIP_PLATFORM_AMD__)) -__launch_bounds__(nwarps*WARP_SIZE, 1) +__launch_bounds__(nwarps*WARP_SIZE, 2) #endif // !(defined(GGML_USE_HIPBLAS) && defined(__HIP_PLATFORM_AMD__)) static __global__ void flash_attn_tile_ext_f32( const char * __restrict__ Q, @@ -80,11 +106,9 @@ static __global__ void flash_attn_tile_ext_f32( const float slope = get_alibi_slope(max_bias, blockIdx.y, n_head_log2, m0, m1); static_assert(D % (2 * WARP_SIZE) == 0, "D not divisible by 2*WARP_SIZE == 64."); - __shared__ float KQ[ncols*FATTN_KQ_STRIDE_TILE_F32]; + __shared__ float KQ[ncols*FATTN_KQ_STRIDE_TILE_F32]; // post-exp fp32 probabilities only - __shared__ float KV_tmp[FATTN_KQ_STRIDE_TILE_F32][D + 1]; // Pad D to avoid memory bank conflicts. - - float2 * KV_tmp2 = (float2 *) KV_tmp; + __shared__ half2 KV_tmp[FATTN_KQ_STRIDE_TILE_F32][D/2 + 1]; // Raw fp16 K/V tile; pad to avoid memory bank conflicts. float kqmax[ncols/nwarps]; #pragma unroll @@ -95,17 +119,21 @@ static __global__ void flash_attn_tile_ext_f32( float2 VKQ[ncols/nwarps][(D/2)/WARP_SIZE] = {{{0.0f, 0.0f}}}; - // Convert Q to half2 and store in registers: - __shared__ float Q_f[ncols][D]; + // Store Q in shared memory, fp32 and pre-scaled (same values as the original kernel; float2-typed + // natural-order layout so the QK loop below pairs it directly with half2 K): + __shared__ float2 Q_f[ncols][D/2]; #pragma unroll for (int j0 = 0; j0 < ncols; j0 += nwarps) { const int j = j0 + threadIdx.y; #pragma unroll - for (int i0 = 0; i0 < D; i0 += 2*WARP_SIZE) { - float2 tmp = ic0 + j < ne01 ? Q_f2[j*(nb01/sizeof(float2)) + i0/2 + threadIdx.x] : make_float2(0.0f, 0.0f); - Q_f[j][i0 + 0*WARP_SIZE + threadIdx.x] = tmp.x * scale; - Q_f[j][i0 + 1*WARP_SIZE + threadIdx.x] = tmp.y * scale; + for (int i0 = 0; i0 < D/2; i0 += WARP_SIZE) { + const int i = i0 + threadIdx.x; + + float2 tmp = ic0 + j < ne01 ? Q_f2[j*(nb01/sizeof(float2)) + i] : make_float2(0.0f, 0.0f); + tmp.x *= scale; + tmp.y *= scale; + Q_f[j][i] = tmp; } } @@ -121,15 +149,16 @@ static __global__ void flash_attn_tile_ext_f32( kqmax_new[j] = kqmax[j]; } + // Stage the K tile as raw half2 (no conversion here; converted in registers at use below): #pragma unroll for (int i_KQ_0 = 0; i_KQ_0 < FATTN_KQ_STRIDE_TILE_F32; i_KQ_0 += nwarps) { const int i_KQ = i_KQ_0 + threadIdx.y; #pragma unroll - for (int k_KQ_0 = 0; k_KQ_0 < D; k_KQ_0 += 2*WARP_SIZE) { - const half2 tmp = K_h2[(k_VKQ_0 + i_KQ)*stride_KV2 + k_KQ_0/2 + threadIdx.x]; - KV_tmp[i_KQ][k_KQ_0 + 0*WARP_SIZE + threadIdx.x] = __low2float(tmp); - KV_tmp[i_KQ][k_KQ_0 + 1*WARP_SIZE + threadIdx.x] = __high2float(tmp); + for (int k_KQ_0 = 0; k_KQ_0 < D/2; k_KQ_0 += WARP_SIZE) { + const int k_KQ = k_KQ_0 + threadIdx.x; + + KV_tmp[i_KQ][k_KQ] = K_h2[(k_VKQ_0 + i_KQ)*stride_KV2 + k_KQ]; } } @@ -138,15 +167,15 @@ static __global__ void flash_attn_tile_ext_f32( float sum[FATTN_KQ_STRIDE_TILE_F32/WARP_SIZE][ncols/nwarps] = {{0.0f}}; #pragma unroll - for (int k_KQ = 0; k_KQ < D; ++k_KQ) { - float K_k[FATTN_KQ_STRIDE_TILE_F32/WARP_SIZE]; - float Q_k[ncols/nwarps]; + for (int k_KQ = 0; k_KQ < D/2; ++k_KQ) { + float2 K_k[FATTN_KQ_STRIDE_TILE_F32/WARP_SIZE]; + float2 Q_k[ncols/nwarps]; #pragma unroll for (int i_KQ_0 = 0; i_KQ_0 < FATTN_KQ_STRIDE_TILE_F32; i_KQ_0 += WARP_SIZE) { const int i_KQ = i_KQ_0 + threadIdx.x; - K_k[i_KQ_0/WARP_SIZE] = KV_tmp[i_KQ][k_KQ]; + K_k[i_KQ_0/WARP_SIZE] = __half22float2(KV_tmp[i_KQ][k_KQ]); // exact fp16->fp32 } #pragma unroll for (int j_KQ_0 = 0; j_KQ_0 < ncols; j_KQ_0 += nwarps) { @@ -159,11 +188,16 @@ static __global__ void flash_attn_tile_ext_f32( for (int i_KQ_0 = 0; i_KQ_0 < FATTN_KQ_STRIDE_TILE_F32; i_KQ_0 += WARP_SIZE) { #pragma unroll for (int j_KQ_0 = 0; j_KQ_0 < ncols; j_KQ_0 += nwarps) { - sum[i_KQ_0/WARP_SIZE][j_KQ_0/nwarps] += K_k[i_KQ_0/WARP_SIZE] * Q_k[j_KQ_0/nwarps]; + // Two ordered fp32 FMAs == two scalar iterations of the original kernel. + sum[i_KQ_0/WARP_SIZE][j_KQ_0/nwarps] += K_k[i_KQ_0/WARP_SIZE].x * Q_k[j_KQ_0/nwarps].x; + sum[i_KQ_0/WARP_SIZE][j_KQ_0/nwarps] += K_k[i_KQ_0/WARP_SIZE].y * Q_k[j_KQ_0/nwarps].y; } } } + // Softcap/mask/running-max on the raw scores. The scores stay in REGISTERS: this pass and the + // softmax pass below use the same (threadIdx.x, threadIdx.y) -> (i_KQ, j_KQ) ownership as the + // QK loop above, so no shared-memory round trip (and no __syncthreads()) is needed in between. #pragma unroll for (int i_KQ_0 = 0; i_KQ_0 < FATTN_KQ_STRIDE_TILE_F32; i_KQ_0 += WARP_SIZE) { const int i_KQ = i_KQ_0 + threadIdx.x; @@ -179,13 +213,9 @@ static __global__ void flash_attn_tile_ext_f32( sum[i_KQ_0/WARP_SIZE][j_KQ_0/nwarps] += mask ? slope*__half2float(maskh[j_KQ*stride_mask + k_VKQ_0 + i_KQ]) : 0.0f; kqmax_new[j_KQ_0/nwarps] = fmaxf(kqmax_new[j_KQ_0/nwarps], sum[i_KQ_0/WARP_SIZE][j_KQ_0/nwarps]); - - KQ[j_KQ*FATTN_KQ_STRIDE_TILE_F32 + i_KQ] = sum[i_KQ_0/WARP_SIZE][j_KQ_0/nwarps]; } } - __syncthreads(); - #pragma unroll for (int j0 = 0; j0 < ncols; j0 += nwarps) { const int j = j0 + threadIdx.y; @@ -199,10 +229,10 @@ static __global__ void flash_attn_tile_ext_f32( for (int i0 = 0; i0 < FATTN_KQ_STRIDE_TILE_F32; i0 += WARP_SIZE) { const int i = i0 + threadIdx.x; - const float diff = KQ[j*FATTN_KQ_STRIDE_TILE_F32 + i] - kqmax[j0/nwarps]; + const float diff = sum[i0/WARP_SIZE][j0/nwarps] - kqmax[j0/nwarps]; const float val = expf(diff); kqsum_add += val; - KQ[j*FATTN_KQ_STRIDE_TILE_F32 + i] = val; + KQ[j*FATTN_KQ_STRIDE_TILE_F32 + i] = val; // fp32 post-exp probability store (exact path) } kqsum[j0/nwarps] = kqsum[j0/nwarps]*KQ_max_scale + kqsum_add; @@ -215,6 +245,9 @@ static __global__ void flash_attn_tile_ext_f32( __syncthreads(); + // Stage the V tile as raw half2. Reusing KV_tmp is safe: the __syncthreads() above orders these + // writes after every thread's last K read (program order puts each thread's K reads before its + // softmax pass, and the barrier puts all softmax passes before any V write). #pragma unroll for (int k0 = 0; k0 < FATTN_KQ_STRIDE_TILE_F32; k0 += nwarps) { const int k = k0 + threadIdx.y; @@ -223,8 +256,7 @@ static __global__ void flash_attn_tile_ext_f32( for (int i0 = 0; i0 < D/2; i0 += WARP_SIZE) { const int i = i0 + threadIdx.x; - KV_tmp2[k*(D/2) + i].x = __low2float(V_h2[(k_VKQ_0 + k)* stride_KV2 + i]); - KV_tmp2[k*(D/2) + i].y = __high2float(V_h2[(k_VKQ_0 + k)* stride_KV2 + i]); + KV_tmp[k][i] = V_h2[(k_VKQ_0 + k)*stride_KV2 + i]; } } @@ -239,7 +271,7 @@ static __global__ void flash_attn_tile_ext_f32( for (int i0 = 0; i0 < D/2; i0 += WARP_SIZE) { const int i = i0 + threadIdx.x; - V_k[i0/WARP_SIZE] = KV_tmp2[k*(D/2) + i]; + V_k[i0/WARP_SIZE] = __half22float2(KV_tmp[k][i]); // exact fp16->fp32 } #pragma unroll for (int j0 = 0; j0 < ncols; j0 += nwarps) {