wylab/llama.cpp
0
0
Fork 1
mirror of https://github.com/ggml-org/llama.cpp.git synced 2026-09-14 18:02:52 +02:00
Code Issues Packages Projects Releases Wiki Activity

add cm2 decode vector support

Browse Source
This commit is contained in:
Ruben Ortlam
2026-09-14 09:51:40 +02:00
parent 4191ad1d88
commit 8d9b18117c
1 changed files with 38 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files
ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm2.comp
+38 -8
View File
@@ -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 {