diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/flash_attention_decode_split_f32_f16_wmma.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/flash_attention_decode_split_f32_f16_wmma.loom index b60cdea28549..a0ca2ecb3afb 100644 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/flash_attention_decode_split_f32_f16_wmma.loom +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/flash_attention_decode_split_f32_f16_wmma.loom @@ -650,8 +650,9 @@ template.def<@ggml.flash_attention.decode_split.reduce_completed.cooperative> de %query_row1 = index.add %query_row0, %c1 : index %query_head0 = index.add %query_head_base, %query_row0 : index %query_head1 = index.add %query_head0, %c1 : index - %query_head_valid0 = index.cmp ult, %query_head0, %query_head_count : index - %query_head_valid1 = index.cmp ult, %query_head1, %query_head_count : index + // Inactive rows must not write query heads owned by the next KV head. + %query_head_valid0 = index.cmp ult, %query_row0, %query_heads_per_key_value_head : index + %query_head_valid1 = index.cmp ult, %query_row1, %query_heads_per_key_value_head : index %safe_query_row0 = scf.select %query_head_valid0, %query_row0, %c0 : index %safe_query_row1 = scf.select %query_head_valid1, %query_row1, %c0 : index %reducer_active0 = scalar.andi %lane_has_block, %query_head_valid0 : i1 diff --git a/tests/test-hrx-ops.cpp b/tests/test-hrx-ops.cpp index 51162048f255..f134e2c9892c 100644 --- a/tests/test-hrx-ops.cpp +++ b/tests/test-hrx-ops.cpp @@ -841,10 +841,10 @@ static QwenFlashAttentionLayoutGraph build_qwen_flash_attention_layout_graph( int64_t key_value_token_count, int64_t qk_head_size = kQwenFlashHeadSize, int64_t value_head_size = kQwenFlashHeadSize, - float scale = 0.0f) { + float scale = 0.0f, + int64_t query_head_count = 32, + int64_t key_value_head_count = 4) { QwenFlashAttentionLayoutGraph graph; - constexpr int64_t query_head_count = 32; - constexpr int64_t key_value_head_count = 4; ggml_tensor * query_storage = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, qk_head_size, query_head_count, query_token_count); @@ -2169,6 +2169,57 @@ static void run_common_flash_attention_cpu_reference_case(int64_t head_size restore_environment_value(kDisableQwenDispatchEnv, had_original_env, original_env_value); } +// Constant V within each KV head gives an exact oracle while distinct heads +// expose writes beyond the query heads owned by a cooperative reducer. +static void run_decode_flash_attention_head_ownership_case(int64_t key_value_token_count, + int64_t key_value_head_count) { + constexpr int64_t query_head_count = 32; + constexpr int64_t head_size = 128; + const int64_t query_heads_per_key_value_head = query_head_count / key_value_head_count; + ggml_backend_t backend = ggml_backend_hrx_init(0); + REQUIRE(backend != nullptr); + + ggml_init_params params = {}; + params.mem_size = 4 * 1024 * 1024; + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + REQUIRE(ctx != nullptr); + QwenFlashAttentionLayoutGraph attention = build_qwen_flash_attention_layout_graph( + ctx, 1, key_value_token_count, head_size, head_size, 0.0f, query_head_count, key_value_head_count); + ggml_cgraph * graph = ggml_new_graph(ctx); + REQUIRE(graph != nullptr); + ggml_build_forward_expand(graph, attention.output); + require_kernel_subsequence(scheduled_kernel_sequence(graph), + { "loom_libs:ggml_flash_attention_decode_split_f32_f16_wmma_next_q8" }); + ggml_backend_buffer_t buffer = ggml_backend_alloc_ctx_tensors(ctx, backend); + REQUIRE(buffer != nullptr); + + const std::vector query = make_flash_query(query_head_count, head_size); + const std::vector key = + make_flash_key_value(key_value_token_count * key_value_head_count, 3, head_size); + const std::vector mask(key_value_token_count, ggml_fp32_to_fp16(0.0f)); + std::vector value(key.size()); + for (size_t i = 0; i < value.size(); ++i) { + const int64_t kv_head = (i / head_size) % key_value_head_count; + value[i] = ggml_fp32_to_fp16(static_cast(kv_head + 1)); + } + std::vector expected(query.size()); + for (size_t i = 0; i < expected.size(); ++i) { + const int64_t query_head = i / head_size; + expected[i] = static_cast(query_head / query_heads_per_key_value_head + 1); + } + set_tensor_bytes(backend, attention.query, query.data(), query.size() * sizeof(float)); + set_tensor_bytes(backend, attention.key, key.data(), key.size() * sizeof(ggml_fp16_t)); + set_tensor_bytes(backend, attention.value, value.data(), value.size() * sizeof(ggml_fp16_t)); + set_tensor_bytes(backend, attention.mask, mask.data(), mask.size() * sizeof(ggml_fp16_t)); + REQUIRE(ggml_backend_graph_compute(backend, graph) == GGML_STATUS_SUCCESS); + require_close(get_f32_tensor(backend, attention.output), expected, 2.0e-2f); + + ggml_backend_buffer_free(buffer); + ggml_free(ctx); + ggml_backend_free(backend); +} + static void run_asymmetric_flash_attention_cpu_reference_case() { static constexpr const char * kDisableQwenDispatchEnv = "GGML_HRX_DISABLE_QWEN_DISPATCH"; const char * original_env = std::getenv(kDisableQwenDispatchEnv); @@ -4698,6 +4749,11 @@ int main() { run_router_projection_case(4); run_router_top8_case(4); run_qwen_flash_attention_case(); + for (const int64_t kv_heads : { 4, 8 }) { + for (const int64_t kv_tokens : { 256, 512 }) { + run_decode_flash_attention_head_ownership_case(kv_tokens, kv_heads); + } + } for (int64_t head_size = 64; head_size <= 512; head_size += 64) { run_common_flash_attention_cpu_reference_case(head_size, 2, 65, 1.0f / std::sqrt(static_cast(head_size)),