hrx: the attention kernel derives the transposed pitch from its key count (fixes a per-key-count recompile that merged with #289)

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=<pad> 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<int64_t>(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<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);
@@ -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 };
 }
 

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<layout>
   %value_view = buffer.view %value_aligned[%c0_offset] : buffer -> view<[%bounded_key_value_token_count]x[%value_width]xf16, %value_layout>