Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
0cc4m merged 12 commits into from
May 30, 2026
Merged

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
12 commits
Select commit Hold shift + click to select a range
b97a6c4
vulkan: add flash attention bf16 kv support
0cc4m May 19, 2026
8f486e9
vulkan: bf16 FA coopmat1 support
0cc4m May 19, 2026
bce8d9e
vulkan: bf16 FA coopmat2 support
0cc4m May 19, 2026
58d2e83
fix FA bf16 f32 fallback
0cc4m May 19, 2026
53356b4
fix FA bf16 coopmat1 shader
0cc4m May 20, 2026
65a14bf
fix FA bf16 coopmat2 shader
0cc4m May 20, 2026
1383a83
code cleanup
0cc4m May 20, 2026
65e493e
cleanup comment change
0cc4m May 20, 2026
141c739
address feedback
0cc4m May 21, 2026
51063b1
add O_TYPE for cm2 FA
0cc4m May 21, 2026
1b71ecb
use O_TYPE for gqaStore function
0cc4m May 21, 2026
48f0c0a
reduce BFLOAT16 ifdefs
0cc4m May 26, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
148 changes: 110 additions & 38 deletions ggml/src/ggml-vulkan/ggml-vulkan.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -691,6 +691,7 @@ struct vk_device_struct {
uint32_t coopmat_int_k;

bool coopmat2;
bool coopmat2_bf16_support {};
bool coopmat2_decode_vector;

bool pipeline_executable_properties_support {};
Expand Down Expand Up @@ -3139,7 +3140,7 @@ struct vk_fa_tuning_params {
};

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);
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);
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);

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) {

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

if (path == FA_COOPMAT2 && k_type == GGML_TYPE_BF16 && !device->coopmat2_bf16_support) {
path = FA_COOPMAT1;
}
if (path == FA_COOPMAT1 && k_type == GGML_TYPE_BF16 && !device->coopmat_bf16_support) {
path = FA_SCALAR;
}
Comment thread
jeffbolznv marked this conversation as resolved.

if (path == FA_COOPMAT1 && device->architecture == vk_device_architecture::NVIDIA_TURING) {
// Nvidia compiler bug, see https://github.com/ggml-org/llama.cpp/pull/19075#issuecomment-3820716090
path = FA_SCALAR;
Expand All @@ -3288,7 +3296,7 @@ static vk_fa_tuning_params get_fa_tuning_params(const vk_device& device, uint32_
bool shape_ok = (f32acc && device->coopmat_support_16x16x16_f32acc) ||
(!f32acc && device->coopmat_support_16x16x16_f16acc);
const vk_fa_tuning_params params = get_fa_tuning_params_coopmat1(device, hsk, hsv, n_rows, n_kv, k_type, v_type, f32acc);
bool shmem_ok = ggml_vk_flash_attn_coopmat_shmem_support(device, params, hsk, hsv, f32acc);
bool shmem_ok = ggml_vk_flash_attn_coopmat_shmem_support(device, params, hsk, hsv, f32acc, k_type);

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

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

const bool bf16_kv = fa.first.k_type == GGML_TYPE_BF16;
const bool use_mmq = ggml_vk_fa_scalar_uses_mmq(device, fa.first.k_type);
const void * spv_data = nullptr;
size_t spv_size = 0;
if (use_mmq) {
const char *name = nullptr;
if (bf16_kv) {
spv_data = flash_attn_f32_f16_fp32_data;
spv_size = flash_attn_f32_f16_fp32_len;
name = aligned ? "flash_attn_f32_bf16_aligned" : "flash_attn_f32_bf16";
} else if (use_mmq) {
#if defined(GGML_VULKAN_INTEGER_DOT_GLSLC_SUPPORT)
if (device->fp16) {
if (f32acc) { spv_data = flash_attn_f32_f16_int8_data; spv_size = flash_attn_f32_f16_int8_len; }
Expand All @@ -3862,6 +3876,7 @@ static void ggml_vk_load_shaders(vk_device& device) {
spv_size = flash_attn_f32_f16_fp32_int8_len;
}
#endif
name = aligned ? "flash_attn_f32_f16_aligned" : "flash_attn_f32_f16";
} else {
if (device->fp16) {
if (f32acc) { spv_data = flash_attn_f32_f16_data; spv_size = flash_attn_f32_f16_len; }
Expand All @@ -3870,8 +3885,8 @@ static void ggml_vk_load_shaders(vk_device& device) {
spv_data = flash_attn_f32_f16_fp32_data;
spv_size = flash_attn_f32_f16_fp32_len;
}
name = aligned ? "flash_attn_f32_f16_aligned" : "flash_attn_f32_f16";
}
const char *name = aligned ? "flash_attn_f32_f16_aligned" : "flash_attn_f32_f16";
ggml_vk_create_pipeline(device, fa.second, name, spv_size, spv_data, "main", 7,
sizeof(vk_flash_attn_push_constants), {Br, 1, 1},
get_fa_spec_constants(fa.first), aligned ? Bc : 1, true,
Expand All @@ -3889,11 +3904,25 @@ static void ggml_vk_load_shaders(vk_device& device) {
const uint32_t fa_sgs = fa.first.subgroup_size;
const bool fa_ds = fa.first.subgroup_size == 0;

const bool bf16_kv = fa.first.k_type == GGML_TYPE_BF16;

const void * spv_data;
size_t spv_size;
if (f32acc) { spv_data = flash_attn_f32_f16_cm1_data; spv_size = flash_attn_f32_f16_cm1_len; }
else { spv_data = flash_attn_f32_f16_f16acc_cm1_data; spv_size = flash_attn_f32_f16_f16acc_cm1_len; }
const char *name = aligned ? "flash_attn_f32_f16_aligned_cm1" : "flash_attn_f32_f16_cm1";
const char *name;
if (bf16_kv) {
#if defined(VK_KHR_shader_bfloat16) && defined(GGML_VULKAN_BFLOAT16_GLSLC_SUPPORT)
if (!device->coopmat_bf16_support) continue;
spv_data = flash_attn_f32_f16_bf16_cm1_data;
spv_size = flash_attn_f32_f16_bf16_cm1_len;
name = aligned ? "flash_attn_f32_bf16_aligned_cm1" : "flash_attn_f32_bf16_cm1";
#else
continue;
#endif
} else {
if (f32acc) { spv_data = flash_attn_f32_f16_cm1_data; spv_size = flash_attn_f32_f16_cm1_len; }
else { spv_data = flash_attn_f32_f16_f16acc_cm1_data; spv_size = flash_attn_f32_f16_f16acc_cm1_len; }
name = aligned ? "flash_attn_f32_f16_aligned_cm1" : "flash_attn_f32_f16_cm1";
}
ggml_vk_create_pipeline(device, fa.second, name, spv_size, spv_data, "main", 7,
sizeof(vk_flash_attn_push_constants), {Br, 1, 1},
get_fa_spec_constants(fa.first), aligned ? Bc : 1, true,
Expand All @@ -3911,10 +3940,20 @@ static void ggml_vk_load_shaders(vk_device& device) {
const bool aligned = fa.first.aligned;
const bool f32acc = fa.first.f32acc;

const bool bf16_kv = fa.first.k_type == GGML_TYPE_BF16;
const void * spv_data;
size_t spv_size;
const char * name;
if (aligned) {
if (bf16_kv) {
#if defined(VK_KHR_shader_bfloat16) && defined(GGML_VULKAN_BFLOAT16_GLSLC_SUPPORT)
if (!device->coopmat2_bf16_support) continue;
spv_data = flash_attn_f32_f16_bf16_cm2_data;
spv_size = flash_attn_f32_f16_bf16_cm2_len;
name = aligned ? "flash_attn_f32_bf16_aligned_cm2" : "flash_attn_f32_bf16_cm2";
#else
continue;
#endif
} else if (aligned) {
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"; }
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"; }
} else {
Expand Down Expand Up @@ -5784,46 +5823,72 @@ static vk_device ggml_vk_get_device(size_t idx) {
found_fp16_256 = false,
found_fp32_128 = false,
found_fp32_256 = false;
bool found_bf16_128 = false,
found_bf16_256 = false;
// need to support fp16*fp16 with fp16/fp32 accumulator, for workgroupsize 128
// with 32x16x16 and 256 with 32x32x16.
for (auto &prop : flexible_dimensions) {
if (prop.saturatingAccumulation == VK_FALSE &&
prop.scope == VK_SCOPE_WORKGROUP_KHR &&
prop.AType == VK_COMPONENT_TYPE_FLOAT16_KHR &&
prop.BType == VK_COMPONENT_TYPE_FLOAT16_KHR) {

if (prop.workgroupInvocations == 128 &&
prop.MGranularity <= 32 &&
prop.NGranularity <= 16 &&
prop.KGranularity <= 16) {
if (prop.CType == VK_COMPONENT_TYPE_FLOAT16_KHR &&
prop.ResultType == VK_COMPONENT_TYPE_FLOAT16_KHR) {
found_fp16_128 = true;
prop.scope == VK_SCOPE_WORKGROUP_KHR) {

if (prop.AType == VK_COMPONENT_TYPE_FLOAT16_KHR &&
prop.BType == VK_COMPONENT_TYPE_FLOAT16_KHR) {

if (prop.workgroupInvocations == 128 &&
prop.MGranularity <= 32 &&
prop.NGranularity <= 16 &&
prop.KGranularity <= 16) {
if (prop.CType == VK_COMPONENT_TYPE_FLOAT16_KHR &&
prop.ResultType == VK_COMPONENT_TYPE_FLOAT16_KHR) {
found_fp16_128 = true;
}
if (prop.CType == VK_COMPONENT_TYPE_FLOAT32_KHR &&
prop.ResultType == VK_COMPONENT_TYPE_FLOAT32_KHR) {
found_fp32_128 = true;
}
}
if (prop.CType == VK_COMPONENT_TYPE_FLOAT32_KHR &&
prop.ResultType == VK_COMPONENT_TYPE_FLOAT32_KHR) {
found_fp32_128 = true;
if (prop.workgroupInvocations == 256 &&
prop.MGranularity <= 32 &&
prop.NGranularity <= 32 &&
prop.KGranularity <= 16) {
if (prop.CType == VK_COMPONENT_TYPE_FLOAT16_KHR &&
prop.ResultType == VK_COMPONENT_TYPE_FLOAT16_KHR) {
found_fp16_256 = true;
}
if (prop.CType == VK_COMPONENT_TYPE_FLOAT32_KHR &&
prop.ResultType == VK_COMPONENT_TYPE_FLOAT32_KHR) {
found_fp32_256 = true;
}
}
}
if (prop.workgroupInvocations == 256 &&
prop.MGranularity <= 32 &&
prop.NGranularity <= 32 &&
prop.KGranularity <= 16) {
if (prop.CType == VK_COMPONENT_TYPE_FLOAT16_KHR &&
prop.ResultType == VK_COMPONENT_TYPE_FLOAT16_KHR) {
found_fp16_256 = true;

#if defined(VK_KHR_shader_bfloat16) && defined(GGML_VULKAN_BFLOAT16_GLSLC_SUPPORT)
if (prop.AType == VK_COMPONENT_TYPE_BFLOAT16_KHR &&
prop.BType == VK_COMPONENT_TYPE_BFLOAT16_KHR &&
prop.CType == VK_COMPONENT_TYPE_FLOAT32_KHR &&
prop.ResultType == VK_COMPONENT_TYPE_FLOAT32_KHR) {

if (prop.workgroupInvocations == 128 &&
prop.MGranularity <= 32 &&
prop.NGranularity <= 16 &&
prop.KGranularity <= 16) {
found_bf16_128 = true;
}
if (prop.CType == VK_COMPONENT_TYPE_FLOAT32_KHR &&
prop.ResultType == VK_COMPONENT_TYPE_FLOAT32_KHR) {
found_fp32_256 = true;
if (prop.workgroupInvocations == 256 &&
prop.MGranularity <= 32 &&
prop.NGranularity <= 32 &&
prop.KGranularity <= 16) {
found_bf16_256 = true;
}
}
#endif
}
}
if (found_fp16_128 && found_fp16_256 &&
found_fp32_128 && found_fp32_256 &&
coopmat2_props.cooperativeMatrixFlexibleDimensionsMaxDimension >= 512) {
device->coopmat2 = true;
device->coopmat2_bf16_support = found_bf16_128 && found_bf16_256;
device->coopmat2_decode_vector = coopmat2_decode_vector_support && coopmat2_decode_vector_features.cooperativeMatrixDecodeVector;
}
}
Expand Down Expand Up @@ -9448,7 +9513,8 @@ static bool ggml_vk_flash_attn_scalar_shmem_support(const vk_device& device, con
const uint32_t Br = params.block_rows;
const uint32_t Bc = params.block_cols;

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

const bool mmq = ggml_vk_fa_scalar_uses_mmq(device, k_type);

Expand Down Expand Up @@ -9489,7 +9555,7 @@ static bool ggml_vk_flash_attn_scalar_shmem_support(const vk_device& device, con
return supported;
}

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) {
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) {
// Needs to be kept up to date on shader changes
const uint32_t Br = params.block_rows;
const uint32_t Bc = params.block_cols;
Expand Down Expand Up @@ -9519,8 +9585,10 @@ static bool ggml_vk_flash_attn_coopmat_shmem_support(const vk_device& device, co
const uint32_t vsh_stride = MatBc / 4 * row_split;
const uint32_t ksh = ((kvshstride >= vsh_stride) ? (Bc * kvshstride) : (Bc * vsh_stride)) * f16vec4;

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

const uint32_t slope = Br * acctype;

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

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

// For scalar/coopmat1 FA, we can use the "large" size to accommodate qga.
// For coopmat2 FA, we always use the small size (which is still pretty large for gqa).
Expand Down Expand Up @@ -16400,6 +16468,7 @@ static bool ggml_backend_vk_device_supports_op(ggml_backend_dev_t dev, const ggm
switch (t) {
case GGML_TYPE_F32:
case GGML_TYPE_F16:
case GGML_TYPE_BF16:
case GGML_TYPE_Q8_0:
case GGML_TYPE_Q5_1:
case GGML_TYPE_Q5_0:
Expand All @@ -16415,6 +16484,9 @@ static bool ggml_backend_vk_device_supports_op(ggml_backend_dev_t dev, const ggm
if (!fa_kv_ok(op->src[1]->type) || !fa_kv_ok(op->src[2]->type)) {
return false;
}
if ((op->src[1]->type == GGML_TYPE_BF16) != (op->src[2]->type == GGML_TYPE_BF16)) {
Comment thread
jeffbolznv marked this conversation as resolved.
return false;
}
if (!coopmat2 && !(device->subgroup_shuffle && device->subgroup_vote)) {
// scalar/coopmat1 FA uses subgroupShuffle/subgroupAll
return false;
Expand Down
12 changes: 11 additions & 1 deletion ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_base.glsl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -97,8 +97,17 @@ layout (binding = 6) readonly buffer MO {uint32_t data_mask_opt[];};
#define FA_TYPE_Q5_0 6u
#define FA_TYPE_Q5_1 7u
#define FA_TYPE_Q8_0 8u
#define FA_TYPE_BF16 30u
#define FA_TYPE_Q1_0 41u

#if defined(BFLOAT16)
#define O_TYPE float
#define O_TYPEV4 vec4
#else
#define O_TYPE FLOAT_TYPE
#define O_TYPEV4 FLOAT_TYPEV4
#endif

// Number of matrix elements per buffer block, derived from the K/V type spec
// constant. F32 is treated as a vec4 "block" of 4 floats. F16 uses block size 1
// and bypasses the dequant path entirely. Quants follow their ggml block sizes.
Expand All @@ -111,6 +120,7 @@ uint fa_block_elems(uint ty) {
case FA_TYPE_Q5_0: return uint(QUANT_K_Q5_0);
case FA_TYPE_Q5_1: return uint(QUANT_K_Q5_1);
case FA_TYPE_Q8_0: return uint(QUANT_K_Q8_0);
case FA_TYPE_BF16: return 1u;
case FA_TYPE_Q1_0: return uint(QUANT_K_Q1_0); // cm2-only, harmless elsewhere
default: return 1u;
}
Expand Down Expand Up @@ -248,7 +258,7 @@ const float FATTN_KQ_MAX_OFFSET = 3.0f*0.6931f;

// Store the output when doing grouped query attention.
// Rows index by Q's dimension 2, and the first N rows are valid.
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)
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)
{
uint32_t offset = (iq2 + r) * HSV / 4 + c;
data_ov4[o_offset + offset] = D_TYPEV4(elems);
Expand Down
Loading
Loading