diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/llm_attention_q_matmul_rope_decode_f32_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/llm_attention_q_matmul_rope_decode_f32_f32.loom index 74435a3bb1a5..0e087e908aff 100644 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/llm_attention_q_matmul_rope_decode_f32_f32.loom +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/llm_attention_q_matmul_rope_decode_f32_f32.loom @@ -30,10 +30,14 @@ template.def<@ggml.mul_mat_f32_f32_decode.publish_pair2> device priority(20) @ll %c0_f32 = scalar.constant 0.0 : f32 %publish = scalar.andi %publish_pair, %row1_valid : i1 %token, %launch_token_count = index.assume %token0, %token_count [lt(%token0, %token_count)] : index, index - %head = index.div %channel, %head_size : index - %head_channel = index.rem %channel, %head_size : index - %theta_channel = index.div %head_channel, %c2 : index + // %channel is the even row of a rotation pair, so index by pair: the remainder is then bounded by + // head_size / 2 and the two-wide store below is in bounds for every head size, not only for a + // power of two (engine #1926: 48, 96 and 160 did not compile). %half_head_size = index.div %head_size, %c2 : index + %pair_channel = index.div %channel, %c2 : index + %head = index.div %pair_channel, %half_head_size : index + %theta_channel = index.rem %pair_channel, %half_head_size : index + %head_channel = index.mul %theta_channel, %c2 : index %positions_noalias, %theta_noalias, %freq_factors_noalias, %output_noalias = buffer.assume.noalias %positions, %theta, %freq_factors, %output : buffer, buffer, buffer, buffer scf.if %publish { %positions_view = buffer.view %positions_noalias[%c0_offset] : buffer -> view<[%launch_token_count]xi32> diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 3f00f132620a..d6d2af1fa53f 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -465,6 +465,12 @@ if (TARGET ggml-hrx) target_include_directories(test-hrx-decode-stride PRIVATE ../ggml/include ../ggml/src ../ggml/src/ggml-hrx) add_test(NAME test-hrx-decode-stride COMMAND test-hrx-decode-stride) + # the fused Q projection at one token on head sizes that are not a power of two (engine #1926); skips without HRX0 + add_executable(test-hrx-attention-q-decode-head-size test-hrx-attention-q-decode-head-size.cpp) + target_link_libraries(test-hrx-attention-q-decode-head-size PRIVATE ggml-hrx ggml hrx::hrx loomc::loomc) + target_include_directories(test-hrx-attention-q-decode-head-size PRIVATE ../ggml/include ../ggml/src ../ggml/src/ggml-hrx) + add_test(NAME test-hrx-attention-q-decode-head-size COMMAND test-hrx-attention-q-decode-head-size) + add_executable(test-hrx-mul-mat-id-k32 test-hrx-mul-mat-id-k32.cpp) target_link_libraries(test-hrx-mul-mat-id-k32 PRIVATE ggml-hrx ggml hrx::hrx loomc::loomc) target_include_directories(test-hrx-mul-mat-id-k32 PRIVATE ../ggml/include ../ggml/src ../ggml/src/ggml-hrx) diff --git a/tests/test-hrx-attention-q-decode-head-size.cpp b/tests/test-hrx-attention-q-decode-head-size.cpp new file mode 100644 index 000000000000..951240c6c0df --- /dev/null +++ b/tests/test-hrx-attention-q-decode-head-size.cpp @@ -0,0 +1,203 @@ +// Copyright 2026 bong-water-water-bong +// SPDX-License-Identifier: Apache-2.0 +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// The fused Q projection at one token (MUL_MAT, the head layout, ROPE) on head sizes that are not a power of two: +// 48 (stories15M), 96 and 160 (Sarashina2.2-3B), with 64 and 128 as the controls, on BF16, F16 and Q8_0 weights - +// the formats that reach llm_attention_q_matmul_rope_decode_f32_f32 at one token. The kernel wrote each rotated +// pair at `channel rem head_size`, which the Loom compiler could bound only for a power of two, so the other head +// sizes failed to compile and the decode returned an error (engine #1926). Each graph runs on the HRX device +// itself (no scheduler), must be planned on that kernel, and is compared with the CPU backend (normalized MSE, as +// test-backend-ops). + +#include "dispatch/dispatch-scheduler.h" +#include "ggml-alloc.h" +#include "ggml-backend.h" +#include "ggml-cpu.h" +#include "ggml.h" +#include "graph/graph.h" +#include "kernel-corpus/kernel-corpus.h" + +#include +#include +#include +#include +#include +#include + +#define REQUIRE(condition) \ + do { \ + if (!(condition)) { \ + std::fprintf(stderr, "%s:%d: requirement failed: %s\n", __FILE__, __LINE__, #condition); \ + std::abort(); \ + } \ + } while (false) + +static constexpr double kMaxNmse = 5e-4; +static constexpr int32_t kPosition = 37; +static const char * kKernel = "loom_libs:llm_attention_q_matmul_rope_decode_f32_f32"; + +static std::string kernel_name_for_id(uint64_t kernel_id) { + const ggml::hrx::KernelResolveResult resolved = + ggml::hrx::resolve_kernel_definition(ggml::hrx::get_qwen_kernel_corpus(), "gfx1151", kernel_id); + REQUIRE(resolved.found()); + return ggml::hrx::kernel_definition_name(*resolved.definition); +} + +// The HRX kernels planned for the graph, or "" when HRX does not take it. +static std::string hrx_plan(ggml_cgraph * graph) { + ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); + REQUIRE(imported.valid()); + ggml::hrx::DispatchScheduler scheduler; + ggml::hrx::DispatchScheduleDiagnostics diagnostics; + if (!scheduler.schedule_graph(imported.graph, { "gfx1151" }, &diagnostics) || !scheduler.plan().valid()) { + return ""; + } + std::string names; + for (const ggml::hrx::Dispatch & dispatch : scheduler.plan().dispatches) { + names += (names.empty() ? "" : "+") + kernel_name_for_id(dispatch.kernel.kernel_id); + } + return names; +} + +struct Case { + ggml_type type; + int64_t head_size; + int64_t head_count; +}; + +struct Data { + std::vector weights; + std::vector x; +}; + +static Data make_data(const Case & c, uint32_t seed) { + std::mt19937 rng(seed); + std::uniform_real_distribution uniform(-1.0f, 1.0f); + const int64_t size = c.head_size * c.head_count; + std::vector w(static_cast(size * size)); + for (float & v : w) { + v = uniform(rng); + } + Data d; + d.weights.resize(ggml_row_size(c.type, size) * size); + ggml_quantize_chunk(c.type, w.data(), d.weights.data(), 0, size, size, nullptr); + d.x.resize(static_cast(size)); + for (float & v : d.x) { + v = uniform(rng); + } + return d; +} + +enum class Run { Done, NotAdmitted, ComputeFailed }; + +static Run run(ggml_backend_t backend, const Case & c, const Data & d, std::vector & result, std::string * plan) { + ggml_init_params params = { 32 * ggml_tensor_overhead() + ggml_graph_overhead(), nullptr, true }; + ggml_context * ctx = ggml_init(params); + REQUIRE(ctx != nullptr); + const int64_t size = c.head_size * c.head_count; + ggml_tensor * weights = ggml_new_tensor_2d(ctx, c.type, size, size); + ggml_tensor * x = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, size, 1); + ggml_tensor * positions = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, 1); + ggml_tensor * q = ggml_reshape_3d(ctx, ggml_mul_mat(ctx, weights, x), c.head_size, c.head_count, 1); + ggml_tensor * out = ggml_rope_ext(ctx, q, positions, nullptr, static_cast(c.head_size), + GGML_ROPE_TYPE_NORMAL, 0, 10000.0f, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f); + ggml_cgraph * graph = ggml_new_graph(ctx); + ggml_build_forward_expand(graph, out); + if (plan != nullptr) { + *plan = ggml_backend_supports_op(backend, out) ? hrx_plan(graph) : ""; + if (plan->empty()) { + ggml_free(ctx); + return Run::NotAdmitted; + } + } + ggml_backend_buffer_t buffer = ggml_backend_alloc_ctx_tensors(ctx, backend); + REQUIRE(buffer != nullptr); + ggml_backend_tensor_set(weights, d.weights.data(), 0, d.weights.size()); + ggml_backend_tensor_set(x, d.x.data(), 0, d.x.size() * sizeof(float)); + ggml_backend_tensor_set(positions, &kPosition, 0, sizeof(kPosition)); + const bool computed = ggml_backend_graph_compute(backend, graph) == GGML_STATUS_SUCCESS; + if (computed) { + result.resize(static_cast(ggml_nelements(out))); + ggml_backend_tensor_get(out, result.data(), 0, result.size() * sizeof(float)); + } + ggml_backend_buffer_free(buffer); + ggml_free(ctx); + return computed ? Run::Done : Run::ComputeFailed; +} + +static double nmse(const std::vector & got, const std::vector & expected) { + double err = 0.0; + double ref = 0.0; + for (size_t i = 0; i < got.size(); ++i) { + const double diff = static_cast(got[i]) - expected[i]; + err += diff * diff; + ref += static_cast(expected[i]) * expected[i]; + } + return ref > 0.0 ? err / ref : err; +} + +int main() { + std::setvbuf(stdout, nullptr, _IOLBF, 0); // the per-case lines must survive the abort in REQUIRE + ggml_backend_dev_t device = ggml_backend_dev_by_name("HRX0"); + if (device == nullptr) { + ggml_backend_load_all(); + device = ggml_backend_dev_by_name("HRX0"); + } + if (device == nullptr) { + std::printf("test-hrx-attention-q-decode-head-size: no HRX0 device, skipped\n"); + return 0; + } + ggml_backend_t hrx = ggml_backend_dev_init(device, nullptr); + ggml_backend_t cpu = ggml_backend_cpu_init(); + REQUIRE(hrx != nullptr && cpu != nullptr); + + std::vector cases; + for (ggml_type type : { GGML_TYPE_BF16, GGML_TYPE_F16, GGML_TYPE_Q8_0 }) { + cases.push_back({ type, 48, 6 }); // stories15M + cases.push_back({ type, 96, 8 }); + cases.push_back({ type, 160, 16 }); // Sarashina2.2-3B + cases.push_back({ type, 64, 16 }); // control + cases.push_back({ type, 128, 8 }); // control + } + int failures = 0; + uint32_t seed = 1; + for (const Case & c : cases) { + const Data d = make_data(c, seed++); + std::string plan; + std::vector got; + std::vector expected; + const Run on_hrx = run(hrx, c, d, got, &plan); + if (on_hrx != Run::Done || plan != kKernel) { + std::printf("%-5s head_size=%3lld heads=%2lld %s\n", ggml_type_name(c.type), (long long) c.head_size, + (long long) c.head_count, + on_hrx == Run::NotAdmitted ? "not admitted by HRX" : + on_hrx == Run::ComputeFailed ? "planned, but the compute failed" : + ("planned on another kernel: " + plan).c_str()); + ++failures; + continue; + } + REQUIRE(run(cpu, c, d, expected, nullptr) == Run::Done); + const double e = nmse(got, expected); + const bool ok = e <= kMaxNmse; + std::printf("%-5s head_size=%3lld heads=%2lld %s nmse=%.3g %s\n", ggml_type_name(c.type), + (long long) c.head_size, (long long) c.head_count, plan.c_str(), e, ok ? "OK" : "FAIL"); + failures += ok ? 0 : 1; + } + ggml_backend_free(cpu); + ggml_backend_free(hrx); + std::printf("test-hrx-attention-q-decode-head-size: %zu cases, %d failures\n", cases.size(), failures); + REQUIRE(failures == 0); + return 0; +}