furyhawk/ik_llama.cpp
1
0
Fork 0
mirror of https://github.com/ikawrakow/ik_llama.cpp.git synced 2026-07-21 02:05:35 +00:00
Code Issues Packages Projects Releases Wiki Activity

DSA: Better way to build the attention mask (#2119)

Browse Source 
* WIP: indexer_topk on CUDA

* Forgot these

* WIP

* WIP

* This seems to work

* Minor

* Fix bug. Fix suggested by @sayap using GLM-5.2

* GLM-DSA: much better PP long context performance (CUDA)

* DSA: Better way to build the attention mask
This commit is contained in:
Kawrakow
2026-07-13 13:28:05 +03:00
committed by GitHub
parent 5d766f536a
commit 593a196f78
9 changed files with 177 additions and 37 deletions
Download Patch File Download Diff File Expand all files Collapse all files
ggml/include/ggml.h
+6
View File
@@ -705,6 +705,7 @@ extern "C" {
GGML_OP_FUSED_RMS_RMS_ADD,
GGML_OP_BLEND,
GGML_OP_INDEXER_TOPK,
GGML_OP_MASK_TOPK,
GGML_OP_SINKHORN,
GGML_OP_COUNT,
@@ -2406,6 +2407,11 @@ extern "C" {
struct ggml_tensor * b,
float c);
GGML_API struct ggml_tensor * ggml_indexer_mask(
struct ggml_context * ctx,
struct ggml_tensor * mask,
struct ggml_tensor * topk);
// sort rows
enum ggml_sort_order {
ggml/src/ggml-cuda.cu
+4
View File
@@ -4137,6 +4137,9 @@ static bool ggml_cuda_compute_forward(ggml_backend_cuda_context & ctx, struct gg
case GGML_OP_INDEXER_TOPK:
ggml_cuda_op_indexer_topk(ctx, dst);
break;
case GGML_OP_MASK_TOPK:
ggml_cuda_op_indexer_mask(ctx, dst);
break;
default:
return false;
}
@@ -5041,6 +5044,7 @@ GGML_CALL static bool ggml_backend_cuda_supports_op(ggml_backend_t backend, cons
op->src[3]->ne[0] == op->src[0]->ne[2];
case GGML_OP_DELTA_NET:
case GGML_OP_INDEXER_TOPK:
case GGML_OP_MASK_TOPK:
return true;
case GGML_OP_SINKHORN: {
const int sink_s = op->op_params[0];
ggml/src/ggml-cuda/indexer_topk.cu
+61
View File
@@ -214,3 +214,64 @@ void ggml_cuda_op_indexer_topk(ggml_backend_cuda_context & ctx, ggml_tensor * ds
}
}
template <typename mask_t>
static __global__ void k_indexer_mask(int ne0, int ne1, int ne2, int ntopk, int ne11,
size_t nb01, size_t nb02, size_t nb03,
size_t nb11, size_t nb12, size_t nb13,
size_t nb1, size_t nb2, size_t nb3,
const mask_t * mask, const int * idx, mask_t * dst) {
int i1 = blockIdx.x;
int i3 = i1 / (ne1*ne2); i1 -= i3*ne1*ne2;
int i2 = i1 / (ne1); i1 -= i2*ne1;
auto m = (const mask_t *)((const char *)mask + i1*nb01 + i2*nb02 + i3*nb03);
auto i = (const int *)((const char *)idx + i1*nb11 + i2*nb12 + i3*nb13);
auto d = (mask_t *)((char *)dst + i1*nb1 + i2*nb2 + i3*nb3);
mask_t inf, zero;
if constexpr (std::is_same_v<mask_t, half>) {
inf = __float2half(-INFINITY);
zero = __float2half(0.0f);
} else {
inf = -INFINITY;
zero = 0.0f;
}
if (i1 < ne11) {
for (int j = threadIdx.x; j < ne0; j += blockDim.x) d[j] = inf;
for (int j = threadIdx.x; j < ntopk; j += blockDim.x) d[i[j]] = zero;
for (int j = threadIdx.x; j < ne0; j += blockDim.x) d[j] += m[j];
} else {
for (int j = threadIdx.x; j < ne0; j += blockDim.x) d[j] = m[j];
}
}
void ggml_cuda_op_indexer_mask(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
auto mask = dst->src[0];
auto topk = dst->src[1];
GGML_ASSERT(mask->ne[0] >= topk->ne[1]);
GGML_ASSERT(mask->ne[1] >= topk->ne[1] && mask->ne[2] == topk->ne[2] && mask->ne[3] == topk->ne[3]);
GGML_ASSERT(ggml_are_same_shape(mask, dst));
GGML_ASSERT(mask->type == GGML_TYPE_F16 || mask->type == GGML_TYPE_F32);
GGML_ASSERT(topk->type == GGML_TYPE_I32);
GGML_ASSERT(mask->type == dst->type);
int nrows = ggml_nrows(dst);
if (dst->type == GGML_TYPE_F16) {
k_indexer_mask<<<nrows, 256, 0, ctx.stream()>>>(dst->ne[0], dst->ne[1], dst->ne[2], topk->ne[0], topk->ne[1],
mask->nb[1], mask->nb[2], mask->nb[3],
topk->nb[1], topk->nb[2], topk->nb[3],
dst->nb[1], dst->nb[2], dst->nb[3],
(const half *)mask->data, (const int *)topk->data, (half *)dst->data);
} else {
k_indexer_mask<<<nrows, 256, 0, ctx.stream()>>>(dst->ne[0], dst->ne[1], dst->ne[2], topk->ne[0], topk->ne[1],
mask->nb[1], mask->nb[2], mask->nb[3],
topk->nb[1], topk->nb[2], topk->nb[3],
dst->nb[1], dst->nb[2], dst->nb[3],
(const float *)mask->data, (const int *)topk->data, (float *)dst->data);
}
}
ggml/src/ggml-cuda/indexer_topk.cuh
+2
View File
@@ -6,3 +6,5 @@
#include "common.cuh"
void ggml_cuda_op_indexer_topk(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
void ggml_cuda_op_indexer_mask(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
ggml/src/ggml.c
+33 -2
View File
@@ -4333,10 +4333,11 @@ static const char * GGML_OP_NAME[GGML_OP_COUNT] = {
"FUSED_RMS_RMS_ADD",
"BLEND",
"INDEXER_TOPK",
"MASK_TOPK",
"SINKHORN",
};
static_assert(GGML_OP_COUNT == 105, "GGML_OP_COUNT != 105");
static_assert(GGML_OP_COUNT == 106, "GGML_OP_COUNT != 106");
static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = {
"none",
@@ -4456,11 +4457,12 @@ static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = {
"rms(x1)+rms(x2)",
"blend(a,b,c)",
"indexer_topk(k, q, w, mask)",
"mask_topk(mask, topk)",
"sinkhorn(x)",
};
static_assert(GGML_OP_COUNT == 105, "GGML_OP_COUNT != 105");
static_assert(GGML_OP_COUNT == 106, "GGML_OP_COUNT != 106");
static_assert(GGML_OP_POOL_COUNT == 2, "GGML_OP_POOL_COUNT != 2");
@@ -10230,6 +10232,29 @@ struct ggml_tensor * ggml_blend(
return result;
}
struct ggml_tensor * ggml_indexer_mask(
struct ggml_context * ctx,
struct ggml_tensor * mask,
struct ggml_tensor * topk) {
GGML_ASSERT(mask->type == GGML_TYPE_F16 || mask->type == GGML_TYPE_F32);
GGML_ASSERT(topk->type == GGML_TYPE_I32);
// The mask may be padded along dim 1, so topk->ne[1] must be <= mask->ne[1]
if (topk->ne[1] > mask->ne[1] || topk->ne[2] != mask->ne[2] || topk->ne[3] != mask->ne[3]) {
printf("%s: Oops. topk is %ld x %ld x %ld x %ld, mask is %ld x %ld x %ld x %ld\n", __func__,
topk->ne[0], topk->ne[1], topk->ne[2], topk->ne[3],
mask->ne[0], mask->ne[1], mask->ne[2], mask->ne[3]);
}
GGML_ASSERT(topk->ne[1] <= mask->ne[1] && topk->ne[2] == mask->ne[2] && topk->ne[3] == mask->ne[3]);
struct ggml_tensor * result = ggml_dup_tensor(ctx, mask);
result->src[0] = mask;
result->src[1] = topk;
result->op = GGML_OP_MASK_TOPK;
return result;
}
// ggml_argsort
struct ggml_tensor * ggml_argsort(
@@ -24943,6 +24968,10 @@ static int ggml_compute_forward(struct ggml_compute_params * params, struct ggml
GGML_ABORT("Fatal error");
}
} break;
case GGML_OP_MASK_TOPK:
{
iqk_mask_topk(tensor, params->ith, params->nth);
} break;
case GGML_OP_WIN_PART:
{
ggml_compute_forward_win_part(params, tensor);
@@ -26005,6 +26034,7 @@ static void ggml_compute_backward(struct ggml_context * ctx, struct ggml_tensor
case GGML_OP_SOLVE_TRI:
case GGML_OP_DELTA_NET:
case GGML_OP_INDEXER_TOPK:
case GGML_OP_MASK_TOPK:
case GGML_OP_SINKHORN:
{
GGML_ABORT("fatal error"); // TODO: not implemented
@@ -26743,6 +26773,7 @@ static int ggml_get_n_tasks(struct ggml_tensor * node, int n_threads) {
case GGML_OP_SOLVE_TRI:
case GGML_OP_DELTA_NET:
case GGML_OP_INDEXER_TOPK:
case GGML_OP_MASK_TOPK:
case GGML_OP_SINKHORN:
{
n_tasks = n_threads;
ggml/src/iqk/iqk_cpu_ops.cpp
+68
View File
@@ -1044,3 +1044,71 @@ void iqk_blend(struct ggml_tensor * dst, int ith, int nth) {
}
}
}
namespace {
inline void iqk_add_f16(int n, const ggml_half * x, ggml_half * y) {
int i = 0;
#if defined __AVX2__ && defined __F16C__
for ( ; i + 7 < n; i += 8) {
auto vx32 = _mm256_cvtph_ps(_mm_loadu_si128((const __m128i *)(x + i)));
auto vy32 = _mm256_cvtph_ps(_mm_loadu_si128((const __m128i *)(y + i)));
vy32 = _mm256_add_ps(vy32, vx32);
_mm_storeu_si128((__m128i *)(y + i), _mm256_cvtps_ph(vy32, _MM_FROUND_TO_NEAREST_INT));
}
#endif
for (; i < n; ++i) {
float y32 = GGML_FP16_TO_FP32(y[i]) + GGML_FP16_TO_FP32(x[i]);
y[i] = GGML_FP32_TO_FP16(y32);
}
}
}
void iqk_mask_topk(struct ggml_tensor * dst, int ith, int nth) {
auto mask = dst->src[0];
auto topk = dst->src[1];
GGML_ASSERT(mask->ne[0] >= topk->ne[1]);
GGML_ASSERT(mask->ne[1] >= topk->ne[1] && mask->ne[2] == topk->ne[2] && mask->ne[3] == topk->ne[3]);
GGML_ASSERT(ggml_are_same_shape(mask, dst));
GGML_ASSERT(mask->type == GGML_TYPE_F16 || mask->type == GGML_TYPE_F32);
GGML_ASSERT(topk->type == GGML_TYPE_I32);
GGML_ASSERT(mask->type == dst->type);
int nrows = ggml_nrows(dst);
int npt = (nrows + nth - 1)/nth;
int first = ith*npt;
int last = std::min(first + npt, nrows);
auto hinf = GGML_FP32_TO_FP16(-INFINITY);
auto hzero = GGML_FP32_TO_FP16(0.0f);
int n = dst->ne[0];
int nidx = topk->ne[0];
for (int ir = first; ir < last; ++ir) {
int i3 = ir/(dst->ne[1]*dst->ne[2]);
int i2 = (ir - i3*dst->ne[1]*dst->ne[2])/dst->ne[1];
int i1 = ir - i3*dst->ne[1]*dst->ne[2] - i2*dst->ne[1];
auto idx = (const int *)((const char *)topk->data + topk->nb[1]*i1 + topk->nb[2]*i2 + topk->nb[3]*i3);
if (dst->type == GGML_TYPE_F16) {
auto x = (const ggml_half *)((const char *)mask->data + mask->nb[1]*i1 + mask->nb[2]*i2 + mask->nb[3]*i3);
auto y = (ggml_half *)((char *)dst->data + i1*dst->nb[1] + i2*dst->nb[2] + i3*dst->nb[3]);
if (i1 < topk->ne[1]) {
for (int j = 0; j < n; ++j) y[j] = hinf;
for (int j = 0; j < nidx; ++j) y[idx[j]] = hzero;
iqk_add_f16(n, x, y);
} else {
for (int j = 0; j < n; ++j) y[j] = x[j];
}
} else {
auto x = (const float *)((const char *)mask->data + mask->nb[1]*i1 + mask->nb[2]*i2 + mask->nb[3]*i3);
auto y = (float *)((char *)dst->data + i1*dst->nb[1] + i2*dst->nb[2] + i3*dst->nb[3]);
if (i1 < topk->ne[1]) {
for (int j = 0; j < n; ++j) y[j] = -INFINITY;
for (int j = 0; j < nidx; ++j) y[idx[j]] = 0.0f;
for (int j = 0; j < n; ++j) y[j] += x[j];
} else {
for (int j = 0; j < n; ++j) y[j] = x[j];
}
}
}
}
ggml/src/iqk/iqk_cpu_ops.h
+2
View File
@@ -43,6 +43,8 @@ void iqk_rms_rms_add(struct ggml_tensor * dst, int ith, int nth);
void iqk_blend(struct ggml_tensor * dst, int ith, int nth);
void iqk_mask_topk(struct ggml_tensor * dst, int ith, int nth);
#ifdef __cplusplus
}
#endif
src/graphs/build_deepseek2.cpp
+1 -34
View File
@@ -670,34 +670,6 @@ ggml_tensor * llm_build_context::build_deepseek2_dsa_sparse_mask(
return sparse;
}
static ggml_tensor * build_deepseek2_dsa_fa_mask(const llama_context & lctx, ggml_context * ctx0, ggml_tensor * KQ_mask, ggml_tensor * sorted) {
GGML_ASSERT(KQ_mask && KQ_mask->type == GGML_TYPE_F16);
GGML_ASSERT(sorted && sorted->type == GGML_TYPE_I32);
GGML_ASSERT(KQ_mask->ne[1] >= sorted->ne[1]);
int n_top_k = (int64_t) lctx.model.hparams.indexer_top_k;
if (lctx.cparams.dsa_top_k >= 0) n_top_k = lctx.cparams.dsa_top_k;
int n_kv_local = KQ_mask->ne[0];
if (n_top_k >= n_kv_local) {
return KQ_mask;
}
GGML_ASSERT(sorted->ne[1] == lctx.inp_mask_inf->ne[1]);
auto top_k = ggml_view_2d(ctx0, sorted, n_top_k, sorted->ne[1], sorted->nb[1], 0);
auto mask32 = ggml_blend(ctx0, lctx.inp_mask_inf, top_k, 0.0f);
if (KQ_mask->ne[1] == mask32->ne[1]) {
auto mask16 = ggml_add(ctx0, KQ_mask, mask32);
return mask16;
}
auto kq1 = ggml_view_2d(ctx0, KQ_mask, KQ_mask->ne[0], mask32->ne[1], KQ_mask->nb[1], 0);
auto kq2 = ggml_view_2d(ctx0, KQ_mask, KQ_mask->ne[0], KQ_mask->ne[1] - mask32->ne[1], KQ_mask->nb[1], mask32->ne[1]*KQ_mask->nb[1]);
kq1 = ggml_add(ctx0, kq1, mask32);
auto mask16 = ggml_concat(ctx0, kq1, kq2, 1);
return mask16;
}
// Adapt the (F32, unpadded {n_kv, n_tokens}) sparse mask for ggml_flash_attn_ext, which on this fork
// requires the mask to be F16, contiguous, and padded in ne[1] to GGML_PAD(n_queries, GGML_KQ_MASK_PAD)
// (build_inp_KQ_mask creates the dense -fa 1 mask exactly that way). We:
@@ -829,7 +801,7 @@ ggml_tensor * llm_build_context::build_deepseek2_layer_attention(
dsa_last_full_sorted = build_deepseek2_dsa_indexer(gf, il, q, cur, KQ_mask, inp_pos);
if (dsa_last_full_sorted && !dsa_fast_path) {
if (lctx.cparams.flash_attn) {
last_sparse_mask_fa = ::build_deepseek2_dsa_fa_mask(lctx, ctx0, KQ_mask, dsa_last_full_sorted);
last_sparse_mask_fa = ggml_indexer_mask(ctx0, KQ_mask, dsa_last_full_sorted);
} else {
last_sparse_mask = build_deepseek2_dsa_sparse_mask(dsa_last_full_sorted, KQ_mask);
}
@@ -1293,11 +1265,6 @@ ggml_cgraph * llm_build_context::build_deepseek2() {
mask = ggml_reshape_2d(ctx0, mask, n_padded, KQ_mask->ne[1]);
ggml_build_forward_expand(gf, mask);
dsa_tg_fast_mask = mask;
} else {
auto minus_inf = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, KQ_mask->ne[0], n_tokens);
minus_inf = ggml_fill_inplace(ctx0, minus_inf, -INFINITY);
ggml_build_forward_expand(gf, minus_inf);
lctx.inp_mask_inf = minus_inf;
}
}
src/llama-context.h
-1
View File
@@ -414,7 +414,6 @@ struct llama_context {
struct ggml_tensor * inp_mtp_states = nullptr;
struct ggml_tensor * inp_mtp_carry = nullptr; // F32 [n_embd, nextn-1] per-head hidden at the last committed position
struct ggml_tensor * inp_dsa_sink = nullptr; // F32 [n_kv, n_tokens] per-sequence attention-sink boost for DSA indexer top-k
struct ggml_tensor * inp_mask_inf = nullptr;
struct openpangu_swa_window_view_state {
bool active = false;