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
@@ -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