167 lines
8.2 KiB
Plaintext
167 lines
8.2 KiB
Plaintext
#include "fattn-vec-f32.cuh"
|
|
#include "fattn-vec-f32-interface.cuh"
|
|
|
|
#define FATTN_VEC_F32_CASE(D, type_K, type_V) \
|
|
if (Q->ne[0] == (D) && K->type == (type_K) && V->type == (type_V)) { \
|
|
ggml_cuda_flash_attn_ext_vec_f32_case<D, D, type_K, type_V>(ctx, dst); \
|
|
return; \
|
|
} \
|
|
|
|
#define FATTN_VEC_F32_CASE_DKDV(Dk, Dv, type_K, type_V) \
|
|
if (Q->ne[0] == (Dk) && V->ne[0] == Dv && K->type == (type_K) && V->type == (type_V)) { \
|
|
ggml_cuda_flash_attn_ext_vec_f32_case<Dk, Dv, type_K, type_V>(ctx, dst); \
|
|
return; \
|
|
}
|
|
|
|
void ggml_cuda_flash_attn_ext_vec_f32(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];
|
|
|
|
#ifdef GGML_CUDA_FA_ALL_QUANTS
|
|
FATTN_VEC_F32_CASE( 64, GGML_TYPE_F16, GGML_TYPE_Q4_0)
|
|
FATTN_VEC_F32_CASE( 64, GGML_TYPE_F16, GGML_TYPE_Q4_1)
|
|
FATTN_VEC_F32_CASE( 64, GGML_TYPE_F16, GGML_TYPE_Q5_0)
|
|
FATTN_VEC_F32_CASE( 64, GGML_TYPE_F16, GGML_TYPE_Q5_1)
|
|
FATTN_VEC_F32_CASE( 64, GGML_TYPE_F16, GGML_TYPE_Q8_0)
|
|
FATTN_VEC_F32_CASE( 64, GGML_TYPE_F16, GGML_TYPE_F16 )
|
|
|
|
FATTN_VEC_F32_CASE(128, GGML_TYPE_Q4_0, GGML_TYPE_Q4_0)
|
|
FATTN_VEC_F32_CASE(128, GGML_TYPE_Q4_1, GGML_TYPE_Q4_0)
|
|
FATTN_VEC_F32_CASE(128, GGML_TYPE_Q5_0, GGML_TYPE_Q4_0)
|
|
FATTN_VEC_F32_CASE(128, GGML_TYPE_Q5_1, GGML_TYPE_Q4_0)
|
|
FATTN_VEC_F32_CASE(128, GGML_TYPE_Q8_0, GGML_TYPE_Q4_0)
|
|
FATTN_VEC_F32_CASE(128, GGML_TYPE_F16, GGML_TYPE_Q4_0)
|
|
|
|
FATTN_VEC_F32_CASE(128, GGML_TYPE_Q4_0, GGML_TYPE_Q4_1)
|
|
FATTN_VEC_F32_CASE(128, GGML_TYPE_Q4_1, GGML_TYPE_Q4_1)
|
|
FATTN_VEC_F32_CASE(128, GGML_TYPE_Q5_0, GGML_TYPE_Q4_1)
|
|
FATTN_VEC_F32_CASE(128, GGML_TYPE_Q5_1, GGML_TYPE_Q4_1)
|
|
FATTN_VEC_F32_CASE(128, GGML_TYPE_Q8_0, GGML_TYPE_Q4_1)
|
|
FATTN_VEC_F32_CASE(128, GGML_TYPE_F16, GGML_TYPE_Q4_1)
|
|
|
|
FATTN_VEC_F32_CASE(128, GGML_TYPE_Q4_0, GGML_TYPE_Q5_0)
|
|
FATTN_VEC_F32_CASE(128, GGML_TYPE_Q4_1, GGML_TYPE_Q5_0)
|
|
FATTN_VEC_F32_CASE(128, GGML_TYPE_Q5_0, GGML_TYPE_Q5_0)
|
|
FATTN_VEC_F32_CASE(128, GGML_TYPE_Q5_1, GGML_TYPE_Q5_0)
|
|
FATTN_VEC_F32_CASE(128, GGML_TYPE_Q8_0, GGML_TYPE_Q5_0)
|
|
FATTN_VEC_F32_CASE(128, GGML_TYPE_F16, GGML_TYPE_Q5_0)
|
|
|
|
FATTN_VEC_F32_CASE(128, GGML_TYPE_Q4_0, GGML_TYPE_Q5_1)
|
|
FATTN_VEC_F32_CASE(128, GGML_TYPE_Q4_1, GGML_TYPE_Q5_1)
|
|
FATTN_VEC_F32_CASE(128, GGML_TYPE_Q5_0, GGML_TYPE_Q5_1)
|
|
FATTN_VEC_F32_CASE(128, GGML_TYPE_Q5_1, GGML_TYPE_Q5_1)
|
|
FATTN_VEC_F32_CASE(128, GGML_TYPE_Q8_0, GGML_TYPE_Q5_1)
|
|
FATTN_VEC_F32_CASE(128, GGML_TYPE_F16, GGML_TYPE_Q5_1)
|
|
|
|
FATTN_VEC_F32_CASE(128, GGML_TYPE_Q4_0, GGML_TYPE_Q8_0)
|
|
FATTN_VEC_F32_CASE(128, GGML_TYPE_Q4_1, GGML_TYPE_Q8_0)
|
|
FATTN_VEC_F32_CASE(128, GGML_TYPE_Q5_0, GGML_TYPE_Q8_0)
|
|
FATTN_VEC_F32_CASE(128, GGML_TYPE_Q5_1, GGML_TYPE_Q8_0)
|
|
FATTN_VEC_F32_CASE(128, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0)
|
|
FATTN_VEC_F32_CASE(128, GGML_TYPE_F16, GGML_TYPE_Q8_0)
|
|
|
|
FATTN_VEC_F32_CASE(128, GGML_TYPE_Q4_0, GGML_TYPE_F16)
|
|
FATTN_VEC_F32_CASE(128, GGML_TYPE_Q4_1, GGML_TYPE_F16)
|
|
FATTN_VEC_F32_CASE(128, GGML_TYPE_Q5_0, GGML_TYPE_F16)
|
|
FATTN_VEC_F32_CASE(128, GGML_TYPE_Q5_1, GGML_TYPE_F16)
|
|
FATTN_VEC_F32_CASE(128, GGML_TYPE_Q8_0, GGML_TYPE_F16)
|
|
FATTN_VEC_F32_CASE(128, GGML_TYPE_F16, GGML_TYPE_F16)
|
|
|
|
FATTN_VEC_F32_CASE(256, GGML_TYPE_F16, GGML_TYPE_F16)
|
|
FATTN_VEC_F32_CASE(256, GGML_TYPE_Q8_0,GGML_TYPE_Q8_0)
|
|
|
|
FATTN_VEC_F32_CASE(128, GGML_TYPE_IQ4_NL, GGML_TYPE_IQ4_NL)
|
|
FATTN_VEC_F32_CASE(128, GGML_TYPE_Q8_0, GGML_TYPE_IQ4_NL)
|
|
FATTN_VEC_F32_CASE(128, GGML_TYPE_Q6_0, GGML_TYPE_Q5_0)
|
|
FATTN_VEC_F32_CASE(128, GGML_TYPE_Q6_0, GGML_TYPE_Q6_0)
|
|
FATTN_VEC_F32_CASE(128, GGML_TYPE_Q8_0, GGML_TYPE_Q6_0)
|
|
|
|
FATTN_VEC_F32_CASE_DKDV(192, 128, GGML_TYPE_F16, GGML_TYPE_F16)
|
|
FATTN_VEC_F32_CASE_DKDV(192, 128, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0)
|
|
|
|
FATTN_VEC_F32_CASE_DKDV(576, 512, GGML_TYPE_F16, GGML_TYPE_F16)
|
|
FATTN_VEC_F32_CASE_DKDV(576, 512, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0)
|
|
#else
|
|
FATTN_VEC_F32_CASE(128, GGML_TYPE_Q4_0, GGML_TYPE_Q4_0)
|
|
|
|
FATTN_VEC_F32_CASE(128, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0)
|
|
|
|
FATTN_VEC_F32_CASE( 64, GGML_TYPE_F16, GGML_TYPE_F16)
|
|
FATTN_VEC_F32_CASE(128, GGML_TYPE_F16, GGML_TYPE_F16)
|
|
FATTN_VEC_F32_CASE(256, GGML_TYPE_F16, GGML_TYPE_F16)
|
|
FATTN_VEC_F32_CASE(256, GGML_TYPE_Q8_0,GGML_TYPE_Q8_0)
|
|
|
|
FATTN_VEC_F32_CASE(128, GGML_TYPE_IQ4_NL, GGML_TYPE_IQ4_NL)
|
|
FATTN_VEC_F32_CASE(128, GGML_TYPE_Q8_0, GGML_TYPE_IQ4_NL)
|
|
FATTN_VEC_F32_CASE(128, GGML_TYPE_Q6_0, GGML_TYPE_Q5_0)
|
|
FATTN_VEC_F32_CASE(128, GGML_TYPE_Q8_0, GGML_TYPE_Q6_0)
|
|
|
|
FATTN_VEC_F32_CASE_DKDV(192, 128, GGML_TYPE_F16, GGML_TYPE_F16)
|
|
FATTN_VEC_F32_CASE_DKDV(192, 128, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0)
|
|
|
|
FATTN_VEC_F32_CASE_DKDV(576, 512, GGML_TYPE_F16, GGML_TYPE_F16)
|
|
FATTN_VEC_F32_CASE_DKDV(576, 512, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0)
|
|
|
|
#endif // GGML_CUDA_FA_ALL_QUANTS
|
|
|
|
on_no_fattn_vec_case(Q->ne[0], V->ne[0]);
|
|
}
|
|
|
|
bool ggml_cuda_fattn_vec_f32_is_supported([[maybe_unused]] ggml_backend_cuda_context & ctx, const ggml_tensor * dst) {
|
|
auto K = dst->src[1];
|
|
auto V = dst->src[2];
|
|
if (K->ne[0] != V->ne[0]) {
|
|
const bool is_192_128 = K->ne[0] == 192 && V->ne[0] == 128;
|
|
const bool is_576_512 = K->ne[0] == 576 && V->ne[0] == 512;
|
|
if (!is_192_128 && !is_576_512) return false;
|
|
if (K->type != V->type) return false;
|
|
// Asymmetric (Dk != Dv) vec_f32 path. F16 KV is always fine (its dot loops stride Dk/2,
|
|
// a clean multiple of 32). Q8_0 KV is now supported: the per-warp Q->q8_1 quantization,
|
|
// the shared->register load, and the quantized-K dot all carry a compile-time-gated
|
|
// per-lane bounds guard (mirroring mainline llama.cpp's fattn-vec ragged-tail idiom) so
|
|
// the ragged final warp iteration for Dk=576 (144 int32 = 4.5*WARP_SIZE) and Dk=192
|
|
// (48 int32 = 1.5*WARP_SIZE) no longer reads Q/K out of bounds. The guard folds to a
|
|
// no-op for aligned Dk (128/256), keeping those instances byte-identical.
|
|
return K->type == GGML_TYPE_F16 || K->type == GGML_TYPE_Q8_0;
|
|
}
|
|
#ifdef GGML_CUDA_FA_ALL_QUANTS
|
|
if (K->ne[0] == 64) {
|
|
return K->type == GGML_TYPE_F16 &&
|
|
(V->type == GGML_TYPE_F16 || V->type == GGML_TYPE_Q4_0 || V->type == GGML_TYPE_Q4_1 ||
|
|
V->type == GGML_TYPE_Q5_0 || V->type == GGML_TYPE_Q5_1 || V->type == GGML_TYPE_Q8_0);
|
|
}
|
|
if (K->ne[0] == 256) {
|
|
return K->type == V->type && (K->type == GGML_TYPE_F16 || K->type == GGML_TYPE_Q8_0);
|
|
}
|
|
if (K->ne[0] != 128 || V->ne[0] != 128) return false;
|
|
if ((K->type == GGML_TYPE_Q4_0 || K->type == GGML_TYPE_Q4_1 || K->type == GGML_TYPE_Q5_0 || K->type == GGML_TYPE_Q5_1 ||
|
|
K->type == GGML_TYPE_Q8_0 || K->type == GGML_TYPE_F16) &&
|
|
(V->type == GGML_TYPE_Q4_0 || V->type == GGML_TYPE_Q4_1 || V->type == GGML_TYPE_Q5_0 || V->type == GGML_TYPE_Q5_1 ||
|
|
V->type == GGML_TYPE_Q8_0 || V->type == GGML_TYPE_F16)) return true;
|
|
return (K->type == GGML_TYPE_Q8_0 && V->type == GGML_TYPE_IQ4_NL) ||
|
|
(K->type == GGML_TYPE_Q6_0 && V->type == GGML_TYPE_Q5_0) ||
|
|
(K->type == GGML_TYPE_Q6_0 && V->type == GGML_TYPE_Q6_0) ||
|
|
(K->type == GGML_TYPE_Q8_0 && V->type == GGML_TYPE_Q6_0) ||
|
|
(K->type == GGML_TYPE_Q8_0 && V->type == GGML_TYPE_IQ4_NL);
|
|
#else
|
|
if (K->ne[0] == 128) {
|
|
if (K->type == V->type) {
|
|
return K->type == GGML_TYPE_Q4_0 || K->type == GGML_TYPE_Q8_0 || K->type == GGML_TYPE_F16 || K->type == GGML_TYPE_IQ4_NL;
|
|
}
|
|
return (K->type == GGML_TYPE_Q8_0 && V->type == GGML_TYPE_IQ4_NL) ||
|
|
(K->type == GGML_TYPE_Q6_0 && V->type == GGML_TYPE_Q5_0) ||
|
|
(K->type == GGML_TYPE_Q8_0 && V->type == GGML_TYPE_Q6_0) ||
|
|
(K->type == GGML_TYPE_Q8_0 && V->type == GGML_TYPE_IQ4_NL);
|
|
}
|
|
if (K->type != V->type) return false;
|
|
if (K->ne[0] == 64) {
|
|
return K->type == GGML_TYPE_F16;
|
|
}
|
|
if (K->ne[0] == 256) {
|
|
return K->type == GGML_TYPE_F16 || K->type == GGML_TYPE_Q8_0;
|
|
}
|
|
return false;
|
|
#endif
|
|
}
|