hrx: prefill attention keeps its rate at long key counts (+26 % prefill at 24k tokens on Qwen3-4B; engine #1989)

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 <algorithm>
 #include <cmath>
 #include <cstdint>
 #include <cstdlib>
@@ -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<int64_t>(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,18 +494,34 @@ static DispatchBinding prepare_flash_attention_value(const DispatchMatchContext
         return { match.value->id, 0, match.value->byte_count };
     }
 
-    const size_t bytes = static_cast<size_t>(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=<pad> 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<int64_t>(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<size_t>(match.key_value_token_count * columns) * sizeof(ggml_fp16_t);
+    const size_t bytes        = static_cast<size_t>(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 });
 
     Dispatch copy;
     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<size_t>(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<size_t>(match.key_value_token_count) * sizeof(ggml_fp16_t) },
               { 4, match.query->nb[1] },

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<layout>
+  %output_view = buffer.view %dst[%zero] : buffer -> view<[%columns]x[%rows]xf16, %output_layout>
   %shared = buffer.alloca<workgroup> align(16) %bytes : buffer
   %shared_view = buffer.view %shared[%zero] : buffer -> view<32x40xf16>
   %transposed = encoding.layout.strided [1, 40] : encoding<layout>
@@ -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
 }
 

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<layout>
   %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>

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)