Skip to content

Commit 6e093b8

Browse files
0cc4m0cc4m
authored
vulkan: add Flash Attention support for BFloat16 KV cache (ggml-org#23420)
* vulkan: add flash attention bf16 kv support * vulkan: bf16 FA coopmat1 support * vulkan: bf16 FA coopmat2 support * fix FA bf16 f32 fallback * fix FA bf16 coopmat1 shader * fix FA bf16 coopmat2 shader * code cleanup * cleanup comment change * address feedback * add O_TYPE for cm2 FA * use O_TYPE for gqaStore function * reduce BFLOAT16 ifdefs
1 parent 3375285 commit 6e093b8

6 files changed

Lines changed: 235 additions & 89 deletions

File tree

‎ggml/src/ggml-vulkan/ggml-vulkan.cpp‎

Lines changed: 110 additions & 38 deletions
Original file line numberDiff line numberDiff line change
@@ -691,6 +691,7 @@ struct vk_device_struct {
691691
uint32_t coopmat_int_k;
692692

693693
bool coopmat2;
694+
bool coopmat2_bf16_support {};
694695
bool coopmat2_decode_vector;
695696

696697
bool pipeline_executable_properties_support {};
@@ -3139,7 +3140,7 @@ struct vk_fa_tuning_params {
31393140
};
31403141

31413142
static bool ggml_vk_flash_attn_scalar_shmem_support(const vk_device& device, const vk_fa_tuning_params& params, uint32_t hsk, uint32_t hsv, bool f32acc, ggml_type k_type, ggml_type v_type);
3142-
static bool ggml_vk_flash_attn_coopmat_shmem_support(const vk_device& device, const vk_fa_tuning_params& params, uint32_t hsk, uint32_t hsv, bool f32acc);
3143+
static bool ggml_vk_flash_attn_coopmat_shmem_support(const vk_device& device, const vk_fa_tuning_params& params, uint32_t hsk, uint32_t hsv, bool f32acc, ggml_type k_type = GGML_TYPE_F16);
31433144

31443145
static vk_fa_tuning_params get_fa_tuning_params_scalar(const vk_device& device, uint32_t hsk, uint32_t hsv, uint32_t n_rows, uint32_t n_kv, ggml_type k_type, ggml_type v_type, bool f32acc) {
31453146

@@ -3279,6 +3280,13 @@ static vk_fa_tuning_params get_fa_tuning_params(const vk_device& device, uint32_
32793280
FaCodePath path = device->coopmat2 ? FA_COOPMAT2 :
32803281
device->coopmat1_fa_support ? FA_COOPMAT1 : FA_SCALAR;
32813282

3283+
if (path == FA_COOPMAT2 && k_type == GGML_TYPE_BF16 && !device->coopmat2_bf16_support) {
3284+
path = FA_COOPMAT1;
3285+
}
3286+
if (path == FA_COOPMAT1 && k_type == GGML_TYPE_BF16 && !device->coopmat_bf16_support) {
3287+
path = FA_SCALAR;
3288+
}
3289+
32823290
if (path == FA_COOPMAT1 && device->architecture == vk_device_architecture::NVIDIA_TURING) {
32833291
// Nvidia compiler bug, see https://github.com/ggml-org/llama.cpp/pull/19075#issuecomment-3820716090
32843292
path = FA_SCALAR;
@@ -3288,7 +3296,7 @@ static vk_fa_tuning_params get_fa_tuning_params(const vk_device& device, uint32_
32883296
bool shape_ok = (f32acc && device->coopmat_support_16x16x16_f32acc) ||
32893297
(!f32acc && device->coopmat_support_16x16x16_f16acc);
32903298
const vk_fa_tuning_params params = get_fa_tuning_params_coopmat1(device, hsk, hsv, n_rows, n_kv, k_type, v_type, f32acc);
3291-
bool shmem_ok = ggml_vk_flash_attn_coopmat_shmem_support(device, params, hsk, hsv, f32acc);
3299+
bool shmem_ok = ggml_vk_flash_attn_coopmat_shmem_support(device, params, hsk, hsv, f32acc, k_type);
32923300

32933301
if (!shape_ok || !shmem_ok) {
32943302
path = FA_SCALAR;
@@ -3334,8 +3342,8 @@ static vk_fa_pipeline_state get_fa_pipeline_state(const vk_device& device, const
33343342

33353343
static std::vector<uint32_t> get_fa_spec_constants(const vk_fa_pipeline_state& state) {
33363344
const auto fa_block_bytes = [](ggml_type t) -> uint32_t {
3337-
// decodeBufF32 uses a block of vec4s for a better memory access pattern.
3338-
return t == GGML_TYPE_F32 ? 16u : (uint32_t) ggml_type_size(t);
3345+
if (t == GGML_TYPE_F32) return 16u;
3346+
return (uint32_t) ggml_type_size(t);
33393347
};
33403348
return {
33413349
/* 0 WorkGroupSize */ state.workgroup_size,
@@ -3849,10 +3857,16 @@ static void ggml_vk_load_shaders(vk_device& device) {
38493857
const uint32_t fa_sgs = fa.first.subgroup_size;
38503858
const bool fa_ds = fa.first.subgroup_size == 0;
38513859

3860+
const bool bf16_kv = fa.first.k_type == GGML_TYPE_BF16;
38523861
const bool use_mmq = ggml_vk_fa_scalar_uses_mmq(device, fa.first.k_type);
38533862
const void * spv_data = nullptr;
38543863
size_t spv_size = 0;
3855-
if (use_mmq) {
3864+
const char *name = nullptr;
3865+
if (bf16_kv) {
3866+
spv_data = flash_attn_f32_f16_fp32_data;
3867+
spv_size = flash_attn_f32_f16_fp32_len;
3868+
name = aligned ? "flash_attn_f32_bf16_aligned" : "flash_attn_f32_bf16";
3869+
} else if (use_mmq) {
38563870
#if defined(GGML_VULKAN_INTEGER_DOT_GLSLC_SUPPORT)
38573871
if (device->fp16) {
38583872
if (f32acc) { spv_data = flash_attn_f32_f16_int8_data; spv_size = flash_attn_f32_f16_int8_len; }
@@ -3862,6 +3876,7 @@ static void ggml_vk_load_shaders(vk_device& device) {
38623876
spv_size = flash_attn_f32_f16_fp32_int8_len;
38633877
}
38643878
#endif
3879+
name = aligned ? "flash_attn_f32_f16_aligned" : "flash_attn_f32_f16";
38653880
} else {
38663881
if (device->fp16) {
38673882
if (f32acc) { spv_data = flash_attn_f32_f16_data; spv_size = flash_attn_f32_f16_len; }
@@ -3870,8 +3885,8 @@ static void ggml_vk_load_shaders(vk_device& device) {
38703885
spv_data = flash_attn_f32_f16_fp32_data;
38713886
spv_size = flash_attn_f32_f16_fp32_len;
38723887
}
3888+
name = aligned ? "flash_attn_f32_f16_aligned" : "flash_attn_f32_f16";
38733889
}
3874-
const char *name = aligned ? "flash_attn_f32_f16_aligned" : "flash_attn_f32_f16";
38753890
ggml_vk_create_pipeline(device, fa.second, name, spv_size, spv_data, "main", 7,
38763891
sizeof(vk_flash_attn_push_constants), {Br, 1, 1},
38773892
get_fa_spec_constants(fa.first), aligned ? Bc : 1, true,
@@ -3889,11 +3904,25 @@ static void ggml_vk_load_shaders(vk_device& device) {
38893904
const uint32_t fa_sgs = fa.first.subgroup_size;
38903905
const bool fa_ds = fa.first.subgroup_size == 0;
38913906

3907+
const bool bf16_kv = fa.first.k_type == GGML_TYPE_BF16;
3908+
38923909
const void * spv_data;
38933910
size_t spv_size;
3894-
if (f32acc) { spv_data = flash_attn_f32_f16_cm1_data; spv_size = flash_attn_f32_f16_cm1_len; }
3895-
else { spv_data = flash_attn_f32_f16_f16acc_cm1_data; spv_size = flash_attn_f32_f16_f16acc_cm1_len; }
3896-
const char *name = aligned ? "flash_attn_f32_f16_aligned_cm1" : "flash_attn_f32_f16_cm1";
3911+
const char *name;
3912+
if (bf16_kv) {
3913+
#if defined(VK_KHR_shader_bfloat16) && defined(GGML_VULKAN_BFLOAT16_GLSLC_SUPPORT)
3914+
if (!device->coopmat_bf16_support) continue;
3915+
spv_data = flash_attn_f32_f16_bf16_cm1_data;
3916+
spv_size = flash_attn_f32_f16_bf16_cm1_len;
3917+
name = aligned ? "flash_attn_f32_bf16_aligned_cm1" : "flash_attn_f32_bf16_cm1";
3918+
#else
3919+
continue;
3920+
#endif
3921+
} else {
3922+
if (f32acc) { spv_data = flash_attn_f32_f16_cm1_data; spv_size = flash_attn_f32_f16_cm1_len; }
3923+
else { spv_data = flash_attn_f32_f16_f16acc_cm1_data; spv_size = flash_attn_f32_f16_f16acc_cm1_len; }
3924+
name = aligned ? "flash_attn_f32_f16_aligned_cm1" : "flash_attn_f32_f16_cm1";
3925+
}
38973926
ggml_vk_create_pipeline(device, fa.second, name, spv_size, spv_data, "main", 7,
38983927
sizeof(vk_flash_attn_push_constants), {Br, 1, 1},
38993928
get_fa_spec_constants(fa.first), aligned ? Bc : 1, true,
@@ -3911,10 +3940,20 @@ static void ggml_vk_load_shaders(vk_device& device) {
39113940
const bool aligned = fa.first.aligned;
39123941
const bool f32acc = fa.first.f32acc;
39133942

3943+
const bool bf16_kv = fa.first.k_type == GGML_TYPE_BF16;
39143944
const void * spv_data;
39153945
size_t spv_size;
39163946
const char * name;
3917-
if (aligned) {
3947+
if (bf16_kv) {
3948+
#if defined(VK_KHR_shader_bfloat16) && defined(GGML_VULKAN_BFLOAT16_GLSLC_SUPPORT)
3949+
if (!device->coopmat2_bf16_support) continue;
3950+
spv_data = flash_attn_f32_f16_bf16_cm2_data;
3951+
spv_size = flash_attn_f32_f16_bf16_cm2_len;
3952+
name = aligned ? "flash_attn_f32_bf16_aligned_cm2" : "flash_attn_f32_bf16_cm2";
3953+
#else
3954+
continue;
3955+
#endif
3956+
} else if (aligned) {
39183957
if (f32acc) { spv_data = flash_attn_f32_f16_cm2_data; spv_size = flash_attn_f32_f16_cm2_len; name = "flash_attn_f32_f16_aligned_f32acc_cm2"; }
39193958
else { spv_data = flash_attn_f32_f16_f16acc_cm2_data; spv_size = flash_attn_f32_f16_f16acc_cm2_len; name = "flash_attn_f32_f16_aligned_f16acc_cm2"; }
39203959
} else {
@@ -5784,46 +5823,72 @@ static vk_device ggml_vk_get_device(size_t idx) {
57845823
found_fp16_256 = false,
57855824
found_fp32_128 = false,
57865825
found_fp32_256 = false;
5826+
bool found_bf16_128 = false,
5827+
found_bf16_256 = false;
57875828
// need to support fp16*fp16 with fp16/fp32 accumulator, for workgroupsize 128
57885829
// with 32x16x16 and 256 with 32x32x16.
57895830
for (auto &prop : flexible_dimensions) {
57905831
if (prop.saturatingAccumulation == VK_FALSE &&
5791-
prop.scope == VK_SCOPE_WORKGROUP_KHR &&
5792-
prop.AType == VK_COMPONENT_TYPE_FLOAT16_KHR &&
5793-
prop.BType == VK_COMPONENT_TYPE_FLOAT16_KHR) {
5794-
5795-
if (prop.workgroupInvocations == 128 &&
5796-
prop.MGranularity <= 32 &&
5797-
prop.NGranularity <= 16 &&
5798-
prop.KGranularity <= 16) {
5799-
if (prop.CType == VK_COMPONENT_TYPE_FLOAT16_KHR &&
5800-
prop.ResultType == VK_COMPONENT_TYPE_FLOAT16_KHR) {
5801-
found_fp16_128 = true;
5832+
prop.scope == VK_SCOPE_WORKGROUP_KHR) {
5833+
5834+
if (prop.AType == VK_COMPONENT_TYPE_FLOAT16_KHR &&
5835+
prop.BType == VK_COMPONENT_TYPE_FLOAT16_KHR) {
5836+
5837+
if (prop.workgroupInvocations == 128 &&
5838+
prop.MGranularity <= 32 &&
5839+
prop.NGranularity <= 16 &&
5840+
prop.KGranularity <= 16) {
5841+
if (prop.CType == VK_COMPONENT_TYPE_FLOAT16_KHR &&
5842+
prop.ResultType == VK_COMPONENT_TYPE_FLOAT16_KHR) {
5843+
found_fp16_128 = true;
5844+
}
5845+
if (prop.CType == VK_COMPONENT_TYPE_FLOAT32_KHR &&
5846+
prop.ResultType == VK_COMPONENT_TYPE_FLOAT32_KHR) {
5847+
found_fp32_128 = true;
5848+
}
58025849
}
5803-
if (prop.CType == VK_COMPONENT_TYPE_FLOAT32_KHR &&
5804-
prop.ResultType == VK_COMPONENT_TYPE_FLOAT32_KHR) {
5805-
found_fp32_128 = true;
5850+
if (prop.workgroupInvocations == 256 &&
5851+
prop.MGranularity <= 32 &&
5852+
prop.NGranularity <= 32 &&
5853+
prop.KGranularity <= 16) {
5854+
if (prop.CType == VK_COMPONENT_TYPE_FLOAT16_KHR &&
5855+
prop.ResultType == VK_COMPONENT_TYPE_FLOAT16_KHR) {
5856+
found_fp16_256 = true;
5857+
}
5858+
if (prop.CType == VK_COMPONENT_TYPE_FLOAT32_KHR &&
5859+
prop.ResultType == VK_COMPONENT_TYPE_FLOAT32_KHR) {
5860+
found_fp32_256 = true;
5861+
}
58065862
}
58075863
}
5808-
if (prop.workgroupInvocations == 256 &&
5809-
prop.MGranularity <= 32 &&
5810-
prop.NGranularity <= 32 &&
5811-
prop.KGranularity <= 16) {
5812-
if (prop.CType == VK_COMPONENT_TYPE_FLOAT16_KHR &&
5813-
prop.ResultType == VK_COMPONENT_TYPE_FLOAT16_KHR) {
5814-
found_fp16_256 = true;
5864+
5865+
#if defined(VK_KHR_shader_bfloat16) && defined(GGML_VULKAN_BFLOAT16_GLSLC_SUPPORT)
5866+
if (prop.AType == VK_COMPONENT_TYPE_BFLOAT16_KHR &&
5867+
prop.BType == VK_COMPONENT_TYPE_BFLOAT16_KHR &&
5868+
prop.CType == VK_COMPONENT_TYPE_FLOAT32_KHR &&
5869+
prop.ResultType == VK_COMPONENT_TYPE_FLOAT32_KHR) {
5870+
5871+
if (prop.workgroupInvocations == 128 &&
5872+
prop.MGranularity <= 32 &&
5873+
prop.NGranularity <= 16 &&
5874+
prop.KGranularity <= 16) {
5875+
found_bf16_128 = true;
58155876
}
5816-
if (prop.CType == VK_COMPONENT_TYPE_FLOAT32_KHR &&
5817-
prop.ResultType == VK_COMPONENT_TYPE_FLOAT32_KHR) {
5818-
found_fp32_256 = true;
5877+
if (prop.workgroupInvocations == 256 &&
5878+
prop.MGranularity <= 32 &&
5879+
prop.NGranularity <= 32 &&
5880+
prop.KGranularity <= 16) {
5881+
found_bf16_256 = true;
58195882
}
58205883
}
5884+
#endif
58215885
}
58225886
}
58235887
if (found_fp16_128 && found_fp16_256 &&
58245888
found_fp32_128 && found_fp32_256 &&
58255889
coopmat2_props.cooperativeMatrixFlexibleDimensionsMaxDimension >= 512) {
58265890
device->coopmat2 = true;
5891+
device->coopmat2_bf16_support = found_bf16_128 && found_bf16_256;
58275892
device->coopmat2_decode_vector = coopmat2_decode_vector_support && coopmat2_decode_vector_features.cooperativeMatrixDecodeVector;
58285893
}
58295894
}
@@ -9448,7 +9513,8 @@ static bool ggml_vk_flash_attn_scalar_shmem_support(const vk_device& device, con
94489513
const uint32_t Br = params.block_rows;
94499514
const uint32_t Bc = params.block_cols;
94509515

9451-
const uint32_t float_type_size = device->fp16 ? sizeof(ggml_fp16_t) : sizeof(float);
9516+
// BF16 uses the fp32 shader (FLOAT_TYPE=float)
9517+
const uint32_t float_type_size = (device->fp16 && k_type != GGML_TYPE_BF16) ? sizeof(ggml_fp16_t) : sizeof(float);
94529518

94539519
const bool mmq = ggml_vk_fa_scalar_uses_mmq(device, k_type);
94549520

@@ -9489,7 +9555,7 @@ static bool ggml_vk_flash_attn_scalar_shmem_support(const vk_device& device, con
94899555
return supported;
94909556
}
94919557

9492-
static bool ggml_vk_flash_attn_coopmat_shmem_support(const vk_device& device, const vk_fa_tuning_params& params, uint32_t hsk, uint32_t hsv, bool f32acc) {
9558+
static bool ggml_vk_flash_attn_coopmat_shmem_support(const vk_device& device, const vk_fa_tuning_params& params, uint32_t hsk, uint32_t hsv, bool f32acc, ggml_type k_type) {
94939559
// Needs to be kept up to date on shader changes
94949560
const uint32_t Br = params.block_rows;
94959561
const uint32_t Bc = params.block_cols;
@@ -9519,8 +9585,10 @@ static bool ggml_vk_flash_attn_coopmat_shmem_support(const vk_device& device, co
95199585
const uint32_t vsh_stride = MatBc / 4 * row_split;
95209586
const uint32_t ksh = ((kvshstride >= vsh_stride) ? (Bc * kvshstride) : (Bc * vsh_stride)) * f16vec4;
95219587

9588+
// BF16 PVMat accumulator is f32 (no bf16 accumulator support), so pvsh is vec4 (16 bytes)
9589+
const uint32_t pvsh_elem_size = (k_type == GGML_TYPE_BF16) ? 16u : f16vec4;
95229590
const uint32_t osh_stride = params.row_split * MatBr / 4;
9523-
const uint32_t pvsh = MatBc * osh_stride * f16vec4;
9591+
const uint32_t pvsh = MatBc * osh_stride * pvsh_elem_size;
95249592

95259593
const uint32_t slope = Br * acctype;
95269594

@@ -9589,7 +9657,7 @@ static void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx
95899657
uint32_t workgroups_y = (uint32_t)neq2;
95909658
uint32_t workgroups_z = (uint32_t)neq3;
95919659

9592-
const bool f32acc = !ctx->device->fp16 || dst->op_params[3] == GGML_PREC_F32;
9660+
const bool f32acc = !ctx->device->fp16 || dst->op_params[3] == GGML_PREC_F32 || k->type == GGML_TYPE_BF16;
95939661

95949662
// For scalar/coopmat1 FA, we can use the "large" size to accommodate qga.
95959663
// For coopmat2 FA, we always use the small size (which is still pretty large for gqa).
@@ -16400,6 +16468,7 @@ static bool ggml_backend_vk_device_supports_op(ggml_backend_dev_t dev, const ggm
1640016468
switch (t) {
1640116469
case GGML_TYPE_F32:
1640216470
case GGML_TYPE_F16:
16471+
case GGML_TYPE_BF16:
1640316472
case GGML_TYPE_Q8_0:
1640416473
case GGML_TYPE_Q5_1:
1640516474
case GGML_TYPE_Q5_0:
@@ -16415,6 +16484,9 @@ static bool ggml_backend_vk_device_supports_op(ggml_backend_dev_t dev, const ggm
1641516484
if (!fa_kv_ok(op->src[1]->type) || !fa_kv_ok(op->src[2]->type)) {
1641616485
return false;
1641716486
}
16487+
if ((op->src[1]->type == GGML_TYPE_BF16) != (op->src[2]->type == GGML_TYPE_BF16)) {
16488+
return false;
16489+
}
1641816490
if (!coopmat2 && !(device->subgroup_shuffle && device->subgroup_vote)) {
1641916491
// scalar/coopmat1 FA uses subgroupShuffle/subgroupAll
1642016492
return false;

‎ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_base.glsl‎

Lines changed: 11 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -97,8 +97,17 @@ layout (binding = 6) readonly buffer MO {uint32_t data_mask_opt[];};
9797
#define FA_TYPE_Q5_0 6u
9898
#define FA_TYPE_Q5_1 7u
9999
#define FA_TYPE_Q8_0 8u
100+
#define FA_TYPE_BF16 30u
100101
#define FA_TYPE_Q1_0 41u
101102

103+
#if defined(BFLOAT16)
104+
#define O_TYPE float
105+
#define O_TYPEV4 vec4
106+
#else
107+
#define O_TYPE FLOAT_TYPE
108+
#define O_TYPEV4 FLOAT_TYPEV4
109+
#endif
110+
102111
// Number of matrix elements per buffer block, derived from the K/V type spec
103112
// constant. F32 is treated as a vec4 "block" of 4 floats. F16 uses block size 1
104113
// and bypasses the dequant path entirely. Quants follow their ggml block sizes.
@@ -111,6 +120,7 @@ uint fa_block_elems(uint ty) {
111120
case FA_TYPE_Q5_0: return uint(QUANT_K_Q5_0);
112121
case FA_TYPE_Q5_1: return uint(QUANT_K_Q5_1);
113122
case FA_TYPE_Q8_0: return uint(QUANT_K_Q8_0);
123+
case FA_TYPE_BF16: return 1u;
114124
case FA_TYPE_Q1_0: return uint(QUANT_K_Q1_0); // cm2-only, harmless elsewhere
115125
default: return 1u;
116126
}
@@ -248,7 +258,7 @@ const float FATTN_KQ_MAX_OFFSET = 3.0f*0.6931f;
248258

249259
// Store the output when doing grouped query attention.
250260
// Rows index by Q's dimension 2, and the first N rows are valid.
251-
void gqaStore(const in uint32_t r, const in uint32_t c, const in FLOAT_TYPEV4 elems, const in uint32_t o_offset, const in uint32_t iq2, const in uint32_t N)
261+
void gqaStore(const in uint32_t r, const in uint32_t c, const in O_TYPEV4 elems, const in uint32_t o_offset, const in uint32_t iq2, const in uint32_t N)
252262
{
253263
uint32_t offset = (iq2 + r) * HSV / 4 + c;
254264
data_ov4[o_offset + offset] = D_TYPEV4(elems);

0 commit comments

Comments
 (0)