From 2cd6c631332377f3d810ba5ad9e6cebd22df7380 Mon Sep 17 00:00:00 2001 From: bong-water-water-bong Date: Tue, 10 Nov 2026 15:31:36 -0400 Subject: [PATCH] hrx: prefill attention keeps its rate at long key counts: a spread pitch for the transposed V and 256-token launches (engine #1989) From 256 query tokens the prefill FLASH_ATTN_EXT dispatch copies V into a transposed buffer and the attention kernel walks it across channels. Two things made that slow at long key counts: - the buffer's rows were 2 x keys bytes apart, so at key counts where that is a multiple of the page size (16,384, 32,768, ...) the channels of a key block shared cache sets. A key count that is a multiple of 64 now gets a pitch of keys + 32; the transpose kernel writes at that pitch and the attention kernel reads at it (two new config values, both defaulting to the packed layout); - one launch over 512 query tokens was slower per token than two over 256 from about 4,096 keys up. From 4,096 keys the dispatch now runs 256-token launches over the one transposed buffer; below that it keeps one launch. GGML_HRX_FA_VALUE_PITCH and GGML_HRX_FA_QUERY_CHUNK are diagnostic overrides (0 and 2048 give the old layout). tests/test-hrx-fa-transposed-value.cpp: values against the CPU backend and the launch plan, 8 cases. Co-Authored-By: Claude Fable 5.2 --- .../common/dispatch-flash-attention.cpp | 42 +++- .../kernels/loom-libs/ops/copy_f32.loom | 12 +- .../ops/flash_attention_f32_f16_wmma.loom | 7 +- tests/CMakeLists.txt | 7 + tests/test-hrx-fa-transposed-value.cpp | 212 ++++++++++++++++++ 5 files changed, 273 insertions(+), 7 deletions(-) create mode 100644 tests/test-hrx-fa-transposed-value.cpp diff --git a/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-flash-attention.cpp b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-flash-attention.cpp index a98f83842ca8..2f060b310bc4 100644 --- a/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-flash-attention.cpp +++ b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-flash-attention.cpp @@ -7,6 +7,7 @@ #include "graph/graph-matcher.h" #include "kernel-corpus/kernel-corpus-catalog-verify.h" +#include #include #include #include @@ -461,6 +462,23 @@ static const GraphNode * single_consumer_with_op(const Graph & graph, ValueId va nullptr; } +// Query tokens per prefill attention dispatch (engine #1989). The kernel takes up to kPrefillMaxQueryTokens, but with +// the transposed V one launch over 512 query tokens is slower per token than two over 256 from about 4,096 keys up, and +// up to 2x slower at long key counts: measured on gfx1151 at Qwen3-4B's geometry, 0.65-0.77 us per key per layer in +// 256-token launches at every key count from 4,096 to 32,768, against 0.72-1.41 us in one launch. The chunks share the +// one transposed V. GGML_HRX_FA_QUERY_CHUNK overrides (diagnostic); GGML_HRX_QUERY_CHUNK_TOKENS still caps it. +static int64_t flash_attention_query_chunk_tokens(int64_t key_value_token_count) { + static const int64_t limit = [] { + const char * text = std::getenv("GGML_HRX_FA_QUERY_CHUNK"); + const int64_t value = text != nullptr && text[0] != '\0' ? static_cast(std::atoll(text)) : 256; + return value >= 16 && value <= kPrefillMaxQueryTokens ? value : kPrefillMaxQueryTokens; + }(); + // Below 4,096 keys one launch is as fast or faster (pp512 on Qwen3-4B: 2,988 tok/s in one launch, 2,929 in two), so + // short prompts keep it. + const int64_t whole = query_chunk_tokens(kPrefillMaxQueryTokens); + return key_value_token_count >= 4096 ? std::min(limit, whole) : whole; +} + static DispatchBinding prepare_flash_attention_value(const DispatchMatchContext & context, const FlashAttentionMatch & match, DispatchMatch & dispatch_match, @@ -476,7 +494,21 @@ static DispatchBinding prepare_flash_attention_value(const DispatchMatchContext return { match.value->id, 0, match.value->byte_count }; } - const size_t bytes = static_cast(match.key_value_token_count * columns) * sizeof(ggml_fp16_t); + // The attention kernel walks the transposed V across channels, i.e. across rows of this buffer. Packed, the rows are + // 2 x keys bytes apart, and when that is a multiple of the page size the channels of a key block share cache sets: + // at 256 query tokens the dispatch took 8.3 ms at 16,384 keys against 5.7 ms at 16,128 and at 16,640 (engine #1989). + // The buffer is private to this dispatch, so a key count that is a multiple of 64 gets a pitch of keys + 32 (an odd + // multiple of 64 bytes): 6.1 ms at 16,384. GGML_HRX_FA_VALUE_PITCH= forces keys + pad (diagnostic; 0 = packed). + static const int64_t forced_pad = [] { + const char * text = std::getenv("GGML_HRX_FA_VALUE_PITCH"); + const int64_t pad = text == nullptr || text[0] == '\0' ? int64_t{ -1 } : static_cast(std::atoll(text)); + return pad >= 0 && pad % 32 == 0 && pad <= 4096 ? pad : int64_t{ -1 }; + }(); + const int64_t pitch = forced_pad >= 0 ? match.key_value_token_count + forced_pad : + match.key_value_token_count % 64 == 0 ? match.key_value_token_count + 32 : + match.key_value_token_count; + const size_t source_bytes = static_cast(match.key_value_token_count * columns) * sizeof(ggml_fp16_t); + const size_t bytes = static_cast(pitch * columns) * sizeof(ggml_fp16_t); const ValueId transposed = match_value(context, dispatch_match, 0); dispatch_match.transients.push_back({ transposed, "common.flash_attention.transposed_value", bytes, 256 }); @@ -484,10 +516,12 @@ static DispatchBinding prepare_flash_attention_value(const DispatchMatchContext copy.kernel = make_kernel_specialization(kCopyTransposeF16Kernel); copy.kernel.compile_parameters.emplace("ggml.copy_transpose_f16.row_count", to_config_value(match.key_value_token_count)); copy.kernel.compile_parameters.emplace("ggml.copy_transpose_f16.column_count", to_config_value(columns)); - copy.bindings.push_back({ match.value->id, 0, bytes }); + copy.kernel.compile_parameters.emplace("ggml.copy_transpose_f16.output_pitch", to_config_value(pitch)); + copy.bindings.push_back({ match.value->id, 0, source_bytes }); copy.bindings.push_back({ transposed, 0, bytes }); dispatch_match.dispatches.push_back(std::move(copy)); attention.compile_parameters.emplace("ggml.flash_attention.value_layout", "1"); + attention.compile_parameters.emplace("ggml.flash_attention.value_transposed_pitch", to_config_value(pitch)); return { transposed, 0, bytes }; } @@ -573,7 +607,7 @@ static bool match_flash_attention_gate_dispatch(const DispatchMatchContext & con dispatch.bindings.push_back({ raw_gate->id, 0, raw_gate->byte_count }); dispatch.bindings.push_back({ output->id, 0, output->byte_count }); return append_query_chunked_dispatches( - dispatch, match.query_token_count, query_chunk_tokens(kPrefillMaxQueryTokens), + dispatch, match.query_token_count, flash_attention_query_chunk_tokens(match.key_value_token_count), { { 0, match.query->nb[1] }, { 3, static_cast(match.key_value_token_count) * sizeof(ggml_fp16_t) }, { 4, raw_gate->nb[2] }, @@ -610,7 +644,7 @@ static bool match_flash_attention_f32_f16_dispatch(const DispatchMatchContext & } } if (!append_query_chunked_dispatches( - dispatch, match.query_token_count, query_chunk_tokens(kPrefillMaxQueryTokens), + dispatch, match.query_token_count, flash_attention_query_chunk_tokens(match.key_value_token_count), { { 0, match.query->nb[1] }, { 3, static_cast(match.key_value_token_count) * sizeof(ggml_fp16_t) }, { 4, match.query->nb[1] }, diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/copy_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/copy_f32.loom index 743075f6658e..5f9d043acadd 100644 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/copy_f32.loom +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/copy_f32.loom @@ -101,6 +101,10 @@ kernel.def target(@ggml_copy_f16_k16_major_target) export("ggml_copy_f16_k16_maj config.decl @ggml.copy_transpose_f16.row_count : %value: index where [range(%value, 32, 32768), multiple_of(%value, 32)] config.decl @ggml.copy_transpose_f16.column_count : %value: index where [range(%value, 32, 32768), multiple_of(%value, 32)] +// Elements between the starts of two output rows; 0 means row_count (packed). A caller that reads the output column by +// column chooses a pitch whose byte size is not a multiple of the page size, so the rows do not share cache sets +// (engine #1989). +config.def @ggml.copy_transpose_f16.output_pitch = 0 : index kernel.def target(@ggml_copy_f32_gfx11_wave64) export("ggml_copy_transpose_f16") @ggml_copy_transpose_f16() { %rows = config.get @ggml.copy_transpose_f16.row_count : index %columns = config.get @ggml.copy_transpose_f16.column_count : index @@ -131,7 +135,11 @@ kernel.def target(@ggml_copy_f32_gfx11_wave64) export("ggml_copy_transpose_f16") %source_col = index.add %col_origin, %col : index %src, %dst = buffer.assume.noalias %input, %output : buffer, buffer %input_view = buffer.view %src[%zero] : buffer -> view<[%rows]x[%columns]xf16> - %output_view = buffer.view %dst[%zero] : buffer -> view<[%columns]x[%rows]xf16> + %output_pitch0 = config.get @ggml.copy_transpose_f16.output_pitch : index + %packed_output = index.cmp eq, %output_pitch0, %c0 : index + %output_pitch = scf.select %packed_output, %rows, %output_pitch0 : index + %output_layout = encoding.layout.strided [%output_pitch, 1] : encoding + %output_view = buffer.view %dst[%zero] : buffer -> view<[%columns]x[%rows]xf16, %output_layout> %shared = buffer.alloca align(16) %bytes : buffer %shared_view = buffer.view %shared[%zero] : buffer -> view<32x40xf16> %transposed = encoding.layout.strided [1, 40] : encoding @@ -154,7 +162,7 @@ kernel.def target(@ggml_copy_f32_gfx11_wave64) export("ggml_copy_transpose_f16") %swapped = vector.from_elements %s0, %s1, %s2, %s3 : vector<4xf16> %target_row = index.add %col_origin, %row : index %target_col = index.add %row_origin, %col : index - vector.store %swapped, %output_view[%target_row, %target_col] : vector<4xf16>, view<[%columns]x[%rows]xf16> + vector.store %swapped, %output_view[%target_row, %target_col] : vector<4xf16>, view<[%columns]x[%rows]xf16, %output_layout> kernel.return } diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/flash_attention_f32_f16_wmma.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/flash_attention_f32_f16_wmma.loom index a4116cecf397..a8bb451f96e4 100644 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/flash_attention_f32_f16_wmma.loom +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/flash_attention_f32_f16_wmma.loom @@ -31,6 +31,8 @@ config.decl @ggml.flash_attention.value_stride : %value: index where [range(%val config.decl @ggml.flash_attention.attention_scale : f32 config.def @ggml.flash_attention.value_layout = 0 : index +// With value_layout 1: elements between two channels of the transposed V; 0 means the key count (packed). +config.def @ggml.flash_attention.value_transposed_pitch = 0 : index config.decl @ggml.flash_attention.apply_gate : %value: index where [range(%value, 0, 1)] @@ -325,7 +327,10 @@ template.def<@ggml.flash_attention.prefill.body> device @ggml_flash_attention_pr %value_stride0 = config.get @ggml.flash_attention.value_stride : index %value_stride = index.assume %value_stride0 [range(%value_stride0, 16, 36864), multiple_of(%value_stride0, 16)] : index %value_row_stride = scf.select %transposed_value, %c1, %value_stride : index - %value_column_stride = scf.select %transposed_value, %bounded_key_value_token_count, %c1 : index + %value_transposed_pitch0 = config.get @ggml.flash_attention.value_transposed_pitch : index + %packed_transposed_value = index.cmp eq, %value_transposed_pitch0, %c0 : index + %value_transposed_pitch = scf.select %packed_transposed_value, %bounded_key_value_token_count, %value_transposed_pitch0 : index + %value_column_stride = scf.select %transposed_value, %value_transposed_pitch, %c1 : index %value_layout = encoding.layout.strided [%value_row_stride, %value_column_stride] : encoding %value_view = buffer.view %value_aligned[%c0_offset] : buffer -> view<[%bounded_key_value_token_count]x[%value_width]xf16, %value_layout> %mask_view = buffer.view %mask_aligned[%c0_offset] : buffer -> view<[%query_token_count]x[%bounded_key_value_token_count]xf16> diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 2ddf39d8e7f3..7111758c8825 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -472,6 +472,13 @@ 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) + # prefill FLASH_ATTN_EXT through the transposed V at 256+ query tokens: values against the CPU and the launch plan + # (engine #1989); skips without HRX0 + add_executable(test-hrx-fa-transposed-value test-hrx-fa-transposed-value.cpp) + target_link_libraries(test-hrx-fa-transposed-value PRIVATE ggml-hrx ggml hrx::hrx loomc::loomc) + target_include_directories(test-hrx-fa-transposed-value PRIVATE ../ggml/include ../ggml/src ../ggml/src/ggml-hrx) + add_test(NAME test-hrx-fa-transposed-value COMMAND test-hrx-fa-transposed-value) + # 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) diff --git a/tests/test-hrx-fa-transposed-value.cpp b/tests/test-hrx-fa-transposed-value.cpp new file mode 100644 index 000000000000..d8791e116254 --- /dev/null +++ b/tests/test-hrx-fa-transposed-value.cpp @@ -0,0 +1,212 @@ +// 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. + +// Prefill FLASH_ATTN_EXT with the transposed V at 256 and 512 query tokens (engine #1989). From 256 query tokens the +// dispatch copies V into a transposed buffer and the attention kernel walks it across channels. That buffer now has a +// pitch of keys + 32 when the key count is a multiple of 64, and from 4,096 keys up 512 query tokens run as two launches +// of 256 over the one buffer. Both are layout and scheduling changes, so the values must not move: each case is compared with the CPU +// backend (normalized MSE, as test-backend-ops), with a causal mask so masked keys and a padded tail are exercised, and +// the plan must be one transpose plus the expected number of attention launches. Key counts cover a padded pitch +// (1,024, 2,048, 4,096), a packed one (1,056, 4,128) and both sides of the 4,096-key launch rule. + +#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 +#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 int64_t kHeadSize = 128; +static constexpr int64_t kKvHeads = 4; +static constexpr int64_t kQueryHeads = 8; +static constexpr double kMaxNmse = 5e-4; + +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); +} + +// Counts the planned dispatches whose kernel name contains `needle`; -1 when HRX does not take the graph. +static int planned(ggml_cgraph * graph, const char * needle) { + ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); + REQUIRE(imported.valid()); + ggml::hrx::DispatchScheduler scheduler; + if (!scheduler.schedule_graph(imported.graph, { "gfx1151" }) || !scheduler.plan().valid()) { + return -1; + } + int count = 0; + for (const ggml::hrx::Dispatch & dispatch : scheduler.plan().dispatches) { + count += kernel_name_for_id(dispatch.kernel.kernel_id).find(needle) != std::string::npos ? 1 : 0; + } + return count; +} + +struct Case { + int64_t queries; + int64_t keys; +}; + +struct Data { + std::vector q; + std::vector k; + std::vector v; + std::vector mask; +}; + +static Data make_data(const Case & c, uint32_t seed) { + std::mt19937 rng(seed); + std::uniform_real_distribution uniform(-1.0f, 1.0f); + Data d; + d.q.resize(static_cast(kHeadSize * kQueryHeads * c.queries)); + for (float & x : d.q) { + x = uniform(rng); + } + d.k.resize(static_cast(kHeadSize * kKvHeads * c.keys)); + d.v.resize(d.k.size()); + for (ggml_fp16_t & x : d.k) { + x = ggml_fp32_to_fp16(uniform(rng)); + } + for (ggml_fp16_t & x : d.v) { + x = ggml_fp32_to_fp16(uniform(rng)); + } + // causal over the last `queries` keys, like a prefill batch appended to earlier context; nothing is fully masked + const ggml_fp16_t hidden = ggml_fp32_to_fp16(-std::numeric_limits::infinity()); + const ggml_fp16_t visible = ggml_fp32_to_fp16(0.0f); + d.mask.resize(static_cast(c.keys * c.queries)); + const int64_t first = c.keys - c.queries; + for (int64_t query = 0; query < c.queries; ++query) { + for (int64_t key = 0; key < c.keys; ++key) { + d.mask[static_cast(query * c.keys + key)] = key <= first + query ? visible : hidden; + } + } + return d; +} + +struct Plan { + int transposes = 0; + int attention = 0; +}; + +static bool run(ggml_backend_t backend, const Case & c, const Data & d, std::vector & result, Plan * plan) { + ggml_init_params params = { 32 * ggml_tensor_overhead() + ggml_graph_overhead(), nullptr, true }; + ggml_context * ctx = ggml_init(params); + REQUIRE(ctx != nullptr); + // the model's layout: token-major rows [head_size, heads, tokens], permuted to [head_size, tokens, heads] + ggml_tensor * q0 = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, kHeadSize, kQueryHeads, c.queries, 1); + ggml_tensor * k0 = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, kHeadSize, kKvHeads, c.keys, 1); + ggml_tensor * v0 = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, kHeadSize, kKvHeads, c.keys, 1); + ggml_tensor * mask = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, c.keys, c.queries, 1, 1); + ggml_tensor * out = ggml_flash_attn_ext(ctx, ggml_permute(ctx, q0, 0, 2, 1, 3), ggml_permute(ctx, k0, 0, 2, 1, 3), + ggml_permute(ctx, v0, 0, 2, 1, 3), mask, + 1.0f / std::sqrt(static_cast(kHeadSize)), 0.0f, 0.0f); + ggml_cgraph * graph = ggml_new_graph(ctx); + ggml_build_forward_expand(graph, out); + if (plan != nullptr) { + plan->transposes = ggml_backend_supports_op(backend, out) ? planned(graph, "copy_transpose_f16") : -1; + plan->attention = plan->transposes < 0 ? -1 : planned(graph, "flash_attention_f32_f16_wmma"); + if (plan->transposes < 0) { + ggml_free(ctx); + return false; + } + } + ggml_backend_buffer_t buffer = ggml_backend_alloc_ctx_tensors(ctx, backend); + REQUIRE(buffer != nullptr); + ggml_backend_tensor_set(q0, d.q.data(), 0, ggml_nbytes(q0)); + ggml_backend_tensor_set(k0, d.k.data(), 0, ggml_nbytes(k0)); + ggml_backend_tensor_set(v0, d.v.data(), 0, ggml_nbytes(v0)); + ggml_backend_tensor_set(mask, d.mask.data(), 0, ggml_nbytes(mask)); + REQUIRE(ggml_backend_graph_compute(backend, graph) == GGML_STATUS_SUCCESS); + 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 true; +} + +int main() { + std::setvbuf(stdout, nullptr, _IOLBF, 0); + 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-fa-transposed-value: 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); + + // With GGML_HRX_FA_QUERY_CHUNK set by hand the launch count is whatever was asked for; only the values are checked. + const bool default_chunk = std::getenv("GGML_HRX_FA_QUERY_CHUNK") == nullptr && + std::getenv("GGML_HRX_QUERY_CHUNK_TOKENS") == nullptr; + const std::vector cases = { { 256, 1024 }, { 256, 1056 }, { 512, 2048 }, { 256, 4096 }, + { 512, 4096 }, { 512, 4128 }, { 768, 4096 }, { 512, 4160 } }; + int failures = 0; + uint32_t seed = 1; + for (const Case & c : cases) { + const Data d = make_data(c, seed++); + Plan plan; + std::vector got; + std::vector expected; + if (!run(hrx, c, d, got, &plan)) { + std::printf("queries=%4lld keys=%5lld not admitted by HRX\n", (long long) c.queries, (long long) c.keys); + ++failures; + continue; + } + REQUIRE(run(cpu, c, d, expected, nullptr)); + 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]; + } + const double nmse = ref > 0.0 ? err / ref : err; + const int want_launches = c.keys >= 4096 ? static_cast((c.queries + 255) / 256) : 1; + const bool plan_ok = plan.transposes == 1 && (!default_chunk || plan.attention == want_launches); + const bool ok = nmse <= kMaxNmse && plan_ok; + std::printf("queries=%4lld keys=%5lld transposes=%d attention launches=%d (want %d) nmse=%.3g %s\n", + (long long) c.queries, (long long) c.keys, plan.transposes, plan.attention, want_launches, nmse, + ok ? "OK" : "FAIL"); + failures += ok ? 0 : 1; + } + ggml_backend_free(cpu); + ggml_backend_free(hrx); + std::printf("test-hrx-fa-transposed-value: %zu cases, %d failures\n", cases.size(), failures); + REQUIRE(failures == 0); + return 0; +}