vulkan/cpu: Support f16 as SET_ROWS src. (#25432)

* vulkan/cpu: Support f16 as SET_ROWS src.

This adds full support for f16 SET_ROWS (equivalent to f32) to vulkan and CPU
backends, and adds more backend tests.

* Set DenormPreserve 16 when supported, to try to fix failures on Intel

* tune error threshold

* update metal supports_op
This commit is contained in:
Jeff Bolz
2026-07-14 08:26:55 -05:00
committed by GitHub
parent 8ff8c4299d
commit 7cbd61002d
7 changed files with 150 additions and 100 deletions
@@ -10,7 +10,7 @@ layout(local_size_x = 32, local_size_y = 1, local_size_z = 1) in;
const uint BLOCK_SIZE = 32;
#endif
layout (binding = 0) readonly buffer S {float data_s[];};
layout (binding = 0) readonly buffer S {S_TYPE data_s[];};
#if defined(SET_ROWS)
#include "generic_binary_head.glsl"
@@ -35,7 +35,7 @@ void quantize(uint dst_idx, uint src_idx)
float vmax = 0.0;
[[unroll]] for (int j = 0; j < QUANT_K_Q4_0; ++j) {
const float v = data_s[src_idx + j];
const float v = float(data_s[src_idx + j]);
if (amax < abs(v)) {
amax = abs(v);
vmax = v;
@@ -48,8 +48,8 @@ void quantize(uint dst_idx, uint src_idx)
data_q[dst_idx].d = float16_t(d);
[[unroll]] for (int j = 0; j < QUANT_K_Q4_0/2; ++j) {
const float x0 = data_s[src_idx + 0 + j]*id;
const float x1 = data_s[src_idx + QUANT_K_Q4_0/2 + j]*id;
const float x0 = float(data_s[src_idx + 0 + j])*id;
const float x1 = float(data_s[src_idx + QUANT_K_Q4_0/2 + j])*id;
const uint xi0 = min(15, int(x0 + 8.5));
const uint xi1 = min(15, int(x1 + 8.5));
@@ -66,7 +66,7 @@ void quantize(uint dst_idx, uint src_idx)
float vmax = -vmin;
[[unroll]] for (int j = 0; j < QUANT_K_Q4_1; ++j) {
const float v = data_s[src_idx + j];
const float v = float(data_s[src_idx + j]);
if (v < vmin) vmin = v;
if (v > vmax) vmax = v;
@@ -79,8 +79,8 @@ void quantize(uint dst_idx, uint src_idx)
data_q[dst_idx].m = float16_t(vmin);
[[unroll]] for (int j = 0; j < QUANT_K_Q4_1/2; ++j) {
const float x0 = (data_s[src_idx + 0 + j] - vmin)*id;
const float x1 = (data_s[src_idx + QUANT_K_Q4_1/2 + j] - vmin)*id;
const float x0 = (float(data_s[src_idx + 0 + j]) - vmin)*id;
const float x1 = (float(data_s[src_idx + QUANT_K_Q4_1/2 + j]) - vmin)*id;
const uint xi0 = min(15, int(x0 + 0.5));
const uint xi1 = min(15, int(x1 + 0.5));
@@ -97,7 +97,7 @@ void quantize(uint dst_idx, uint src_idx)
float vmax = 0.0;
[[unroll]] for (int j = 0; j < QUANT_K_Q5_0; ++j) {
const float v = data_s[src_idx + j];
const float v = float(data_s[src_idx + j]);
if (amax < abs(v)) {
amax = abs(v);
vmax = v;
@@ -111,8 +111,8 @@ void quantize(uint dst_idx, uint src_idx)
uint32_t qh = 0;
[[unroll]] for (int j = 0; j < QUANT_K_Q5_0/2; ++j) {
const float x0 = data_s[src_idx + 0 + j]*id;
const float x1 = data_s[src_idx + QUANT_K_Q5_0/2 + j]*id;
const float x0 = float(data_s[src_idx + 0 + j])*id;
const float x1 = float(data_s[src_idx + QUANT_K_Q5_0/2 + j])*id;
const uint xi0 = min(31, int(x0 + 16.5));
const uint xi1 = min(31, int(x1 + 16.5));
@@ -129,11 +129,11 @@ void quantize(uint dst_idx, uint src_idx)
#if defined(DATA_A_Q5_1)
void quantize(uint dst_idx, uint src_idx)
{
float min = data_s[src_idx + 0];
float min = float(data_s[src_idx + 0]);
float max = min;
[[unroll]] for (int j = 1; j < QUANT_K_Q5_1; ++j) {
const float v = data_s[src_idx + j];
const float v = float(data_s[src_idx + j]);
min = v < min ? v : min;
max = v > max ? v : max;
}
@@ -146,8 +146,8 @@ void quantize(uint dst_idx, uint src_idx)
uint32_t qh = 0;
[[unroll]] for (int j = 0; j < QUANT_K_Q5_1/2; ++j) {
const float x0 = (data_s[src_idx + 0 + j] - min)*id;
const float x1 = (data_s[src_idx + QUANT_K_Q5_1/2 + j] - min)*id;
const float x0 = (float(data_s[src_idx + 0 + j]) - min)*id;
const float x1 = (float(data_s[src_idx + QUANT_K_Q5_1/2 + j]) - min)*id;
const uint xi0 = uint(x0 + 0.5);
const uint xi1 = uint(x1 + 0.5);
@@ -166,7 +166,7 @@ void quantize(uint dst_idx, uint src_idx)
float amax = 0.0; // absolute max
[[unroll]] for (int j = 0; j < QUANT_K_Q8_0; j++) {
const float v = data_s[src_idx + j];
const float v = float(data_s[src_idx + j]);
amax = max(amax, abs(v));
}
@@ -176,7 +176,7 @@ void quantize(uint dst_idx, uint src_idx)
data_q[dst_idx].d = float16_t(d);
[[unroll]] for (int j = 0; j < QUANT_K_Q8_0; ++j) {
const float x0 = data_s[src_idx + j]*id;
const float x0 = float(data_s[src_idx + j])*id;
data_q[dst_idx].qs[j] = int8_t(round(x0));
}
@@ -189,7 +189,7 @@ void quantize(uint dst_idx, uint src_idx)
float sum_abs = 0.0;
[[unroll]] for (int j = 0; j < QUANT_K_Q1_0; j++) {
sum_abs += abs(data_s[src_idx + j]);
sum_abs += abs(float(data_s[src_idx + j]));
}
const float d = sum_abs / QUANT_K_Q1_0;
@@ -201,7 +201,7 @@ void quantize(uint dst_idx, uint src_idx)
}
[[unroll]] for (int j = 0; j < QUANT_K_Q1_0; ++j) {
if (data_s[src_idx + j] >= 0.0) {
if (float(data_s[src_idx + j]) >= 0.0) {
data_q[dst_idx].qs[j / 8] |= uint8_t(1 << (j % 8));
}
}
@@ -226,7 +226,7 @@ void quantize(uint dst_idx, uint src_idx)
float vmax = 0.0;
[[unroll]] for (int j = 0; j < QUANT_K_IQ4_NL; ++j) {
const float v = data_s[src_idx + j];
const float v = float(data_s[src_idx + j]);
if (amax < abs(v)) {
amax = abs(v);
vmax = v;
@@ -238,16 +238,16 @@ void quantize(uint dst_idx, uint src_idx)
float sumqx = 0, sumq2 = 0;
[[unroll]] for (int j = 0; j < QUANT_K_IQ4_NL/2; ++j) {
const float x0 = data_s[src_idx + 0 + j]*id;
const float x1 = data_s[src_idx + QUANT_K_IQ4_NL/2 + j]*id;
const float x0 = float(data_s[src_idx + 0 + j])*id;
const float x1 = float(data_s[src_idx + QUANT_K_IQ4_NL/2 + j])*id;
const uint xi0 = best_index(x0);
const uint xi1 = best_index(x1);
data_q[dst_idx].qs[j] = uint8_t(xi0 | (xi1 << 4));
const float v0 = kvalues_iq4nl[xi0];
const float v1 = kvalues_iq4nl[xi1];
const float w0 = data_s[src_idx + 0 + j]*data_s[src_idx + 0 + j];
const float w1 = data_s[src_idx + QUANT_K_IQ4_NL/2 + j]*data_s[src_idx + QUANT_K_IQ4_NL/2 + j];
sumqx += w0*v0*data_s[src_idx + j] + w1*v1*data_s[src_idx + QUANT_K_IQ4_NL/2 + j];
const float w0 = float(data_s[src_idx + 0 + j])*float(data_s[src_idx + 0 + j]);
const float w1 = float(data_s[src_idx + QUANT_K_IQ4_NL/2 + j])*float(data_s[src_idx + QUANT_K_IQ4_NL/2 + j]);
sumqx += w0*v0*float(data_s[src_idx + j]) + w1*v1*float(data_s[src_idx + QUANT_K_IQ4_NL/2 + j]);
sumq2 += w0*v0*v0 + w1*v1*v1;
}
@@ -259,14 +259,14 @@ void quantize(uint dst_idx, uint src_idx)
#if defined(DATA_A_F32) || defined(DATA_A_F16)
void quantize(uint dst_idx, uint src_idx)
{
data_q[dst_idx] = A_TYPE(data_s[src_idx]);
data_q[dst_idx] = A_TYPE(float(data_s[src_idx]));
}
#endif
#if defined(DATA_A_BF16)
void quantize(uint dst_idx, uint src_idx)
{
data_q[dst_idx] = A_TYPE(fp32_to_bf16(data_s[src_idx]));
data_q[dst_idx] = A_TYPE(fp32_to_bf16(float(data_s[src_idx])));
}
#endif
@@ -824,13 +824,15 @@ void process_shaders() {
string_to_spv("cpy_transpose_32", "copy_transpose.comp", {{"A_TYPE", "uint"}, {"D_TYPE", "uint"}});
for (std::string t : {"q1_0", "q4_0", "q4_1", "q5_0", "q5_1", "q8_0", "iq4_nl"}) {
string_to_spv("cpy_f32_" + t, "copy_to_quant.comp", {{"DATA_A_" + to_uppercase(t), "1"}, {"D_TYPE", "float"}, {"FLOAT_TYPE", "float"}});
string_to_spv("cpy_f32_" + t, "copy_to_quant.comp", {{"DATA_A_" + to_uppercase(t), "1"}, {"S_TYPE", "float"}, {"D_TYPE", "float"}, {"FLOAT_TYPE", "float"}});
string_to_spv("cpy_" + t + "_f32", "copy_from_quant.comp", {{"DATA_A_" + to_uppercase(t), "1"}, {"D_TYPE", "float"}, {"FLOAT_TYPE", "float"}});
}
for (std::string t : {"f32", "f16", "bf16", "q1_0", "q4_0", "q4_1", "q5_0", "q5_1", "q8_0", "iq4_nl"}) {
string_to_spv("set_rows_" + t + "_i32", "copy_to_quant.comp", {{"SET_ROWS", "1"}, {"DATA_A_" + to_uppercase(t), "1"}, {"B_TYPE", "uint"}, {"B_SIZE", "32"}, {"D_TYPE", "float"}, {"FLOAT_TYPE", "float"}});
string_to_spv("set_rows_" + t + "_i64", "copy_to_quant.comp", {{"SET_ROWS", "1"}, {"DATA_A_" + to_uppercase(t), "1"}, {"B_TYPE", "uvec2"}, {"B_SIZE", "64"}, {"D_TYPE", "float"}, {"FLOAT_TYPE", "float"}});
for (auto src : {std::pair{"f32", "float"}, std::pair{"f16", "float16_t"}}) {
for (std::string dst : {"f32", "f16", "bf16", "q1_0", "q4_0", "q4_1", "q5_0", "q5_1", "q8_0", "iq4_nl"}) {
string_to_spv("set_rows_" + std::string(src.first) + "_" + dst + "_i32", "copy_to_quant.comp", {{"SET_ROWS", "1"}, {"DATA_A_" + to_uppercase(dst), "1"}, {"B_TYPE", "uint"}, {"B_SIZE", "32"}, {"S_TYPE", src.second}, {"D_TYPE", "float"}, {"FLOAT_TYPE", "float"}});
string_to_spv("set_rows_" + std::string(src.first) + "_" + dst + "_i64", "copy_to_quant.comp", {{"SET_ROWS", "1"}, {"DATA_A_" + to_uppercase(dst), "1"}, {"B_TYPE", "uvec2"}, {"B_SIZE", "64"}, {"S_TYPE", src.second}, {"D_TYPE", "float"}, {"FLOAT_TYPE", "float"}});
}
}
auto get_type_str = [](bool f16) {