Lumpiasty/llama.cpp
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

vulkan: add Flash Attention support for BFloat16 KV cache (#23420)

Browse Source 
* 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
This commit is contained in:
Ruben Ortlam
2026-05-30 10:39:31 +02:00
committed by GitHub
parent 337528571d
commit 6e093b80ea
6 changed files with 234 additions and 88 deletions
Download Patch File Download Diff File Expand all files Collapse all files
ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp
+22
View File
@@ -662,6 +662,28 @@ void process_shaders() {
}
}
const std::map<std::string, std::string> fa_bf16_dict = {
{"FLOAT_TYPE", "bfloat16_t"},
{"FLOAT_TYPEV2", "bf16vec2"},
{"FLOAT_TYPEV4", "bf16vec4"},
{"ACC_TYPE", "float"},
{"ACC_TYPEV2", "vec2"},
{"ACC_TYPEV4", "vec4"},
{"BFLOAT16", "1"},
};
#if defined(GGML_VULKAN_BFLOAT16_GLSLC_SUPPORT) && defined(GGML_VULKAN_COOPMAT_GLSLC_SUPPORT)
string_to_spv("flash_attn_f32_f16_bf16", "flash_attn_cm1.comp",
merge_maps(fa_bf16_dict, {{"Q_TYPE", "float"}, {"D_TYPE", "float"}, {"D_TYPEV4", "vec4"}, {"COOPMAT", "1"}}),
true, true, false, false);
#endif
#if defined(GGML_VULKAN_BFLOAT16_GLSLC_SUPPORT) && defined(GGML_VULKAN_COOPMAT2_GLSLC_SUPPORT)
string_to_spv("flash_attn_f32_f16_bf16", "flash_attn_cm2.comp",
merge_maps(fa_bf16_dict, {{"Q_TYPE", "float"}, {"D_TYPE", "float"}, {"D_TYPEV4", "vec4"}}),
true, false, true, false);
#endif
std::map<std::string, std::string> base_dict = {{"FLOAT_TYPE", "float"}, {"FLOAT_TYPEV2", "vec2"}};
for (const auto& tname : type_names) {