From 5c1e95c901bc4e48e3cf12acb4ab1686b181283f Mon Sep 17 00:00:00 2001
From: Ruben Ortlam <rortlam@redhat.com>
Date: Tue, 24 Mar 2026 13:57:40 +0100
Subject: [PATCH] add coopmat2 support

---
 .../vulkan-shaders/dequant_funcs_cm2.glsl        | 16 ++++++++++++++--
 .../ggml-vulkan/vulkan-shaders/mul_mm_cm2.comp   | 13 +++++++++++--
 2 files changed, 25 insertions(+), 4 deletions(-)

diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/dequant_funcs_cm2.glsl b/ggml/src/ggml-vulkan/vulkan-shaders/dequant_funcs_cm2.glsl
index c582aba87d..2c485bc842 100644
--- a/ggml/src/ggml-vulkan/vulkan-shaders/dequant_funcs_cm2.glsl
+++ b/ggml/src/ggml-vulkan/vulkan-shaders/dequant_funcs_cm2.glsl
@@ -25,16 +25,28 @@ float16_t dequantFuncQ1_0(const in decodeBufQ1_0 bl, const in uint blockCoords[2
     return bit != 0u ? d : -d;
 }
 
+#ifdef A_TYPE_REPACKED
+layout(buffer_reference, std430, buffer_reference_align = 16) buffer decodeBufQ4_0 {
+   uint16_t qs[8];
+};
+#else
 layout(buffer_reference, std430, buffer_reference_align = 2) buffer decodeBufQ4_0 {
    block_q4_0_packed16 block;
 };
+#endif
 
 float16_t dequantFuncQ4_0(const in decodeBufQ4_0 bl, const in uint blockCoords[2], const in uint coordInBlock[2])
 {
-    const float16_t d = bl.block.d;
     const uint idx = coordInBlock[1];
-    const uint shift = (idx & 0x10) >> 2;
+#ifdef A_TYPE_REPACKED
+    const uint ib = pos_a + blockCoords[0] * (p.stride_a / QUANT_K) + blockCoords[1];
+    const float16_t d = data_a_deltas[p.deltas_offset + ib];
+    uint32_t qs = uint32_t(bl.qs[(idx & 0xE) >> 1]);
+#else
+    const float16_t d = bl.block.d;
     uint32_t qs = uint32_t(bl.block.qs[(idx & 0xE) >> 1]);
+#endif
+    const uint shift = (idx & 0x10) >> 2;
     qs >>= shift;
     qs &= 0x0F0F;
     qs = unpack8(qs)[idx & 1];
diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm_cm2.comp b/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm_cm2.comp
index 497a18ff8a..3bb3c4af72 100644
--- a/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm_cm2.comp
+++ b/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm_cm2.comp
@@ -63,13 +63,22 @@ layout (push_constant) uniform parameter
 #endif
     // N dimension for the B matrix can be >= p.N
     uint padded_N;
+    uint deltas_offset;
 } p;
 
 
+#ifdef A_TYPE_REPACKED
+struct block_q4_0_quants { uint16_t qs[8]; };
+layout (binding = 0) readonly buffer A {block_q4_0_quants data_a[];};
+layout (binding = 0) readonly buffer A_DELTAS {float16_t data_a_deltas[];};
+#else
 layout (binding = 0) readonly buffer A {A_TYPE data_a[];};
+#endif
 layout (binding = 1) readonly buffer B {B_TYPE data_b[];};
 layout (binding = 2) writeonly buffer D {D_TYPE data_d[];};
 
+uint pos_a;
+
 #if QUANT_K > 1
 #define DECODEFUNCA , dequantFuncA
 
@@ -254,10 +263,10 @@ void main() {
 #endif
 
 #ifdef MUL_MAT_ID
-    uint pos_a = expert_idx * (p.batch_stride_a / QUANT_K);
+    pos_a = expert_idx * (p.batch_stride_a / QUANT_K);
     uint pos_b = 0;
 #else
-    uint pos_a = batch_idx_a * (p.batch_stride_a / QUANT_K);
+    pos_a = batch_idx_a * (p.batch_stride_a / QUANT_K);
     uint pos_b = batch_idx * p.batch_stride_b;
     uint pos_d = batch_idx * p.batch_stride_d + ik * p.batch_stride_d * p.num_batches;
 #endif