From 6e3ce2f0bc4561fa2b8c47274f7d0d913d700875 Mon Sep 17 00:00:00 2001 From: bong-water-water-bong Date: Wed, 11 Nov 2026 07:19:04 -0400 Subject: [PATCH] hrx: the routed gate/up SwiGLU narrows its input at 2^-6, so an input above 65504 stays finite (engine #2127) The four projection kernels of the fused routed gate/up SwiGLU (Q4_K, Q5_K, Q6_K and the generic one) loaded the f32 input and narrowed it to f16 unscaled. One input above 65504 became inf and every output of the routed feed-forward was non-finite. They now narrow at 2^-6, and the SwiGLU body scales the gate and up sums back by 2^6 in f32 before the activation. Both factors are powers of two, so an input that fit before is rounded as before. Test: test-hrx-ops routed_moe.{q4_K,q6_K}.large_input.{60000,100000}. One input is the spike and its gate/up weight columns are exactly zero, so the device's own result without the spike is the reference (the CPU is not: it quantises the activation per block, and a spike costs the rest of its block their precision). Co-Authored-By: Claude Fable 5.2 --- .../mul_mat_id_q4k_f16_wmma_projection.loom | 6 +- .../mul_mat_id_q5k_f16_wmma_projection.loom | 6 +- .../mul_mat_id_q6k_f16_wmma_projection.loom | 6 +- .../mul_mat_id_swiglu_f16_f16_accumulate.loom | 6 +- .../mul_mat_id_swiglu_f16_f16_wmma_body.loom | 8 +- tests/test-hrx-ops.cpp | 137 +++++++++++++++++- 6 files changed, 160 insertions(+), 9 deletions(-) diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/mul_mat_id_q4k_f16_wmma_projection.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/mul_mat_id_q4k_f16_wmma_projection.loom index c744d20f85a0..0359c23fd81b 100644 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/mul_mat_id_q4k_f16_wmma_projection.loom +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/mul_mat_id_q4k_f16_wmma_projection.loom @@ -181,7 +181,11 @@ template.def<@ggml.mul_mat_id_q4k_f16_wmma.project_vector16> device @ggml_mul_ma %token, %input_token_count = index.assume %token0, %token_count [lt(%token0, %token_count)] : index, index %input_k = index.add %k_origin, %load_k : index %loaded = vector.load %input_view[%token, %input_k] : view<[%token_count]x[%input_size]xf32> -> vector<4xf32> - %converted = vector.fptrunc %loaded : vector<4xf32> to vector<4xf16> + // engine #2127: narrow the routed gate/up input at 2^-6, so an input above 65504 stays finite through the f16 + // stage; the SwiGLU body scales the gate and up sums back by 2^6 (mul_mat_id_swiglu_f16_f16_wmma_body.loom). + %input_prescale = vector.constant 0.015625 : vector<4xf32> + %prescaled = vector.mulf %loaded, %input_prescale : vector<4xf32> + %converted = vector.fptrunc %prescaled : vector<4xf32> to vector<4xf16> scf.yield %converted : vector<4xf16> } else { scf.yield %c0_f16x4 : vector<4xf16> diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/mul_mat_id_q5k_f16_wmma_projection.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/mul_mat_id_q5k_f16_wmma_projection.loom index 34714bf1633a..b3ec0bebc96f 100644 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/mul_mat_id_q5k_f16_wmma_projection.loom +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/mul_mat_id_q5k_f16_wmma_projection.loom @@ -202,7 +202,11 @@ template.def<@ggml.mul_mat_id_q5k_f16_wmma.project_vector16> device @ggml_mul_ma %token, %input_token_count = index.assume %token0, %token_count [lt(%token0, %token_count)] : index, index %input_k = index.add %k_origin, %load_k : index %loaded = vector.load %input_view[%token, %input_k] : view<[%token_count]x[%input_size]xf32> -> vector<4xf32> - %converted = vector.fptrunc %loaded : vector<4xf32> to vector<4xf16> + // engine #2127: narrow the routed gate/up input at 2^-6, so an input above 65504 stays finite through the f16 + // stage; the SwiGLU body scales the gate and up sums back by 2^6 (mul_mat_id_swiglu_f16_f16_wmma_body.loom). + %input_prescale = vector.constant 0.015625 : vector<4xf32> + %prescaled = vector.mulf %loaded, %input_prescale : vector<4xf32> + %converted = vector.fptrunc %prescaled : vector<4xf32> to vector<4xf16> scf.yield %converted : vector<4xf16> } else { scf.yield %c0_f16x4 : vector<4xf16> diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/mul_mat_id_q6k_f16_wmma_projection.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/mul_mat_id_q6k_f16_wmma_projection.loom index 284577982a6d..87d4b6dc4c28 100644 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/mul_mat_id_q6k_f16_wmma_projection.loom +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/mul_mat_id_q6k_f16_wmma_projection.loom @@ -150,7 +150,11 @@ template.def<@ggml.mul_mat_id_q6k_f16_wmma.project_vector16> device @ggml_mul_ma %token, %input_token_count = index.assume %token0, %token_count [lt(%token0, %token_count)] : index, index %input_k = index.add %k_origin, %load_k : index %loaded = vector.load %input_view[%token, %input_k] : view<[%token_count]x[%input_size]xf32> -> vector<4xf32> - %converted = vector.fptrunc %loaded : vector<4xf32> to vector<4xf16> + // engine #2127: narrow the routed gate/up input at 2^-6, so an input above 65504 stays finite through the f16 + // stage; the SwiGLU body scales the gate and up sums back by 2^6 (mul_mat_id_swiglu_f16_f16_wmma_body.loom). + %input_prescale = vector.constant 0.015625 : vector<4xf32> + %prescaled = vector.mulf %loaded, %input_prescale : vector<4xf32> + %converted = vector.fptrunc %prescaled : vector<4xf32> to vector<4xf16> scf.yield %converted : vector<4xf16> } else { scf.yield %c0_f16x4 : vector<4xf16> diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/mul_mat_id_swiglu_f16_f16_accumulate.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/mul_mat_id_swiglu_f16_f16_accumulate.loom index 874bd565ae5a..a018580ac304 100644 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/mul_mat_id_swiglu_f16_f16_accumulate.loom +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/mul_mat_id_swiglu_f16_f16_accumulate.loom @@ -117,7 +117,11 @@ template.def<@ggml.mul_mat_id_swiglu_f16_f16.project_generic_vector16> device @g %token, %input_token_count = index.assume %token0, %token_count [lt(%token0, %token_count)] : index, index %input_k = index.add %k_origin, %load_k : index %loaded = vector.load %input_view[%token, %input_k] : view<[%token_count]x[%input_size]xf32> -> vector<4xf32> - %converted = vector.fptrunc %loaded : vector<4xf32> to vector<4xf16> + // engine #2127: narrow the routed gate/up input at 2^-6, so an input above 65504 stays finite through the f16 + // stage; the SwiGLU body scales the gate and up sums back by 2^6 (mul_mat_id_swiglu_f16_f16_wmma_body.loom). + %input_prescale = vector.constant 0.015625 : vector<4xf32> + %prescaled = vector.mulf %loaded, %input_prescale : vector<4xf32> + %converted = vector.fptrunc %prescaled : vector<4xf32> to vector<4xf16> scf.yield %converted : vector<4xf16> } else { scf.yield %c0_f16x4 : vector<4xf16> diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/mul_mat_id_swiglu_f16_f16_wmma_body.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/mul_mat_id_swiglu_f16_f16_wmma_body.loom index 21169146fca9..eefa942ef547 100644 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/mul_mat_id_swiglu_f16_f16_wmma_body.loom +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/mul_mat_id_swiglu_f16_f16_wmma_body.loom @@ -175,13 +175,17 @@ template.def<@ggml.mul_mat_id_swiglu_f16_f16.body> device @ggml_mul_mat_id_swigl vector.fragment.store %gate_result, %result_fragment_view[%c0, %c0] shape [%m, %n] : vector<16xf16>, view<16x16xf16, %result_fragment_layout> kernel.barrier scope(subgroup) ordering(acq_rel) %gate_values = vector.load %result_physical_view[%publish_route, %publish_channel_add] : view<16x16xf16> -> vector<8xf16> - %gate_wide = vector.extf %gate_values : vector<8xf16> to vector<8xf32> + // engine #2127: the projections narrowed the input at 2^-6, so the gate and up sums are at 2^-6 too; undo it in f32 + %sum_rescale = vector.constant 64.0 : vector<8xf32> + %gate_prescaled = vector.extf %gate_values : vector<8xf16> to vector<8xf32> + %gate_wide = vector.mulf %gate_prescaled, %sum_rescale : vector<8xf32> kernel.barrier scope(subgroup) ordering(acq_rel) vector.fragment.store %up_result, %result_fragment_view[%c0, %c0] shape [%m, %n] : vector<16xf16>, view<16x16xf16, %result_fragment_layout> kernel.barrier scope(subgroup) ordering(acq_rel) scf.if %writes { %up_values = vector.load %result_physical_view[%publish_route, %publish_channel_add] : view<16x16xf16> -> vector<8xf16> - %up_wide = vector.extf %up_values : vector<8xf16> to vector<8xf32> + %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> %mask = vector.mask.range [%channel to %bounded_output_size step %c1] : index -> vector<8xi1> diff --git a/tests/test-hrx-ops.cpp b/tests/test-hrx-ops.cpp index 8b56790ec347..ae4a73919f15 100644 --- a/tests/test-hrx-ops.cpp +++ b/tests/test-hrx-ops.cpp @@ -950,7 +950,8 @@ struct RoutedMoeGraph { static RoutedMoeGraph build_routed_moe_graph(ggml_context * ctx, ggml_type down_weight_type, - bool include_next_rmsnorm) { + bool include_next_rmsnorm, + ggml_type gate_up_weight_type = GGML_TYPE_Q4_K) { RoutedMoeGraph graph; constexpr int64_t token_count = 1; graph.logits = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, kQwenRouterExpertCount, token_count); @@ -962,9 +963,9 @@ static RoutedMoeGraph build_routed_moe_graph(ggml_context * ctx, graph.input = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, kQwenHiddenSize, 1, token_count); graph.gate_weight = - ggml_new_tensor_3d(ctx, GGML_TYPE_Q4_K, kQwenHiddenSize, kQwenMoeIntermediate, kQwenRouterExpertCount); + ggml_new_tensor_3d(ctx, gate_up_weight_type, kQwenHiddenSize, kQwenMoeIntermediate, kQwenRouterExpertCount); graph.up_weight = - ggml_new_tensor_3d(ctx, GGML_TYPE_Q4_K, kQwenHiddenSize, kQwenMoeIntermediate, kQwenRouterExpertCount); + ggml_new_tensor_3d(ctx, gate_up_weight_type, kQwenHiddenSize, kQwenMoeIntermediate, kQwenRouterExpertCount); graph.down_weight = ggml_new_tensor_3d(ctx, down_weight_type, kQwenMoeIntermediate, kQwenHiddenSize, kQwenRouterExpertCount); graph.hidden_state = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, kQwenHiddenSize, token_count); @@ -4721,6 +4722,128 @@ static void run_routed_moe_cpu_reference_case(ggml_type down_weight_type, bool i ggml_backend_free(hrx_backend); } +// engine #2127: the routed gate/up SwiGLU with ONE input above the f16 maximum. The gate and up weights are the usual +// pattern with their first 32 columns exactly zero, and the spike is input 5. On the reference the spike contributes +// nothing, so the projections, the gated product and the output are those of an ordinary input; on the device a spike +// that became inf when narrowed gives inf * 0 = NaN. So the only value out of range is the one the gate/up projection +// kernels narrow. `spike` is that input. +static void run_routed_moe_large_input_case(ggml_type gate_up_weight_type, float spike) { + 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_Q4_K, false, gate_up_weight_type); + 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); + + const std::vector logits = make_router_logits(1); + std::vector input = make_pattern_f32(kQwenHiddenSize, 15, 0.01f); + input[5] = 0.25f; + const std::vector 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)); + + // make_quantized_rows' pattern with the first 256 columns zero: a whole quantisation block of zeros decodes to exact + // zeros in every format used here (the plain-pass comparison below would show it if one did not) + const auto small_rows = [gate_up_weight_type](int64_t row_length, int64_t row_count, int seed) { + const ggml_type_traits * traits = ggml_get_type_traits(gate_up_weight_type); + REQUIRE(traits != nullptr); + REQUIRE(traits->from_float_ref != nullptr); + ggml_quantize_init(gate_up_weight_type); + const size_t row_size = ggml_row_size(gate_up_weight_type, row_length); + std::vector data(static_cast(row_count) * row_size); + std::vector row(static_cast(row_length)); + for (int64_t r = 0; r < row_count; ++r) { + for (int64_t c = 0; c < row_length; ++c) { + const int value = static_cast((r * 13 + c * 7 + seed * 31) % 101) - 50; + row[static_cast(c)] = c < 256 ? 0.0f : static_cast(value) * 0.005f; + } + traits->from_float_ref(row.data(), data.data() + static_cast(r) * row_size, row_length); + } + return data; + }; + { + const std::vector gate = small_rows(kQwenHiddenSize, kQwenMoeIntermediate * kQwenRouterExpertCount, 17); + set_tensor_pair_bytes(cpu_backend, cpu.gate_weight, hrx_backend, hrx.gate_weight, gate.data(), gate.size()); + } + { + const std::vector up = small_rows(kQwenHiddenSize, kQwenMoeIntermediate * kQwenRouterExpertCount, 18); + set_tensor_pair_bytes(cpu_backend, cpu.up_weight, hrx_backend, hrx.up_weight, up.data(), up.size()); + } + { + const std::vector down = + make_quantized_rows(GGML_TYPE_Q4_K, kQwenMoeIntermediate, kQwenHiddenSize * kQwenRouterExpertCount, 19); + set_tensor_pair_bytes(cpu_backend, cpu.down_weight, hrx_backend, hrx.down_weight, down.data(), down.size()); + } + + // First pass: input 5 is an ordinary value. This is the reference for the spike pass, on both backends, because + // the spike's weight column is exactly zero. (The CPU is not a usable reference for the spike pass: for a K-quant + // weight it quantises the activation to 8 bits per block, so a spike costs the rest of its block their precision.) + 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 cpu_plain = get_f32_tensor(cpu_backend, cpu.output); + const std::vector hrx_plain = get_f32_tensor(hrx_backend, hrx.output); + require_close(hrx_plain, cpu_plain, 2.5f, 1.0e-1f); + + // Second pass on the device: the same graph with the spike. + input[5] = spike; + ggml_backend_tensor_set(hrx.input, input.data(), 0, input.size() * sizeof(float)); + REQUIRE(ggml_backend_graph_compute(hrx_backend, hrx_graph) == GGML_STATUS_SUCCESS); + ggml_backend_synchronize(hrx_backend); + const std::vector hrx_spike = get_f32_tensor(hrx_backend, hrx.output); + REQUIRE(hrx_spike.size() == hrx_plain.size()); + + double error = 0.0; + double norm = 0.0; + size_t non_finite = 0; + for (size_t i = 0; i < hrx_plain.size(); ++i) { + if (!std::isfinite(hrx_spike[i])) { + ++non_finite; + continue; + } + const double diff = static_cast(hrx_spike[i]) - static_cast(hrx_plain[i]); + error += diff * diff; + norm += static_cast(hrx_plain[i]) * static_cast(hrx_plain[i]); + } + std::fprintf(stderr, "routed large input: spike=%g non-finite=%zu of %zu nmse(spike pass vs plain pass, device)=%g\n", + static_cast(spike), non_finite, hrx_plain.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; @@ -5805,6 +5928,14 @@ static void add_cases(test_runner::Suite & suite) { [=] { run_asymmetric_flash_attention_cpu_reference_case(); }); add(suite, "routed_moe.q6_k", [=] { run_routed_moe_cpu_reference_case(GGML_TYPE_Q6_K, false); }); + // engine #2127: 6.0e4 is inside f16 (the control); 1.0e5 is not + // Q4_K and Q6_K gate/up weights reach the fused routed SwiGLU in this graph. Q5_K and IQ4_XS do not (the scheduler + // does not pick ggml_mul_mat_id_swiglu_f16_f16_wmma for them here), so their projection kernels have no arm. + for (const ggml_type gate_up_type : { GGML_TYPE_Q4_K, GGML_TYPE_Q6_K }) { + const std::string name = std::string("routed_moe.") + ggml_type_name(gate_up_type) + ".large_input."; + add(suite, name + "60000", [=] { run_routed_moe_large_input_case(gate_up_type, 6.0e4f); }); + add(suite, name + "100000", [=] { run_routed_moe_large_input_case(gate_up_type, 1.0e5f); }); + } add(suite, "decode_attention_qkv_scheduling", [=] { run_decode_attention_qkv_scheduling_case(); }); add(suite, "routed_moe_decode.q4_k",