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:
@@ -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
|
||||
|
||||
|
||||
Reference in New Issue
Block a user