From 8d9b18117ce919af15e355fa8e90fab2d4efd6ff Mon Sep 17 00:00:00 2001
From: Ruben Ortlam <rortlam@redhat.com>
Date: Mon, 14 Sep 2026 09:51:40 +0200
Subject: [PATCH] add cm2 decode vector support

---
 .../vulkan-shaders/flash_attn_cm2.comp        | 46 +++++++++++++++----
 1 file changed, 38 insertions(+), 8 deletions(-)

diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm2.comp b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm2.comp
index b708325c89..293dcaecd3 100644
--- a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm2.comp
+++ b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm2.comp
@@ -29,6 +29,12 @@
 #include "dequant_funcs_cm2.glsl"
 #endif
 
+#ifdef GL_NV_cooperative_matrix_decode_vector
+#define FA_GATHER_BS 4u
+#else
+#define FA_GATHER_BS 1u
+#endif
+
 // buffer_reference stride = sizeof(struct) = FaBlockBytesK/V.
 layout(buffer_reference, std430, buffer_reference_align = 1) buffer decodeBufFA_K {
     uint8_t raw[FaBlockBytesK];
@@ -116,19 +122,43 @@ layout (binding = 3) readonly buffer MF16 {float16_t data_mf16[];};
 uint32_t g_k_off_elem, g_v_off_elem, g_m_off_elem;
 
 #if !defined(BFLOAT16)
-// Gather decode: ignore the pre-resolved block and read the selected KV row via the
-// index list. blockCoords[0] = KV slot in [0,n_kv_max), [1] = head dim.
+// blockCoords are in block units: KV slot = blockCoords[0],
+// head dim = blockCoords[1]*FA_GATHER_BS + coordInBlock[1].
 float16_t faGatherK(const decodeBufFA_K unused, const uint32_t blockCoords[2], const uint32_t coordInBlock[2]) {
     if (blockCoords[0] >= p.split_kv) { return float16_t(0); }
     const int r = data_sparse[sparse_base + blockCoords[0]];
-    return r < 0 ? float16_t(0) : data_kf16[g_k_off_elem + uint(r) * k_stride + blockCoords[1]];
+    return r < 0 ? float16_t(0) : data_kf16[g_k_off_elem + uint(r) * k_stride + blockCoords[1] * FA_GATHER_BS + coordInBlock[1]];
 }
 
 float16_t faGatherV(const decodeBufFA_V unused, const uint32_t blockCoords[2], const uint32_t coordInBlock[2]) {
     if (blockCoords[0] >= p.split_kv) { return float16_t(0); }
     const int r = data_sparse[sparse_base + blockCoords[0]];
-    return r < 0 ? float16_t(0) : data_vf16[g_v_off_elem + uint(r) * v_stride + blockCoords[1]];
+    return r < 0 ? float16_t(0) : data_vf16[g_v_off_elem + uint(r) * v_stride + blockCoords[1] * FA_GATHER_BS + coordInBlock[1]];
 }
+
+#ifdef GL_NV_cooperative_matrix_decode_vector
+f16vec4 faGatherKVector(const decodeBufFA_K unused, const uint32_t blockCoords[2], const uint32_t coordInBlock[2]) {
+    if (blockCoords[0] >= p.split_kv) { return f16vec4(0); }
+    const int r = data_sparse[sparse_base + blockCoords[0]];
+    if (r < 0) { return f16vec4(0); }
+    const uint32_t o = g_k_off_elem + uint(r) * k_stride + blockCoords[1] * FA_GATHER_BS + coordInBlock[1];
+    return f16vec4(data_kf16[o], data_kf16[o + 1], data_kf16[o + 2], data_kf16[o + 3]);
+}
+
+f16vec4 faGatherVVector(const decodeBufFA_V unused, const uint32_t blockCoords[2], const uint32_t coordInBlock[2]) {
+    if (blockCoords[0] >= p.split_kv) { return f16vec4(0); }
+    const int r = data_sparse[sparse_base + blockCoords[0]];
+    if (r < 0) { return f16vec4(0); }
+    const uint32_t o = g_v_off_elem + uint(r) * v_stride + blockCoords[1] * FA_GATHER_BS + coordInBlock[1];
+    return f16vec4(data_vf16[o], data_vf16[o + 1], data_vf16[o + 2], data_vf16[o + 3]);
+}
+
+#define FAGATHERK , faGatherK, faGatherKVector
+#define FAGATHERV , faGatherV, faGatherVVector
+#else
+#define FAGATHERK , faGatherK
+#define FAGATHERV , faGatherV
+#endif
 #endif
 
 // Add gathered mask to S (slope==1 since sparse requires max_bias==0). col = slot in block jblk.
@@ -218,8 +248,8 @@ void main() {
 
     tensorViewNV<2, false, 1, 0> tensorViewTranspose = createTensorViewNV(2, false, 1, 0);
 
-    const uint bs_k = fa_block_elems(FaTypeK);
-    const uint bs_v = fa_block_elems(FaTypeV);
+    const uint bs_k = USE_SPARSE ? FA_GATHER_BS : fa_block_elems(FaTypeK);
+    const uint bs_v = USE_SPARSE ? FA_GATHER_BS : fa_block_elems(FaTypeV);
     tensorLayoutK = setTensorLayoutBlockSizeNV(tensorLayoutK, 1, bs_k);
     tensorLayoutV = setTensorLayoutBlockSizeNV(tensorLayoutV, 1, bs_v);
 
@@ -355,7 +385,7 @@ void main() {
 #else
         const bool k_use_decode = (bs_k > 1u);
         if (USE_SPARSE) {
-            coopMatLoadTensorNV(K_T, data_k, k_offset, sliceTensorLayoutNV(tensorLayoutK, j * Bc, Bc, 0, HSK_pad), tensorViewTranspose, faGatherK);
+            coopMatLoadTensorNV(K_T, data_k, k_offset, sliceTensorLayoutNV(tensorLayoutK, j * Bc, Bc, 0, HSK_pad), tensorViewTranspose FAGATHERK);
         } else if (k_use_decode) {
             coopMatLoadTensorNV(K_T, data_k, k_offset, sliceTensorLayoutNV(tensorLayoutK, j * Bc, Bc, 0, HSK_pad), tensorViewTranspose FADECODEK);
         } else {
@@ -429,7 +459,7 @@ void main() {
 #else
         const bool v_use_decode = (bs_v > 1u);
         if (USE_SPARSE) {
-            coopMatLoadTensorNV(V, data_v, v_offset, sliceTensorLayoutNV(tensorLayoutV, j * Bc, Bc, 0, HSV_pad), faGatherV);
+            coopMatLoadTensorNV(V, data_v, v_offset, sliceTensorLayoutNV(tensorLayoutV, j * Bc, Bc, 0, HSV_pad) FAGATHERV);
         } else if (v_use_decode) {
             coopMatLoadTensorNV(V, data_v, v_offset, sliceTensorLayoutNV(tensorLayoutV, j * Bc, Bc, 0, HSV_pad) FADECODEV);
         } else {