Skip to content

Commit e920c52

Browse files
jeffbolznvjeffbolznv
authored
vulkan: Use native e2m1 and e4m3 conversions for mxfp4/nvfp4 (ggml-org#25338)
This uses the new VK_EXT_shader_ocp_microscaling_types extension to do fp4 type promotions, and also uses the float8 extension to do ue4m3 promotions for nvfp4. It's reasonable to assume that an implementation that supports fp4 will also support fp8, so we don't need to handle all possible combinations of support.
1 parent 259ae1d commit e920c52

10 files changed

Lines changed: 345 additions & 64 deletions

File tree

‎ggml/src/ggml-vulkan/CMakeLists.txt‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -97,6 +97,18 @@ if (Vulkan_FOUND)
9797
"GGML_VULKAN_BFLOAT16_GLSLC_SUPPORT"
9898
)
9999

100+
test_shader_extension_support(
101+
"GL_EXT_float_e2m1"
102+
"${CMAKE_CURRENT_SOURCE_DIR}/vulkan-shaders/feature-tests/float_e2m1.comp"
103+
"GGML_VULKAN_FLOAT_E2M1_GLSLC_SUPPORT"
104+
)
105+
106+
test_shader_extension_support(
107+
"GL_EXT_float_e4m3"
108+
"${CMAKE_CURRENT_SOURCE_DIR}/vulkan-shaders/feature-tests/float_e4m3.comp"
109+
"GGML_VULKAN_FLOAT_E4M3_GLSLC_SUPPORT"
110+
)
111+
100112
target_link_libraries(ggml-vulkan PRIVATE Vulkan::Vulkan)
101113
target_include_directories(ggml-vulkan PRIVATE ${CMAKE_CURRENT_BINARY_DIR})
102114

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

Lines changed: 170 additions & 60 deletions
Large diffs are not rendered by default.

