add coopmat2 support

This commit is contained in:
Ruben Ortlam
2026-03-24 13:57:40 +01:00
parent c285bb9838
commit 5c1e95c901
2 changed files with 25 additions and 4 deletions
@@ -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];
@@ -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