diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/mul_mat_f32_f32_wmma_core.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/mul_mat_f32_f32_wmma_core.loom index f7f667fa5dc5..8a17547146eb 100644 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/mul_mat_f32_f32_wmma_core.loom +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/mul_mat_f32_f32_wmma_core.loom @@ -60,6 +60,11 @@ func.def inline @ggml_mul_mat_f32_f32_wmma_tiled(%weight_format: index, %token_c %result_stage_bytes = index.constant 2048 : offset %c0_f32x4 = vector.constant 0.0 : vector<4xf32> %c0_f16x4 = vector.constant 0.0 : vector<4xf16> + // engine #1962 (the dense form of #1539): the activations go through f16 for the matrix multiply and f16 ends at + // 65504, so an unnormalised activation above that became inf and the product NaN. Narrow at 2^-6 and scale the f32 + // accumulator back by 2^6 on publish: exact while the scaled value is a normal f16, window [3.9e-3, 4.19e6]. + %activation_prescale = vector.constant 0.015625 : vector<4xf32> + %result_unscale = vector.constant 64.0 : vector<4xf32> %zero_accumulator = vector.constant 0.0 : vector<4xf32> %m = index.constant 16 : index %n = index.constant 16 : index @@ -136,7 +141,8 @@ func.def inline @ggml_mul_mat_f32_f32_wmma_tiled(%weight_format: index, %token_c %bounded_token, %input_token_count = index.assume %token, %bounded_token_count [lt(%token, %bounded_token_count)] : index, index %mask = vector.mask.range [%input_k to %bounded_input_size step %c1] : index -> vector<4xi1> %loaded = vector.load.mask %input_view[%bounded_token, %input_k], %mask, %c0_f32x4 : view<[%bounded_token_count]x[%bounded_input_size]xf32>, vector<4xi1>, vector<4xf32> - %converted = vector.fptrunc %loaded : vector<4xf32> to vector<4xf16> + %prescaled = vector.mulf %loaded, %activation_prescale : vector<4xf32> + %converted = vector.fptrunc %prescaled : vector<4xf32> to vector<4xf16> scf.yield %converted : vector<4xf16> } else { scf.yield %c0_f16x4 : vector<4xf16> @@ -184,28 +190,32 @@ func.def inline @ggml_mul_mat_f32_f32_wmma_tiled(%weight_format: index, %token_c vector.fragment.store %result00, %result_fragment_view[%c0, %c0] shape [%m, %n] : vector<4xf32>, view<16x16xf32, %result_fragment_layout> kernel.barrier scope(subgroup) ordering(acq_rel) scf.if %writes00 { - %values = vector.load %result_physical_view[%publish_token, %publish_channel_add] : view<16x16xf32> -> vector<4xf32> + %scaled_values = vector.load %result_physical_view[%publish_token, %publish_channel_add] : view<16x16xf32> -> vector<4xf32> + %values = vector.mulf %scaled_values, %result_unscale : vector<4xf32> template.apply<@ggml.mul_mat_f32_f32_wmma.publish_vector4>(%writes00, %token0, %channel0, %bounded_token_count, %bounded_output_size, %output_accumulation, %output_unary_op, %values, %bias, %output, %residual_input, %residual_output, %norm_weight, %normalized_output, %completion_counters) : (i1, index, index, index, index, index, index, vector<4xf32>, buffer, buffer, buffer, buffer, buffer, buffer, buffer) } kernel.barrier scope(subgroup) ordering(acq_rel) vector.fragment.store %result01, %result_fragment_view[%c0, %c0] shape [%m, %n] : vector<4xf32>, view<16x16xf32, %result_fragment_layout> kernel.barrier scope(subgroup) ordering(acq_rel) scf.if %writes01 { - %values = vector.load %result_physical_view[%publish_token, %publish_channel_add] : view<16x16xf32> -> vector<4xf32> + %scaled_values = vector.load %result_physical_view[%publish_token, %publish_channel_add] : view<16x16xf32> -> vector<4xf32> + %values = vector.mulf %scaled_values, %result_unscale : vector<4xf32> template.apply<@ggml.mul_mat_f32_f32_wmma.publish_vector4>(%writes01, %token1, %channel0, %bounded_token_count, %bounded_output_size, %output_accumulation, %output_unary_op, %values, %bias, %output, %residual_input, %residual_output, %norm_weight, %normalized_output, %completion_counters) : (i1, index, index, index, index, index, index, vector<4xf32>, buffer, buffer, buffer, buffer, buffer, buffer, buffer) } kernel.barrier scope(subgroup) ordering(acq_rel) vector.fragment.store %result10, %result_fragment_view[%c0, %c0] shape [%m, %n] : vector<4xf32>, view<16x16xf32, %result_fragment_layout> kernel.barrier scope(subgroup) ordering(acq_rel) scf.if %writes10 { - %values = vector.load %result_physical_view[%publish_token, %publish_channel_add] : view<16x16xf32> -> vector<4xf32> + %scaled_values = vector.load %result_physical_view[%publish_token, %publish_channel_add] : view<16x16xf32> -> vector<4xf32> + %values = vector.mulf %scaled_values, %result_unscale : vector<4xf32> template.apply<@ggml.mul_mat_f32_f32_wmma.publish_vector4>(%writes10, %token0, %channel1, %bounded_token_count, %bounded_output_size, %output_accumulation, %output_unary_op, %values, %bias, %output, %residual_input, %residual_output, %norm_weight, %normalized_output, %completion_counters) : (i1, index, index, index, index, index, index, vector<4xf32>, buffer, buffer, buffer, buffer, buffer, buffer, buffer) } kernel.barrier scope(subgroup) ordering(acq_rel) vector.fragment.store %result11, %result_fragment_view[%c0, %c0] shape [%m, %n] : vector<4xf32>, view<16x16xf32, %result_fragment_layout> kernel.barrier scope(subgroup) ordering(acq_rel) scf.if %writes11 { - %values = vector.load %result_physical_view[%publish_token, %publish_channel_add] : view<16x16xf32> -> vector<4xf32> + %scaled_values = vector.load %result_physical_view[%publish_token, %publish_channel_add] : view<16x16xf32> -> vector<4xf32> + %values = vector.mulf %scaled_values, %result_unscale : vector<4xf32> template.apply<@ggml.mul_mat_f32_f32_wmma.publish_vector4>(%writes11, %token1, %channel1, %bounded_token_count, %bounded_output_size, %output_accumulation, %output_unary_op, %values, %bias, %output, %residual_input, %residual_output, %norm_weight, %normalized_output, %completion_counters) : (i1, index, index, index, index, index, index, vector<4xf32>, buffer, buffer, buffer, buffer, buffer, buffer, buffer) } template.apply<@ggml.mul_mat_f32_f32_wmma.finish_tile>(%bounded_token_count, %bounded_output_size, %output_tile_count, %token_tile_count, %output_accumulation, %output_unary_op, %epsilon, %channel_tile, %token_tile, %bias, %output, %residual_input, %residual_output, %norm_weight, %normalized_output, %completion_counters) : (index, index, index, index, index, index, f32, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer) diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/mul_mat_id_f32_f32_wmma_core.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/mul_mat_id_f32_f32_wmma_core.loom index 4be48c72cfed..d5abb157250b 100644 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/mul_mat_id_f32_f32_wmma_core.loom +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/mul_mat_id_f32_f32_wmma_core.loom @@ -76,6 +76,16 @@ template.def<@ggml.mul_mat_id_f32_f32_wmma.core> device @ggml_mul_mat_id_f32_f32 %c0_i32 = scalar.constant 0 : i32 %cn1_i32 = scalar.constant -1 : i32 %c0_f16x4 = vector.constant 0.0 : vector<4xf16> + // engine #1539: the activations go through f16 for the matrix multiply, and f16 ends at 65504. An expert's + // down-projection reads the SwiGLU output, which is not normalised: flint's layer 46 reaches +-9.1e4 on an ordinary + // prompt, the narrowing made that inf, the route sum inf - inf = NaN, and the NaN was stored as that token's keys and + // values. So the activations are narrowed at 2^-6 of their value and the f32 accumulator is scaled back by 2^6 when a + // tile is published. Both factors are powers of two, so nothing is rounded differently while the scaled value is a + // normal f16: the representable window moves from [6.1e-5, 6.55e4] to [3.9e-3, 4.19e6]. Below 3.9e-3 an element is + // rounded to a multiple of 3.8e-6 instead of keeping 11 bits, which is far under the rounding already applied to the + // order-1 elements it is summed with (4.9e-4 at 1.0). + %activation_prescale = vector.constant 0.015625 : vector<4xf32> + %result_unscale = vector.constant 64.0 : vector<4xf32> %zero_accumulator = vector.constant 0.0 : vector<4xf32> %m = index.constant 16 : index %n = index.constant 16 : index @@ -209,7 +219,8 @@ template.def<@ggml.mul_mat_id_f32_f32_wmma.core> device @ggml_mul_mat_id_f32_f32 %input_k = index.add %k_origin, %load_k : index %mask = vector.mask.range [%input_k to %input_size step %c1] : index -> vector<4xi1> %loaded = vector.load.mask %input_view[%token, %input_route, %input_k], %mask, %c0_f32x4 : view<[%bounded_token_count]x[%bounded_input_route_count]x[%input_size]xf32>, vector<4xi1>, vector<4xf32> - %converted = vector.fptrunc %loaded : vector<4xf32> to vector<4xf16> + %prescaled = vector.mulf %loaded, %activation_prescale : vector<4xf32> + %converted = vector.fptrunc %prescaled : vector<4xf32> to vector<4xf16> scf.yield %converted : vector<4xf16> } else { scf.yield %c0_f16x4 : vector<4xf16> @@ -277,28 +288,32 @@ template.def<@ggml.mul_mat_id_f32_f32_wmma.core> device @ggml_mul_mat_id_f32_f32 vector.fragment.store %result00, %result_fragment_view[%c0, %c0] shape [%m, %n] : vector<4xf32>, view<16x16xf32, %result_fragment_layout> kernel.barrier scope(subgroup) ordering(acq_rel) scf.if %writes00 { - %values = vector.load %result_physical_view[%publish_route, %publish_channel_add] : view<16x16xf32> -> vector<4xf32> + %scaled_values = vector.load %result_physical_view[%publish_route, %publish_channel_add] : view<16x16xf32> -> vector<4xf32> + %values = vector.mulf %scaled_values, %result_unscale : vector<4xf32> template.apply<@ggml.mul_mat_id_f32_f32_wmma.publish_vector4>(%writes00, %bounded_assignment0, %assignment_token0, %assignment_route0, %channel0, %bounded_token_count, %bounded_route_count, %bounded_output_size, %values, %bias, %output, %residual_input, %residual_output, %norm_weight, %normalized_output, %completion_counters) : (i1, index, index, index, index, index, index, index, vector<4xf32>, buffer, buffer, buffer, buffer, buffer, buffer, buffer) } kernel.barrier scope(subgroup) ordering(acq_rel) vector.fragment.store %result01, %result_fragment_view[%c0, %c0] shape [%m, %n] : vector<4xf32>, view<16x16xf32, %result_fragment_layout> kernel.barrier scope(subgroup) ordering(acq_rel) scf.if %writes01 { - %values = vector.load %result_physical_view[%publish_route, %publish_channel_add] : view<16x16xf32> -> vector<4xf32> + %scaled_values = vector.load %result_physical_view[%publish_route, %publish_channel_add] : view<16x16xf32> -> vector<4xf32> + %values = vector.mulf %scaled_values, %result_unscale : vector<4xf32> template.apply<@ggml.mul_mat_id_f32_f32_wmma.publish_vector4>(%writes01, %bounded_assignment1, %assignment_token1, %assignment_route1, %channel0, %bounded_token_count, %bounded_route_count, %bounded_output_size, %values, %bias, %output, %residual_input, %residual_output, %norm_weight, %normalized_output, %completion_counters) : (i1, index, index, index, index, index, index, index, vector<4xf32>, buffer, buffer, buffer, buffer, buffer, buffer, buffer) } kernel.barrier scope(subgroup) ordering(acq_rel) vector.fragment.store %result10, %result_fragment_view[%c0, %c0] shape [%m, %n] : vector<4xf32>, view<16x16xf32, %result_fragment_layout> kernel.barrier scope(subgroup) ordering(acq_rel) scf.if %writes10 { - %values = vector.load %result_physical_view[%publish_route, %publish_channel_add] : view<16x16xf32> -> vector<4xf32> + %scaled_values = vector.load %result_physical_view[%publish_route, %publish_channel_add] : view<16x16xf32> -> vector<4xf32> + %values = vector.mulf %scaled_values, %result_unscale : vector<4xf32> template.apply<@ggml.mul_mat_id_f32_f32_wmma.publish_vector4>(%writes10, %bounded_assignment0, %assignment_token0, %assignment_route0, %channel1, %bounded_token_count, %bounded_route_count, %bounded_output_size, %values, %bias, %output, %residual_input, %residual_output, %norm_weight, %normalized_output, %completion_counters) : (i1, index, index, index, index, index, index, index, vector<4xf32>, buffer, buffer, buffer, buffer, buffer, buffer, buffer) } kernel.barrier scope(subgroup) ordering(acq_rel) vector.fragment.store %result11, %result_fragment_view[%c0, %c0] shape [%m, %n] : vector<4xf32>, view<16x16xf32, %result_fragment_layout> kernel.barrier scope(subgroup) ordering(acq_rel) scf.if %writes11 { - %values = vector.load %result_physical_view[%publish_route, %publish_channel_add] : view<16x16xf32> -> vector<4xf32> + %scaled_values = vector.load %result_physical_view[%publish_route, %publish_channel_add] : view<16x16xf32> -> vector<4xf32> + %values = vector.mulf %scaled_values, %result_unscale : vector<4xf32> template.apply<@ggml.mul_mat_id_f32_f32_wmma.publish_vector4>(%writes11, %bounded_assignment1, %assignment_token1, %assignment_route1, %channel1, %bounded_token_count, %bounded_route_count, %bounded_output_size, %values, %bias, %output, %residual_input, %residual_output, %norm_weight, %normalized_output, %completion_counters) : (i1, index, index, index, index, index, index, index, vector<4xf32>, buffer, buffer, buffer, buffer, buffer, buffer, buffer) } kernel.barrier scope(workgroup) ordering(acq_rel) diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index 306217b7841a..629a707f936e 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -4592,6 +4592,88 @@ struct test_mul_mat_id : public test_case { } }; +// GGML_OP_MUL_MAT whose activations exceed the f16 range: the dense form of the engine #1539 case below. +struct test_mul_mat_large_activation : public test_mul_mat { + const float magnitude; + + std::string vars() override { + return test_mul_mat::vars() + ",magnitude=" + var_to_str(magnitude); + } + + test_mul_mat_large_activation(ggml_type type_a, int64_t m, int64_t n, int64_t k, float magnitude) + : test_mul_mat(type_a, GGML_TYPE_F32, m, n, k, {1, 1}, {1, 1}), magnitude(magnitude) {} + + void initialize_tensors(ggml_context * ctx) override { + test_case::initialize_tensors(ctx); + for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != NULL; t = ggml_get_next_tensor(ctx, t)) { + if (strcmp(t->name, "b") != 0) { + continue; + } + std::vector data(ggml_nelements(t)); + ggml_backend_tensor_get(t, data.data(), 0, ggml_nbytes(t)); + if (magnitude < 1.0f) { + // the lower edge of the same window: every activation scaled down (see the MUL_MAT_ID case below) + for (float & value : data) { + value *= magnitude; + } + ggml_backend_tensor_set(t, data.data(), 0, ggml_nbytes(t)); + continue; + } + for (int64_t token = 0; token < t->ne[1]; token += 4) { + for (int64_t i = 0; i < t->ne[0]; i += 64) { + data[token * t->ne[0] + i] = ((i / 64 + token / 4) % 2 == 0) ? magnitude : -magnitude; + } + } + ggml_backend_tensor_set(t, data.data(), 0, ggml_nbytes(t)); + } + } +}; + +// GGML_OP_MUL_MAT_ID whose activations exceed the f16 range (engine #1539). +// An expert's down-projection reads the SwiGLU output, which nothing normalises: flint-40b's layer 46 reaches +-9.1e4 on +// an ordinary 22-token prompt. A backend that narrows the activations to f16 for the multiply turns anything above 65504 +// into inf, and the sum over the routes into NaN - finite on the CPU, non-finite on the device, and invisible to every +// other MUL_MAT_ID case here because init_tensor_uniform() stays inside [-1, 1]. One activation in 64 of every fourth +// token is set to +-magnitude; the rest keep their order-1 values, as in the model. +// A magnitude below 1 is the other end of the same window: every activation is scaled by it, so a narrowing that +// trades the top of the f16 range for the bottom has a case pulling back. 3.2e-2 is the smallest column maximum +// measured on flint's expert down-projections (37278 columns over 141 layer executions; median 4.1). +struct test_mul_mat_id_large_activation : public test_mul_mat_id { + const float magnitude; + + std::string vars() override { + return test_mul_mat_id::vars() + ",magnitude=" + var_to_str(magnitude); + } + + test_mul_mat_id_large_activation(ggml_type type_a, int n_mats, int n_used, int64_t m, int64_t n, int64_t k, float magnitude) + : test_mul_mat_id(type_a, GGML_TYPE_F32, n_mats, n_used, false, m, n, k), magnitude(magnitude) {} + + void initialize_tensors(ggml_context * ctx) override { + init_mul_mat_id_tensors(ctx, n_mats); + for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != NULL; t = ggml_get_next_tensor(ctx, t)) { + if (strcmp(t->name, "b") != 0) { + continue; + } + std::vector data(ggml_nelements(t)); + ggml_backend_tensor_get(t, data.data(), 0, ggml_nbytes(t)); + const int64_t per_token = t->ne[0] * t->ne[1]; + if (magnitude < 1.0f) { + for (float & value : data) { + value *= magnitude; + } + ggml_backend_tensor_set(t, data.data(), 0, ggml_nbytes(t)); + continue; + } + for (int64_t token = 0; token < t->ne[2]; token += 4) { + for (int64_t i = 0; i < per_token; i += 64) { + data[token * per_token + i] = ((i / 64 + token / 4) % 2 == 0) ? magnitude : -magnitude; + } + } + ggml_backend_tensor_set(t, data.data(), 0, ggml_nbytes(t)); + } + } +}; + // GGML_OP_MUL_MAT_ID + GGML_OP_ADD or GGML_OP_MUL struct test_mul_mat_id_fusion : public test_case { const ggml_type type_a; @@ -9321,6 +9403,36 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_mul_mat_row_padded_activation(GGML_TYPE_BF16, GGML_TYPE_F32, 3712, 512, 232, 264)); test_cases.emplace_back(new test_mul_mat_id(GGML_TYPE_BF16, GGML_TYPE_F32, 8, 2, false, 384, 8, 1480)); + // engine #1539: flint-40b's expert down-projection (256 experts, 6 used, 512 -> 2048, q5_K) at the 22-token batch + // that produced inf, with activations of the measured magnitude; then one decode-sized batch of the same contraction. + test_cases.emplace_back(new test_mul_mat_id_large_activation(GGML_TYPE_Q5_K, 256, 6, 2048, 22, 512, 9.1e4f)); + test_cases.emplace_back(new test_mul_mat_id_large_activation(GGML_TYPE_Q5_K, 256, 6, 2048, 1, 512, 9.1e4f)); + test_cases.emplace_back(new test_mul_mat_id_large_activation(GGML_TYPE_Q5_K, 256, 6, 2048, 22, 512, 3.2e-2f)); + // engine #1962: the dense form. Weights [k, 2048], k = 4096 and 512, at decode, small-batch and prefill token counts. + // Only the cases this backend answers correctly today are listed, so the suite stays a gate: f16 weights are absent + // because the CPU reference itself returns NaN there, and q4_K/q5_K at 512 tokens and q6_K at 1, 8 (k = 4096) and 512 + // tokens still overflow in kernel families this change does not touch (they are listed on the issue). + for (ggml_type type_a : { GGML_TYPE_Q8_0, GGML_TYPE_IQ4_XS, GGML_TYPE_Q4_0, GGML_TYPE_BF16 }) { + for (int64_t n : { 1, 8, 22, 512 }) { + test_cases.emplace_back(new test_mul_mat_large_activation(type_a, 2048, n, 4096, 9.1e4f)); + test_cases.emplace_back(new test_mul_mat_large_activation(type_a, 2048, n, 512, 9.1e4f)); + } + } + for (ggml_type type_a : { GGML_TYPE_Q4_K, GGML_TYPE_Q5_K }) { + for (int64_t n : { 1, 8, 22 }) { + test_cases.emplace_back(new test_mul_mat_large_activation(type_a, 2048, n, 4096, 9.1e4f)); + test_cases.emplace_back(new test_mul_mat_large_activation(type_a, 2048, n, 512, 9.1e4f)); + } + } + test_cases.emplace_back(new test_mul_mat_large_activation(GGML_TYPE_Q6_K, 2048, 22, 4096, 9.1e4f)); + test_cases.emplace_back(new test_mul_mat_large_activation(GGML_TYPE_Q6_K, 2048, 22, 512, 9.1e4f)); + test_cases.emplace_back(new test_mul_mat_large_activation(GGML_TYPE_Q6_K, 2048, 8, 512, 9.1e4f)); + // and the lower edge: every activation at 2.4e-2, the smallest column maximum measured on the dense matmul inputs of + // Qwen3-0.6B, Qwen3-4B, ZAYA1-8B and flint-40b (26k columns; an all-zero column is exact at any scale and is not counted). + for (ggml_type type_a : { GGML_TYPE_Q4_K, GGML_TYPE_Q8_0, GGML_TYPE_BF16 }) { + test_cases.emplace_back(new test_mul_mat_large_activation(type_a, 2048, 22, 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})); for (ggml_type type_a : all_types) {