diff --git a/ggml/src/ggml-cuda/fattn-mma-f16.cuh b/ggml/src/ggml-cuda/fattn-mma-f16.cuh index 2a558bce..c82647ee 100644 --- a/ggml/src/ggml-cuda/fattn-mma-f16.cuh +++ b/ggml/src/ggml-cuda/fattn-mma-f16.cuh @@ -580,17 +580,18 @@ static __device__ __forceinline__ void flash_attn_ext_f16_load_tile_q( } } -// Prompt-sized tiles over a quantized cache are software pipelined: the raw rows of the next K or V tile and the next -// mask tile are copied into shared memory with cp.async while the tensor cores work on the current tile, then -// dequantized shared -> shared. Needs whole 16-byte aligned rows (head size 256: 272 bytes of q8_0, 144 of q4_0) -// and no OOB rows (ncols2 > 1 implies a padded cache); the staging buffers fit in the Q tile's shared memory. +// Tiles of 32 or more columns over a quantized cache (prompt tiles and the 4x8 decode/verify tile) are software +// pipelined: the raw rows of the next K or V tile and the next mask tile are copied into shared memory with cp.async +// while the tensor cores work on the current tile, then dequantized shared -> shared. Needs whole 16-byte aligned +// rows (head size 256: 272 bytes of q8_0, 144 of q4_0) and no OOB rows (ncols2 > 1 implies a padded cache). The +// staging buffers keep the 2 blocks per SM of these tiles. static constexpr __host__ __device__ int ggml_cuda_fattn_mma_q_row_bytes(const ggml_type type_KV, const int D) { return D/32 * (type_KV == GGML_TYPE_Q4_0 ? 18 : 34); } static constexpr __host__ __device__ bool ggml_cuda_fattn_mma_kv_prefetch_shape( const ggml_type type_KV, const int DKQ, const int DV, const int ncols1, const int ncols2) { - return ggml_cuda_fattn_mma_kv_native(type_KV) && ncols1*ncols2 >= 64 && ncols2 > 1 && + return ggml_cuda_fattn_mma_kv_native(type_KV) && ncols1*ncols2 >= 32 && ncols2 > 1 && ggml_cuda_fattn_mma_q_row_bytes(type_KV, DKQ) % 16 == 0 && ggml_cuda_fattn_mma_q_row_bytes(type_KV, DV) % 16 == 0; } @@ -609,6 +610,14 @@ static constexpr __device__ bool ggml_cuda_fattn_mma_kv_prefetch( #endif // CP_ASYNC_AVAILABLE } +// Decode/verify tiles (32 columns) have little math to hide a load behind, so they start the next tile's raw K right +// after its staging buffer is consumed and raw V one phase early (same buffers, same barriers). Prompt tiles keep the +// later start: on them the longer-lived addresses cost registers and spills. +static constexpr __host__ __device__ bool ggml_cuda_fattn_mma_kv_lead( + const ggml_type type_KV, const int DKQ, const int DV, const int ncols1, const int ncols2) { + return ggml_cuda_fattn_mma_kv_prefetch_shape(type_KV, DKQ, DV, ncols1, ncols2) && ncols1*ncols2 <= 32; +} + // Starts copying the raw rows of a [nbatch_fa x row_bytes] quantized K/V tile into shared memory, row after row. template static __device__ __forceinline__ void flash_attn_ext_f16_load_raw_q( @@ -794,6 +803,7 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter( constexpr bool kv_prefetch = ggml_cuda_fattn_mma_kv_prefetch(type_KV, DKQ, DV, ncols1, ncols2); static_assert(!kv_prefetch || (nstages == 0 && nbatch_K2 == DKQ/2 && nbatch_V2 == DV/2 && !oob_check), "K/V prefetch needs whole rows per pass and no OOB rows"); + constexpr bool kv_lead = ggml_cuda_fattn_mma_kv_lead(type_KV, DKQ, DV, ncols1, ncols2); constexpr int row_bytes_K = ggml_cuda_fattn_mma_q_row_bytes(type_KV, DKQ); constexpr int row_bytes_V = ggml_cuda_fattn_mma_q_row_bytes(type_KV, DV); // Raw staging tiles behind the mask tile, see flash_attn_ext_f16_process_tile. @@ -819,12 +829,18 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter( flash_attn_ext_f16_load_tile (V_h2 + int64_t(k_VKQ_0)*stride_V, tile_V, nbatch_V2, stride_V, k_VKQ_sup); } else if constexpr (kv_prefetch) { - // Raw K and the mask of this tile were started by the previous iteration (or before the first one); - // the barrier also retires every warp's VKQ reads of tile_V before tile_K (the same memory) is written. - cp_async_wait_all(); - __syncthreads(); - flash_attn_ext_f16_load_raw_q - ((const char *) V_h2 + int64_t(k_VKQ_0)*stride_V, raw_V, stride_V); + // Raw K and the mask of this tile were started by the previous iteration (or before the first one). With kv_lead + // raw V was started after them, so all but the newest cp.async group must be done, otherwise raw V starts here. + // The barrier also retires every warp's VKQ reads of tile_V before tile_K (the same memory) is written. + if constexpr (kv_lead) { + cp_async_wait_group<1>(); + __syncthreads(); + } else { + cp_async_wait_all(); + __syncthreads(); + flash_attn_ext_f16_load_raw_q + ((const char *) V_h2 + int64_t(k_VKQ_0)*stride_V, raw_V, stride_V); + } } else { constexpr bool use_cp_async = nstages == 1; if (ncols2 > 1 || mask_h) { @@ -855,6 +871,12 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter( ((const char *) K_h2 + int64_t(k_VKQ_0)*stride_K, tile_K, 2*k0_start, 2*k0_diff, stride_K, k_VKQ_sup); } __syncthreads(); + if constexpr (kv_lead && !last_iter) { + // raw_K is free again: start the next tile's K now, so that it has the whole iteration to arrive. + flash_attn_ext_f16_load_raw_q + ((const char *) K_h2 + int64_t(k_VKQ_0 + nbatch_fa)*stride_K, raw_K, stride_K); + cp_async_commit_group(); + } } // Calculate tile of KQ: @@ -1207,6 +1229,21 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter( if (use_cp_async) { cp_async_wait_all(); } + } else if constexpr (kv_lead) { + // Raw V of this tile is the oldest pending group (the next raw K may still be in flight). Every + // warp is past KQ and the softmax, so the mask tile can take the next tile. + if constexpr (last_iter) { + cp_async_wait_group<0>(); + } else { + cp_async_wait_group<1>(); + } + __syncthreads(); + if constexpr (!last_iter) { + flash_attn_ext_f16_load_mask + (mask_h + k_VKQ_0 + nbatch_fa, tile_mask, stride_mask, k_VKQ_sup, jt*ncols1, ne01); + cp_async_commit_group(); + } + flash_attn_ext_f16_dequantize_raw_q(raw_V, tile_V); } else if constexpr (kv_prefetch) { // Every warp is past KQ and the softmax, so raw K and the mask tile can take the next tile. cp_async_wait_all(); @@ -1223,6 +1260,12 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter( ((const char *) V_h2 + int64_t(k_VKQ_0)*stride_V, tile_V, i0_start, i0_diff, stride_V, k_VKQ_sup); } __syncthreads(); + if constexpr (kv_lead && !last_iter) { + // raw_V is free again: start the next tile's V. + flash_attn_ext_f16_load_raw_q + ((const char *) V_h2 + int64_t(k_VKQ_0 + nbatch_fa)*stride_V, raw_V, stride_V); + cp_async_commit_group(); + } } } const half2 * tile_V_i = !V_is_K_view || i0_stop > 2*nbatch_K2 ? tile_V : tile_V + i0_start/2; @@ -1514,13 +1557,20 @@ static __device__ __forceinline__ void flash_attn_ext_f16_process_tile( int kb0 = kb0_start; - // Quantized prompt tiles: start the raw K and mask tiles of the first iteration, see flash_attn_ext_f16_iter. + // Quantized K/V tiles: start the raw K and mask tiles of the first iteration (with kv_lead also raw V, as a second + // cp.async group), see flash_attn_ext_f16_iter. if constexpr (ggml_cuda_fattn_mma_kv_prefetch(type_KV, DKQ, DV, ncols1, ncols2)) { char * raw_K = (char *) (tile_mask + ncols1*(nbatch_fa + 8)); flash_attn_ext_f16_load_mask (mask_h + kb0*nbatch_fa, tile_mask, stride_mask, nbatch_fa, jt*ncols1, ne01); flash_attn_ext_f16_load_raw_q ((const char *) K_h2 + int64_t(kb0)*nbatch_fa*stride_K, raw_K, stride_K); + if constexpr (ggml_cuda_fattn_mma_kv_lead(type_KV, DKQ, DV, ncols1, ncols2)) { + cp_async_commit_group(); + flash_attn_ext_f16_load_raw_q + ((const char *) V_h2 + int64_t(kb0)*nbatch_fa*stride_V, raw_K + nbatch_fa*ggml_cuda_fattn_mma_q_row_bytes(type_KV, DKQ), stride_V); + cp_async_commit_group(); + } } // Preload mask and K data for first iteration when using cp_async with multiple stages: