hrx: the fused routed feed-forward keeps its f16 planes at 2^-6, so large SwiGLU products and down sums stay finite (engine #2186)

ggml/src/ggml-hrx/dispatch_registration/qwen/dispatch-routed-ffn.cpp

@@ -44,6 +44,10 @@ static constexpr int64_t                       kRoutedFfnRouteCount
 static constexpr size_t                        kRoutedFfnPlanTransientAlignment  = 256;
 static constexpr const char *                  kRoutedFfnF16GateUpOutputName     = "qwen.moe.gate_up_swiglu_f16";
 static constexpr const char *                  kRoutedFfnF16RoutedDownOutputName = "qwen.moe.routed_down_f16";
+// engine #2186: both f16 planes of the fused routed feed-forward hold their values at 2^-6. The gate/up kernel
+// publishes the SwiGLU product that way, the routed down projection is linear and carries the factor, and the weighted
+// reduce multiplies it back in f32. Recorded on the plan's alternate values so a reader has to ask for it.
+static constexpr int                           kRoutedFfnF16PlaneScaleLog2       = 6;
 static constexpr const char *                  kRoutedFfnQ8GateUpOutputName      = "qwen.decode.moe.gate_up_swiglu_q8";
 static constexpr const char *                  kRoutedFfnQ8HiddenOutputName      = "qwen.decode.moe.hidden_q8";
 static constexpr const char *                  kRoutedFfnQ8NextRmsNormOutputName = "qwen.prefill.moe.next_rmsnorm_q8";
@@ -387,7 +391,8 @@ static RoutedDownMatch match_routed_ffn_down_grouped(const DispatchMatchContext
     }
 
     const CommandPlanAlternateValue * input_alternate =
-        find_alternate_value(context.plan, input->id, GGML_TYPE_F16, f16_gate_up_output_size(token_count));
+        find_alternate_value(context.plan, input->id, GGML_TYPE_F16, f16_gate_up_output_size(token_count),
+                             kRoutedFfnF16PlaneScaleLog2);
     if (input_alternate == nullptr) {
         return {};
     }
@@ -873,7 +878,8 @@ static WeightedReduceMatch match_routed_ffn_down_weighted_reduce(const DispatchM
     }
     const int64_t                     token_count = match.token_count;
     const CommandPlanAlternateValue * routed_alternate =
-        find_alternate_value(context.plan, routed_output->id, GGML_TYPE_F16, f16_routed_down_output_size(token_count));
+        find_alternate_value(context.plan, routed_output->id, GGML_TYPE_F16, f16_routed_down_output_size(token_count),
+                             kRoutedFfnF16PlaneScaleLog2);
     if (routed_alternate == nullptr) {
         return {};
     }
@@ -917,7 +923,8 @@ static bool match_routed_ffn_gate_up_swiglu_f16_wmma_dispatch(const DispatchMatc
         { f16_output, kRoutedFfnF16GateUpOutputName, f16_output_bytes, kRoutedFfnPlanTransientAlignment });
     Status metadata_status;
     if (!dispatch_match.metadata.append_alternate_value(
-            { match.glu_output->id, f16_output, GGML_TYPE_F16, f16_output_bytes, kRoutedFfnF16GateUpOutputName },
+            { match.glu_output->id, f16_output, GGML_TYPE_F16, f16_output_bytes, kRoutedFfnF16GateUpOutputName,
+              kRoutedFfnF16PlaneScaleLog2 },
             metadata_status)) {
         return false;
     }
@@ -1186,7 +1193,8 @@ static bool build_routed_ffn_down_grouped_dispatch(const DispatchMatchContext &
         { f16_output, kRoutedFfnF16RoutedDownOutputName, f16_output_bytes, kRoutedFfnPlanTransientAlignment });
     Status metadata_status;
     if (!dispatch_match.metadata.append_alternate_value(
-            { match.output->id, f16_output, GGML_TYPE_F16, f16_output_bytes, kRoutedFfnF16RoutedDownOutputName },
+            { match.output->id, f16_output, GGML_TYPE_F16, f16_output_bytes, kRoutedFfnF16RoutedDownOutputName,
+              kRoutedFfnF16PlaneScaleLog2 },
             metadata_status)) {
         return false;
     }

ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/mul_mat_id_swiglu_f16_f16_wmma_body.loom

@@ -187,7 +187,14 @@ template.def<@ggml.mul_mat_id_swiglu_f16_f16.body> device @ggml_mul_mat_id_swigl
       %up_prescaled = vector.extf %up_values : vector<8xf16> to vector<8xf32>
       %up_wide = vector.mulf %up_prescaled, %sum_rescale : vector<8xf32>
       %wide_values = template.apply<@ggml.mul_mat_id_swiglu_f16_f16.apply_vector8>(%gate_wide, %up_wide) : (vector<8xf32>, vector<8xf32>) -> (vector<8xf32>)
-      %values = vector.fptrunc %wide_values : vector<8xf32> to vector<8xf16>
+      // engine #2186: the product goes to an f16 plane, and the routed down projection that reads it accumulates and
+      // publishes in f16 as well, so the product and every partial down sum had to stay under 65504 (measured: with
+      // products up to 2.6e4 the down sums already overflowed and 770 of 2048 outputs were inf). The plane is
+      // published at 2^-6; the down projection is linear and carries the factor; the weighted reduce, which widens to
+      // f32, multiplies it back. The plan records the factor on both planes (scale_log2).
+      %product_prescale = vector.constant 0.015625 : vector<8xf32>
+      %scaled_values = vector.mulf %wide_values, %product_prescale : vector<8xf32>
+      %values = vector.fptrunc %scaled_values : vector<8xf32> to vector<8xf16>
       %mask = vector.mask.range [%channel to %bounded_output_size step %c1] : index -> vector<8xi1>
       vector.store.mask %values, %output_view[%bounded_assignment, %channel], %mask : vector<8xf16>, view<[%output_row_count]x[%bounded_output_size]xf16>, vector<8xi1>
     }

ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/routed_down_quantized_f16_wmma.loom

@@ -592,7 +592,10 @@ kernel.def target(@qwen3_moe_routed_down_gfx11_wave64) @qwen3_moe_routed_down_we
     %weighted_sum = scf.for %route = [%c0 to %bounded_route_count step %c1](%sum = %c0_f32x4 : vector<4xf32>) -> (vector<4xf32>) unroll {
       %assignment = index.madd %token, %bounded_route_count, %route : index
       %routed = vector.load.mask %routed_output_view[%assignment, %channel], %mask, %c0_f16x4 : view<[%assignment_count]x[%bounded_output_size]xf16>, vector<4xi1>, vector<4xf16>
-      %wide = vector.extf %routed : vector<4xf16> to vector<4xf32>
+      // engine #2186: the routed down plane is stored at 2^-6 (see mul_mat_id_swiglu_f16_f16_wmma_body.loom)
+      %stored = vector.extf %routed : vector<4xf16> to vector<4xf32>
+      %plane_unscale = vector.constant 64.0 : vector<4xf32>
+      %wide = vector.mulf %stored, %plane_unscale : vector<4xf32>
       %route_weight = view.load %route_weight_view[%token, %route] : view<[%token_count]x[%bounded_route_count]xf32> -> f32
       %route_weight_x4 = vector.splat %route_weight : vector<4xf32>
       %weighted = vector.mulf %wide, %route_weight_x4 : vector<4xf32>

ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/routed_down_weighted_reduce_next_rmsnorm_f32.loom

@@ -65,7 +65,10 @@ kernel.def target(@qwen3_moe_routed_down_next_norm_gfx11_wave32) @qwen3_moe_rout
     %weighted_sum = scf.for %route = [%c0 to %route_count step %c1](%sum = %c0_f32x8 : vector<8xf32>) -> (vector<8xf32>) unroll {
       %assignment = index.madd %token, %route_count, %route : index
       %routed_f16 = vector.load %routed_output_view[%assignment, %channel] : view<[%assignment_count]x[%output_size]xf16> -> vector<8xf16>
-      %routed = vector.extf %routed_f16 : vector<8xf16> to vector<8xf32>
+      // engine #2186: the routed down plane is stored at 2^-6 (see mul_mat_id_swiglu_f16_f16_wmma_body.loom)
+      %routed_stored = vector.extf %routed_f16 : vector<8xf16> to vector<8xf32>
+      %plane_unscale = vector.constant 64.0 : vector<8xf32>
+      %routed = vector.mulf %routed_stored, %plane_unscale : vector<8xf32>
       %route_weight = view.load %route_weights_view[%token, %route] : view<[%launch_token_count]x[%route_count]xf32> -> f32
       %route_weight_x8 = vector.splat %route_weight : vector<8xf32>
       %weighted = vector.mulf %routed, %route_weight_x8 : vector<8xf32>

ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/routed_down_weighted_reduce_next_rmsnorm_q8_1_x4.loom

@@ -89,7 +89,10 @@ kernel.def target(@qwen3_moe_routed_down_next_norm_q8_gfx11_wave32) @qwen3_moe_r
     %weighted_sum = scf.for %route = [%c0 to %route_count step %c1](%sum = %c0_f32x8 : vector<8xf32>) -> (vector<8xf32>) unroll {
       %assignment = index.madd %token, %route_count, %route : index
       %routed_f16 = vector.load %routed_output_view[%assignment, %channel] : view<[%assignment_count]x[%output_size]xf16> -> vector<8xf16>
-      %routed = vector.extf %routed_f16 : vector<8xf16> to vector<8xf32>
+      // engine #2186: the routed down plane is stored at 2^-6 (see mul_mat_id_swiglu_f16_f16_wmma_body.loom)
+      %routed_stored = vector.extf %routed_f16 : vector<8xf16> to vector<8xf32>
+      %plane_unscale = vector.constant 64.0 : vector<8xf32>
+      %routed = vector.mulf %routed_stored, %plane_unscale : vector<8xf32>
       %route_weight = view.load %route_weights_view[%token, %route] : view<[%launch_token_count]x[%route_count]xf32> -> f32
       %route_weight_x8 = vector.splat %route_weight : vector<8xf32>
       %weighted = vector.mulf %routed, %route_weight_x8 : vector<8xf32>

tests/hrx-backend-test.cpp

@@ -7151,10 +7151,16 @@ static void append_qwen_routed_down_for_graph(ggml::hrx::Graph &              gr
     REQUIRE(output_value != nullptr);
     const ggml::hrx::CommandPlanMoeRoutingBundle * routing_bundle =
         plan.metadata.find_moe_routing_bundle(route_ids_value->id);
+    // engine #2186: both f16 planes of the fused routed feed-forward are stored at 2^-6, and the plan says so: a reader
+    // that does not ask for that scale must not find them.
     const ggml::hrx::CommandPlanAlternateValue * gate_up_alternate = plan.metadata.find_alternate_value(
-        glu_value->id, GGML_TYPE_F16, qwen_routed_gate_up_f16_output_size(tensors.output->ne[2]));
+        glu_value->id, GGML_TYPE_F16, qwen_routed_gate_up_f16_output_size(tensors.output->ne[2]), 6);
     const ggml::hrx::CommandPlanAlternateValue * routed_down_alternate = plan.metadata.find_alternate_value(
-        output_value->id, GGML_TYPE_F16, qwen_routed_down_f16_output_size(tensors.output->ne[2]));
+        output_value->id, GGML_TYPE_F16, qwen_routed_down_f16_output_size(tensors.output->ne[2]), 6);
+    REQUIRE(plan.metadata.find_alternate_value(glu_value->id, GGML_TYPE_F16,
+                                               qwen_routed_gate_up_f16_output_size(tensors.output->ne[2])) == nullptr);
+    REQUIRE(plan.metadata.find_alternate_value(output_value->id, GGML_TYPE_F16,
+                                               qwen_routed_down_f16_output_size(tensors.output->ne[2])) == nullptr);
     REQUIRE(routing_bundle != nullptr);
     REQUIRE(gate_up_alternate != nullptr);
     REQUIRE(routed_down_alternate != nullptr);
@@ -8275,6 +8281,8 @@ static void run_qwen_routed_gate_up_dispatch_checks() {
         // the f32 GLU output the down projection reads.
         REQUIRE(plan.metadata.find_alternate_value(imported.graph.values().find_tensor(tensors.glu)->id, GGML_TYPE_F16,
                                                    qwen_routed_gate_up_f16_output_size(4)) == nullptr);
+        REQUIRE(plan.metadata.find_alternate_value(imported.graph.values().find_tensor(tensors.glu)->id, GGML_TYPE_F16,
+                                                   qwen_routed_gate_up_f16_output_size(4), 6) == nullptr);
         const size_t             down_index = producer_index_for_tensor(imported.graph, tensors.output);
         ggml::hrx::DispatchMatch down_match;
         REQUIRE(match_dispatch_at_index(imported.graph, plan, covered_nodes, down_index, down_match));

tests/test-hrx-ops.cpp

@@ -946,6 +946,7 @@ struct RoutedMoeGraph {
     ggml_tensor * hidden_state = nullptr;
     ggml_tensor * norm_weight  = nullptr;
     ggml_tensor * output       = nullptr;
+    ggml_tensor * glu          = nullptr;  // the SwiGLU output, readable on a backend that keeps intermediates
 };
 
 static RoutedMoeGraph build_routed_moe_graph(ggml_context * ctx,
@@ -981,6 +982,7 @@ static RoutedMoeGraph build_routed_moe_graph(ggml_context * ctx,
     REQUIRE(up != nullptr);
     ggml_tensor * glu = ggml_glu_split(ctx, gate, up, GGML_GLU_OP_SWIGLU);
     REQUIRE(glu != nullptr);
+    graph.glu          = glu;
     ggml_tensor * down = ggml_mul_mat_id(ctx, graph.down_weight, glu, route_ids);
     REQUIRE(down != nullptr);
     ggml_tensor * weighted = ggml_mul(ctx, down, route_weights);
@@ -4844,6 +4846,121 @@ static void run_routed_moe_large_input_case(ggml_type gate_up_weight_type, float
     ggml_backend_free(hrx_backend);
 }
 
+// engine #2186: the fused routed feed-forward keeps the SwiGLU product and the routed down sums in f16. With every
+// input scaled by `scale` the gate and up sums grow with it and the product with its square, so above some scale they
+// leave the f16 range (65504) while the result stays far inside single precision.
+// The reference is the CPU on the same weights held as f32 (the quantised rows, dequantised): with K-quant weights the
+// CPU quantises the activation to 8 bits per block, which costs it more precision than the device's f16 does.
+static void run_routed_moe_large_product_case(ggml_type gate_up_weight_type, float scale) {
+    ggml_backend_t cpu_backend = init_cpu_backend();
+    ggml_backend_t hrx_backend = ggml_backend_hrx_init(0);
+    REQUIRE(hrx_backend != nullptr);
+
+    ggml_init_params params = {};
+    params.mem_size         = 64 * 1024 * 1024;
+    params.no_alloc         = true;
+    ggml_context * cpu_ctx  = ggml_init(params);
+    ggml_context * hrx_ctx  = ggml_init(params);
+    REQUIRE(cpu_ctx != nullptr);
+    REQUIRE(hrx_ctx != nullptr);
+
+    RoutedMoeGraph cpu = build_routed_moe_graph(cpu_ctx, GGML_TYPE_F32, false, GGML_TYPE_F32);
+    RoutedMoeGraph hrx = build_routed_moe_graph(hrx_ctx, GGML_TYPE_Q4_K, false, gate_up_weight_type);
+
+    ggml_cgraph * cpu_graph = ggml_new_graph(cpu_ctx);
+    ggml_cgraph * hrx_graph = ggml_new_graph(hrx_ctx);
+    REQUIRE(cpu_graph != nullptr);
+    REQUIRE(hrx_graph != nullptr);
+    ggml_build_forward_expand(cpu_graph, cpu.output);
+    ggml_build_forward_expand(hrx_graph, hrx.output);
+
+    REQUIRE(test_sequence_has(scheduled_kernel_sequence(hrx_graph), "loom_libs:ggml_mul_mat_id_swiglu_f16_f16_wmma"));
+
+    ggml_backend_buffer_t cpu_buffer = ggml_backend_alloc_ctx_tensors(cpu_ctx, cpu_backend);
+    ggml_backend_buffer_t hrx_buffer = ggml_backend_alloc_ctx_tensors(hrx_ctx, hrx_backend);
+    REQUIRE(cpu_buffer != nullptr);
+    REQUIRE(hrx_buffer != nullptr);
+
+    // distinct router logits: make_router_logits() repeats eight values across the 128 experts, and with ties the CPU
+    // and the device may each pick a different valid top eight, which is a different graph, not an error
+    std::vector<float> logits(static_cast<size_t>(kQwenRouterExpertCount));
+    for (int64_t expert = 0; expert < kQwenRouterExpertCount; ++expert) {
+        logits[static_cast<size_t>(expert)] = 0.05f * static_cast<float>((expert * 37) % kQwenRouterExpertCount);
+    }
+    std::vector<float>       input  = make_pattern_f32(kQwenHiddenSize, 15, 0.01f);
+    for (float & value : input) {
+        value *= scale;
+    }
+    const std::vector<float> hidden = make_pattern_f32(kQwenHiddenSize, 16, 0.02f);
+    set_tensor_pair_bytes(cpu_backend, cpu.logits, hrx_backend, hrx.logits, logits.data(),
+                          logits.size() * sizeof(float));
+    set_tensor_pair_bytes(cpu_backend, cpu.input, hrx_backend, hrx.input, input.data(), input.size() * sizeof(float));
+    set_tensor_pair_bytes(cpu_backend, cpu.hidden_state, hrx_backend, hrx.hidden_state, hidden.data(),
+                          hidden.size() * sizeof(float));
+    // the device gets the quantised rows, the CPU the same rows dequantised to f32
+    const auto set_weight_pair = [&](ggml_tensor * cpu_weight, ggml_tensor * hrx_weight, ggml_type type,
+                                     int64_t row_length, int64_t row_count, int seed) {
+        const std::vector<uint8_t> rows   = make_quantized_rows(type, row_length, row_count, seed);
+        const ggml_type_traits *   traits = ggml_get_type_traits(type);
+        REQUIRE(traits != nullptr);
+        REQUIRE(traits->to_float != nullptr);
+        std::vector<float> wide(static_cast<size_t>(row_length * row_count));
+        const size_t       row_size = ggml_row_size(type, row_length);
+        for (int64_t r = 0; r < row_count; ++r) {
+            traits->to_float(rows.data() + static_cast<size_t>(r) * row_size,
+                             wide.data() + static_cast<size_t>(r * row_length), row_length);
+        }
+        ggml_backend_tensor_set(hrx_weight, rows.data(), 0, rows.size());
+        ggml_backend_tensor_set(cpu_weight, wide.data(), 0, wide.size() * sizeof(float));
+    };
+    set_weight_pair(cpu.gate_weight, hrx.gate_weight, gate_up_weight_type, kQwenHiddenSize,
+                    kQwenMoeIntermediate * kQwenRouterExpertCount, 17);
+    set_weight_pair(cpu.up_weight, hrx.up_weight, gate_up_weight_type, kQwenHiddenSize,
+                    kQwenMoeIntermediate * kQwenRouterExpertCount, 18);
+    set_weight_pair(cpu.down_weight, hrx.down_weight, GGML_TYPE_Q4_K, kQwenMoeIntermediate,
+                    kQwenHiddenSize * kQwenRouterExpertCount, 19);
+
+    REQUIRE(ggml_backend_graph_compute(cpu_backend, cpu_graph) == GGML_STATUS_SUCCESS);
+    REQUIRE(ggml_backend_graph_compute(hrx_backend, hrx_graph) == GGML_STATUS_SUCCESS);
+    ggml_backend_synchronize(cpu_backend);
+    ggml_backend_synchronize(hrx_backend);
+    const std::vector<float> cpu_out = get_f32_tensor(cpu_backend, cpu.output);
+    const std::vector<float> hrx_out = get_f32_tensor(hrx_backend, hrx.output);
+    const std::vector<float> cpu_glu = get_f32_tensor(cpu_backend, cpu.glu);
+    REQUIRE(cpu_out.size() == hrx_out.size());
+
+    float glu_max = 0.0f;
+    for (float value : cpu_glu) {
+        REQUIRE(std::isfinite(value));
+        glu_max = std::max(glu_max, std::fabs(value));
+    }
+    double error      = 0.0;
+    double norm       = 0.0;
+    size_t non_finite = 0;
+    for (size_t i = 0; i < cpu_out.size(); ++i) {
+        REQUIRE(std::isfinite(cpu_out[i]));
+        if (!std::isfinite(hrx_out[i])) {
+            ++non_finite;
+            continue;
+        }
+        const double diff = static_cast<double>(hrx_out[i]) - static_cast<double>(cpu_out[i]);
+        error += diff * diff;
+        norm  += static_cast<double>(cpu_out[i]) * static_cast<double>(cpu_out[i]);
+    }
+    std::fprintf(stderr, "routed large product: scale=%g max|silu(gate)*up| (cpu)=%g non-finite=%zu of %zu nmse(device vs cpu)=%g\n",
+                 static_cast<double>(scale), static_cast<double>(glu_max), non_finite, cpu_out.size(),
+                 norm > 0.0 ? error / norm : -1.0);
+    REQUIRE(non_finite == 0);
+    REQUIRE(error / norm < 1.0e-3);
+
+    ggml_backend_buffer_free(cpu_buffer);
+    ggml_backend_buffer_free(hrx_buffer);
+    ggml_free(cpu_ctx);
+    ggml_free(hrx_ctx);
+    ggml_backend_free(cpu_backend);
+    ggml_backend_free(hrx_backend);
+}
+
 static void run_decode_routed_moe_scheduling_case(ggml_type down_weight_type, bool alias_gate_input = false) {
     ggml_init_params params = {};
     params.mem_size         = 128 * 1024 * 1024;
@@ -5938,6 +6055,11 @@ static void add_cases(test_runner::Suite & suite) {
     }
     add(suite, "decode_attention_qkv_scheduling",
         [=] { run_decode_attention_qkv_scheduling_case(); });
+    // engine #2186: the gate/up product, at the ordinary scale and at scales that take it past the f16 range
+    for (float scale : { 1.0f, 30.0f, 100.0f, 300.0f, 1000.0f }) {
+        add(suite, std::string("routed_moe.q4_K.large_product.") + std::to_string(static_cast<int>(scale)),
+            [=] { run_routed_moe_large_product_case(GGML_TYPE_Q4_K, scale); });
+    }
     add(suite, "routed_moe_decode.q4_k",
         [=] { run_decode_routed_moe_scheduling_case(GGML_TYPE_Q4_K); });
     add(suite, "routed_moe_decode.q6_k",