‎ggml/src/ggml-vulkan/vulkan-shaders/CMakeLists.txt‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,14 @@ if (GGML_VULKAN_BFLOAT16_GLSLC_SUPPORT)
2323
add_compile_definitions(GGML_VULKAN_BFLOAT16_GLSLC_SUPPORT)
2424
message(STATUS "Enabling bfloat16 glslc support")
2525
endif()
26+
if (GGML_VULKAN_FLOAT_E2M1_GLSLC_SUPPORT)
27+
add_compile_definitions(GGML_VULKAN_FLOAT_E2M1_GLSLC_SUPPORT)
28+
message(STATUS "Enabling E2M1 glslc support")
29+
endif()
30+
if (GGML_VULKAN_FLOAT_E4M3_GLSLC_SUPPORT)
31+
add_compile_definitions(GGML_VULKAN_FLOAT_E4M3_GLSLC_SUPPORT)
32+
message(STATUS "Enabling E4M3 glslc support")
33+
endif()
2634
if (GGML_VULKAN_SHADER_DEBUG_INFO)
2735
add_compile_definitions(GGML_VULKAN_SHADER_DEBUG_INFO)
2836
message(STATUS "Enabling shader debug info")

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

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -480,12 +480,22 @@ vec4 dequantize4(uint ib, uint iqs, uint a_offset) {
480480
#if defined(DATA_A_MXFP4)
481481
vec2 dequantize(uint ib, uint iqs, uint a_offset) {
482482
const uint vui = uint(data_a[a_offset + ib].qs[iqs]);
483+
#ifdef USE_OCP_FP4
484+
return vec2(unpackFloat2xfe2m1EXT(uint8_t(vui)));
485+
#else
483486
return vec2(kvalues_mxfp4[vui & 0xF], kvalues_mxfp4[vui >> 4]) * 0.5;
487+
#endif
484488
}
485489
vec4 dequantize4(uint ib, uint iqs, uint a_offset) {
490+
#ifdef USE_OCP_FP4
491+
const uint16_t vui = uint16_t(uint(data_a[a_offset + ib].qs[iqs]) |
492+
uint(data_a[a_offset + ib].qs[iqs + 1]) << 8);
493+
return vec4(unpackFloat4xfe2m1EXT(vui));
494+
#else
486495
vec2 v0 = dequantize(ib, iqs, a_offset);
487496
vec2 v1 = dequantize(ib, iqs + 1, a_offset);
488497
return vec4(v0.x, v0.y, v1.x, v1.y);
498+
#endif
489499
}
490500
#endif
491501

@@ -495,16 +505,30 @@ vec2 dequantize(uint ib, uint iqs, uint a_offset) {
495505
const float d = ue4m3_to_fp32(data_a[a_offset + ib].d[sub]);
496506
const uint j = iqs & 7;
497507
const uint shift = (iqs & 8) >> 1; // 0 or 4
508+
#ifdef USE_OCP_FP4
509+
const uint vui = uint(data_a_packed16[a_offset + ib].qs[(sub * 8u + j) / 2u]);
510+
return vec2(bitcastExtractfe2m1EXT(unpack8(vui).xy, shift)) * d;
511+
#else
498512
const uint vui0 = uint(data_a[a_offset + ib].qs[sub * 8u + j]);
499513
const uint vui1 = uint(data_a[a_offset + ib].qs[sub * 8u + j + 1]);
500514
const uint qs0 = (vui0 >> shift) & 0xF;
501515
const uint qs1 = (vui1 >> shift) & 0xF;
502516
return vec2(float(kvalues_mxfp4[qs0]), float(kvalues_mxfp4[qs1])) * d * 0.5;
517+
#endif
503518
}
504519
vec4 dequantize4(uint ib, uint iqs, uint a_offset) {
520+
#ifdef USE_OCP_FP4
521+
const uint sub = iqs >> 4;
522+
const float d = ue4m3_to_fp32(data_a[a_offset + ib].d[sub]);
523+
const uint j = iqs & 7;
524+
const uint shift = (iqs & 8) >> 1; // 0 or 4
525+
const uint vui = data_a_packed32[a_offset + ib].qs[(sub * 8u + j) / 4u];
526+
return vec4(bitcastExtractfe2m1EXT(unpack8(vui), shift)) * d;
527+
#else
505528
const vec2 v0 = dequantize(ib, iqs, a_offset);
506529
const vec2 v1 = dequantize(ib, iqs + 2u, a_offset);
507530
return vec4(v0.x, v0.y, v1.x, v1.y);
531+
#endif
508532
}
509533
#endif
510534

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

Lines changed: 27 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1232,11 +1232,15 @@ float16_t dequantFuncMXFP4(const in decodeBufMXFP4 bl, const in uint blockCoords
12321232
const uint idx = coordInBlock[1];
12331233
const uint iqs = idx & 0xF;
12341234
const uint shift = (idx & 0x10) >> 2;
1235+
#ifdef USE_OCP_FP4
1236+
return float16_t(bitcastExtractfe2m1EXT(bl.block.qs[iqs], shift)) * float16_t(d);
1237+
#else
12351238
uint32_t qs = bl.block.qs[iqs];
12361239
qs >>= shift;
12371240
qs &= 0xF;
12381241
float16_t ret = float16_t(kvalues_mxfp4[qs] * d * 0.5);
12391242
return ret;
1243+
#endif
12401244
}
12411245

12421246
f16vec4 dequantFuncMXFP4_v(const in decodeBufMXFP4 bl, const in uint blockCoords[2], const in uint coordInBlock[2])
@@ -1245,6 +1249,16 @@ f16vec4 dequantFuncMXFP4_v(const in decodeBufMXFP4 bl, const in uint blockCoords
12451249
const uint idx = coordInBlock[1];
12461250
const uint iqs = idx & 0xF;
12471251
const uint shift = (idx & 0x10) >> 2;
1252+
#ifdef USE_OCP_FP4
1253+
const fe2m1vec4 qv = bitcastExtractfe2m1EXT(
1254+
u8vec4(
1255+
bl.block.qs[iqs],
1256+
bl.block.qs[iqs + 1u],
1257+
bl.block.qs[iqs + 2u],
1258+
bl.block.qs[iqs + 3u]),
1259+
shift);
1260+
return f16vec4(qv) * float16_t(d);
1261+
#else
12481262
uvec4 qv = uvec4(
12491263
uint(bl.block.qs[iqs]),
12501264
uint(bl.block.qs[iqs + 1u]),
@@ -1257,6 +1271,7 @@ f16vec4 dequantFuncMXFP4_v(const in decodeBufMXFP4 bl, const in uint blockCoords
12571271
float(kvalues_mxfp4[qv.z]),
12581272
float(kvalues_mxfp4[qv.w])) * d * 0.5f;
12591273
return f16vec4(ret);
1274+
#endif
12601275
}
12611276
#endif
12621277

@@ -1275,10 +1290,15 @@ float16_t dequantFuncNVFP4(const in decodeBufNVFP4 bl, const in uint blockCoords
12751290
const uint sub = (idx & 0x30) >> 4;
12761291
const uint iqs = ((idx & 0x30) >> 1) + (idx & 0x7);
12771292
const uint shift = (idx & 0x8) >> 1;
1293+
#ifdef USE_OCP_FP4
1294+
const float16_t d = float16_t(ue4m3_from_bits(bl.block.d[sub]));
1295+
return float16_t(bitcastExtractfe2m1EXT(bl.block.qs[iqs], shift)) * d;
1296+
#else
12781297
const float d = ue4m3_to_fp32(bl.block.d[sub]);
12791298
uint qs = uint(bl.block.qs[iqs]);
12801299
qs = (qs >> shift) & 0xF;
12811300
return float16_t(kvalues_mxfp4[qs] * d * 0.5);
1301+
#endif
12821302
}
12831303

12841304
f16vec4 dequantFuncNVFP4_v(const in decodeBufNVFP4 bl, const in uint blockCoords[2], const in uint coordInBlock[2])
@@ -1288,16 +1308,22 @@ f16vec4 dequantFuncNVFP4_v(const in decodeBufNVFP4 bl, const in uint blockCoords
12881308
const uint sub = idx >> 4;
12891309
const uint qs_w = ((idx & 0x30) >> 3) + ((idx & 0x4u) >> 2); // iqs / 4, in [0,8)
12901310
const uint shift = (idx & 0x8) >> 1;
1291-
const float d = ue4m3_to_fp32(bl.block.d[sub]);
12921311

12931312
const uint qsw = uint32_t(bl32.block.qs[qs_w]);
1313+
#ifdef USE_OCP_FP4
1314+
const float16_t d = float16_t(ue4m3_from_bits(bl.block.d[sub]));
1315+
const fe2m1vec4 qv = bitcastExtractfe2m1EXT(unpack8(qsw), shift);
1316+
return f16vec4(qv) * d;
1317+
#else
1318+
const float d = ue4m3_to_fp32(bl.block.d[sub]);
12941319
const u8vec4 qv = unpack8((qsw >> shift) & 0x0F0F0F0Fu);
12951320
const vec4 ret = vec4(
12961321
float(kvalues_mxfp4[qv.x]),
12971322
float(kvalues_mxfp4[qv.y]),
12981323
float(kvalues_mxfp4[qv.z]),
12991324
float(kvalues_mxfp4[qv.w])) * d * 0.5f;
13001325
return f16vec4(ret);
1326+
#endif
13011327
}
13021328
#endif
13031329

‎ggml/src/ggml-vulkan/vulkan-shaders/feature-tests/float_e2m1.comp‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,7 @@
1+
#version 460
2+
3+
#extension GL_EXT_float_e2m1 : require
4+
5+
void main()
6+
{
7+
}

‎ggml/src/ggml-vulkan/vulkan-shaders/feature-tests/float_e4m3.comp‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,7 @@
1+
#version 460
2+
3+
#extension GL_EXT_float_e4m3 : require
4+
5+
void main()
6+
{
7+
}

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

Lines changed: 16 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -502,14 +502,21 @@ void load_a_to_shmem(const uint pos_a, const uint row, const uint col, const uin
502502
const uint ib = idx / 8;
503503
const uint iqs = (idx & 0x07) * 2;
504504

505-
const float d = e8m0_to_fp32(data_a[ib].e) * 0.5;
506505
const uint vui = uint(data_a[ib].qs[iqs]);
507506
const uint vui2 = uint(data_a[ib].qs[iqs+1]);
508507

508+
#ifdef USE_OCP_FP4
509+
const float d = e8m0_to_fp32(data_a[ib].e);
510+
const u8vec2 packed = u8vec2(vui, vui2);
511+
buf_a[buf_idx ] = FLOAT_TYPEV2(bitcastExtractfe2m1EXT(packed, 0u)) * FLOAT_TYPE(d);
512+
buf_a[buf_idx + 8] = FLOAT_TYPEV2(bitcastExtractfe2m1EXT(packed, 4u)) * FLOAT_TYPE(d);
513+
#else
514+
const float d = e8m0_to_fp32(data_a[ib].e) * 0.5;
509515
buf_a[buf_idx ] = FLOAT_TYPEV2(kvalues_mxfp4[vui & 0xF] * d,
510516
kvalues_mxfp4[vui2 & 0xF] * d);
511517
buf_a[buf_idx + 8] = FLOAT_TYPEV2(kvalues_mxfp4[vui >> 4] * d,
512518
kvalues_mxfp4[vui2 >> 4] * d);
519+
#endif
513520
#elif defined(DATA_A_NVFP4)
514521
const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row;
515522
// lo and hi nibbles are 8 elements apart, which doesn't quite line up with
@@ -519,15 +526,22 @@ void load_a_to_shmem(const uint pos_a, const uint row, const uint col, const uin
519526
const uint ib = idx / 16u;
520527
const uint sub = (idx & 0xC) >> 2;
521528
const uint iqs = (idx & 0xF) * 2;
522-
const float d = ue4m3_to_fp32(data_a[ib].d[sub]) * 0.5;
523529
const uint vui = uint(data_a[ib].qs[iqs]);
524530
const uint vui2 = uint(data_a[ib].qs[iqs+1]);
525531

532+
#ifdef USE_OCP_FP4
533+
const FLOAT_TYPE d = FLOAT_TYPE(ue4m3_from_bits(data_a[ib].d[sub]));
534+
const u8vec2 packed = u8vec2(vui, vui2);
535+
buf_a[buf_idx ] = FLOAT_TYPEV2(bitcastExtractfe2m1EXT(packed, 0u)) * d;
536+
buf_a[buf_idx + 4] = FLOAT_TYPEV2(bitcastExtractfe2m1EXT(packed, 4u)) * d;
537+
#else
538+
const float d = ue4m3_to_fp32(data_a[ib].d[sub]) * 0.5;
526539
buf_a[buf_idx ] = FLOAT_TYPEV2(kvalues_mxfp4[vui & 0xF] * d,
527540
kvalues_mxfp4[vui2 & 0xF] * d);
528541
buf_a[buf_idx + 4] = FLOAT_TYPEV2(kvalues_mxfp4[vui >> 4] * d,
529542
kvalues_mxfp4[vui2 >> 4] * d);
530543
#endif
544+
#endif
531545
}
532546

533547
#if !defined(MUL_MAT_ID)

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

Lines changed: 30 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,11 @@
77
#extension GL_EXT_shader_explicit_arithmetic_types_int8 : require
88
#extension GL_EXT_shader_16bit_storage : require
99

10+
#ifdef USE_OCP_FP4
11+
#extension GL_EXT_float_e2m1 : require
12+
#extension GL_EXT_float_e4m3 : require
13+
#endif
14+
1015
#if defined(DATA_A_F32)
1116
#define QUANT_K 1
1217
#define QUANT_R 1
@@ -1730,6 +1735,12 @@ struct block_nvfp4
17301735
uint8_t qs[QUANT_K_NVFP4 / 2];
17311736
};
17321737

1738+
struct block_nvfp4_packed16
1739+
{
1740+
uint16_t d[QUANT_K_NVFP4 / 16 / 2];
1741+
uint16_t qs[QUANT_K_NVFP4 / 2 / 2];
1742+
};
1743+
17331744
struct block_nvfp4_packed32
17341745
{
17351746
uint32_t d[QUANT_K_NVFP4 / 16 / 4];
@@ -1741,6 +1752,7 @@ struct block_nvfp4_packed32
17411752
#define QUANT_R QUANT_R_NVFP4
17421753
#define QUANT_AUXF 1
17431754
#define A_TYPE block_nvfp4
1755+
#define A_TYPE_PACKED16 block_nvfp4_packed16
17441756
#define A_TYPE_PACKED32 block_nvfp4_packed32
17451757
#endif
17461758

@@ -1764,14 +1776,16 @@ void init_iq_shmem(uvec3 wgsize)
17641776
#endif
17651777

17661778
#if defined(DATA_A_MXFP4) || defined(DATA_A_NVFP4)
1779+
#if !defined(USE_OCP_FP4)
17671780
const int8_t kvalues_mxfp4_const[16] = {
17681781
int8_t(0), int8_t(1), int8_t(2), int8_t(3), int8_t(4), int8_t(6), int8_t(8), int8_t(12),
17691782
int8_t(0), int8_t(-1), int8_t(-2), int8_t(-3), int8_t(-4), int8_t(-6), int8_t(-8), int8_t(-12),
17701783
};
17711784

17721785
shared int8_t kvalues_mxfp4[16];
1786+
#endif
17731787

1774-
#if defined(DATA_A_NVFP4)
1788+
#if defined(DATA_A_NVFP4) && !defined(USE_OCP_FP4)
17751789
// UE4M3 scale in NVFP4 blocks use only 7 bits; sign (bit 7) is always zero.
17761790
shared float ue4m3_fp32_lut[128];
17771791

@@ -1789,6 +1803,7 @@ float ue4m3_to_fp32_build(uint u) {
17891803
}
17901804
#endif
17911805

1806+
#if !defined(USE_OCP_FP4)
17921807
#define NEEDS_INIT_IQ_SHMEM
17931808
void init_iq_shmem(uvec3 wgsize)
17941809
{
@@ -1804,6 +1819,7 @@ void init_iq_shmem(uvec3 wgsize)
18041819
barrier();
18051820
}
18061821
#endif
1822+
#endif
18071823

18081824
// returns the bfloat value in the low 16b.
18091825
// See ggml_compute_fp32_to_bf16
@@ -1838,8 +1854,21 @@ float e8m0_to_fp32(uint8_t x) {
18381854
}
18391855

18401856
#if defined(DATA_A_NVFP4)
1857+
#if defined(USE_OCP_FP4)
1858+
floate4m3_t ue4m3_from_bits(uint8_t x) {
1859+
if (x == uint8_t(0x7F)) {
1860+
return floate4m3_t(0.0);
1861+
}
1862+
return uintBitsToFloate4m3EXT(x);
1863+
}
1864+
#endif
1865+
18411866
float ue4m3_to_fp32(uint8_t x) {
1867+
#if defined(USE_OCP_FP4)
1868+
return float(ue4m3_from_bits(x));
1869+
#else
18421870
return ue4m3_fp32_lut[uint(x)];
1871+
#endif
18431872
}
18441873
#endif
18451874

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

Lines changed: 44 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -610,6 +610,15 @@ void matmul_shaders(bool fp16, MatMulIdType matmul_id_type, bool coopmat, bool c
610610
string_to_spv(shader_name + "_" + tname + "_f16" + dot2_sfx, source_name, merge_maps(merge_maps(base_dict, float_type_dict), {{data_a_key, "1"}, {"LOAD_VEC_A", load_vec_a}, {"LOAD_VEC_B", load_vec}, {"B_TYPE", aligned_b_type_f16}, {"B_TYPE_SCALAR", "float16_t"}, {"B_TYPEV4", "f16vec4"}, {"D_TYPE", "float"}}), fp16, coopmat, coopmat2, f16acc);
611611
}
612612

613+
#if defined(GGML_VULKAN_FLOAT_E2M1_GLSLC_SUPPORT) && defined(GGML_VULKAN_FLOAT_E4M3_GLSLC_SUPPORT)
614+
if ((coopmat || coopmat2) && (tname == "mxfp4" || tname == "nvfp4")) {
615+
if (!coopmat2) {
616+
string_to_spv(shader_name + "_" + tname + "_f32_ocp" + dot2_sfx, source_name, merge_maps(merge_maps(base_dict, float_type_dict), {{data_a_key, "1"}, {"USE_OCP_FP4", "1"}, {"LOAD_VEC_A", load_vec_a}, {"LOAD_VEC_B", load_vec}, {"B_TYPE", aligned_b_type_f32}, {"B_TYPE_SCALAR", "float"}, {"B_TYPEV4", "vec4"}, {"D_TYPE", "float"}}), fp16, coopmat, coopmat2, f16acc);
617+
}
618+
string_to_spv(shader_name + "_" + tname + "_f16_ocp" + dot2_sfx, source_name, merge_maps(merge_maps(base_dict, float_type_dict), {{data_a_key, "1"}, {"USE_OCP_FP4", "1"}, {"LOAD_VEC_A", load_vec_a}, {"LOAD_VEC_B", load_vec}, {"B_TYPE", aligned_b_type_f16}, {"B_TYPE_SCALAR", "float16_t"}, {"B_TYPEV4", "f16vec4"}, {"D_TYPE", "float"}}), fp16, coopmat, coopmat2, f16acc);
619+
}
620+
#endif
621+
613622
#if defined(GGML_VULKAN_INTEGER_DOT_GLSLC_SUPPORT)
614623
// Integer dot mmq performs better with f32 accumulators (different shader, skip for dot2)
615624
if (!f16acc && !coopmat && !coopmat2 && !dot2 && (is_legacy_quant(tname) || is_k_quant(tname) || tname == "mxfp4")) {
@@ -732,6 +741,20 @@ void process_shaders() {
732741
string_to_spv("mul_mat_vec_" + tname + "_f32_f32_subgroup_no_shmem", shader, merge_maps(base_dict, {{data_a_key, "1"}, {"B_TYPE", "float"}, {"B_TYPEV2", "vec2"}, {"B_TYPEV4", "vec4"}, {"D_TYPE", "float"}, {"USE_SUBGROUP_ADD_NO_SHMEM", "1"}}));
733742
string_to_spv("mul_mat_vec_" + tname + "_f16_f32_subgroup_no_shmem", shader, merge_maps(base_dict, {{data_a_key, "1"}, {"B_TYPE", "float16_t"}, {"B_TYPEV2", "f16vec2"}, {"B_TYPEV4", "f16vec4"}, {"D_TYPE", "float"}, {"USE_SUBGROUP_ADD_NO_SHMEM", "1"}}));
734743

744+
#if defined(GGML_VULKAN_FLOAT_E2M1_GLSLC_SUPPORT) && defined(GGML_VULKAN_FLOAT_E4M3_GLSLC_SUPPORT)
745+
if (tname == "mxfp4" || tname == "nvfp4") {
746+
string_to_spv("mul_mat_vec_" + tname + "_f32_f32_ocp", shader, merge_maps(base_dict, {{data_a_key, "1"}, {"USE_OCP_FP4", "1"}, {"B_TYPE", "float"}, {"B_TYPEV2", "vec2"}, {"B_TYPEV4", "vec4"}, {"D_TYPE", "float"}}));
747+
string_to_spv("mul_mat_vec_" + tname + "_f16_f32_ocp", shader, merge_maps(base_dict, {{data_a_key, "1"}, {"USE_OCP_FP4", "1"}, {"B_TYPE", "float16_t"}, {"B_TYPEV2", "f16vec2"}, {"B_TYPEV4", "f16vec4"}, {"D_TYPE", "float"}}));
748+
string_to_spv("mul_mat_vec_" + tname + "_f32_f32_ocp_subgroup", shader, merge_maps(base_dict, {{data_a_key, "1"}, {"USE_OCP_FP4", "1"}, {"B_TYPE", "float"}, {"B_TYPEV2", "vec2"}, {"B_TYPEV4", "vec4"}, {"D_TYPE", "float"}, {"USE_SUBGROUP_ADD", "1"}}));
749+
string_to_spv("mul_mat_vec_" + tname + "_f16_f32_ocp_subgroup", shader, merge_maps(base_dict, {{data_a_key, "1"}, {"USE_OCP_FP4", "1"}, {"B_TYPE", "float16_t"}, {"B_TYPEV2", "f16vec2"}, {"B_TYPEV4", "f16vec4"}, {"D_TYPE", "float"}, {"USE_SUBGROUP_ADD", "1"}}));
750+
string_to_spv("mul_mat_vec_" + tname + "_f32_f32_ocp_subgroup_no_shmem", shader, merge_maps(base_dict, {{data_a_key, "1"}, {"USE_OCP_FP4", "1"}, {"B_TYPE", "float"}, {"B_TYPEV2", "vec2"}, {"B_TYPEV4", "vec4"}, {"D_TYPE", "float"}, {"USE_SUBGROUP_ADD_NO_SHMEM", "1"}}));
751+
string_to_spv("mul_mat_vec_" + tname + "_f16_f32_ocp_subgroup_no_shmem", shader, merge_maps(base_dict, {{data_a_key, "1"}, {"USE_OCP_FP4", "1"}, {"B_TYPE", "float16_t"}, {"B_TYPEV2", "f16vec2"}, {"B_TYPEV4", "f16vec4"}, {"D_TYPE", "float"}, {"USE_SUBGROUP_ADD_NO_SHMEM", "1"}}));
752+
string_to_spv("mul_mat_vec_id_" + tname + "_f32_f32_ocp", shader, merge_maps(base_dict, {{"MUL_MAT_ID", "1"}, {data_a_key, "1"}, {"USE_OCP_FP4", "1"}, {"B_TYPE", "float"}, {"B_TYPEV2", "vec2"}, {"B_TYPEV4", "vec4"}, {"D_TYPE", "float"}}));
753+
string_to_spv("mul_mat_vec_id_" + tname + "_f32_f32_ocp_subgroup", shader, merge_maps(base_dict, {{"MUL_MAT_ID", "1"}, {data_a_key, "1"}, {"USE_OCP_FP4", "1"}, {"B_TYPE", "float"}, {"B_TYPEV2", "vec2"}, {"B_TYPEV4", "vec4"}, {"D_TYPE", "float"}, {"USE_SUBGROUP_ADD", "1"}}));
754+
string_to_spv("mul_mat_vec_id_" + tname + "_f32_f32_ocp_subgroup_no_shmem", shader, merge_maps(base_dict, {{"MUL_MAT_ID", "1"}, {data_a_key, "1"}, {"USE_OCP_FP4", "1"}, {"B_TYPE", "float"}, {"B_TYPEV2", "vec2"}, {"B_TYPEV4", "vec4"}, {"D_TYPE", "float"}, {"USE_SUBGROUP_ADD_NO_SHMEM", "1"}}));
755+
}
756+
#endif
757+
735758
string_to_spv("mul_mat_vec_id_" + tname + "_f32_f32", shader, merge_maps(base_dict, {{"MUL_MAT_ID", "1"}, {data_a_key, "1"}, {"B_TYPE", "float"}, {"B_TYPEV2", "vec2"}, {"B_TYPEV4", "vec4"}, {"D_TYPE", "float"}}));
736759
string_to_spv("mul_mat_vec_id_" + tname + "_f32_f32_subgroup", shader, merge_maps(base_dict, {{"MUL_MAT_ID", "1"}, {data_a_key, "1"}, {"B_TYPE", "float"}, {"B_TYPEV2", "vec2"}, {"B_TYPEV4", "vec4"}, {"D_TYPE", "float"}, {"USE_SUBGROUP_ADD", "1"}}));
737760
string_to_spv("mul_mat_vec_id_" + tname + "_f32_f32_subgroup_no_shmem", shader, merge_maps(base_dict, {{"MUL_MAT_ID", "1"}, {data_a_key, "1"}, {"B_TYPE", "float"}, {"B_TYPEV2", "vec2"}, {"B_TYPEV4", "vec4"}, {"D_TYPE", "float"}, {"USE_SUBGROUP_ADD_NO_SHMEM", "1"}}));
@@ -1233,6 +1256,27 @@ void write_output_files() {
12331256
}
12341257
}
12351258

1259+
#if defined(GGML_VULKAN_FLOAT_E2M1_GLSLC_SUPPORT) && defined(GGML_VULKAN_FLOAT_E4M3_GLSLC_SUPPORT)
1260+
for (const std::string& btype : {"f16", "f32"}) {
1261+
for (const std::string& tname : {"mxfp4", "nvfp4"}) {
1262+
hdr << "extern const void * arr_dmmv_" << tname << "_" << btype << "_f32_ocp_data[3];\n";
1263+
hdr << "extern const uint64_t arr_dmmv_" << tname << "_" << btype << "_f32_ocp_len[3];\n";
1264+
if (basename(input_filepath) == "mul_mat_vec.comp") {
1265+
src << "const void * arr_dmmv_" << tname << "_" << btype << "_f32_ocp_data[3] = {mul_mat_vec_" << tname << "_" << btype << "_f32_ocp_data, mul_mat_vec_" << tname << "_" << btype << "_f32_ocp_subgroup_data, mul_mat_vec_" << tname << "_" << btype << "_f32_ocp_subgroup_no_shmem_data};\n";
1266+
src << "const uint64_t arr_dmmv_" << tname << "_" << btype << "_f32_ocp_len[3] = {mul_mat_vec_" << tname << "_" << btype << "_f32_ocp_len, mul_mat_vec_" << tname << "_" << btype << "_f32_ocp_subgroup_len, mul_mat_vec_" << tname << "_" << btype << "_f32_ocp_subgroup_no_shmem_len};\n";
1267+
}
1268+
}
1269+
}
1270+
for (const std::string& tname : {"mxfp4", "nvfp4"}) {
1271+
hdr << "extern const void * arr_dmmv_id_" << tname << "_f32_f32_ocp_data[3];\n";
1272+
hdr << "extern const uint64_t arr_dmmv_id_" << tname << "_f32_f32_ocp_len[3];\n";
1273+
if (basename(input_filepath) == "mul_mat_vec.comp") {
1274+
src << "const void * arr_dmmv_id_" << tname << "_f32_f32_ocp_data[3] = {mul_mat_vec_id_" << tname << "_f32_f32_ocp_data, mul_mat_vec_id_" << tname << "_f32_f32_ocp_subgroup_data, mul_mat_vec_id_" << tname << "_f32_f32_ocp_subgroup_no_shmem_data};\n";
1275+
src << "const uint64_t arr_dmmv_id_" << tname << "_f32_f32_ocp_len[3] = {mul_mat_vec_id_" << tname << "_f32_f32_ocp_len, mul_mat_vec_id_" << tname << "_f32_f32_ocp_subgroup_len, mul_mat_vec_id_" << tname << "_f32_f32_ocp_subgroup_no_shmem_len};\n";
1276+
}
1277+
}
1278+
#endif
1279+
12361280
if (input_filepath == "") {
12371281
write_file_if_changed(target_hpp, hdr.str());
12381282
}

0 commit comments

Comments
 (0)