Skip to content

Commit 0f53c83

Browse files
jeffbolznvbaramofme
jeffbolznv
authored and
baramofme
committed
vulkan: Support asymmetric FA in scalar/mmq/coopmat1 paths (ggml-org#22589)
1 parent bd2ad91 commit 0f53c83

4 files changed

Lines changed: 39 additions & 120 deletions

File tree

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

Lines changed: 24 additions & 44 deletions
Original file line numberDiff line numberDiff line change
@@ -3098,12 +3098,6 @@ static vk_fa_tuning_params get_fa_tuning_params_coopmat2(const vk_device& device
30983098
}
30993099

31003100
static vk_fa_tuning_params get_fa_tuning_params(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) {
3101-
// Mixed K/V is only implemented on the coopmat2 (flash_attn_cm2) path; never use scalar/cm1.
3102-
if (k_type != v_type) {
3103-
GGML_ASSERT(device->coopmat2);
3104-
return get_fa_tuning_params_coopmat2(device, hsk, hsv, n_rows, n_kv, k_type, v_type, f32acc);
3105-
}
3106-
31073101
FaCodePath path = device->coopmat2 ? FA_COOPMAT2 :
31083102
device->coopmat1_fa_support ? FA_COOPMAT1 : FA_SCALAR;
31093103

@@ -3115,7 +3109,7 @@ static vk_fa_tuning_params get_fa_tuning_params(const vk_device& device, uint32_
31153109
if (path == FA_COOPMAT1) {
31163110
bool shape_ok = (f32acc && device->coopmat_support_16x16x16_f32acc) ||
31173111
(!f32acc && device->coopmat_support_16x16x16_f16acc);
3118-
const vk_fa_tuning_params params = get_fa_tuning_params_coopmat1(device, hsk, hsv, n_rows, n_kv, k_type, f32acc);
3112+
const vk_fa_tuning_params params = get_fa_tuning_params_coopmat1(device, hsk, hsv, n_rows, n_kv, k_type, v_type, f32acc);
31193113
bool shmem_ok = ggml_vk_flash_attn_coopmat_shmem_support(device, params, hsk, hsv, f32acc);
31203114

31213115
if (!shape_ok || !shmem_ok) {
@@ -3135,9 +3129,9 @@ static vk_fa_tuning_params get_fa_tuning_params(const vk_device& device, uint32_
31353129

31363130
switch (path) {
31373131
case FA_SCALAR:
3138-
return get_fa_tuning_params_scalar(device, hsk, hsv, n_rows, n_kv, k_type, f32acc);
3132+
return get_fa_tuning_params_scalar(device, hsk, hsv, n_rows, n_kv, k_type, v_type, f32acc);
31393133
case FA_COOPMAT1:
3140-
return get_fa_tuning_params_coopmat1(device, hsk, hsv, n_rows, n_kv, k_type, f32acc);
3134+
return get_fa_tuning_params_coopmat1(device, hsk, hsv, n_rows, n_kv, k_type, v_type, f32acc);
31413135
case FA_COOPMAT2:
31423136
return get_fa_tuning_params_coopmat2(device, hsk, hsv, n_rows, n_kv, k_type, v_type, f32acc);
31433137
default:
@@ -3731,35 +3725,29 @@ static void ggml_vk_load_shaders(vk_device& device) {
37313725
#endif
37323726

37333727
#if defined(VK_NV_cooperative_matrix2) && defined(GGML_VULKAN_COOPMAT2_GLSLC_SUPPORT)
3734-
#define CREATE_FA_CM2_MIXED() \
3735-
for (int fa_k_ty = 0; fa_k_ty < (int)GGML_TYPE_COUNT; ++fa_k_ty) { \
3736-
for (auto &fa : device->pipeline_flash_attn_f32_f16[fa_k_ty]) { \
3737-
FaCodePath path = fa.first.path; \
3738-
uint32_t Br = fa.first.Br; \
3739-
uint32_t Bc = fa.first.Bc; \
3740-
bool aligned = fa.first.aligned; \
3741-
bool f32acc = fa.first.f32acc; \
3742-
if (path == FA_COOPMAT2) { \
3743-
if (aligned) { \
3744-
if (f32acc) { \
3745-
ggml_vk_create_pipeline(device, fa.second, "flash_attn_f32_f16_mixed_aligned_f32acc_cm2", flash_attn_f32_f16_mixed_cm2_len, flash_attn_f32_f16_mixed_cm2_data, "main", 7, sizeof(vk_flash_attn_push_constants), {Br, 1, 1}, get_fa_spec_constants(fa.first), Bc, true, false, 0); \
3746-
} else { \
3747-
ggml_vk_create_pipeline(device, fa.second, "flash_attn_f32_f16_mixed_aligned_f16acc_cm2", flash_attn_f32_f16_mixed_f16acc_cm2_len, flash_attn_f32_f16_mixed_f16acc_cm2_data, "main", 7, sizeof(vk_flash_attn_push_constants), {Br, 1, 1}, get_fa_spec_constants(fa.first), Bc, true, false, 0); \
3748-
} \
3749-
} else { \
3750-
if (f32acc) { \
3751-
ggml_vk_create_pipeline(device, fa.second, "flash_attn_f32_f16_mixed_f32acc_cm2", flash_attn_f32_f16_mixed_cm2_len, flash_attn_f32_f16_mixed_cm2_data, "main", 7, sizeof(vk_flash_attn_push_constants), {Br, 1, 1}, get_fa_spec_constants(fa.first), 1, true, false, 0); \
3752-
} else { \
3753-
ggml_vk_create_pipeline(device, fa.second, "flash_attn_f32_f16_mixed_f16acc_cm2", flash_attn_f32_f16_mixed_f16acc_cm2_len, flash_attn_f32_f16_mixed_f16acc_cm2_data, "main", 7, sizeof(vk_flash_attn_push_constants), {Br, 1, 1}, get_fa_spec_constants(fa.first), 1, true, false, 0); \
3754-
} \
3755-
} \
3756-
} \
3757-
} \
3758-
}
37593728
if (device->coopmat2) {
3760-
CREATE_FA_CM2_MIXED();
3729+
for (auto &fa : device->pipeline_flash_attn_f32_f16) {
3730+
if (fa.first.path != FA_COOPMAT2) continue;
3731+
const uint32_t Br = fa.first.Br;
3732+
const uint32_t Bc = fa.first.Bc;
3733+
const bool aligned = fa.first.aligned;
3734+
const bool f32acc = fa.first.f32acc;
3735+
3736+
const void * spv_data;
3737+
size_t spv_size;
3738+
const char * name;
3739+
if (aligned) {
3740+
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"; }
3741+
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"; }
3742+
} else {
3743+
if (f32acc) { spv_data = flash_attn_f32_f16_cm2_data; spv_size = flash_attn_f32_f16_cm2_len; name = "flash_attn_f32_f16_f32acc_cm2"; }
3744+
else { spv_data = flash_attn_f32_f16_f16acc_cm2_data; spv_size = flash_attn_f32_f16_f16acc_cm2_len; name = "flash_attn_f32_f16_f16acc_cm2"; }
3745+
}
3746+
ggml_vk_create_pipeline(device, fa.second, name, spv_size, spv_data, "main", 7,
3747+
sizeof(vk_flash_attn_push_constants), {Br, 1, 1},
3748+
get_fa_spec_constants(fa.first), aligned ? Bc : 1, true, false, 0);
3749+
}
37613750
}
3762-
#undef CREATE_FA_CM2_MIXED
37633751
#endif
37643752

37653753
const int mul_mat_id_param_count = 5;
@@ -9313,10 +9301,6 @@ static void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx
93139301

93149302
tuning_params = get_fa_tuning_params(ctx->device, HSK, HSV, N, KV, k->type, v->type, f32acc);
93159303

9316-
if (tuning_params.path != FA_COOPMAT2) {
9317-
GGML_ASSERT(k->type == v->type);
9318-
}
9319-
93209304
const uint32_t q_stride = (uint32_t)(nbq1 / ggml_type_size(q->type));
93219305
uint32_t k_stride = (uint32_t)(nbk1 / ggml_type_size(k->type));
93229306
uint32_t v_stride = (uint32_t)(nbv1 / ggml_type_size(v->type));
@@ -16108,10 +16092,6 @@ static bool ggml_backend_vk_device_supports_op(ggml_backend_dev_t dev, const ggm
1610816092
if (op->src[3] && op->src[3]->type != GGML_TYPE_F16) {
1610916093
return false;
1611016094
}
16111-
// mismatching K/V type is currently supported for coopmat2 only.
16112-
if (op->src[1]->type != op->src[2]->type && !coopmat2) {
16113-
return false;
16114-
}
1611516095
auto fa_kv_ok = [coopmat2](ggml_type t) {
1611616096
switch (t) {
1611716097
case GGML_TYPE_F32:

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

Lines changed: 0 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -140,35 +140,6 @@ uint fa_quant_r_mmq(uint ty) {
140140
#define USE_DECODE_K (FaTypeK != FA_TYPE_F16)
141141
#define USE_DECODE_V (FaTypeV != FA_TYPE_F16)
142142

143-
#if defined(DATA_A_TURBO3_0)
144-
const float T3C[8] = float[8](
145-
-0.190685, -0.117832, -0.065717, -0.021460,
146-
0.021460, 0.065717, 0.117832, 0.190685
147-
);
148-
FLOAT_TYPEV4 dequantize4(uint ib, uint iqs, uint a_offset, uint binding_idx) {
149-
FLOAT_TYPEV4 r;
150-
for (int k = 0; k < 4; k++) {
151-
uint j = iqs + uint(k);
152-
float nm;
153-
uint qb;
154-
uint sb;
155-
if (binding_idx == BINDING_IDX_K) {
156-
nm = float(data_k_t3[a_offset + ib].norm);
157-
qb = uint(data_k_t3[a_offset + ib].qs[j / 4]);
158-
sb = uint(data_k_t3[a_offset + ib].signs[j / 8]);
159-
} else {
160-
nm = float(data_v_t3[a_offset + ib].norm);
161-
qb = uint(data_v_t3[a_offset + ib].qs[j / 4]);
162-
sb = uint(data_v_t3[a_offset + ib].signs[j / 8]);
163-
}
164-
uint lo = (qb >> ((j % 4) * 2)) & 0x3;
165-
uint hi = (sb >> (j % 8)) & 0x1;
166-
r[k] = FLOAT_TYPE(T3C[lo | (hi << 2)] * nm);
167-
}
168-
return r;
169-
}
170-
#endif
171-
172143
#define CEIL_DIV(a, b) (((a) + (b) - 1) / (b))
173144

174145

‎ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm2.comp‎

Lines changed: 14 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -28,43 +28,28 @@ layout(buffer_reference, std430, buffer_reference_align = 1) buffer decodeBufFA_
2828
uint8_t raw[FaBlockBytesV];
2929
};
3030

31-
uint fa_block_elems(uint ty) {
32-
switch (ty) {
33-
case 0u: return 4u; // GGML_TYPE_F32: vec4 block (matches decodeBufF32 / dequantFuncF32)
34-
case 1u: return 1u; // GGML_TYPE_F16
35-
case 2u: return uint(QUANT_K_Q4_0);
36-
case 3u: return uint(QUANT_K_Q4_1);
37-
case 6u: return uint(QUANT_K_Q5_0);
38-
case 7u: return uint(QUANT_K_Q5_1);
39-
case 8u: return uint(QUANT_K_Q8_0);
40-
case 41u: return uint(QUANT_K_Q1_0);
41-
default:
42-
return 1u;
43-
}
44-
}
45-
4631
float16_t faDecodeK(const decodeBufFA_K bl_in, const uint blockCoords[2], const uint coordInBlock[2]) {
4732
switch (FaTypeK) {
48-
case 0u: return dequantFuncF32(decodeBufF32(bl_in), blockCoords, coordInBlock);
49-
case 2u: return dequantFuncQ4_0(decodeBufQ4_0(bl_in), blockCoords, coordInBlock);
50-
case 3u: return dequantFuncQ4_1(decodeBufQ4_1(bl_in), blockCoords, coordInBlock);
51-
case 6u: return dequantFuncQ5_0(decodeBufQ5_0(bl_in), blockCoords, coordInBlock);
52-
case 7u: return dequantFuncQ5_1(decodeBufQ5_1(bl_in), blockCoords, coordInBlock);
53-
case 8u: return dequantFuncQ8_0(decodeBufQ8_0(bl_in), blockCoords, coordInBlock);
54-
case 41u: return dequantFuncQ1_0(decodeBufQ1_0(bl_in), blockCoords, coordInBlock);
33+
case FA_TYPE_F32: return dequantFuncF32 (decodeBufF32 (bl_in), blockCoords, coordInBlock);
34+
case FA_TYPE_Q4_0: return dequantFuncQ4_0(decodeBufQ4_0(bl_in), blockCoords, coordInBlock);
35+
case FA_TYPE_Q4_1: return dequantFuncQ4_1(decodeBufQ4_1(bl_in), blockCoords, coordInBlock);
36+
case FA_TYPE_Q5_0: return dequantFuncQ5_0(decodeBufQ5_0(bl_in), blockCoords, coordInBlock);
37+
case FA_TYPE_Q5_1: return dequantFuncQ5_1(decodeBufQ5_1(bl_in), blockCoords, coordInBlock);
38+
case FA_TYPE_Q8_0: return dequantFuncQ8_0(decodeBufQ8_0(bl_in), blockCoords, coordInBlock);
39+
case FA_TYPE_Q1_0: return dequantFuncQ1_0(decodeBufQ1_0(bl_in), blockCoords, coordInBlock);
5540
default: return float16_t(0);
5641
}
5742
}
5843

5944
float16_t faDecodeV(const decodeBufFA_V bl_in, const uint blockCoords[2], const uint coordInBlock[2]) {
6045
switch (FaTypeV) {
61-
case 0u: return dequantFuncF32(decodeBufF32(bl_in), blockCoords, coordInBlock);
62-
case 2u: return dequantFuncQ4_0(decodeBufQ4_0(bl_in), blockCoords, coordInBlock);
63-
case 3u: return dequantFuncQ4_1(decodeBufQ4_1(bl_in), blockCoords, coordInBlock);
64-
case 6u: return dequantFuncQ5_0(decodeBufQ5_0(bl_in), blockCoords, coordInBlock);
65-
case 7u: return dequantFuncQ5_1(decodeBufQ5_1(bl_in), blockCoords, coordInBlock);
66-
case 8u: return dequantFuncQ8_0(decodeBufQ8_0(bl_in), blockCoords, coordInBlock);
67-
case 41u: return dequantFuncQ1_0(decodeBufQ1_0(bl_in), blockCoords, coordInBlock);
46+
case FA_TYPE_F32: return dequantFuncF32 (decodeBufF32 (bl_in), blockCoords, coordInBlock);
47+
case FA_TYPE_Q4_0: return dequantFuncQ4_0(decodeBufQ4_0(bl_in), blockCoords, coordInBlock);
48+
case FA_TYPE_Q4_1: return dequantFuncQ4_1(decodeBufQ4_1(bl_in), blockCoords, coordInBlock);
49+
case FA_TYPE_Q5_0: return dequantFuncQ5_0(decodeBufQ5_0(bl_in), blockCoords, coordInBlock);
50+
case FA_TYPE_Q5_1: return dequantFuncQ5_1(decodeBufQ5_1(bl_in), blockCoords, coordInBlock);
51+
case FA_TYPE_Q8_0: return dequantFuncQ8_0(decodeBufQ8_0(bl_in), blockCoords, coordInBlock);
52+
case FA_TYPE_Q1_0: return dequantFuncQ1_0(decodeBufQ1_0(bl_in), blockCoords, coordInBlock);
6853
default: return float16_t(0);
6954
}
7055
}

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

Lines changed: 1 addition & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -650,26 +650,9 @@ void process_shaders() {
650650

651651
if (fp16) {
652652
#if defined(GGML_VULKAN_COOPMAT2_GLSLC_SUPPORT)
653-
string_to_spv("flash_attn_f32_f16_mixed", "flash_attn_cm2.comp",
653+
string_to_spv("flash_attn_f32_f16", "flash_attn_cm2.comp",
654654
merge_maps(fa_base_dict, {{"Q_TYPE", "float"}, {"D_TYPE", "float"}, {"D_TYPEV4", "vec4"}}), fp16, false, true, f16acc);
655655
#endif
656-
}
657-
658-
for (const auto& tname : type_names) {
659-
if (tname == "bf16") continue;
660-
661-
if (fp16) {
662-
#if defined(GGML_VULKAN_COOPMAT_GLSLC_SUPPORT)
663-
if (tname == "f16") {
664-
string_to_spv("flash_attn_f32_f16_" + tname, "flash_attn_cm1.comp",
665-
merge_maps(fa_base_dict, {{"Q_TYPE", "float"}, {"D_TYPE", "float"}, {"D_TYPEV4", "vec4"}, {"COOPMAT", "1"}}), fp16, true, false, f16acc);
666-
} else if (tname == "q4_0" || tname == "q4_1" || tname == "q5_0" || tname == "q5_1" || tname == "iq4_nl" || tname == "q8_0" || tname == "f32") {
667-
std::string data_a_key = "DATA_A_" + to_uppercase(tname);
668-
string_to_spv("flash_attn_f32_f16_" + tname, "flash_attn_cm1.comp",
669-
merge_maps(fa_base_dict, {{data_a_key, "1"}, {"Q_TYPE", "float"}, {"D_TYPE", "float"}, {"D_TYPEV4", "vec4"}, {"BLOCK_SIZE", "QUANT_K_"+to_uppercase(tname)}, {"COOPMAT", "1"}}), fp16, true, false, f16acc);
670-
}
671-
#endif
672-
}
673656

674657
#if defined(GGML_VULKAN_COOPMAT_GLSLC_SUPPORT)
675658
string_to_spv("flash_attn_f32_f16", "flash_attn_cm1.comp",

0 commit comments

Comments
 (0)