vulkan: Use native e2m1 and e4m3 conversions for mxfp4/nvfp4 (#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.
This commit is contained in:
Jeff Bolz
2026-07-13 08:44:17 -05:00
committed by GitHub
parent 259ae1df8b
commit e920c523e3
10 changed files with 344 additions and 63 deletions
@@ -23,6 +23,14 @@ if (GGML_VULKAN_BFLOAT16_GLSLC_SUPPORT)
add_compile_definitions(GGML_VULKAN_BFLOAT16_GLSLC_SUPPORT)
message(STATUS "Enabling bfloat16 glslc support")
endif()
if (GGML_VULKAN_FLOAT_E2M1_GLSLC_SUPPORT)
add_compile_definitions(GGML_VULKAN_FLOAT_E2M1_GLSLC_SUPPORT)
message(STATUS "Enabling E2M1 glslc support")
endif()
if (GGML_VULKAN_FLOAT_E4M3_GLSLC_SUPPORT)
add_compile_definitions(GGML_VULKAN_FLOAT_E4M3_GLSLC_SUPPORT)
message(STATUS "Enabling E4M3 glslc support")
endif()
if (GGML_VULKAN_SHADER_DEBUG_INFO)
add_compile_definitions(GGML_VULKAN_SHADER_DEBUG_INFO)
message(STATUS "Enabling shader debug info")
@@ -480,12 +480,22 @@ vec4 dequantize4(uint ib, uint iqs, uint a_offset) {
#if defined(DATA_A_MXFP4)
vec2 dequantize(uint ib, uint iqs, uint a_offset) {
const uint vui = uint(data_a[a_offset + ib].qs[iqs]);
#ifdef USE_OCP_FP4
return vec2(unpackFloat2xfe2m1EXT(uint8_t(vui)));
#else
return vec2(kvalues_mxfp4[vui & 0xF], kvalues_mxfp4[vui >> 4]) * 0.5;
#endif
}
vec4 dequantize4(uint ib, uint iqs, uint a_offset) {
#ifdef USE_OCP_FP4
const uint16_t vui = uint16_t(uint(data_a[a_offset + ib].qs[iqs]) |
uint(data_a[a_offset + ib].qs[iqs + 1]) << 8);
return vec4(unpackFloat4xfe2m1EXT(vui));
#else
vec2 v0 = dequantize(ib, iqs, a_offset);
vec2 v1 = dequantize(ib, iqs + 1, a_offset);
return vec4(v0.x, v0.y, v1.x, v1.y);
#endif
}
#endif
@@ -495,16 +505,30 @@ vec2 dequantize(uint ib, uint iqs, uint a_offset) {
const float d = ue4m3_to_fp32(data_a[a_offset + ib].d[sub]);
const uint j = iqs & 7;
const uint shift = (iqs & 8) >> 1; // 0 or 4
#ifdef USE_OCP_FP4
const uint vui = uint(data_a_packed16[a_offset + ib].qs[(sub * 8u + j) / 2u]);
return vec2(bitcastExtractfe2m1EXT(unpack8(vui).xy, shift)) * d;
#else
const uint vui0 = uint(data_a[a_offset + ib].qs[sub * 8u + j]);
const uint vui1 = uint(data_a[a_offset + ib].qs[sub * 8u + j + 1]);
const uint qs0 = (vui0 >> shift) & 0xF;
const uint qs1 = (vui1 >> shift) & 0xF;
return vec2(float(kvalues_mxfp4[qs0]), float(kvalues_mxfp4[qs1])) * d * 0.5;
#endif
}
vec4 dequantize4(uint ib, uint iqs, uint a_offset) {
#ifdef USE_OCP_FP4
const uint sub = iqs >> 4;
const float d = ue4m3_to_fp32(data_a[a_offset + ib].d[sub]);
const uint j = iqs & 7;
const uint shift = (iqs & 8) >> 1; // 0 or 4
const uint vui = data_a_packed32[a_offset + ib].qs[(sub * 8u + j) / 4u];
return vec4(bitcastExtractfe2m1EXT(unpack8(vui), shift)) * d;
#else
const vec2 v0 = dequantize(ib, iqs, a_offset);
const vec2 v1 = dequantize(ib, iqs + 2u, a_offset);
return vec4(v0.x, v0.y, v1.x, v1.y);
#endif
}
#endif
@@ -1232,11 +1232,15 @@ float16_t dequantFuncMXFP4(const in decodeBufMXFP4 bl, const in uint blockCoords
const uint idx = coordInBlock[1];
const uint iqs = idx & 0xF;
const uint shift = (idx & 0x10) >> 2;
#ifdef USE_OCP_FP4
return float16_t(bitcastExtractfe2m1EXT(bl.block.qs[iqs], shift)) * float16_t(d);
#else
uint32_t qs = bl.block.qs[iqs];
qs >>= shift;
qs &= 0xF;
float16_t ret = float16_t(kvalues_mxfp4[qs] * d * 0.5);
return ret;
#endif
}
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
const uint idx = coordInBlock[1];
const uint iqs = idx & 0xF;
const uint shift = (idx & 0x10) >> 2;
#ifdef USE_OCP_FP4
const fe2m1vec4 qv = bitcastExtractfe2m1EXT(
u8vec4(
bl.block.qs[iqs],
bl.block.qs[iqs + 1u],
bl.block.qs[iqs + 2u],
bl.block.qs[iqs + 3u]),
shift);
return f16vec4(qv) * float16_t(d);
#else
uvec4 qv = uvec4(
uint(bl.block.qs[iqs]),
uint(bl.block.qs[iqs + 1u]),
@@ -1257,6 +1271,7 @@ f16vec4 dequantFuncMXFP4_v(const in decodeBufMXFP4 bl, const in uint blockCoords
float(kvalues_mxfp4[qv.z]),
float(kvalues_mxfp4[qv.w])) * d * 0.5f;
return f16vec4(ret);
#endif
}
#endif
@@ -1275,10 +1290,15 @@ float16_t dequantFuncNVFP4(const in decodeBufNVFP4 bl, const in uint blockCoords
const uint sub = (idx & 0x30) >> 4;
const uint iqs = ((idx & 0x30) >> 1) + (idx & 0x7);
const uint shift = (idx & 0x8) >> 1;
#ifdef USE_OCP_FP4
const float16_t d = float16_t(ue4m3_from_bits(bl.block.d[sub]));
return float16_t(bitcastExtractfe2m1EXT(bl.block.qs[iqs], shift)) * d;
#else
const float d = ue4m3_to_fp32(bl.block.d[sub]);
uint qs = uint(bl.block.qs[iqs]);
qs = (qs >> shift) & 0xF;
return float16_t(kvalues_mxfp4[qs] * d * 0.5);
#endif
}
f16vec4 dequantFuncNVFP4_v(const in decodeBufNVFP4 bl, const in uint blockCoords[2], const in uint coordInBlock[2])
@@ -1288,9 +1308,14 @@ f16vec4 dequantFuncNVFP4_v(const in decodeBufNVFP4 bl, const in uint blockCoords
const uint sub = idx >> 4;
const uint qs_w = ((idx & 0x30) >> 3) + ((idx & 0x4u) >> 2); // iqs / 4, in [0,8)
const uint shift = (idx & 0x8) >> 1;
const float d = ue4m3_to_fp32(bl.block.d[sub]);
const uint qsw = uint32_t(bl32.block.qs[qs_w]);
#ifdef USE_OCP_FP4
const float16_t d = float16_t(ue4m3_from_bits(bl.block.d[sub]));
const fe2m1vec4 qv = bitcastExtractfe2m1EXT(unpack8(qsw), shift);
return f16vec4(qv) * d;
#else
const float d = ue4m3_to_fp32(bl.block.d[sub]);
const u8vec4 qv = unpack8((qsw >> shift) & 0x0F0F0F0Fu);
const vec4 ret = vec4(
float(kvalues_mxfp4[qv.x]),
@@ -1298,6 +1323,7 @@ f16vec4 dequantFuncNVFP4_v(const in decodeBufNVFP4 bl, const in uint blockCoords
float(kvalues_mxfp4[qv.z]),
float(kvalues_mxfp4[qv.w])) * d * 0.5f;
return f16vec4(ret);
#endif
}
#endif
@@ -0,0 +1,7 @@
#version 460
#extension GL_EXT_float_e2m1 : require
void main()
{
}
@@ -0,0 +1,7 @@
#version 460
#extension GL_EXT_float_e4m3 : require
void main()
{
}
@@ -502,14 +502,21 @@ void load_a_to_shmem(const uint pos_a, const uint row, const uint col, const uin
const uint ib = idx / 8;
const uint iqs = (idx & 0x07) * 2;
const float d = e8m0_to_fp32(data_a[ib].e) * 0.5;
const uint vui = uint(data_a[ib].qs[iqs]);
const uint vui2 = uint(data_a[ib].qs[iqs+1]);
#ifdef USE_OCP_FP4
const float d = e8m0_to_fp32(data_a[ib].e);
const u8vec2 packed = u8vec2(vui, vui2);
buf_a[buf_idx ] = FLOAT_TYPEV2(bitcastExtractfe2m1EXT(packed, 0u)) * FLOAT_TYPE(d);
buf_a[buf_idx + 8] = FLOAT_TYPEV2(bitcastExtractfe2m1EXT(packed, 4u)) * FLOAT_TYPE(d);
#else
const float d = e8m0_to_fp32(data_a[ib].e) * 0.5;
buf_a[buf_idx ] = FLOAT_TYPEV2(kvalues_mxfp4[vui & 0xF] * d,
kvalues_mxfp4[vui2 & 0xF] * d);
buf_a[buf_idx + 8] = FLOAT_TYPEV2(kvalues_mxfp4[vui >> 4] * d,
kvalues_mxfp4[vui2 >> 4] * d);
#endif
#elif defined(DATA_A_NVFP4)
const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row;
// 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
const uint ib = idx / 16u;
const uint sub = (idx & 0xC) >> 2;
const uint iqs = (idx & 0xF) * 2;
const float d = ue4m3_to_fp32(data_a[ib].d[sub]) * 0.5;
const uint vui = uint(data_a[ib].qs[iqs]);
const uint vui2 = uint(data_a[ib].qs[iqs+1]);
#ifdef USE_OCP_FP4
const FLOAT_TYPE d = FLOAT_TYPE(ue4m3_from_bits(data_a[ib].d[sub]));
const u8vec2 packed = u8vec2(vui, vui2);
buf_a[buf_idx ] = FLOAT_TYPEV2(bitcastExtractfe2m1EXT(packed, 0u)) * d;
buf_a[buf_idx + 4] = FLOAT_TYPEV2(bitcastExtractfe2m1EXT(packed, 4u)) * d;
#else
const float d = ue4m3_to_fp32(data_a[ib].d[sub]) * 0.5;
buf_a[buf_idx ] = FLOAT_TYPEV2(kvalues_mxfp4[vui & 0xF] * d,
kvalues_mxfp4[vui2 & 0xF] * d);
buf_a[buf_idx + 4] = FLOAT_TYPEV2(kvalues_mxfp4[vui >> 4] * d,
kvalues_mxfp4[vui2 >> 4] * d);
#endif
#endif
}
#if !defined(MUL_MAT_ID)
+30 -1
View File
@@ -7,6 +7,11 @@
#extension GL_EXT_shader_explicit_arithmetic_types_int8 : require
#extension GL_EXT_shader_16bit_storage : require
#ifdef USE_OCP_FP4
#extension GL_EXT_float_e2m1 : require
#extension GL_EXT_float_e4m3 : require
#endif
#if defined(DATA_A_F32)
#define QUANT_K 1
#define QUANT_R 1
@@ -1730,6 +1735,12 @@ struct block_nvfp4
uint8_t qs[QUANT_K_NVFP4 / 2];
};
struct block_nvfp4_packed16
{
uint16_t d[QUANT_K_NVFP4 / 16 / 2];
uint16_t qs[QUANT_K_NVFP4 / 2 / 2];
};
struct block_nvfp4_packed32
{
uint32_t d[QUANT_K_NVFP4 / 16 / 4];
@@ -1741,6 +1752,7 @@ struct block_nvfp4_packed32
#define QUANT_R QUANT_R_NVFP4
#define QUANT_AUXF 1
#define A_TYPE block_nvfp4
#define A_TYPE_PACKED16 block_nvfp4_packed16
#define A_TYPE_PACKED32 block_nvfp4_packed32
#endif
@@ -1764,14 +1776,16 @@ void init_iq_shmem(uvec3 wgsize)
#endif
#if defined(DATA_A_MXFP4) || defined(DATA_A_NVFP4)
#if !defined(USE_OCP_FP4)
const int8_t kvalues_mxfp4_const[16] = {
int8_t(0), int8_t(1), int8_t(2), int8_t(3), int8_t(4), int8_t(6), int8_t(8), int8_t(12),
int8_t(0), int8_t(-1), int8_t(-2), int8_t(-3), int8_t(-4), int8_t(-6), int8_t(-8), int8_t(-12),
};
shared int8_t kvalues_mxfp4[16];
#endif
#if defined(DATA_A_NVFP4)
#if defined(DATA_A_NVFP4) && !defined(USE_OCP_FP4)
// UE4M3 scale in NVFP4 blocks use only 7 bits; sign (bit 7) is always zero.
shared float ue4m3_fp32_lut[128];
@@ -1789,6 +1803,7 @@ float ue4m3_to_fp32_build(uint u) {
}
#endif
#if !defined(USE_OCP_FP4)
#define NEEDS_INIT_IQ_SHMEM
void init_iq_shmem(uvec3 wgsize)
{
@@ -1804,6 +1819,7 @@ void init_iq_shmem(uvec3 wgsize)
barrier();
}
#endif
#endif
// returns the bfloat value in the low 16b.
// See ggml_compute_fp32_to_bf16
@@ -1838,8 +1854,21 @@ float e8m0_to_fp32(uint8_t x) {
}
#if defined(DATA_A_NVFP4)
#if defined(USE_OCP_FP4)
floate4m3_t ue4m3_from_bits(uint8_t x) {
if (x == uint8_t(0x7F)) {
return floate4m3_t(0.0);
}
return uintBitsToFloate4m3EXT(x);
}
#endif
float ue4m3_to_fp32(uint8_t x) {
#if defined(USE_OCP_FP4)
return float(ue4m3_from_bits(x));
#else
return ue4m3_fp32_lut[uint(x)];
#endif
}
#endif
@@ -610,6 +610,15 @@ void matmul_shaders(bool fp16, MatMulIdType matmul_id_type, bool coopmat, bool c
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);
}
#if defined(GGML_VULKAN_FLOAT_E2M1_GLSLC_SUPPORT) && defined(GGML_VULKAN_FLOAT_E4M3_GLSLC_SUPPORT)
if ((coopmat || coopmat2) && (tname == "mxfp4" || tname == "nvfp4")) {
if (!coopmat2) {
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);
}
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);
}
#endif
#if defined(GGML_VULKAN_INTEGER_DOT_GLSLC_SUPPORT)
// Integer dot mmq performs better with f32 accumulators (different shader, skip for dot2)
if (!f16acc && !coopmat && !coopmat2 && !dot2 && (is_legacy_quant(tname) || is_k_quant(tname) || tname == "mxfp4")) {
@@ -732,6 +741,20 @@ void process_shaders() {
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"}}));
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"}}));
#if defined(GGML_VULKAN_FLOAT_E2M1_GLSLC_SUPPORT) && defined(GGML_VULKAN_FLOAT_E4M3_GLSLC_SUPPORT)
if (tname == "mxfp4" || tname == "nvfp4") {
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"}}));
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"}}));
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"}}));
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"}}));
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"}}));
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"}}));
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"}}));
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"}}));
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"}}));
}
#endif
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"}}));
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"}}));
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() {
}
}
#if defined(GGML_VULKAN_FLOAT_E2M1_GLSLC_SUPPORT) && defined(GGML_VULKAN_FLOAT_E4M3_GLSLC_SUPPORT)
for (const std::string& btype : {"f16", "f32"}) {
for (const std::string& tname : {"mxfp4", "nvfp4"}) {
hdr << "extern const void * arr_dmmv_" << tname << "_" << btype << "_f32_ocp_data[3];\n";
hdr << "extern const uint64_t arr_dmmv_" << tname << "_" << btype << "_f32_ocp_len[3];\n";
if (basename(input_filepath) == "mul_mat_vec.comp") {
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";
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";
}
}
}
for (const std::string& tname : {"mxfp4", "nvfp4"}) {
hdr << "extern const void * arr_dmmv_id_" << tname << "_f32_f32_ocp_data[3];\n";
hdr << "extern const uint64_t arr_dmmv_id_" << tname << "_f32_f32_ocp_len[3];\n";
if (basename(input_filepath) == "mul_mat_vec.comp") {
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";
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";
}
}
#endif
if (input_filepath == "") {
write_file_if_changed(target_hpp, hdr.str());
}