diff --git a/tests/hrx-backend-test.cpp b/tests/hrx-backend-test.cpp index 8fffb7d6f07f..38abc3c83b63 100644 --- a/tests/hrx-backend-test.cpp +++ b/tests/hrx-backend-test.cpp @@ -4119,8 +4119,15 @@ static void schedule_flash_attention_command(ggml_context * ctx, REQUIRE(copy.bindings.size() == 2); REQUIRE(copy.bindings[1].value == dispatch.bindings[2].value); REQUIRE(copy.bindings[0].value != copy.bindings[1].value); - REQUIRE(copy.bindings[1].length == static_cast(key_value_token_count * key_value_head_count * value_head_size) * sizeof(ggml_fp16_t)); - require_compile_parameter(dispatch, "ggml.flash_attention.value_layout", "1"); + // engine #1989: the transposed V has the spread pitch - keys + 32 when the key count is a multiple of 64 - which + // the transpose kernel is told and the attention kernel derives from its key count (value_layout 2). The source + // binding stays the packed V. + const int64_t transposed_pitch = key_value_token_count % 64 == 0 ? key_value_token_count + 32 : key_value_token_count; + REQUIRE(copy.bindings[0].length == static_cast(key_value_token_count * key_value_head_count * value_head_size) * sizeof(ggml_fp16_t)); + REQUIRE(copy.bindings[1].length == static_cast(transposed_pitch * key_value_head_count * value_head_size) * sizeof(ggml_fp16_t)); + require_compile_parameter(copy, "ggml.copy_transpose_f16.output_pitch", std::to_string(transposed_pitch)); + require_compile_parameter(dispatch, "ggml.flash_attention.value_layout", "2"); + REQUIRE(dispatch.kernel.compile_parameters.count("ggml.flash_attention.value_transposed_pitch") == 0); } else { REQUIRE(dispatch.kernel.compile_parameters.count("ggml.flash_attention.value_layout") == 0); } @@ -8375,7 +8382,7 @@ static void run_alias_value_import_checks() { value_map .add_snapshot_value({ ggml::hrx::ValueId(0), ggml::hrx::ValueKind::External, ggml::hrx::ValueStorageId(0), ggml::hrx::ValueId(0), ggml::hrx::ValueId(), 0, 8 * sizeof(float), GGML_TYPE_F32, ne, - nb, 8, 8 * sizeof(float), true, nullptr, std::nullopt }) + nb, 8, 8 * sizeof(float), true, nullptr, ggml::hrx::AliasMiss::None, std::nullopt }) .success()); REQUIRE(value_map.add_snapshot_storage({ ggml::hrx::ValueStorageId(1), ggml::hrx::ValueId(1), 8 * sizeof(float) }) .success()); @@ -8383,7 +8390,7 @@ static void run_alias_value_import_checks() { value_map .add_snapshot_value({ ggml::hrx::ValueId(1), ggml::hrx::ValueKind::Transient, ggml::hrx::ValueStorageId(1), ggml::hrx::ValueId(1), ggml::hrx::ValueId(), 0, 8 * sizeof(float), GGML_TYPE_F32, ne, - nb, 8, 8 * sizeof(float), true, nullptr, std::nullopt }) + nb, 8, 8 * sizeof(float), true, nullptr, ggml::hrx::AliasMiss::None, std::nullopt }) .success()); REQUIRE(value_map.alias_storage(ggml::hrx::ValueId(1), ggml::hrx::ValueId(0)).success()); const ggml::hrx::Value * aliased_value = value_map.find(ggml::hrx::ValueId(1));