diff --git a/ggml/include/ggml.h b/ggml/include/ggml.h index f408f91817e7..d47c24592027 100644 --- a/ggml/include/ggml.h +++ b/ggml/include/ggml.h @@ -2370,6 +2370,19 @@ extern "C" { float scale, float alpha); + // DDTree/tree-causal flash attention carrier. Uses FLASH_ATTN_SPARSE with + // extra sources so CUDA/HIP backends can specialize; dense masked FA fallback + // remains exact when no tree kernel is available. + GGML_API struct ggml_tensor * ggml_flash_attn_tree( + struct ggml_context * ctx, + struct ggml_tensor * q, + struct ggml_tensor * k, + struct ggml_tensor * v, + struct ggml_tensor * mask, + struct ggml_tensor * parent_ids, + struct ggml_tensor * positions, + float scale); + // TurboQuant FWHT rotation. direction: 0 = forward, 1 = inverse. // Applies signs1 -> FWHT -> signs2 (forward) or signs2 -> FWHT -> signs1 (inverse). // Used for KV cache rotation in TurboQuant quantization types (TQ3_0). diff --git a/ggml/src/ggml-cuda/fattn-common.cuh b/ggml/src/ggml-cuda/fattn-common.cuh index 832423c8db97..514be9e785ad 100644 --- a/ggml/src/ggml-cuda/fattn-common.cuh +++ b/ggml/src/ggml-cuda/fattn-common.cuh @@ -5,6 +5,7 @@ #include "vecdotq.cuh" #include "tq3-quant.cuh" +#include #include #define FATTN_KQ_STRIDE 256 @@ -26,6 +27,7 @@ typedef void (* fattn_kernel_t)( const char * __restrict__ mask, const char * __restrict__ sinks, const int * __restrict__ KV_max, + const int * __restrict__ KV_min, float * __restrict__ dst, float2 * __restrict__ dst_meta, const float scale, @@ -766,6 +768,54 @@ static __global__ void flash_attn_mask_to_KV_max( KV_max[sequence*ne31 + jt] = KV_max_sj; } +template +__launch_bounds__(FATTN_KQ_STRIDE/2, 1) +static __global__ void flash_attn_mask_to_KV_min( + const half2 * __restrict__ mask, int * __restrict__ KV_min, const int ne30, const int s31, const int s33) { + const int ne31 = gridDim.x; + const int tid = threadIdx.x; + const int sequence = blockIdx.y; + const int jt = blockIdx.x; + + mask += sequence*s33 + jt*ncols1*s31; + + __shared__ int buf_iw[WARP_SIZE]; + if (tid < WARP_SIZE) { + buf_iw[tid] = 1; + } + __syncthreads(); + + int KV_min_sj = 0; + for (; KV_min_sj < ne30 * FATTN_KQ_STRIDE; KV_min_sj += FATTN_KQ_STRIDE) { + int all_inf = 1; + +#pragma unroll + for (int j = 0; j < ncols1; ++j) { + const float2 tmp = __half22float2(mask[j*s31 + KV_min_sj/2 + tid]); + all_inf = all_inf && int(isinf(tmp.x)) && int(isinf(tmp.y)); + } + + all_inf = warp_reduce_all(all_inf); + if (tid % WARP_SIZE == 0) { + buf_iw[tid / WARP_SIZE] = all_inf; + } + __syncthreads(); + all_inf = buf_iw[tid % WARP_SIZE]; + __syncthreads(); + all_inf = warp_reduce_all(all_inf); + + if (!all_inf) { + break; + } + } + + if (threadIdx.x != 0) { + return; + } + + KV_min[sequence*ne31 + jt] = KV_min_sj; +} + template // D == head size __launch_bounds__(D, 1) static __global__ void flash_attn_stream_k_fixup_uniform( @@ -1041,6 +1091,7 @@ void launch_fattn( ggml_cuda_pool_alloc K_f16(pool); ggml_cuda_pool_alloc V_f16(pool); ggml_cuda_pool_alloc KV_max(pool); + ggml_cuda_pool_alloc KV_min(pool); ggml_cuda_pool_alloc dst_tmp(pool); ggml_cuda_pool_alloc dst_tmp_meta(pool); @@ -1121,10 +1172,18 @@ void launch_fattn( const int ntiles_z_gqa = ((gqa_ratio + ncols2 - 1) / ncols2); const int ntiles_dst = ntiles_x * ntiles_z_gqa * K->ne[2] * Q->ne[3]; + static const bool kv_min_enabled = + (std::getenv("GGML_CUDA_FA_KV_MIN") != nullptr || + std::getenv("DFLASH_LAGUNA_FA_KV_MIN") != nullptr) && + std::getenv("GGML_CUDA_FA_KV_MIN_DISABLE") == nullptr && + std::getenv("DFLASH_LAGUNA_FA_KV_MIN_DISABLE") == nullptr; + // Optional optimization where the mask is scanned to determine whether part of the calculation can be skipped. - // Only worth the overhead if there is at lease one FATTN_KQ_STRIDE x FATTN_KQ_STRIDE square to be skipped or - // multiple sequences of possibly different lengths. - if (mask && K->ne[1] % FATTN_KQ_STRIDE == 0 && (Q->ne[1] >= 1024 || Q->ne[3] > 1)) { + // KV_max skips trailing all-masked tiles. KV_min, when explicitly enabled, skips leading all-masked tiles. + // Both bounds are conservative at FATTN_KQ_STRIDE granularity; the mask remains authoritative. + const bool can_scan_kv_bounds = mask && K->ne[1] % FATTN_KQ_STRIDE == 0; + const bool kv_max_enabled = can_scan_kv_bounds && (Q->ne[1] >= 1024 || Q->ne[3] > 1); + if (can_scan_kv_bounds && (kv_max_enabled || kv_min_enabled)) { const int s31 = mask->nb[1] / sizeof(half2); const int s33 = mask->nb[3] / sizeof(half2); @@ -1134,10 +1193,18 @@ void launch_fattn( const int ne_KV_max = blocks_num_KV_max.x*blocks_num_KV_max.y; const int iter_k = K->ne[1] / FATTN_KQ_STRIDE; - KV_max.alloc(ne_KV_max); - flash_attn_mask_to_KV_max<<>> - ((const half2 *) mask->data, KV_max.ptr, iter_k, s31, s33); - CUDA_CHECK(cudaGetLastError()); + if (kv_max_enabled) { + KV_max.alloc(ne_KV_max); + flash_attn_mask_to_KV_max<<>> + ((const half2 *) mask->data, KV_max.ptr, iter_k, s31, s33); + CUDA_CHECK(cudaGetLastError()); + } + if (kv_min_enabled) { + KV_min.alloc(ne_KV_max); + flash_attn_mask_to_KV_min<<>> + ((const half2 *) mask->data, KV_min.ptr, iter_k, s31, s33); + CUDA_CHECK(cudaGetLastError()); + } } const dim3 block_dim(warp_size, nwarps, 1); @@ -1245,6 +1312,7 @@ void launch_fattn( mask ? ((const char *) mask->data) : nullptr, sinks ? ((const char *) sinks->data) : nullptr, KV_max.ptr, + KV_min.ptr, !stream_k && parallel_blocks > 1 ? dst_tmp.ptr : (float *) KQV->data, dst_tmp_meta.ptr, scale, max_bias, m0, m1, n_head_log2, logit_softcap, Q->ne[0], ne01, Q->ne[2], Q->ne[3], Q->nb[1], Q->nb[2], Q->nb[3], diff --git a/ggml/src/ggml-cuda/fattn-mma-f16.cuh b/ggml/src/ggml-cuda/fattn-mma-f16.cuh index b40e47a456ae..c20baf65a4f5 100644 --- a/ggml/src/ggml-cuda/fattn-mma-f16.cuh +++ b/ggml/src/ggml-cuda/fattn-mma-f16.cuh @@ -1553,6 +1553,7 @@ static __global__ void flash_attn_ext_f16( const char * __restrict__ mask, const char * __restrict__ sinks, const int * __restrict__ KV_max, + const int * __restrict__ KV_min, float * __restrict__ dst, float2 * __restrict__ dst_meta, const float scale, @@ -1569,6 +1570,7 @@ static __global__ void flash_attn_ext_f16( const int32_t ne31, const int32_t ne32, const int32_t ne33, const int32_t nb31, const int32_t nb32, const int64_t nb33) { #if defined(FLASH_ATTN_AVAILABLE) && (defined(VOLTA_MMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || (defined(AMD_WMMA_AVAILABLE) && defined(RDNA4)) || defined(AMD_MFMA_AVAILABLE)) + GGML_UNUSED(KV_min); // Skip unused kernel variants for faster compilation: if (use_logit_softcap && !(DKQ == 128 || DKQ == 256 || DKQ == 512)) { @@ -1710,7 +1712,7 @@ static __global__ void flash_attn_ext_f16( (Q_f2, K_h2, V_h2, mask_h, sinks_f, dstk, dst_meta, scale, slope, logit_softcap, ne01, ne02, gqa_ratio, ne11, stride_Q1, stride_Q2, stride_K, stride_V, stride_mask, jt, zt_gqa, kb0_start, kb0_stop); #else - GGML_UNUSED_VARS(Q, K, V, mask, sinks, KV_max, dst, dst_meta, scale, + GGML_UNUSED_VARS(Q, K, V, mask, sinks, KV_max, KV_min, dst, dst_meta, scale, max_bias, m0, m1, n_head_log2, logit_softcap, ne00, ne01, ne02, ne03, nb01, nb02, nb03, diff --git a/ggml/src/ggml-cuda/fattn-sparse.cu b/ggml/src/ggml-cuda/fattn-sparse.cu index 151578542f20..479e5cccb280 100644 --- a/ggml/src/ggml-cuda/fattn-sparse.cu +++ b/ggml/src/ggml-cuda/fattn-sparse.cu @@ -29,6 +29,7 @@ #include #include #include +#include static ggml_cuda_sparse_attn_fn_t s_sparse_kernel = nullptr; @@ -81,17 +82,427 @@ __global__ void k_bf16_to_f32_flat( dst[idx] = __bfloat162float(src[idx]); } + +static __device__ __forceinline__ float tree_q8_0_load_d128(const char * row, const int d) { + const block_q8_0 * blocks = (const block_q8_0 *) row; + const block_q8_0 & b = blocks[d / QK8_0]; + return __half2float(b.d) * (float)b.qs[d % QK8_0]; +} + +__global__ void k_flash_attn_tree_q8_0_d128( + const char * __restrict__ Q, + const char * __restrict__ K, + const char * __restrict__ V, + const char * __restrict__ mask, + char * __restrict__ dst, + const float scale, + const int32_t n_q, + const int32_t n_head, + const int32_t n_head_kv, + const int32_t n_kv, + const int32_t n_batch, + const int64_t q_nb1, + const int64_t q_nb2, + const int64_t q_nb3, + const int64_t k_nb1, + const int64_t k_nb2, + const int64_t k_nb3, + const int64_t v_nb1, + const int64_t v_nb2, + const int64_t v_nb3, + const int64_t m_nb1, + const int64_t m_nb3, + const int64_t d_nb1, + const int64_t d_nb2, + const int64_t d_nb3) { + constexpr int D = 128; + const int q = blockIdx.x; + const int h = blockIdx.y; + const int b = blockIdx.z; + const int d = threadIdx.x; + if (q >= n_q || h >= n_head || b >= n_batch || d >= D) { + return; + } + + const int gqa = n_head / n_head_kv; + const int hk = h / gqa; + + const char * q_row = Q + (int64_t)b*q_nb3 + (int64_t)h*q_nb2 + (int64_t)q*q_nb1; + const char * k_base = K + (int64_t)b*k_nb3 + (int64_t)hk*k_nb2; + const char * v_base = V + (int64_t)b*v_nb3 + (int64_t)hk*v_nb2; + const char * m_row = mask + (int64_t)(b % n_batch)*m_nb3 + (int64_t)q*m_nb1; + + const float qd = ((const float *)q_row)[d]; + float acc = 0.0f; + float m = -3.4028234663852886e38f; + float ss = 0.0f; + + __shared__ float red[D]; + for (int k = 0; k < n_kv; ++k) { + const half mh = ((const half *)m_row)[k]; + const float mask_v = __half2float(mh); + if (mask_v <= -60000.0f) { + continue; + } + + const char * k_row = k_base + (int64_t)k*k_nb1; + const float kd = tree_q8_0_load_d128(k_row, d); + red[d] = qd * kd; + __syncthreads(); + + for (int stride = D/2; stride > 0; stride >>= 1) { + if (d < stride) { + red[d] += red[d + stride]; + } + __syncthreads(); + } + + const float score = red[0] * scale + mask_v; + const float m_new = fmaxf(m, score); + const float old_scale = ss == 0.0f ? 0.0f : expf(m - m_new); + const float p = expf(score - m_new); + const char * v_row = v_base + (int64_t)k*v_nb1; + const float vd = tree_q8_0_load_d128(v_row, d); + acc = acc * old_scale + p * vd; + ss = ss * old_scale + p; + m = m_new; + __syncthreads(); + } + + float * dst_row = (float *)(dst + (int64_t)b*d_nb3 + (int64_t)q*d_nb2 + (int64_t)h*d_nb1); + dst_row[d] = ss > 0.0f ? acc / ss : 0.0f; +} + +static bool ggml_cuda_flash_attn_tree_q8_0_supported(const ggml_tensor * dst) { + const ggml_tensor * Q = dst->src[0]; + const ggml_tensor * K = dst->src[1]; + const ggml_tensor * V = dst->src[2]; + const ggml_tensor * M = dst->src[3]; + if (!Q || !K || !V || !M || dst->type != GGML_TYPE_F32) { + return false; + } + if (Q->type != GGML_TYPE_F32 || K->type != GGML_TYPE_Q8_0 || V->type != GGML_TYPE_Q8_0 || M->type != GGML_TYPE_F16) { + return false; + } + if (Q->ne[0] != 128 || K->ne[0] != 128 || V->ne[0] != 128 || Q->ne[3] != 1) { + return false; + } + if (Q->ne[2] % K->ne[2] != 0 || V->ne[2] != K->ne[2] || M->ne[0] < K->ne[1] || M->ne[1] < Q->ne[1]) { + return false; + } + if (Q->nb[0] != (int64_t)sizeof(float) || dst->nb[0] != (int64_t)sizeof(float)) { + return false; + } + return Q->ne[1] <= 65; +} + +static void ggml_cuda_flash_attn_tree_q8_0(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { + ggml_cuda_set_device(ctx.device); + GGML_ASSERT(ggml_cuda_flash_attn_tree_q8_0_supported(dst)); + const ggml_tensor * Q = dst->src[0]; + const ggml_tensor * K = dst->src[1]; + const ggml_tensor * V = dst->src[2]; + const ggml_tensor * M = dst->src[3]; + + float op_params[2]; + memcpy(op_params, dst->op_params, sizeof(op_params)); + const float scale = op_params[0]; + + const dim3 grid((unsigned)Q->ne[1], (unsigned)Q->ne[2], (unsigned)Q->ne[3]); + constexpr int block = 128; + k_flash_attn_tree_q8_0_d128<<>>( + (const char *)Q->data, + (const char *)K->data, + (const char *)V->data, + (const char *)M->data, + (char *)dst->data, + scale, + (int32_t)Q->ne[1], + (int32_t)Q->ne[2], + (int32_t)K->ne[2], + (int32_t)K->ne[1], + (int32_t)Q->ne[3], + Q->nb[1], Q->nb[2], Q->nb[3], + K->nb[1], K->nb[2], K->nb[3], + V->nb[1], V->nb[2], V->nb[3], + M->nb[1], M->nb[3], + dst->nb[1], dst->nb[2], dst->nb[3]); + CUDA_CHECK(cudaGetLastError()); +} + +__global__ void k_flash_attn_tree_q8_0_d128_partials( + const char * __restrict__ Q, + const char * __restrict__ K, + const char * __restrict__ V, + const char * __restrict__ mask, + float * __restrict__ partial_acc, + float2 * __restrict__ partial_meta, + const float scale, + const int32_t n_q, + const int32_t n_head, + const int32_t n_head_kv, + const int32_t n_kv, + const int32_t n_batch, + const int32_t n_chunks, + const int32_t chunk_size, + const int64_t q_nb1, + const int64_t q_nb2, + const int64_t q_nb3, + const int64_t k_nb1, + const int64_t k_nb2, + const int64_t k_nb3, + const int64_t v_nb1, + const int64_t v_nb2, + const int64_t v_nb3, + const int64_t m_nb1, + const int64_t m_nb3) { + constexpr int D = 128; + const int q = blockIdx.x; + const int h = blockIdx.y; + const int z = blockIdx.z; + const int c = z % n_chunks; + const int b = z / n_chunks; + const int d = threadIdx.x; + if (q >= n_q || h >= n_head || b >= n_batch || d >= D) { + return; + } + + const int gqa = n_head / n_head_kv; + const int hk = h / gqa; + + const char * q_row = Q + (int64_t)b*q_nb3 + (int64_t)h*q_nb2 + (int64_t)q*q_nb1; + const char * k_base = K + (int64_t)b*k_nb3 + (int64_t)hk*k_nb2; + const char * v_base = V + (int64_t)b*v_nb3 + (int64_t)hk*v_nb2; + const char * m_row = mask + (int64_t)(b % n_batch)*m_nb3 + (int64_t)q*m_nb1; + + const float qd = ((const float *)q_row)[d]; + float acc = 0.0f; + float m = -3.4028234663852886e38f; + float ss = 0.0f; + + __shared__ float red[D]; + const int k_begin = c * chunk_size; + const int k_end = (k_begin + chunk_size < n_kv) ? k_begin + chunk_size : n_kv; + for (int k = k_begin; k < k_end; ++k) { + const float mask_v = __half2float(((const half *)m_row)[k]); + if (mask_v <= -60000.0f) { + continue; + } + + const char * k_row = k_base + (int64_t)k*k_nb1; + const float kd = tree_q8_0_load_d128(k_row, d); + red[d] = qd * kd; + __syncthreads(); + + for (int stride = D/2; stride > 0; stride >>= 1) { + if (d < stride) { + red[d] += red[d + stride]; + } + __syncthreads(); + } + + const float score = red[0] * scale + mask_v; + const float m_new = fmaxf(m, score); + const float old_scale = ss == 0.0f ? 0.0f : expf(m - m_new); + const float p = expf(score - m_new); + const char * v_row = v_base + (int64_t)k*v_nb1; + const float vd = tree_q8_0_load_d128(v_row, d); + acc = acc * old_scale + p * vd; + ss = ss * old_scale + p; + m = m_new; + __syncthreads(); + } + + const int64_t partial_row = ((((int64_t)b*n_head + h)*n_q + q)*n_chunks + c); + partial_acc[partial_row*D + d] = acc; + if (d == 0) { + partial_meta[partial_row] = make_float2(m, ss); + } +} + +__global__ void k_flash_attn_tree_q8_0_d128_combine( + const float * __restrict__ partial_acc, + const float2 * __restrict__ partial_meta, + char * __restrict__ dst, + const int32_t n_q, + const int32_t n_head, + const int32_t n_batch, + const int32_t n_chunks, + const int64_t d_nb1, + const int64_t d_nb2, + const int64_t d_nb3) { + constexpr int D = 128; + const int q = blockIdx.x; + const int h = blockIdx.y; + const int b = blockIdx.z; + const int d = threadIdx.x; + if (q >= n_q || h >= n_head || b >= n_batch || d >= D) { + return; + } + + float m = -3.4028234663852886e38f; + for (int c = 0; c < n_chunks; ++c) { + const int64_t row = ((((int64_t)b*n_head + h)*n_q + q)*n_chunks + c); + const float2 ms = partial_meta[row]; + if (ms.y > 0.0f) { + m = fmaxf(m, ms.x); + } + } + + float ss = 0.0f; + float acc = 0.0f; + for (int c = 0; c < n_chunks; ++c) { + const int64_t row = ((((int64_t)b*n_head + h)*n_q + q)*n_chunks + c); + const float2 ms = partial_meta[row]; + if (ms.y <= 0.0f) { + continue; + } + const float s = expf(ms.x - m); + ss += ms.y * s; + acc += partial_acc[row*D + d] * s; + } + + float * dst_row = (float *)(dst + (int64_t)b*d_nb3 + (int64_t)q*d_nb2 + (int64_t)h*d_nb1); + dst_row[d] = ss > 0.0f ? acc / ss : 0.0f; +} + +static bool ggml_cuda_flash_attn_tree_q8_0_parallel_supported(const ggml_tensor * dst) { + if (!ggml_cuda_flash_attn_tree_q8_0_supported(dst)) { + return false; + } + const ggml_tensor * Q = dst->src[0]; + const ggml_tensor * K = dst->src[1]; + return Q->ne[3] == 1 && K->ne[1] > 64; +} + +static void ggml_cuda_flash_attn_tree_q8_0_parallel(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { + ggml_cuda_set_device(ctx.device); + GGML_ASSERT(ggml_cuda_flash_attn_tree_q8_0_parallel_supported(dst)); + const ggml_tensor * Q = dst->src[0]; + const ggml_tensor * K = dst->src[1]; + const ggml_tensor * V = dst->src[2]; + const ggml_tensor * M = dst->src[3]; + + float op_params[2]; + memcpy(op_params, dst->op_params, sizeof(op_params)); + const float scale = op_params[0]; + + int chunk_size = 64; + if (const char * env = std::getenv("DFLASH_LAGUNA_TREE_ATTN_Q8_PAR_CHUNK")) { + const int v = std::atoi(env); + if (v == 32 || v == 64 || v == 128 || v == 256) { + chunk_size = v; + } + } + if (const char * env = std::getenv("GGML_CUDA_TREE_ATTN_Q8_PAR_CHUNK")) { + const int v = std::atoi(env); + if (v == 32 || v == 64 || v == 128 || v == 256) { + chunk_size = v; + } + } + + const int n_q = (int)Q->ne[1]; + const int n_head = (int)Q->ne[2]; + const int n_batch = (int)Q->ne[3]; + const int n_kv = (int)K->ne[1]; + const int n_chunks = (n_kv + chunk_size - 1) / chunk_size; + const size_t n_rows = (size_t)n_batch * (size_t)n_head * (size_t)n_q * (size_t)n_chunks; + + float * partial_acc = nullptr; + float2 * partial_meta = nullptr; + CUDA_CHECK(cudaMallocAsync(&partial_acc, n_rows * 128 * sizeof(float), ctx.stream())); + CUDA_CHECK(cudaMallocAsync(&partial_meta, n_rows * sizeof(float2), ctx.stream())); + + const dim3 grid_part((unsigned)n_q, (unsigned)n_head, (unsigned)(n_chunks*n_batch)); + constexpr int block = 128; + k_flash_attn_tree_q8_0_d128_partials<<>>( + (const char *)Q->data, + (const char *)K->data, + (const char *)V->data, + (const char *)M->data, + partial_acc, + partial_meta, + scale, + n_q, + n_head, + (int32_t)K->ne[2], + n_kv, + n_batch, + n_chunks, + chunk_size, + Q->nb[1], Q->nb[2], Q->nb[3], + K->nb[1], K->nb[2], K->nb[3], + V->nb[1], V->nb[2], V->nb[3], + M->nb[1], M->nb[3]); + CUDA_CHECK(cudaGetLastError()); + + const dim3 grid_combine((unsigned)n_q, (unsigned)n_head, (unsigned)n_batch); + k_flash_attn_tree_q8_0_d128_combine<<>>( + partial_acc, + partial_meta, + (char *)dst->data, + n_q, + n_head, + n_batch, + n_chunks, + dst->nb[1], dst->nb[2], dst->nb[3]); + CUDA_CHECK(cudaGetLastError()); + + CUDA_CHECK(cudaFreeAsync(partial_acc, ctx.stream())); + CUDA_CHECK(cudaFreeAsync(partial_meta, ctx.stream())); +} + void ggml_cuda_flash_attn_sparse(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { // K/V may be quantized (Q8_0, Q4_0, etc.). If so, we dequantize them to F16 // into temporary buffers before the S<->H transpose into BF16 pFlash layout. + // Tree-attention carrier: src[3] is the exact dense mask fallback, while + // src[4]/src[5] carry parent ids / positions for a future specialized + // backend kernel. Until that kernel is installed, run dense masked FA so the + // graph path is bit-equivalent to ggml_flash_attn_ext. + const bool tree_mode = dst->src[4] != nullptr; + const bool tree_vec_enabled = tree_mode && + (std::getenv("GGML_CUDA_TREE_ATTN_VEC") != nullptr || + std::getenv("DFLASH_LAGUNA_TREE_ATTN_VEC") != nullptr) && + std::getenv("GGML_CUDA_TREE_ATTN_VEC_DISABLE") == nullptr && + std::getenv("DFLASH_LAGUNA_TREE_ATTN_VEC_DISABLE") == nullptr && + ggml_cuda_flash_attn_ext_vec_supported(dst); + const bool tree_vec_masked_enabled = tree_mode && + (std::getenv("GGML_CUDA_TREE_ATTN_VEC_MASKED") != nullptr || + std::getenv("DFLASH_LAGUNA_TREE_ATTN_VEC_MASKED") != nullptr) && + std::getenv("GGML_CUDA_TREE_ATTN_VEC_MASKED_DISABLE") == nullptr && + std::getenv("DFLASH_LAGUNA_TREE_ATTN_VEC_MASKED_DISABLE") == nullptr && + ggml_cuda_flash_attn_ext_vec_supported(dst); + const bool tree_q8_enabled = tree_mode && + (std::getenv("GGML_CUDA_TREE_ATTN_Q8") != nullptr || + std::getenv("DFLASH_LAGUNA_TREE_ATTN_Q8") != nullptr) && + std::getenv("GGML_CUDA_TREE_ATTN_Q8_DISABLE") == nullptr && + std::getenv("DFLASH_LAGUNA_TREE_ATTN_Q8_DISABLE") == nullptr && + ggml_cuda_flash_attn_tree_q8_0_supported(dst); + const bool tree_q8_parallel_enabled = tree_mode && + (std::getenv("GGML_CUDA_TREE_ATTN_Q8_PAR") != nullptr || + std::getenv("DFLASH_LAGUNA_TREE_ATTN_Q8_PAR") != nullptr) && + std::getenv("GGML_CUDA_TREE_ATTN_Q8_PAR_DISABLE") == nullptr && + std::getenv("DFLASH_LAGUNA_TREE_ATTN_Q8_PAR_DISABLE") == nullptr && + ggml_cuda_flash_attn_tree_q8_0_parallel_supported(dst); + + if (tree_q8_parallel_enabled) { + ggml_cuda_flash_attn_tree_q8_0_parallel(ctx, dst); + return; + } + if (tree_q8_enabled) { + ggml_cuda_flash_attn_tree_q8_0(ctx, dst); + return; + } + // When no sparse kernel is registered, fall back to dense FA. - if (!s_sparse_kernel) { + if (tree_mode || !s_sparse_kernel) { const enum ggml_op saved_op = dst->op; float saved_params[GGML_MAX_OP_PARAMS / sizeof(float)]; memcpy(saved_params, dst->op_params, GGML_MAX_OP_PARAMS); ggml_tensor * saved_src3 = dst->src[3]; ggml_tensor * saved_src4 = dst->src[4]; + ggml_tensor * saved_src5 = dst->src[5]; float op_params[2]; memcpy(op_params, dst->op_params, sizeof(op_params)); @@ -101,15 +512,23 @@ void ggml_cuda_flash_attn_sparse(ggml_backend_cuda_context & ctx, ggml_tensor * dst->op = GGML_OP_FLASH_ATTN_EXT; memset(dst->op_params, 0, GGML_MAX_OP_PARAMS); memcpy(dst->op_params, ext_params, sizeof(ext_params)); - dst->src[3] = nullptr; + dst->src[3] = tree_mode ? saved_src3 : nullptr; dst->src[4] = nullptr; + dst->src[5] = nullptr; - ggml_cuda_flash_attn_ext(ctx, dst); + if (tree_vec_masked_enabled) { + ggml_cuda_flash_attn_ext_vec_masked_force(ctx, dst); + } else if (tree_vec_enabled) { + ggml_cuda_flash_attn_ext_vec_force(ctx, dst); + } else { + ggml_cuda_flash_attn_ext(ctx, dst); + } dst->op = saved_op; memcpy(dst->op_params, saved_params, GGML_MAX_OP_PARAMS); dst->src[3] = saved_src3; dst->src[4] = saved_src4; + dst->src[5] = saved_src5; return; } diff --git a/ggml/src/ggml-cuda/fattn-tile.cuh b/ggml/src/ggml-cuda/fattn-tile.cuh index 26721cc4c7de..0c5743b619c8 100644 --- a/ggml/src/ggml-cuda/fattn-tile.cuh +++ b/ggml/src/ggml-cuda/fattn-tile.cuh @@ -762,6 +762,7 @@ static __global__ void flash_attn_tile( const char * __restrict__ mask, const char * __restrict__ sinks, const int * __restrict__ KV_max, + const int * __restrict__ KV_min, float * __restrict__ dst, float2 * __restrict__ dst_meta, const float scale, @@ -787,7 +788,7 @@ static __global__ void flash_attn_tile( #endif // GGML_USE_WMMA_FATTN (use_logit_softcap && !(DV == 128 || DV == 256 || DV == 512)) ) { - GGML_UNUSED_VARS(Q, K, V, mask, sinks, KV_max, dst, dst_meta, scale, + GGML_UNUSED_VARS(Q, K, V, mask, sinks, KV_max, KV_min, dst, dst_meta, scale, max_bias, m0, m1, n_head_log2, logit_softcap, ne00, ne01, ne02, ne03, nb01, nb02, nb03, @@ -911,10 +912,11 @@ static __global__ void flash_attn_tile( __syncthreads(); // Main loop over KV cache: + const int k_VKQ_min = KV_min ? KV_min[sequence*gridDim.x + blockIdx.x] : 0; const int k_VKQ_max = KV_max ? KV_max[sequence*gridDim.x + blockIdx.x] : ne11; if (ncols2 == 1) { // Branch with out-of-bounds checks. - int k_VKQ_0 = blockIdx.y*nbatch_fa; + int k_VKQ_0 = k_VKQ_min + blockIdx.y*nbatch_fa; while (k_VKQ_0 < k_VKQ_max - nbatch_fa) { constexpr bool oob_check = false; flash_attn_tile_iter @@ -930,7 +932,7 @@ static __global__ void flash_attn_tile( } } else { // Branch without out-of-bounds checks. - for (int k_VKQ_0 = blockIdx.y*nbatch_fa; k_VKQ_0 < k_VKQ_max; k_VKQ_0 += gridDim.y*nbatch_fa) { + for (int k_VKQ_0 = k_VKQ_min + blockIdx.y*nbatch_fa; k_VKQ_0 < k_VKQ_max; k_VKQ_0 += gridDim.y*nbatch_fa) { constexpr bool oob_check = false; flash_attn_tile_iter (Q_tmp, K_h2, V_h2, maskh, ne01, logit_softcap, slope, KQ, KV_tmp, @@ -1092,7 +1094,7 @@ static __global__ void flash_attn_tile( } } #else - GGML_UNUSED_VARS(Q, K, V, mask, sinks, KV_max, dst, dst_meta, scale, + GGML_UNUSED_VARS(Q, K, V, mask, sinks, KV_max, KV_min, dst, dst_meta, scale, max_bias, m0, m1, n_head_log2, logit_softcap, ne00, ne01, ne02, ne03, nb01, nb02, nb03, diff --git a/ggml/src/ggml-cuda/fattn-vec.cuh b/ggml/src/ggml-cuda/fattn-vec.cuh index bcf1dd804420..2324e6ef28ef 100644 --- a/ggml/src/ggml-cuda/fattn-vec.cuh +++ b/ggml/src/ggml-cuda/fattn-vec.cuh @@ -16,7 +16,7 @@ static constexpr __device__ int ggml_cuda_fattn_vec_get_nthreads_device() { #pragma clang diagnostic push #pragma clang diagnostic ignored "-Wpass-failed" #endif // __clang__ -template // D == head size +template // D == head size __launch_bounds__(ggml_cuda_fattn_vec_get_nthreads_device(), 1) static __global__ void flash_attn_ext_vec( const char * __restrict__ Q, @@ -25,6 +25,7 @@ static __global__ void flash_attn_ext_vec( const char * __restrict__ mask, const char * __restrict__ sinks, const int * __restrict__ KV_max, + const int * __restrict__ KV_min, float * __restrict__ dst, float2 * __restrict__ dst_meta, const float scale, @@ -44,7 +45,7 @@ static __global__ void flash_attn_ext_vec( // Skip unused kernel variants for faster compilation: if (use_logit_softcap && !(D == 128 || D == 256)) { - GGML_UNUSED_VARS(Q, K, V, mask, sinks, KV_max, dst, dst_meta, scale, + GGML_UNUSED_VARS(Q, K, V, mask, sinks, KV_max, KV_min, dst, dst_meta, scale, max_bias, m0, m1, n_head_log2, logit_softcap, ne00, ne01, ne02, ne03, nb01, nb02, nb03, @@ -241,21 +242,24 @@ static __global__ void flash_attn_ext_vec( #endif // V_DOT2_F32_F16_AVAILABLE } + const int k_VKQ_min = KV_min ? KV_min[sequence*gridDim.x + blockIdx.x] : 0; const int k_VKQ_max = KV_max ? KV_max[sequence*gridDim.x + blockIdx.x] : ne11; - K += blockIdx.y*nthreads * nb11; - V += blockIdx.y*nthreads * nb21; - maskh += blockIdx.y*nthreads; - for (int k_VKQ_0 = blockIdx.y*nthreads; k_VKQ_0 < k_VKQ_max; k_VKQ_0 += gridDim.y*nthreads, + K += (k_VKQ_min + blockIdx.y*nthreads) * nb11; + V += (k_VKQ_min + blockIdx.y*nthreads) * nb21; + maskh += k_VKQ_min + blockIdx.y*nthreads; + for (int k_VKQ_0 = k_VKQ_min + blockIdx.y*nthreads; k_VKQ_0 < k_VKQ_max; k_VKQ_0 += gridDim.y*nthreads, // Increment pointers after each loop: K += gridDim.y*nthreads*nb11, V += gridDim.y*nthreads*nb21, maskh += gridDim.y*nthreads) { // Calculate KQ tile and keep track of new maximum KQ values: float KQ_reg[ncols]; // KQ in registers. + bool KQ_masked[ncols]; float KQ_max_new[ncols]; #pragma unroll for (int j = 0; j < ncols; ++j) { KQ_max_new[j] = KQ_max[j]; + KQ_masked[j] = false; } #pragma unroll @@ -264,21 +268,34 @@ static __global__ void flash_attn_ext_vec( #pragma unroll for (int j = 0; j < ncols; ++j) { - float sum = vec_dot_KQ(K + i_KQ*nb11, Q_reg[j], Q_i32[j], Q_ds[j]); + const bool q_in_bounds = ncols == 1 || ic0 + j < int(ne01.z); + float mask_v = 0.0f; + bool masked = false; + if (mask && q_in_bounds) { + mask_v = slope*__half2float(maskh[j*ne11 + i_KQ]); + if constexpr (skip_masked) { + masked = mask_v <= -60000.0f; + } + } + + float sum = masked ? 0.0f : vec_dot_KQ(K + i_KQ*nb11, Q_reg[j], Q_i32[j], Q_ds[j]); sum = warp_reduce_sum(sum); - if (use_logit_softcap) { - sum = logit_softcap*tanhf(sum); - } + if (!masked) { + if (use_logit_softcap) { + sum = logit_softcap*tanhf(sum); + } - if (mask && (ncols == 1 || ic0 + j < int(ne01.z))) { - sum += slope*__half2float(maskh[j*ne11 + i_KQ]); - } + if (mask && q_in_bounds) { + sum += mask_v; + } - KQ_max_new[j] = fmaxf(KQ_max_new[j], sum + FATTN_KQ_MAX_OFFSET); + KQ_max_new[j] = fmaxf(KQ_max_new[j], sum + FATTN_KQ_MAX_OFFSET); + } if ((nthreads_KQ == WARP_SIZE ? threadIdx.x : threadIdx.x % nthreads_KQ) == uint32_t(i_KQ_0)) { KQ_reg[j] = sum; + KQ_masked[j] = masked; } } } @@ -292,7 +309,11 @@ static __global__ void flash_attn_ext_vec( const float KQ_max_scale = expf(KQ_max[j] - KQ_max_new[j]); KQ_max[j] = KQ_max_new[j]; - KQ_reg[j] = expf(KQ_reg[j] - KQ_max[j]); + if constexpr (skip_masked) { + KQ_reg[j] = KQ_masked[j] ? 0.0f : expf(KQ_reg[j] - KQ_max[j]); + } else { + KQ_reg[j] = expf(KQ_reg[j] - KQ_max[j]); + } KQ_sum[j] = KQ_sum[j]*KQ_max_scale + KQ_reg[j]; KQ[j*nthreads + tid] = KQ_reg[j]; @@ -321,9 +342,17 @@ static __global__ void flash_attn_ext_vec( #ifdef V_DOT2_F32_F16_AVAILABLE half2 KQ_k[ncols]; + bool any_KQ_k = false; #pragma unroll for (int j = 0; j < ncols; ++j) { - KQ_k[j] = __half2half2(KQ[j*nthreads + k]); + const half kq = KQ[j*nthreads + k]; + KQ_k[j] = __half2half2(kq); + any_KQ_k = any_KQ_k || __half2float(kq) != 0.0f; + } + if constexpr (skip_masked) { + if (!any_KQ_k) { + continue; + } } #pragma unroll for (int i_VKQ_0 = 0; i_VKQ_0 < D/2; i_VKQ_0 += nthreads_V*V_rows_per_thread/2) { @@ -350,9 +379,16 @@ static __global__ void flash_attn_ext_vec( } #else float KQ_k[ncols]; + bool any_KQ_k = false; #pragma unroll for (int j = 0; j < ncols; ++j) { KQ_k[j] = KQ[j*nthreads + k]; + any_KQ_k = any_KQ_k || KQ_k[j] != 0.0f; + } + if constexpr (skip_masked) { + if (!any_KQ_k) { + continue; + } } #pragma unroll for (int i_VKQ_0 = 0; i_VKQ_0 < D/2; i_VKQ_0 += nthreads_V*V_rows_per_thread/2) { @@ -508,7 +544,7 @@ static __global__ void flash_attn_ext_vec( dst_meta[((sequence*int(ne01.z) + ic0 + tid)*ne02 + head)*gridDim.y + blockIdx.y] = make_float2(KQ_max[tid], KQ_sum[tid]); } #else - GGML_UNUSED_VARS(Q, K, V, mask, sinks, KV_max, dst, dst_meta, scale, + GGML_UNUSED_VARS(Q, K, V, mask, sinks, KV_max, KV_min, dst, dst_meta, scale, max_bias, m0, m1, n_head_log2, logit_softcap, ne00, ne01, ne02, ne03, nb01, nb02, nb03, @@ -524,13 +560,13 @@ static __global__ void flash_attn_ext_vec( #pragma clang diagnostic pop #endif // __clang__ -template +template void ggml_cuda_flash_attn_ext_vec_case_impl(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { const int cc = ggml_cuda_info().devices[ggml_cuda_get_device()].cc; const int nthreads = ggml_cuda_fattn_vec_get_nthreads_host(cc); const int nwarps = nthreads / WARP_SIZE; - fattn_kernel_t fattn_kernel = flash_attn_ext_vec; + fattn_kernel_t fattn_kernel = flash_attn_ext_vec; const bool need_f16_K = type_K == GGML_TYPE_F16; const bool need_f16_V = type_V == GGML_TYPE_F16; constexpr size_t nbytes_shared = 0; @@ -549,10 +585,46 @@ void ggml_cuda_flash_attn_ext_vec_case(ggml_backend_cuda_context & ctx, ggml_ten constexpr int cols_per_block = 1; if (logit_softcap == 0.0f) { constexpr bool use_logit_softcap = false; - ggml_cuda_flash_attn_ext_vec_case_impl(ctx, dst); + constexpr bool skip_masked = false; + ggml_cuda_flash_attn_ext_vec_case_impl(ctx, dst); + } else { + constexpr bool use_logit_softcap = true; + constexpr bool skip_masked = false; + ggml_cuda_flash_attn_ext_vec_case_impl(ctx, dst); + } + return; + } + + constexpr int cols_per_block = 2; + if (logit_softcap == 0.0f) { + constexpr bool use_logit_softcap = false; + constexpr bool skip_masked = false; + ggml_cuda_flash_attn_ext_vec_case_impl(ctx, dst); + } else { + constexpr bool use_logit_softcap = true; + constexpr bool skip_masked = false; + ggml_cuda_flash_attn_ext_vec_case_impl(ctx, dst); + } +} + +template +void ggml_cuda_flash_attn_ext_vec_masked_case(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { + const ggml_tensor * KQV = dst; + const ggml_tensor * Q = dst->src[0]; + + float logit_softcap; + memcpy(&logit_softcap, (const float *) KQV->op_params + 2, sizeof(float)); + + if (Q->ne[1] == 1) { + constexpr int cols_per_block = 1; + if (logit_softcap == 0.0f) { + constexpr bool use_logit_softcap = false; + constexpr bool skip_masked = true; + ggml_cuda_flash_attn_ext_vec_case_impl(ctx, dst); } else { constexpr bool use_logit_softcap = true; - ggml_cuda_flash_attn_ext_vec_case_impl(ctx, dst); + constexpr bool skip_masked = true; + ggml_cuda_flash_attn_ext_vec_case_impl(ctx, dst); } return; } @@ -560,10 +632,12 @@ void ggml_cuda_flash_attn_ext_vec_case(ggml_backend_cuda_context & ctx, ggml_ten constexpr int cols_per_block = 2; if (logit_softcap == 0.0f) { constexpr bool use_logit_softcap = false; - ggml_cuda_flash_attn_ext_vec_case_impl(ctx, dst); + constexpr bool skip_masked = true; + ggml_cuda_flash_attn_ext_vec_case_impl(ctx, dst); } else { constexpr bool use_logit_softcap = true; - ggml_cuda_flash_attn_ext_vec_case_impl(ctx, dst); + constexpr bool skip_masked = true; + ggml_cuda_flash_attn_ext_vec_case_impl(ctx, dst); } } diff --git a/ggml/src/ggml-cuda/fattn-wmma-f16.cu b/ggml/src/ggml-cuda/fattn-wmma-f16.cu index f19defbff939..5aa145592c4c 100644 --- a/ggml/src/ggml-cuda/fattn-wmma-f16.cu +++ b/ggml/src/ggml-cuda/fattn-wmma-f16.cu @@ -30,6 +30,7 @@ static __global__ void flash_attn_ext_f16( const char * __restrict__ mask, const char * __restrict__ sinks, const int * __restrict__ KV_max, + const int * __restrict__ KV_min, float * __restrict__ dst, float2 * __restrict__ dst_meta, const float scale, @@ -188,8 +189,9 @@ static __global__ void flash_attn_ext_f16( __syncthreads(); // Iterate over ne11 == previous tokens: + const int k_VKQ_min = KV_min ? KV_min[sequence*gridDim.x + blockIdx.x] : 0; const int k_VKQ_max = KV_max ? KV_max[sequence*gridDim.x + blockIdx.x] : ne11; - for (int k_VKQ_0 = blockIdx.y*FATTN_KQ_STRIDE; k_VKQ_0 < k_VKQ_max; k_VKQ_0 += gridDim.y*FATTN_KQ_STRIDE) { + for (int k_VKQ_0 = k_VKQ_min + blockIdx.y*FATTN_KQ_STRIDE; k_VKQ_0 < k_VKQ_max; k_VKQ_0 += gridDim.y*FATTN_KQ_STRIDE) { // Calculate tile of KQ: #pragma unroll for (int i_KQ_0 = 0; i_KQ_0 < FATTN_KQ_STRIDE; i_KQ_0 += KQ_stride_tc) { @@ -493,7 +495,7 @@ static __global__ void flash_attn_ext_f16( dst_meta[j_dst_unrolled] = dst_meta_val; } #else - GGML_UNUSED_VARS(Q, K, V, mask, sinks, KV_max, dst, dst_meta, scale, + GGML_UNUSED_VARS(Q, K, V, mask, sinks, KV_max, KV_min, dst, dst_meta, scale, max_bias, m0, m1, n_head_log2, logit_softcap, ne00, ne01, ne02, ne03, nb01, nb02, nb03, diff --git a/ggml/src/ggml-cuda/fattn.cu b/ggml/src/ggml-cuda/fattn.cu index 0eff4f1b7c6a..0d6fd3e98887 100644 --- a/ggml/src/ggml-cuda/fattn.cu +++ b/ggml/src/ggml-cuda/fattn.cu @@ -220,11 +220,26 @@ static void ggml_cuda_flash_attn_ext_mma_f16(ggml_backend_cuda_context & ctx, gg } \ } \ +#define FATTN_VEC_MASKED_CASE(D, type_K, type_V) \ + { \ + const bool type_K_okay = K->type == (type_K) || (K->type == GGML_TYPE_F32 && (type_K) == GGML_TYPE_F16);\ + const bool type_V_okay = V->type == (type_V) || (V->type == GGML_TYPE_F32 && (type_V) == GGML_TYPE_F16);\ + if (Q->ne[0] == (D) && type_K_okay && type_V_okay) { \ + ggml_cuda_flash_attn_ext_vec_masked_case(ctx, dst); \ + return; \ + } \ + } \ + #define FATTN_VEC_CASES_ALL_D(type_K, type_V) \ FATTN_VEC_CASE( 64, type_K, type_V) \ FATTN_VEC_CASE(128, type_K, type_V) \ FATTN_VEC_CASE(256, type_K, type_V) \ +#define FATTN_VEC_MASKED_CASES_ALL_D(type_K, type_V) \ + FATTN_VEC_MASKED_CASE( 64, type_K, type_V) \ + FATTN_VEC_MASKED_CASE(128, type_K, type_V) \ + FATTN_VEC_MASKED_CASE(256, type_K, type_V) \ + static void ggml_cuda_flash_attn_ext_vec(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { ggml_tensor * Q = dst->src[0]; ggml_tensor * K = dst->src[1]; @@ -312,6 +327,82 @@ static void ggml_cuda_flash_attn_ext_vec(ggml_backend_cuda_context & ctx, ggml_t GGML_ABORT("fatal error"); } +static void ggml_cuda_flash_attn_ext_vec_masked(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { + ggml_tensor * Q = dst->src[0]; + ggml_tensor * K = dst->src[1]; + ggml_tensor * V = dst->src[2]; + + FATTN_VEC_MASKED_CASE(128, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0) + + GGML_ABORT("fatal error"); +} + +static bool ggml_cuda_flash_attn_ext_vec_kv_type_supported(const ggml_tensor * K, const ggml_tensor * V) { + if (K->type != V->type) { + return false; + } + + switch (K->type) { + case GGML_TYPE_F32: + case GGML_TYPE_F16: + case GGML_TYPE_Q4_0: + case GGML_TYPE_Q8_0: + case GGML_TYPE_BF16: + return true; +#ifdef GGML_CUDA_FA_ALL_QUANTS + case GGML_TYPE_Q4_1: + case GGML_TYPE_Q5_0: + case GGML_TYPE_Q5_1: + return true; +#endif // GGML_CUDA_FA_ALL_QUANTS +#ifndef GGML_USE_HIP + case GGML_TYPE_TQ3_0: + return true; +#endif // GGML_USE_HIP + default: + return false; + } +} + +bool ggml_cuda_flash_attn_ext_vec_supported(const ggml_tensor * dst) { +#ifndef FLASH_ATTN_AVAILABLE + GGML_UNUSED(dst); + return false; +#else + const ggml_tensor * Q = dst->src[0]; + const ggml_tensor * K = dst->src[1]; + const ggml_tensor * V = dst->src[2]; + const ggml_tensor * mask = dst->src[3]; + + if (!Q || !K || !V || Q->type != GGML_TYPE_F32 || dst->type != GGML_TYPE_F32) { + return false; + } + if (Q->ne[2] % K->ne[2] != 0 || K->ne[1] % FATTN_KQ_STRIDE != 0) { + return false; + } + if (mask && mask->ne[2] != 1) { + return false; + } + if (!(Q->ne[0] == 64 || Q->ne[0] == 128 || Q->ne[0] == 256) || V->ne[0] != K->ne[0]) { + return false; + } + return ggml_cuda_flash_attn_ext_vec_kv_type_supported(K, V); +#endif // FLASH_ATTN_AVAILABLE +} + +void ggml_cuda_flash_attn_ext_vec_force(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { + ggml_cuda_set_device(ctx.device); + GGML_ASSERT(ggml_cuda_flash_attn_ext_vec_supported(dst)); + ggml_cuda_flash_attn_ext_vec(ctx, dst); +} + +void ggml_cuda_flash_attn_ext_vec_masked_force(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { + ggml_cuda_set_device(ctx.device); + GGML_ASSERT(ggml_cuda_flash_attn_ext_vec_supported(dst)); + GGML_ASSERT(dst->src[3] != nullptr); + ggml_cuda_flash_attn_ext_vec_masked(ctx, dst); +} + // Best FlashAttention kernel for a specific GPU: enum best_fattn_kernel { BEST_FATTN_KERNEL_NONE = 0, @@ -466,6 +557,15 @@ static best_fattn_kernel ggml_cuda_get_best_fattn_kernel(const int device, const // For small batch sizes the vector kernel may be preferable over the kernels optimized for large batch sizes: const bool can_use_vector_kernel = Q->ne[0] <= 256 && Q->ne[0] % 64 == 0 && K->ne[1] % FATTN_KQ_STRIDE == 0; + static const bool force_tile_kernel = + (std::getenv("GGML_CUDA_FA_FORCE_TILE") != nullptr || + std::getenv("DFLASH_LAGUNA_FA_FORCE_TILE") != nullptr) && + std::getenv("GGML_CUDA_FA_FORCE_TILE_DISABLE") == nullptr && + std::getenv("DFLASH_LAGUNA_FA_FORCE_TILE_DISABLE") == nullptr; + if (force_tile_kernel) { + return BEST_FATTN_KERNEL_TILE; + } + // If Turing tensor cores are available, use them: if (turing_mma_available(cc) && Q->ne[0] != 40 && Q->ne[0] != 72) { if (can_use_vector_kernel) { diff --git a/ggml/src/ggml-cuda/fattn.cuh b/ggml/src/ggml-cuda/fattn.cuh index 78705d59951c..3c14f9045d6b 100644 --- a/ggml/src/ggml-cuda/fattn.cuh +++ b/ggml/src/ggml-cuda/fattn.cuh @@ -2,4 +2,10 @@ void ggml_cuda_flash_attn_ext(ggml_backend_cuda_context & ctx, ggml_tensor * dst); +bool ggml_cuda_flash_attn_ext_vec_supported(const ggml_tensor * dst); + +void ggml_cuda_flash_attn_ext_vec_force(ggml_backend_cuda_context & ctx, ggml_tensor * dst); + +void ggml_cuda_flash_attn_ext_vec_masked_force(ggml_backend_cuda_context & ctx, ggml_tensor * dst); + bool ggml_cuda_flash_attn_ext_supported(int device, const ggml_tensor * dst); diff --git a/ggml/src/ggml.c b/ggml/src/ggml.c index 10c95a999df2..08055e28a4fa 100644 --- a/ggml/src/ggml.c +++ b/ggml/src/ggml.c @@ -5421,6 +5421,37 @@ struct ggml_tensor * ggml_flash_attn_sparse( return result; } +struct ggml_tensor * ggml_flash_attn_tree( + struct ggml_context * ctx, + struct ggml_tensor * q, + struct ggml_tensor * k, + struct ggml_tensor * v, + struct ggml_tensor * mask, + struct ggml_tensor * parent_ids, + struct ggml_tensor * positions, + float scale) { + + GGML_ASSERT(parent_ids != NULL); + GGML_ASSERT(parent_ids->type == GGML_TYPE_I32); + GGML_ASSERT(positions != NULL); + GGML_ASSERT(positions->type == GGML_TYPE_I32); + + if (mask) { + GGML_ASSERT(mask->type == GGML_TYPE_F16); + GGML_ASSERT(ggml_is_contiguous(mask)); + GGML_ASSERT(q->ne[2] % mask->ne[2] == 0); + GGML_ASSERT(q->ne[3] % mask->ne[3] == 0); + } + + struct ggml_tensor * result = ggml_flash_attn_sparse( + ctx, q, k, v, scale, /*alpha=*/-1.0f); + result->src[3] = mask; + result->src[4] = parent_ids; + result->src[5] = positions; + + return result; +} + // ggml_flash_attn_back struct ggml_tensor * ggml_flash_attn_back(