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 2f060b310bc4..4af9148a7551 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 @@ -498,15 +498,15 @@ static DispatchBinding prepare_flash_attention_value(const DispatchMatchContext // 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 = [] { + // multiple of 64 bytes): 6.1 ms at 16,384. The attention kernel derives the same pitch from its key count + // (value_layout 2), so the pitch is not one of its compile parameters. GGML_HRX_FA_VALUE_PITCH=0 keeps the packed + // layout (diagnostic). + static const bool spread_pitch = [] { 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 }; + return text == nullptr || text[0] == '\0' || std::atoi(text) != 0; }(); - 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 int64_t pitch = spread_pitch && 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); @@ -520,8 +520,7 @@ static DispatchBinding prepare_flash_attention_value(const DispatchMatchContext 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)); + attention.compile_parameters.emplace("ggml.flash_attention.value_layout", spread_pitch ? "2" : "1"); return { transposed, 0, bytes }; } 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 a8bb451f96e4..ed07c1980198 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,8 +31,10 @@ 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 +// value_layout: 0 = row-major V; 1 = transposed V, packed (channels `keys` elements apart); 2 = transposed V with the +// spread pitch: keys + 32 when the key count is a multiple of 64, else keys. The pitch is derived from the launch's key +// count here, by the same rule the dispatcher uses for the buffer it allocates, so it is not a compile parameter and a +// new key count does not compile a new kernel (engine #1989). config.decl @ggml.flash_attention.apply_gate : %value: index where [range(%value, 0, 1)] @@ -323,13 +325,19 @@ template.def<@ggml.flash_attention.prefill.body> device @ggml_flash_attention_pr %output_aligned = buffer.assume.alignment %output_noalias {minimum_alignment = 16} : buffer %key_view = buffer.view %key_aligned[%c0_offset] : buffer -> view<[%bounded_key_value_token_count]x[%key_width]xf16> %value_layout_tag = config.get @ggml.flash_attention.value_layout : index - %transposed_value = index.cmp eq, %value_layout_tag, %c1 : index + %transposed_value = index.cmp uge, %value_layout_tag, %c1 : index %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_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 + %c2_layout = index.constant 2 : index + %c32_pitch = index.constant 32 : index + %c64_pitch = index.constant 64 : index + %spread_value = index.cmp eq, %value_layout_tag, %c2_layout : index + %key_count_rem64 = index.rem %bounded_key_value_token_count, %c64_pitch : index + %key_count_is_multiple_of_64 = index.cmp eq, %key_count_rem64, %c0 : index + %spread_applies = scalar.andi %spread_value, %key_count_is_multiple_of_64 : i1 + %spread_pitch = index.add %bounded_key_value_token_count, %c32_pitch : index + %value_transposed_pitch = scf.select %spread_applies, %spread_pitch, %bounded_key_value_token_count : 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>