diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/mul_mat_q5_k_q8_plane_wmma.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/mul_mat_q5_k_q8_plane_wmma.loom index 4df29247a9d3..b4eb46e34a0b 100644 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/mul_mat_q5_k_q8_plane_wmma.loom +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/mul_mat_q5_k_q8_plane_wmma.loom @@ -466,9 +466,37 @@ template.def<@ggml.mul_mat_q5_k_q8_plane.wmmai8.body> device @ggml_mul_mat_q5_k_ %q8_scale_dst0_0 = index.add %q8_scale_tile2, %q8_token_vector : index %q8_scale_dst1_0 = index.add %q8_scale_dst0_0, %c16 : index %q8_d0 = scalar.extf %q8_d0_f16 : f16 to f32 - %q8_scaled_sum0 = scalar.extf %q8_s0_f16 : f16 to f32 + %q8_stored_sum0 = scalar.extf %q8_s0_f16 : f16 to f32 %q8_d1 = scalar.extf %q8_d1_f16 : f16 to f32 - %q8_scaled_sum1 = scalar.extf %q8_s1_f16 : f16 to f32 + %q8_stored_sum1 = scalar.extf %q8_s1_f16 : f16 to f32 + // engine #1962: the block's scaled sum (scale times the sum of its 32 codes) is stored as f16, and it leaves the + // f16 range as soon as one activation in the block does, although the scale itself still fits up to 8.3e6. A + // stored sum that is not finite is rebuilt here from the codes this lane already holds; a finite one is used as + // stored, so ordinary data takes the path it always took. + // The rebuilt value cannot overflow: |sum of 32 codes| <= 32 * 127 = 4064 and the scale is an f16 (<= 65504), so + // the product is at most 2.7e8 in f32. The compare is false for inf and for NaN, so both are rebuilt. + %q8_f16_max = scalar.constant 65504.0 : f32 + %q8_sum_zero = scalar.constant 0.0 : f32 + %q8_stored_abs0 = scalar.absf %q8_stored_sum0 : f32 + %q8_stored_abs1 = scalar.absf %q8_stored_sum1 : f32 + %q8_sum_fits0 = scalar.cmpf ole, %q8_stored_abs0, %q8_f16_max : f32 + %q8_sum_fits1 = scalar.cmpf ole, %q8_stored_abs1, %q8_f16_max : f32 + %q8_scaled_sum0 = scf.if %q8_sum_fits0 -> (f32) { + scf.yield %q8_stored_sum0 : f32 + } else { + %q8_codes0 = vector.sitofp %q8_values0 : vector<32xi8> to vector<32xf32> + %q8_code_sum0 = vector.reduce %q8_codes0, %q8_sum_zero : vector<32xf32>, f32 + %q8_rebuilt0 = scalar.mulf %q8_code_sum0, %q8_d0 : f32 + scf.yield %q8_rebuilt0 : f32 + } + %q8_scaled_sum1 = scf.if %q8_sum_fits1 -> (f32) { + scf.yield %q8_stored_sum1 : f32 + } else { + %q8_codes1 = vector.sitofp %q8_values1 : vector<32xi8> to vector<32xf32> + %q8_code_sum1 = vector.reduce %q8_codes1, %q8_sum_zero : vector<32xf32>, f32 + %q8_rebuilt1 = scalar.mulf %q8_code_sum1, %q8_d1 : f32 + scf.yield %q8_rebuilt1 : f32 + } func.call @ggml_q5_plane_dense_q8_wmmai8_stage_metadata(%token256, %scratch, %as_off, %asum_off, %q8_scale_dst0_0, %q8_scale_dst1_0, %q8_d0, %q8_d1, %q8_scaled_sum0, %q8_scaled_sum1) : (i1, buffer, offset, offset, index, index, f32, f32, f32, f32) } // One lane decodes two adjacent Q5_K groups; each selects its bits from the shared 32-byte fifth-bit plane. diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index ffbc25c30411..83e90abdf59b 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -9528,11 +9528,20 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_mul_mat_resident_weight(GGML_TYPE_Q5_K, 2048, 512, 4096)); // engine #1962: the two small-batch Q6_K kernels with activations of the magnitude measured in flint-40b, and the // lower edge (the smallest dense-input column maximum measured on four models). The 512-token prefill kernels for - // Q4_K, Q5_K and Q6_K still overflow above 65504 and are not listed until they pass. + // Q4_K and Q6_K still overflow above 65504 and are not listed until they pass. test_cases.emplace_back(new test_mul_mat_resident_weight_large_activation(GGML_TYPE_Q6_K, 2048, 8, 4096, 9.1e4f)); test_cases.emplace_back(new test_mul_mat_resident_weight_large_activation(GGML_TYPE_Q6_K, 4096, 8, 512, 9.1e4f)); test_cases.emplace_back(new test_mul_mat_resident_weight_large_activation(GGML_TYPE_Q6_K, 2048, 8, 4096, 2.4e-2f)); test_cases.emplace_back(new test_mul_mat_resident_weight_large_activation(GGML_TYPE_Q6_K, 4096, 8, 512, 2.4e-2f)); + // engine #1962: the Q5_K prefill reads the activations as 8-bit blocks (block_q8_1_x4). Each block carries its scale + // and its sum; a sum stored already multiplied by the scale left the f16 range as soon as one activation did. + // 6.5e4 is the last step of the ladder where the stored sum is still finite (the stored path); 6.6e4 is the first + // where it is not (the rebuilt path). + test_cases.emplace_back(new test_mul_mat_resident_weight_large_activation(GGML_TYPE_Q5_K, 2048, 512, 4096, 6.5e4f)); + test_cases.emplace_back(new test_mul_mat_resident_weight_large_activation(GGML_TYPE_Q5_K, 2048, 512, 4096, 6.6e4f)); + test_cases.emplace_back(new test_mul_mat_resident_weight_large_activation(GGML_TYPE_Q5_K, 2048, 512, 4096, 9.1e4f)); + test_cases.emplace_back(new test_mul_mat_resident_weight_large_activation(GGML_TYPE_Q5_K, 2048, 512, 4096, 1.0e6f)); + test_cases.emplace_back(new test_mul_mat_resident_weight_large_activation(GGML_TYPE_Q5_K, 2048, 512, 4096, 2.4e-2f)); test_cases.emplace_back(new test_mul_mat(GGML_TYPE_Q4_0, GGML_TYPE_F32, 576, 512, 576, {1,1}, {1,1})); test_cases.emplace_back(new test_mul_mat(GGML_TYPE_Q4_0, GGML_TYPE_F32, 1, 2048, 8192, {1, 1}, {1, 1}));