diff --git a/common/common.cpp b/common/common.cpp index 5dc820ca2..fc6779fb3 100644 --- a/common/common.cpp +++ b/common/common.cpp @@ -1892,6 +1892,15 @@ bool gpt_params_find_arg(int argc, char ** argv, const std::string & arg, gpt_pa params.mla_attn = std::stoi(argv[i]); return true; } + if (arg == "-dsa" || arg == "--dsa") { + params.dsa = true; + return true; + } + if (arg == "-dsatk" || arg == "--dsa-top-k") { + CHECK_ARG + params.dsa_top_k = std::stoi(argv[i]); + return true; + } if (arg == "-amb" || arg == "--attention-max-batch") { CHECK_ARG params.attn_max_batch = std::stoi(argv[i]); @@ -3011,6 +3020,8 @@ void gpt_params_print_usage(int /*argc*/, char ** argv, const gpt_params & param options.push_back({ "*", "-no-fa, --no-flash-attn", "disable Flash Attention (default: %s)", params.flash_attn ? "enabled" : "disabled" }); options.push_back({ "*", "-fa, --flash-attn (auto|on|off|0|1)", "set Flash Attention (default: %s)", params.flash_attn ? "on" : "off" }); options.push_back({ "*", "-mla, --mla-use", "enable MLA (default: %d)", params.mla_attn }); + options.push_back({ "*", "-dsa, --dsa", "enable GLM DSA sparse attention (GLM-DSA arch only; default: %s)", params.dsa ? "enabled" : "disabled" }); + options.push_back({ "*", "-dsatk, --dsa-top-k", "DSA top-k override; <0 uses the model's configured indexer_top_k (default: %d)", params.dsa_top_k }); options.push_back({ "*", "-amb, --attention-max-batch", "max batch size for attention computations (default: %d)", params.attn_max_batch}); options.push_back({ "*", "-no-fmoe, --no-fused-moe", "disable fused MoE (default: %s)", params.fused_moe_up_gate ? "enabled" : "disabled" }); options.push_back({ "*", "-ger, --grouped-expert-routing", "enable grouped expert routing (default: %s)", params.grouped_expert_routing ? "enabled" : "disabled" }); @@ -4256,6 +4267,8 @@ struct llama_context_params common_context_params_to_llama(const gpt_params & pa cparams.fused_mmad = params.fused_mmad; cparams.rope_cache = params.rope_cache; cparams.graph_reuse = params.graph_reuse; + cparams.dsa = params.dsa; + cparams.dsa_top_k = params.dsa_top_k; cparams.k_cache_hadamard = params.k_cache_hadamard; cparams.v_cache_hadamard = params.v_cache_hadamard; cparams.split_mode_graph_scheduling = params.split_mode_graph_scheduling; diff --git a/common/common.h b/common/common.h index b99848d17..21ef179b3 100644 --- a/common/common.h +++ b/common/common.h @@ -417,6 +417,8 @@ struct gpt_params { bool grouped_expert_routing = false; // if to use grouped expert routing (BailingMoeV2 arch) bool rope_cache = false; // if to use RoPE cache (for supported models) bool graph_reuse = true; // if to reuse compute graphs + bool dsa = false; // enable GLM DSA sparse attention (off by default; opt-in via --dsa) + int dsa_top_k = -1; // DSA top-k override (<0 => use the model's configured indexer_top_k) int min_experts = -1; float thresh_experts = 0; diff --git a/ggml/include/ggml.h b/ggml/include/ggml.h index e3b212cc1..a39d61a36 100644 --- a/ggml/include/ggml.h +++ b/ggml/include/ggml.h @@ -703,6 +703,7 @@ extern "C" { GGML_OP_FAKE_CPY, GGML_OP_FUSED_NORM, GGML_OP_FUSED_RMS_RMS_ADD, + GGML_OP_BLEND, GGML_OP_COUNT, }; @@ -2393,6 +2394,14 @@ extern "C" { struct ggml_tensor * a, float c); + // Overwrite values in a with c for the indeces stored in b + GGML_API struct ggml_tensor * ggml_blend( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + float c); + + // sort rows enum ggml_sort_order { GGML_SORT_ORDER_ASC, diff --git a/ggml/src/ggml-cuda.cu b/ggml/src/ggml-cuda.cu index 8536179c7..e006e661c 100644 --- a/ggml/src/ggml-cuda.cu +++ b/ggml/src/ggml-cuda.cu @@ -56,6 +56,7 @@ #include "ggml-cuda/reduce.cuh" #include "ggml-cuda/tri.cuh" #include "ggml-cuda/delta-net.cuh" +#include "ggml-cuda/blend.cuh" #include #include @@ -3640,6 +3641,9 @@ static bool ggml_cuda_compute_forward(ggml_backend_cuda_context & ctx, struct gg case GGML_OP_REDUCE: ggml_cuda_op_reduce(ctx, dst); break; + case GGML_OP_BLEND: + ggml_cuda_op_blend(ctx, dst); + break; case GGML_OP_FAKE_CPY: break; case GGML_OP_ARGMAX: @@ -4865,6 +4869,10 @@ GGML_CALL static bool ggml_backend_cuda_supports_op(ggml_backend_t backend, cons if (src0_type == GGML_TYPE_F16 && src1_type == GGML_TYPE_F32) { return true; } + if (src0_type == GGML_TYPE_I32 && src1_type == GGML_TYPE_I32) { + // DSA lightning-indexer top_k indices (I32) copy/cont. + return true; + } if (ggml_is_quantized(src0_type) && (src1_type == GGML_TYPE_F16 || src1_type == GGML_TYPE_F32)) { return true; } @@ -4879,6 +4887,7 @@ GGML_CALL static bool ggml_backend_cuda_supports_op(ggml_backend_t backend, cons return false; } break; case GGML_OP_REDUCE: + case GGML_OP_BLEND: case GGML_OP_FAKE_CPY: case GGML_OP_ARGMAX: return true; @@ -4960,11 +4969,23 @@ GGML_CALL static bool ggml_backend_cuda_supports_op(ggml_backend_t backend, cons return ggml_is_contiguous(op->src[0]); //case GGML_OP_ROPE: // return ggml_is_contiguous(op->src[0]); + case GGML_OP_ARGSORT: + return true; + case GGML_OP_ARGSORT_THRESH: + // The CUDA bitonic argsort launches one thread per (padded) column, so the + // row width rounded up to a power of 2 must fit in a single CUDA block (<=1024 + // threads). Wider rows (e.g. the DSA lightning-indexer scoring over a large + // n_kv) would fail at launch with "invalid configuration argument"; report them + // as unsupported so the scheduler falls back to the CPU argsort (no size limit). + { + int64_t ncols = op->src[0]->ne[0]; + int64_t ncols_pad = 1; + while (ncols_pad < ncols) ncols_pad *= 2; + return ncols_pad <= 1024; + } case GGML_OP_IM2COL: case GGML_OP_POOL_2D: case GGML_OP_SUM_ROWS: - case GGML_OP_ARGSORT: - case GGML_OP_ARGSORT_THRESH: case GGML_OP_GROUPED_TOPK: case GGML_OP_ACC: case GGML_OP_GROUP_NORM: diff --git a/ggml/src/ggml-cuda/argsort.cu b/ggml/src/ggml-cuda/argsort.cu index d589ad3f8..28531cce0 100644 --- a/ggml/src/ggml-cuda/argsort.cu +++ b/ggml/src/ggml-cuda/argsort.cu @@ -426,6 +426,151 @@ static void argsort_openai_f32_f32_i32_cuda(const float * x, float * weights, in } } +#if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) && CUDART_VERSION >= 11070 +# define GGML_CUDA_USE_CUB +#endif // !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) && CUDART_VERSION >= 11070 + +#ifdef GGML_CUDA_USE_CUB +# include +# if (CCCL_MAJOR_VERSION >= 3 && CCCL_MINOR_VERSION >= 1) +# define STRIDED_ITERATOR_AVAILABLE +# include +# endif +using namespace cub; +#endif // GGML_CUDA_USE_CUB + +#ifndef STRIDED_ITERATOR_AVAILABLE +static __global__ void init_offsets(int * offsets, const int ncols, const int nrows) { + const int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx <= nrows) { + offsets[idx] = idx * ncols; + } +} +#endif // STRIDED_ITERATOR_AVAILABLE + +#ifdef GGML_CUDA_USE_CUB +static __global__ void init_indices(int * indices, const int ncols, const int nrows) { + const int col = blockIdx.x * blockDim.x + threadIdx.x; + const int row = blockIdx.y; + + if (col < ncols && row < nrows) { + indices[row * ncols + col] = col; + } +} + +void argsort_f32_i32_cuda_cub(ggml_cuda_pool & pool, + const float * x, + int * dst, + const int ncols, + const int nrows, + ggml_sort_order order, + cudaStream_t stream) { + ggml_cuda_pool_alloc temp_indices_alloc(pool, ncols * nrows); + ggml_cuda_pool_alloc temp_keys_alloc(pool, ncols * nrows); + + int * temp_indices = temp_indices_alloc.get(); + float * temp_keys = temp_keys_alloc.get(); + + static const int block_size = 256; + const dim3 grid_size((ncols + block_size - 1) / block_size, nrows); + init_indices<<>>(temp_indices, ncols, nrows); + +#ifdef STRIDED_ITERATOR_AVAILABLE + auto offset_iterator = cuda::make_strided_iterator(cuda::make_counting_iterator(0), ncols); +#else + // offset_iterator needs to populate nrows + 1 elements, so we also have to ceildiv nrows + 1 by block_size + const int nrows_offset = nrows + 1; + ggml_cuda_pool_alloc offsets_alloc(pool, nrows_offset); + int * offset_iterator = offsets_alloc.get(); + const dim3 offset_grid((nrows_offset + block_size - 1) / block_size); + init_offsets<<>>(offset_iterator, ncols, nrows); +#endif + CUDA_CHECK(cudaMemcpyAsync(temp_keys, x, ncols * nrows * sizeof(float), cudaMemcpyDeviceToDevice, stream)); + + size_t temp_storage_bytes = 0; + + bool is_capturing = false; +#ifdef USE_CUDA_GRAPH + // Currently (confirmed for CCCL <= 3.2) DeviceSegmentedSort does not support stream capture, while DeviceSegmentedRadixSort does. + // See https://github.com/NVIDIA/cccl/issues/5661#issuecomment-3229037149 + // TODO: constrain this to the CCCL versions that have this issue once it's resolved in a future CCCL release. + cudaStreamCaptureStatus capture_status; + CUDA_CHECK(cudaStreamIsCapturing(stream, &capture_status)); + is_capturing = (capture_status != cudaStreamCaptureStatusNone); +#endif // USE_CUDA_GRAPH + + if (order == GGML_SORT_ORDER_ASC) { + if (nrows == 1) { + CUDA_CHECK(DeviceRadixSort::SortPairs(nullptr, temp_storage_bytes, temp_keys, temp_keys, // keys (in-place) + temp_indices, dst, // values (indices) + ncols, 0, sizeof(float) * 8, stream)); + } else if (is_capturing) { + CUDA_CHECK(DeviceSegmentedRadixSort::SortPairs( + nullptr, temp_storage_bytes, temp_keys, temp_keys, // keys (in-place) + temp_indices, dst, // values (indices) + ncols * nrows, nrows, // num items, num segments + offset_iterator, offset_iterator + 1, 0, sizeof(float) * 8, stream)); + } else { + CUDA_CHECK(DeviceSegmentedSort::SortPairs(nullptr, temp_storage_bytes, temp_keys, + temp_keys, // keys (in-place) + temp_indices, dst, // values (indices) + ncols * nrows, nrows, // num items, num segments + offset_iterator, offset_iterator + 1, stream)); + } + } else { + if (nrows == 1) { + CUDA_CHECK(DeviceRadixSort::SortPairsDescending(nullptr, temp_storage_bytes, temp_keys, + temp_keys, // keys (in-place) + temp_indices, dst, // values (indices) + ncols, 0, sizeof(float) * 8, stream)); + } else if (is_capturing) { + CUDA_CHECK(DeviceSegmentedRadixSort::SortPairsDescending( + nullptr, temp_storage_bytes, temp_keys, temp_keys, temp_indices, dst, ncols * nrows, nrows, + offset_iterator, offset_iterator + 1, 0, sizeof(float) * 8, stream)); + } else { + CUDA_CHECK(DeviceSegmentedSort::SortPairsDescending(nullptr, temp_storage_bytes, temp_keys, temp_keys, + temp_indices, dst, ncols * nrows, nrows, + offset_iterator, offset_iterator + 1, stream)); + } + } + + ggml_cuda_pool_alloc temp_storage_alloc(pool, temp_storage_bytes); + void * d_temp_storage = temp_storage_alloc.get(); + + if (order == GGML_SORT_ORDER_ASC) { + if (nrows == 1) { + CUDA_CHECK(DeviceRadixSort::SortPairs(d_temp_storage, temp_storage_bytes, temp_keys, + temp_keys, // keys (in-place) + temp_indices, dst, // values (indices) + ncols, 0, sizeof(float) * 8, stream)); + } else if (is_capturing) { + CUDA_CHECK(DeviceSegmentedRadixSort::SortPairs(d_temp_storage, temp_storage_bytes, temp_keys, temp_keys, + temp_indices, dst, ncols * nrows, nrows, offset_iterator, + offset_iterator + 1, 0, sizeof(float) * 8, stream)); + } else { + CUDA_CHECK(DeviceSegmentedSort::SortPairs(d_temp_storage, temp_storage_bytes, temp_keys, temp_keys, + temp_indices, dst, ncols * nrows, nrows, offset_iterator, + offset_iterator + 1, stream)); + } + } else { + if (nrows == 1) { + CUDA_CHECK(DeviceRadixSort::SortPairsDescending(d_temp_storage, temp_storage_bytes, temp_keys, + temp_keys, // keys (in-place) + temp_indices, dst, // values (indices) + ncols, 0, sizeof(float) * 8, stream)); + } else if (is_capturing) { + CUDA_CHECK(DeviceSegmentedRadixSort::SortPairsDescending( + d_temp_storage, temp_storage_bytes, temp_keys, temp_keys, temp_indices, dst, ncols * nrows, nrows, + offset_iterator, offset_iterator + 1, 0, sizeof(float) * 8, stream)); + } else { + CUDA_CHECK(DeviceSegmentedSort::SortPairsDescending(d_temp_storage, temp_storage_bytes, temp_keys, + temp_keys, temp_indices, dst, ncols * nrows, nrows, + offset_iterator, offset_iterator + 1, stream)); + } + } +} +#endif // GGML_CUDA_USE_CUB + void ggml_cuda_op_argsort(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { const ggml_tensor * src0 = dst->src[0]; const float * src0_d = (const float *)src0->data; @@ -439,8 +584,20 @@ void ggml_cuda_op_argsort(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { const int64_t ncols = src0->ne[0]; const int64_t nrows = ggml_nrows(src0); +#ifdef GGML_CUDA_USE_CUB + const int ncols_pad = next_power_of_2(ncols); + const size_t shared_mem = ncols_pad * sizeof(int); + const size_t max_shared_mem = ggml_cuda_info().devices[ggml_cuda_get_device()].smpb; + enum ggml_sort_order order = (enum ggml_sort_order) dst->op_params[0]; + if (shared_mem > max_shared_mem || ncols > 1024) { + ggml_cuda_pool & pool = ctx.pool(); + argsort_f32_i32_cuda_cub(pool, src0_d, (int *) dst_d, ncols, nrows, order, stream); + return; + } +#endif + argsort_f32_T_cuda(src0_d, (int *)dst_d, ncols, nrows, ncols, order, -1, 0.f, stream); } diff --git a/ggml/src/ggml-cuda/blend.cu b/ggml/src/ggml-cuda/blend.cu new file mode 100644 index 000000000..8c14a05f8 --- /dev/null +++ b/ggml/src/ggml-cuda/blend.cu @@ -0,0 +1,101 @@ +#include "blend.cuh" + +#define CUDA_BLEND_BLOCK_SIZE 256 + +template +static __global__ void kernel_blend(int n, int nidx, const Data * x, const Idx * idx, Data * y, float c, + int ne1, int ne2, + 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) { + Data b; + if constexpr (std::is_same_v) { + b = __float2bfloat16(c); + } else { + b = (Data)c; + } + int ii = blockIdx.x; + int i3 = ii / (ne1*ne2); ii -= i3*ne1*ne2; + int i2 = ii / ne1; + int i1 = ii - i2*ne1; + auto x_row = x + i1*nb01 + i2*nb02 + i3*nb03; + auto y_row = y + i1*nb1 + i2*nb2 + i3*nb3; + auto idx_row = idx + i1*nb11 + i2*nb12 + i3*nb13; + + if (x_row != y_row) { + for (int i = threadIdx.x; i < n; i += blockDim.x) { + y_row[i] = x_row[i]; + } + __syncthreads(); + } + for (int i = threadIdx.x; i < nidx; i += blockDim.x) { + y_row[idx_row[i]] = b; + } +} + +void ggml_cuda_op_blend(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { + const ggml_tensor * src0 = dst->src[0]; + const ggml_tensor * src1 = dst->src[1]; + + GGML_ASSERT(src0->type == GGML_TYPE_F32 || src0->type == GGML_TYPE_F16 || src0->type == GGML_TYPE_BF16); + GGML_ASSERT(src0->type == dst->type); + GGML_ASSERT(ggml_are_same_shape(src0, dst)); + GGML_ASSERT(src1->type == GGML_TYPE_I32 || src1->type == GGML_TYPE_I64); + GGML_ASSERT(src1->ne[1] == src0->ne[1] && src1->ne[2] == src0->ne[2] && src1->ne[3] == src0->ne[3]); + GGML_ASSERT(src1->ne[0] <= src0->ne[0]); + + float c; + memcpy(&c, dst->op_params, sizeof(c)); + + auto nrows = ggml_nrows(dst); + dim3 grid_dims(nrows, 1, 1); + dim3 block_size(CUDA_BLEND_BLOCK_SIZE, 1, 1); + if (src1->type == GGML_TYPE_I32) { + auto idx = (const int32_t *)src1->data; + if (src0->type == GGML_TYPE_F32) { + kernel_blend<<>>(src0->ne[0], src1->ne[0], + (const float *)src0->data, idx, (float *)dst->data, c, src0->ne[1], src0->ne[2], + src0->nb[1]/sizeof(float), src0->nb[2]/sizeof(float), src0->nb[3]/sizeof(float), + src1->nb[1]/sizeof(int32_t), src1->nb[2]/sizeof(int32_t), src1->nb[3]/sizeof(int32_t), + dst->nb[1]/sizeof(float), dst->nb[2]/sizeof(float), dst->nb[3]/sizeof(float)); + } + else if (src0->type == GGML_TYPE_F16) { + kernel_blend<<>>(src0->ne[0], src1->ne[0], + (const half *)src0->data, idx, (half *)dst->data, c, src0->ne[1], src0->ne[2], + src0->nb[1]/sizeof(half), src0->nb[2]/sizeof(half), src0->nb[3]/sizeof(half), + src1->nb[1]/sizeof(int32_t), src1->nb[2]/sizeof(int32_t), src1->nb[3]/sizeof(int32_t), + dst->nb[1]/sizeof(half), dst->nb[2]/sizeof(half), dst->nb[3]/sizeof(half)); + } + else { + kernel_blend<<>>(src0->ne[0], src1->ne[0], + (const nv_bfloat16 *)src0->data, idx, (nv_bfloat16 *)dst->data, c, src0->ne[1], src0->ne[2], + src0->nb[1]/sizeof(nv_bfloat16), src0->nb[2]/sizeof(nv_bfloat16), src0->nb[3]/sizeof(nv_bfloat16), + src1->nb[1]/sizeof(int32_t), src1->nb[2]/sizeof(int32_t), src1->nb[3]/sizeof(int32_t), + dst->nb[1]/sizeof(nv_bfloat16), dst->nb[2]/sizeof(nv_bfloat16), dst->nb[3]/sizeof(nv_bfloat16)); + } + } else { + auto idx = (const int64_t *)src1->data; + if (src0->type == GGML_TYPE_F32) { + kernel_blend<<>>(src0->ne[0], src1->ne[0], + (const float *)src0->data, idx, (float *)dst->data, c, src0->ne[1], src0->ne[2], + src0->nb[1]/sizeof(float), src0->nb[2]/sizeof(float), src0->nb[3]/sizeof(float), + src1->nb[1]/sizeof(int64_t), src1->nb[2]/sizeof(int64_t), src1->nb[3]/sizeof(int64_t), + dst->nb[1]/sizeof(float), dst->nb[2]/sizeof(float), dst->nb[3]/sizeof(float)); + } + else if (src0->type == GGML_TYPE_F16) { + kernel_blend<<>>(src0->ne[0], src1->ne[0], + (const half *)src0->data, idx, (half *)dst->data, c, src0->ne[1], src0->ne[2], + src0->nb[1]/sizeof(half), src0->nb[2]/sizeof(half), src0->nb[3]/sizeof(half), + src1->nb[1]/sizeof(int64_t), src1->nb[2]/sizeof(int64_t), src1->nb[3]/sizeof(int64_t), + dst->nb[1]/sizeof(half), dst->nb[2]/sizeof(half), dst->nb[3]/sizeof(half)); + } + else { + kernel_blend<<>>(src0->ne[0], src1->ne[0], + (const nv_bfloat16 *)src0->data, idx, (nv_bfloat16 *)dst->data, c, src0->ne[1], src0->ne[2], + src0->nb[1]/sizeof(nv_bfloat16), src0->nb[2]/sizeof(nv_bfloat16), src0->nb[3]/sizeof(nv_bfloat16), + src1->nb[1]/sizeof(int64_t), src1->nb[2]/sizeof(int64_t), src1->nb[3]/sizeof(int64_t), + dst->nb[1]/sizeof(nv_bfloat16), dst->nb[2]/sizeof(nv_bfloat16), dst->nb[3]/sizeof(nv_bfloat16)); + } + } + +} diff --git a/ggml/src/ggml-cuda/blend.cuh b/ggml/src/ggml-cuda/blend.cuh new file mode 100644 index 000000000..85a904a3c --- /dev/null +++ b/ggml/src/ggml-cuda/blend.cuh @@ -0,0 +1,3 @@ +#include "common.cuh" + +void ggml_cuda_op_blend(ggml_backend_cuda_context & ctx, ggml_tensor * dst); diff --git a/ggml/src/ggml-cuda/cpy.cu b/ggml/src/ggml-cuda/cpy.cu index a0f303309..14e47b16d 100644 --- a/ggml/src/ggml-cuda/cpy.cu +++ b/ggml/src/ggml-cuda/cpy.cu @@ -643,6 +643,10 @@ void ggml_cuda_cpy(ggml_backend_cuda_context & ctx, const ggml_tensor * src0, gg ggml_cpy_flt_cuda (src0_ddc, src1_ddc, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb13, main_stream, dest_ptrs_d, graph_cpynode_index); } else if (src0->type == GGML_TYPE_I32 && src1->type == GGML_TYPE_F32) { ggml_cpy_flt_cuda (src0_ddc, src1_ddc, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb13, main_stream, dest_ptrs_d, graph_cpynode_index); + } else if (src0->type == GGML_TYPE_I32 && src1->type == GGML_TYPE_I32) { + // Needed for the DSA lightning-indexer: top_k (argsort) produces I32 indices that may be + // cont'd / copied / moved across backends (CPU argsort fallback for large n_kv). + ggml_cpy_flt_cuda (src0_ddc, src1_ddc, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb13, main_stream, dest_ptrs_d, graph_cpynode_index); } else if (ggml_are_same_shape(src0, src1) && src0->type == GGML_TYPE_Q8_0 && src1->type == GGML_TYPE_Q8_0) { // This is needed for MLA with mla=2 when using q8_0 cache. transpose_q8_0(ctx, src0, src1); diff --git a/ggml/src/ggml.c b/ggml/src/ggml.c index 19ba0436c..53aef9869 100644 --- a/ggml/src/ggml.c +++ b/ggml/src/ggml.c @@ -4320,9 +4320,10 @@ static const char * GGML_OP_NAME[GGML_OP_COUNT] = { "FAKE_CPY", "FUSED_NORM", "FUSED_RMS_RMS_ADD", + "BLEND", }; -static_assert(GGML_OP_COUNT == 102, "GGML_OP_COUNT != 102"); +static_assert(GGML_OP_COUNT == 103, "GGML_OP_COUNT != 103"); static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = { "none", @@ -4440,10 +4441,11 @@ static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = { "fake_cpy(x,y)", "norm(x,y)", "rms(x1)+rms(x2)", + "blend(a,b,c)", }; -static_assert(GGML_OP_COUNT == 102, "GGML_OP_COUNT != 102"); +static_assert(GGML_OP_COUNT == 103, "GGML_OP_COUNT != 103"); static_assert(GGML_OP_POOL_COUNT == 2, "GGML_OP_POOL_COUNT != 2"); @@ -10137,6 +10139,28 @@ struct ggml_tensor * ggml_fill_inplace( return ggml_fill_impl(ctx, a, c, true); } +// ggml blend +struct ggml_tensor * ggml_blend( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + float c) { + GGML_ASSERT(a->type == GGML_TYPE_F32 || a->type == GGML_TYPE_F16 || a->type == GGML_TYPE_BF16); + GGML_ASSERT(b->type == GGML_TYPE_I32 || b->type == GGML_TYPE_I64); + GGML_ASSERT(b->ne[0] <= a->ne[0]); + for (int dim = 1; dim < GGML_MAX_DIMS; ++dim) { + GGML_ASSERT(a->ne[dim] == b->ne[dim]); + } + + struct ggml_tensor * result = ggml_dup_tensor(ctx, a); + result->src[0] = a; + result->src[1] = b; + memcpy(result->op_params, &c, sizeof(c)); + result->op = GGML_OP_BLEND; + + return result; +} + // ggml_argsort struct ggml_tensor * ggml_argsort( @@ -14974,8 +14998,6 @@ static void ggml_compute_forward_concat_any( GGML_ASSERT(src0->type == src1->type && src0->type == dst->type); const int32_t dim = ggml_get_op_params_i32(dst, 0); - // Let's do it for dim = 0 only for now - GGML_ASSERT(dim == 0); int ith = params->ith; int nth = params->nth; @@ -14983,21 +15005,46 @@ static void ggml_compute_forward_concat_any( int64_t nrows = ggml_nrows(dst); int64_t nrows_per_thread = (nrows + nth - 1)/nth; int64_t first_row = ith*nrows_per_thread; - if (first_row >= nrows) return; int64_t last_row = MIN(first_row + nrows_per_thread, nrows); + if (first_row >= last_row) return; int64_t src0_row_size = ggml_row_size(src0->type, src0->ne[0]); int64_t src1_row_size = ggml_row_size(src1->type, src1->ne[0]); - for (int64_t row = first_row; row < last_row; ++row) { - int64_t i3 = row/(dst->ne[1]*dst->ne[2]); - int64_t i2 = (row - i3*dst->ne[1]*dst->ne[2])/dst->ne[1]; - int64_t i1 = row - i3*dst->ne[1]*dst->ne[2] - i2*dst->ne[1]; - char * y = (char *)dst->data + i1*dst->nb[1] + i2*dst->nb[2] + i3*dst->nb[3]; - const char * x0 = (const char *)src0->data + i1*src0->nb[1] + i2*src0->nb[2] + i3*src0->nb[3]; - const char * x1 = (const char *)src1->data + i1*src1->nb[1] + i2*src1->nb[2] + i3*src1->nb[3]; - memcpy(y, x0, src0_row_size); - memcpy(y + src0_row_size, x1, src1_row_size); + if (dim == 0) { + for (int64_t row = first_row; row < last_row; ++row) { + int64_t i3 = row/(dst->ne[1]*dst->ne[2]); + int64_t i2 = (row - i3*dst->ne[1]*dst->ne[2])/dst->ne[1]; + int64_t i1 = row - i3*dst->ne[1]*dst->ne[2] - i2*dst->ne[1]; + char * y = (char *)dst->data + i1*dst->nb[1] + i2*dst->nb[2] + i3*dst->nb[3]; + const char * x0 = (const char *)src0->data + i1*src0->nb[1] + i2*src0->nb[2] + i3*src0->nb[3]; + const char * x1 = (const char *)src1->data + i1*src1->nb[1] + i2*src1->nb[2] + i3*src1->nb[3]; + memcpy(y, x0, src0_row_size); + memcpy(y + src0_row_size, x1, src1_row_size); + } + } + else { + GGML_ASSERT(src0_row_size == src1_row_size); + for (int64_t row = first_row; row < last_row; ++row) { + int64_t i3 = row/(dst->ne[1]*dst->ne[2]); + int64_t i2 = (row - i3*dst->ne[1]*dst->ne[2])/dst->ne[1]; + int64_t i1 = row - i3*dst->ne[1]*dst->ne[2] - i2*dst->ne[1]; + char * y = (char *)dst->data + i1*dst->nb[1] + i2*dst->nb[2] + i3*dst->nb[3]; + const char * x; + if (dim == 1) { + x = i1 < src0->ne[1] ? (const char *)src0->data + i1*src0->nb[1] + i2*src0->nb[2] + i3*src0->nb[3] + : (const char *)src1->data + (i1 - src0->ne[1])*src1->nb[1] + i2*src1->nb[2] + i3*src1->nb[3]; + } + else if (dim == 2) { + x = i2 < src0->ne[2] ? (const char *)src0->data + i1*src0->nb[1] + i2*src0->nb[2] + i3*src0->nb[3] + : (const char *)src1->data + i1*src1->nb[1] + (i2 - src0->ne[2])*src1->nb[2] + i3*src1->nb[3]; + } + else { + x = i3 < src0->ne[3] ? (const char *)src0->data + i1*src0->nb[1] + i2*src0->nb[2] + i3*src0->nb[3] + : (const char *)src1->data + i1*src1->nb[1] + i2*src1->nb[2] + (i3 - src0->ne[3])*src1->nb[3]; + } + memcpy(y, x, src0_row_size); + } } } @@ -24404,6 +24451,10 @@ static int ggml_compute_forward(struct ggml_compute_params * params, struct ggml { iqk_rms_rms_add(tensor, params->ith, params->nth); } break; + case GGML_OP_BLEND: + { + iqk_blend(tensor, params->ith, params->nth); + } break; case GGML_OP_FUSED_NORM: { ggml_compute_forward_fused_norm(params, tensor); @@ -25243,6 +25294,7 @@ static void ggml_compute_backward(struct ggml_context * ctx, struct ggml_tensor case GGML_OP_FUSED_RMS_NORM: case GGML_OP_FUSED_RMS_RMS_ADD: case GGML_OP_FUSED_NORM: + case GGML_OP_BLEND: { GGML_ABORT("fatal error"); // TODO: not implemented } @@ -26443,6 +26495,7 @@ static int ggml_get_n_tasks(struct ggml_tensor * node, int n_threads) { case GGML_OP_RMS_NORM: case GGML_OP_FUSED_RMS_NORM: case GGML_OP_FUSED_RMS_RMS_ADD: + case GGML_OP_BLEND: case GGML_OP_FUSED_NORM: case GGML_OP_RMS_NORM_BACK: case GGML_OP_GROUP_NORM: diff --git a/ggml/src/iqk/iqk_cpu_ops.cpp b/ggml/src/iqk/iqk_cpu_ops.cpp index 56dedc423..0102f8cb3 100644 --- a/ggml/src/iqk/iqk_cpu_ops.cpp +++ b/ggml/src/iqk/iqk_cpu_ops.cpp @@ -974,3 +974,73 @@ void iqk_rms_rms_add(struct ggml_tensor * dst, int ith, int nth) { } } } + +namespace { +template +inline void iqk_blend_row(int n, int nidx, const Data * x, const Idx * idx, Data * y, float c) { + Data b; + if constexpr (std::is_same_v) { + b = c; + } + else if constexpr (std::is_same_v) { + b = GGML_FP32_TO_FP16(c); + } + else { + b = GGML_FP32_TO_BF16(c); + } + if (y != x) { + for (int j = 0; j < n; ++j) y[j] = x[j]; + } + for (int j = 0; j < nidx; ++j) y[idx[j]] = b; +} +} + +void iqk_blend(struct ggml_tensor * dst, int ith, int nth) { + auto src0 = dst->src[0]; + auto src1 = dst->src[1]; + GGML_ASSERT(src0->type == GGML_TYPE_F32 || src0->type == GGML_TYPE_F16 || src0->type == GGML_TYPE_BF16); + GGML_ASSERT(src1->type == GGML_TYPE_I32 || src1->type == GGML_TYPE_I64); + GGML_ASSERT(src1->ne[0] <= src0->ne[0]); + GGML_ASSERT(src0->type == dst->type); + GGML_ASSERT(src0->ne[0] == dst->ne[0]); + for (int dim = 1; dim < GGML_MAX_DIMS; ++dim) { + GGML_ASSERT(src0->ne[dim] == src1->ne[dim]); + GGML_ASSERT(src0->ne[dim] == dst->ne[dim]); + } + float c; + std::memcpy(&c, dst->op_params, sizeof(c)); + int nrows = ggml_nrows(src0); + int npt = (nrows + nth - 1)/nth; + int first = ith*npt; + int last = std::min(nrows, first + npt); + for (int ir = first; ir < last; ++ir) { + int ii = ir; + int i3 = ii/(src0->ne[1]*src0->ne[2]); ii -= i3*src0->ne[1]*src0->ne[2]; + int i2 = ii/(src0->ne[1]); ii -= i2*src0->ne[1]; + int i1 = ii; + auto x = (const char *)src0->data + i1*src0->nb[1] + i2*src0->nb[2] + i3*src0->nb[3]; + auto y = ( char *) dst->data + i1* dst->nb[1] + i2* dst->nb[2] + i3* dst->nb[3]; + auto idx = (const char *)src1->data + i1*src1->nb[1] + i2*src1->nb[2] + i3*src1->nb[3]; + if (src1->type == GGML_TYPE_I32) { + if (src0->type == GGML_TYPE_F32) { + iqk_blend_row(src0->ne[0], src1->ne[0], (const float *)x, (const int32_t *)idx, (float *)y, c); + } + else if (src0->type == GGML_TYPE_F16) { + iqk_blend_row(src0->ne[0], src1->ne[0], (const ggml_half *)x, (const int32_t *)idx, (ggml_half *)y, c); + } + else { + iqk_blend_row(src0->ne[0], src1->ne[0], (const ggml_bf16_t *)x, (const int32_t *)idx, (ggml_bf16_t *)y, c); + } + } else { + if (src0->type == GGML_TYPE_F32) { + iqk_blend_row(src0->ne[0], src1->ne[0], (const float *)x, (const int64_t *)idx, (float *)y, c); + } + else if (src0->type == GGML_TYPE_F16) { + iqk_blend_row(src0->ne[0], src1->ne[0], (const ggml_half *)x, (const int64_t *)idx, (ggml_half *)y, c); + } + else { + iqk_blend_row(src0->ne[0], src1->ne[0], (const ggml_bf16_t *)x, (const int64_t *)idx, (ggml_bf16_t *)y, c); + } + } + } +} diff --git a/ggml/src/iqk/iqk_cpu_ops.h b/ggml/src/iqk/iqk_cpu_ops.h index 0a049caff..7d6eed745 100644 --- a/ggml/src/iqk/iqk_cpu_ops.h +++ b/ggml/src/iqk/iqk_cpu_ops.h @@ -41,6 +41,8 @@ bool iqk_ssm_conv4(int nr, int nc, int nt, void iqk_rms_rms_add(struct ggml_tensor * dst, int ith, int nth); +void iqk_blend(struct ggml_tensor * dst, int ith, int nth); + #ifdef __cplusplus } #endif diff --git a/include/llama.h b/include/llama.h index c0902faa9..ff6915f1d 100644 --- a/include/llama.h +++ b/include/llama.h @@ -489,6 +489,8 @@ extern "C" { bool fused_mmad; // whether to use fused mul+multi_add op [EXPERIMENTAL] bool rope_cache; // whether to use RoPE cache [EXPERIMENTAL] bool graph_reuse; // whether to reuse graphs when possible [EXPERIMENTAL] + bool dsa; // enable GLM DSA sparse attention (off by default) [EXPERIMENTAL] + int dsa_top_k; // DSA top-k override (<0 => model's configured indexer_top_k) [EXPERIMENTAL] int min_experts; float thresh_experts; bool only_active_experts; diff --git a/src/graphs/build_deepseek2.cpp b/src/graphs/build_deepseek2.cpp index c5129b8f9..ed057177e 100644 --- a/src/graphs/build_deepseek2.cpp +++ b/src/graphs/build_deepseek2.cpp @@ -323,6 +323,370 @@ ggml_tensor * llm_build_context::build_deepseek2_tp_attention( return combined; } +// DSA lightning indexer (GLM-5.2 / DeepSeek-V3.2). CACHE-BACKED: the batch's indexer keys are +// (Hadamard-rotated and) written to a persistent per-layer indexer-key cache (kv_self.kr_l[il]) at +// kv_head, then the FULL [head_size, n_kv] cached key set is read back and scored against the current +// batch's indexer queries. This makes DECODE correct: a generated token (kv_head>0) scores against +// ALL past indexer keys, not just itself. Returns top_k [n_top_k, n_tokens] over the n_kv key axis. +// +// The Walsh-Hadamard rotation H (orthonormal, H^2==I) is applied to both indexer_q and indexer_k. +// (H q)*(H k) == q*k so it is score-preserving; its purpose is to improve the precision of the keys +// we store in the F16 indexer cache (matches the reference). Gated by cparams.dsa_indexer_hadamard. +// +// MULTI-SEQUENCE & FA: both are now handled. +// - Multi-sequence batches (n_seq>1) are correct: each token is written to its own cache cell at +// kv_self.head+i (per-sequence), the base block-diagonal KQ_mask drives cross-sequence keys to +// -inf before argsort, and the attention sink is anchored per-sequence (inp_dsa_sink). Validated +// n_seq=4 == n_seq=1 chunk-for-chunk (UPDATE 5). +// - The FA path (-fa 1) consumes the sparse top-k mask too: build_deepseek2_dsa_fa_mask converts +// the F32 sparse mask into the padded F16 mask the flash-attention kernel reads (UPDATE 4). It +// does NOT fall back to the dense KQ_mask. +// - Serving (context-shift / defrag / multi-turn seq_rm): the persistent indexer-key cache kr_l is +// maintained by build_k_shift (delta-RoPE around the Hadamard), build_defrag (row move), and the +// seq ops are metadata-only so kr_l rows stay matched to their cells (UPDATE 6). +ggml_tensor * llm_build_context::build_deepseek2_dsa_indexer( + ggml_cgraph * gf, + int il, + ggml_tensor * qr, + ggml_tensor * cur, + ggml_tensor * KQ_mask, + ggml_tensor * inp_pos) { + const auto & layer = model.layers[il]; + + const int64_t n_ihead = hparams.indexer_n_head; + const int64_t head_size = hparams.indexer_head_size; + const int64_t rope_dim = n_rot; // n_embd_head_qk_rope + const int64_t nope_dim = head_size - rope_dim; + + // ---- indexer_q : {head_size * n_ihead, n_tokens} ---- + ggml_tensor * indexer_q = ggml_mul_mat(ctx0, layer.indexer_attn_q_b, qr); + cb(indexer_q, "dsa_indexer_q", il); + + // split rope/nope along dim0, per head + ggml_tensor * indexer_q_pe = ggml_view_3d(ctx0, indexer_q, rope_dim, n_ihead, n_tokens, + ggml_row_size(indexer_q->type, head_size), + ggml_row_size(indexer_q->type, head_size) * n_ihead, 0); + ggml_tensor * indexer_q_nope = ggml_view_3d(ctx0, indexer_q, nope_dim, n_ihead, n_tokens, + ggml_row_size(indexer_q->type, head_size), + ggml_row_size(indexer_q->type, head_size) * n_ihead, + ggml_row_size(indexer_q->type, rope_dim)); + + indexer_q_pe = ggml_rope_ext(ctx0, indexer_q_pe, inp_pos, nullptr, n_rot, + LLAMA_ROPE_TYPE_NEOX, n_ctx_orig, freq_base, freq_scale, + ext_factor, attn_factor, beta_fast, beta_slow); + + // {head_size, n_ihead, n_tokens} + indexer_q = ggml_concat(ctx0, indexer_q_pe, indexer_q_nope, 0); + cb(indexer_q, "dsa_indexer_q_cat", il); + + // ---- indexer_k : {head_size, n_tokens} (single key head, MQA) ---- + ggml_tensor * indexer_k = ggml_mul_mat(ctx0, layer.indexer_attn_k, cur); + // LayerNorm (with weight + bias) over head_size + indexer_k = llm_build_norm(ctx0, indexer_k, hparams, layer.indexer_k_norm, layer.indexer_k_norm_b, LLM_NORM, cb, il); + cb(indexer_k, "dsa_indexer_k", il); + + ggml_tensor * indexer_k_pe = ggml_view_3d(ctx0, indexer_k, rope_dim, 1, n_tokens, + ggml_row_size(indexer_k->type, head_size), + ggml_row_size(indexer_k->type, head_size), 0); + ggml_tensor * indexer_k_nope = ggml_view_3d(ctx0, indexer_k, nope_dim, 1, n_tokens, + ggml_row_size(indexer_k->type, head_size), + ggml_row_size(indexer_k->type, head_size), + ggml_row_size(indexer_k->type, rope_dim)); + + indexer_k_pe = ggml_rope_ext(ctx0, indexer_k_pe, inp_pos, nullptr, n_rot, + LLAMA_ROPE_TYPE_NEOX, n_ctx_orig, freq_base, freq_scale, + ext_factor, attn_factor, beta_fast, beta_slow); + + // {head_size, 1, n_tokens} + indexer_k = ggml_concat(ctx0, indexer_k_pe, indexer_k_nope, 0); + cb(indexer_k, "dsa_indexer_k_cat", il); + + // ---- Walsh-Hadamard rotation (score-preserving; improves cached-K F16 precision) ---- + // nrot = largest power of 2 dividing head_size (== head_size for head_size = 128). + // DSA_HADAMARD_DISABLE: DEBUG-ONLY env knob (no CLI surface). Default = rotation enabled. + static const bool dsa_had_disable = getenv("DSA_HADAMARD_DISABLE") != nullptr; + if (lctx.cparams.dsa_indexer_hadamard && !dsa_had_disable) { + GGML_ASSERT((head_size & ~(head_size - 1)) == head_size); + indexer_q = ggml_hadamard(ctx0, indexer_q, head_size); + indexer_k = ggml_hadamard(ctx0, indexer_k, head_size); + } + + // ---- write the batch's indexer keys into the persistent indexer-key cache at kv_head ---- + // kr_l[il] is [head_size, kv_size] (F16, MQA single head). Store {head_size, n_tokens}. + ggml_tensor * kr_cache = kv_self.kr_l[il]; + GGML_ASSERT(kr_cache && "DSA indexer key cache not allocated"); + { + ggml_tensor * indexer_k_2d = ggml_reshape_2d(ctx0, indexer_k, head_size, n_tokens); + ggml_tensor * kr_view = ggml_view_2d(ctx0, kr_cache, head_size, n_tokens, + ggml_row_size(kr_cache->type, head_size), + ggml_row_size(kr_cache->type, head_size) * kv_head); + ggml_tensor * kr_cpy = ggml_cpy(ctx0, indexer_k_2d, kr_view); + // GRAPH-REUSE FIXUP REGISTRATION: the K/V cache_copies fixup in update_cache_copies() + // re-points the K/V cache writes to the current kv_head when a graph is reused, but it + // does NOT touch this indexer-key (kr_l) write, whose view bakes kv_head at build time. + // Under FA the cache pads to 256, so consecutive decode ubatches keep the SAME n_kv and + // can_reuse_graph() reuses the graph -- without this registration the kr_l write stays + // baked at the first ubatch's kv_head, so later ubatches never write their recent index + // keys (those cells read uninitialized -> block-max-pool/top-k drops the genuinely + // attended recent block -> degraded/NaN sparse-FA decode). Register it like K/V so + // update_cache_copies() patches view_offs = kv_head * step each reuse. + // step = one index-key row = head_size * F16 = kr_cache->nb[1]. + if ((size_t) il < lctx.dsa_cache_copies.size()) { + lctx.dsa_cache_copies[il].cpy = kr_cpy; + lctx.dsa_cache_copies[il].step = kr_cache->nb[1]; + } + ggml_build_forward_expand(gf, kr_cpy); + } + + // ---- read back the full cached key set: {head_size, n_kv} ---- + ggml_tensor * cached_k = ggml_view_2d(ctx0, kr_cache, head_size, n_kv, + ggml_row_size(kr_cache->type, head_size), 0); + cb(cached_k, "dsa_cached_k", il); + + // ---- indexer weights : {n_ihead, n_tokens} ---- + ggml_tensor * indexer_weights = ggml_mul_mat(ctx0, layer.indexer_proj, cur); + indexer_weights = ggml_scale(ctx0, indexer_weights, 1.0f / sqrtf(float(head_size * n_ihead))); + cb(indexer_weights, "dsa_indexer_weights", il); + + // ---- scores ---- + // indexer_q : {head_size, n_ihead, n_tokens} -> {head_size, n_tokens, n_ihead} + // cached_k : {head_size, n_kv} -> {head_size, n_kv, 1}; broadcasts over q's n_ihead dim. + ggml_tensor * indexer_k_b = ggml_reshape_3d(ctx0, cached_k, head_size, n_kv, 1); + + ggml_tensor * indexer_score = ggml_view_2d(ctx0, KQ_mask, n_kv, n_tokens, KQ_mask->nb[1], 0); + if (indexer_score->type != GGML_TYPE_F32) { + indexer_score = ggml_cast(ctx0, indexer_score, GGML_TYPE_F32); + cb(indexer_score, "indexer_score_f32", il); + } + for (int head = 0; head < n_ihead; ++head) { + int il_cb = 1000*(il + 1) + head; + // [1, n_tokens] + auto w = ggml_cont(ctx0, ggml_view_2d(ctx0, indexer_weights, 1, indexer_weights->ne[1], indexer_weights->nb[1], indexer_weights->nb[0]*head)); + cb(w, "iweights", il_cb); + // [head_size, n_tokens] + auto q = ggml_view_2d(ctx0, indexer_q, indexer_q->ne[0], indexer_q->ne[2], indexer_q->nb[2], indexer_q->nb[1]*head); + // [n_kv, n_tokens] + auto kq = ggml_mul_mat(ctx0, indexer_k_b, q); + cb(kq, "ikq", il_cb); + // [n_kv, n_tokens] + kq = ggml_relu(ctx0, kq); + cb(kq, "ikq_relu", il_cb); + // [n_kv, n_tokens] + auto score = ggml_mul(ctx0, kq, w); + cb(score, "score", il_cb); + indexer_score = ggml_add_inplace(ctx0, indexer_score, score); + cb(indexer_score, "indexer_score", il_cb); + ggml_build_forward_expand(gf, indexer_score); + } + + //// {n_kv(keys), n_tokens(q), n_ihead} (k's head dim broadcasts over n_ihead) + //ggml_tensor * indexer_kq = ggml_mul_mat(ctx0, indexer_k_b, indexer_q); + //cb(indexer_kq, "dsa_indexer_kq", il); + + //// -> {n_ihead, n_tokens(q), n_kv(keys)} for per-head weighting + //indexer_kq = ggml_cont(ctx0, ggml_permute(ctx0, indexer_kq, 2, 1, 0, 3)); + + //ggml_tensor * indexer_score = ggml_relu(ctx0, indexer_kq); + + //// weights {n_ihead, n_tokens} -> {n_ihead, n_tokens, 1} broadcast over keys + //indexer_weights = ggml_reshape_3d(ctx0, indexer_weights, n_ihead, n_tokens, 1); + //indexer_score = ggml_mul(ctx0, indexer_score, indexer_weights); + + //// sum over heads -> {1, n_tokens(q), n_kv(keys)} + //indexer_score = ggml_sum_rows(ctx0, indexer_score); + + //// -> {n_kv(keys), n_tokens(q), 1} + //indexer_score = ggml_cont(ctx0, ggml_permute(ctx0, indexer_score, 2, 1, 0, 3)); + //cb(indexer_score, "dsa_indexer_score", il); + + //// add base causal mask over the n_kv keys: first n_tokens query columns of KQ_mask {n_kv, n_tokens_pad}. + //ggml_tensor * causal = ggml_view_2d(ctx0, KQ_mask, n_kv, n_tokens, KQ_mask->nb[1], 0); + //// Under -fa 1 the dense KQ_mask is F16; CPU ggml_add only supports F32+F16 when src0 is F16, + //// not F32(score)+F16(mask) (it aborts). Cast the causal mask view to F32 so the add is valid on + //// CPU. (CUDA add accepts mixed types, so this only bit the CPU build.) + //if (causal->type != GGML_TYPE_F32) causal = ggml_cast(ctx0, ggml_cont(ctx0, causal), GGML_TYPE_F32); + //indexer_score = ggml_add(ctx0, indexer_score, causal); + //cb(indexer_score, "dsa_indexer_score_masked", il); + + // Attention-sink force-inclusion: add a finite positive boost to each query's OWN SEQUENCE's + // first n_sink present tokens so the sink token(s) always survive the top-k selection. Masking + // the sink collapses most transformers; a heavily-quantized (IQ2) indexer does not reliably rank + // it high on its own. The boost is finite, so it cannot un-mask future/causal -inf positions + // (-inf + boost = -inf). + // + // MULTI-SEQUENCE: the boost MUST be per-(key,query), not a global per-key vector. With several + // sequences packed into one ubatch (seq 0 at cache cells [0,n0), seq 1 at [n0,n1), ...), a global + // "key index < n_sink" boost only protects sequence 0's sink; sequence 1's sink (at cell n0, not + // cell 0) is left unprotected and gets dropped from top-k once the mask bites, collapsing it. + // We therefore use a per-graph input tensor inp_dsa_sink {n_kv, n_tokens} (filled on the CPU from + // kv_self.cells like the KQ_mask, in llama_set_inputs): inp_dsa_sink[j,i] = 1e20 iff key cell i + // belongs to query j's sequence AND its pos is within [min present pos of that seq, +n_sink). + // + // SERVING: the anchor is each sequence's FIRST PRESENT pos, not absolute pos < n_sink. After + // multi-turn seq_rm drops a sequence's early tokens its earliest survivor has pos >= n_sink; an + // absolute test would then protect nothing and let the (now-)sink be masked out. For a fresh + // sequence starting at pos 0, min(pos)==0 so the boosted set is exactly the old "cell pos < + // n_sink" set with the same 1e20 magnitude — n_seq==1 from pos 0 stays byte-identical. + // DSA_SINK: DEBUG-ONLY env knob (no CLI surface). Default = 1 (protect each sequence's first + // present token from being masked out of top-k). Must stay in sync with the two fill sites. + if (lctx.inp_dsa_sink) { + indexer_score = ggml_add(ctx0, indexer_score, lctx.inp_dsa_sink); + cb(indexer_score, "dsa_indexer_score_sink", il); + ggml_build_forward_expand(gf, indexer_score); + } + + // FULL descending argsort of the per-query scores over the n_kv axis: {n_kv, n_tokens} (I32). + // We return the full ranking (not just the top-k view): the sparse-mask builder writes a value + // into EVERY key slot keyed by its rank, which avoids relying on ggml_set_rows preserving an + // uninitialized base for partially-written destinations (a CUDA in-place quirk that corrupted + // decode when n_kv > top_k). + //ggml_tensor * sorted = ggml_cont(ctx0, ggml_argsort(ctx0, indexer_score, GGML_SORT_ORDER_DESC)); + //ggml_tensor * sorted = ggml_argsort(ctx0, indexer_score, GGML_SORT_ORDER_DESC); + ggml_tensor * sorted; + if (cparams.flash_attn) { + int64_t n_top_k = (int64_t) hparams.indexer_top_k; + if (lctx.cparams.dsa_top_k >= 0) n_top_k = lctx.cparams.dsa_top_k; + if (n_top_k > indexer_score->ne[0]) n_top_k = indexer_score->ne[0]; + sorted = ggml_top_k(ctx0, indexer_score, n_top_k); + sorted = ggml_cont(ctx0, sorted); + } else { + sorted = ggml_argsort(ctx0, indexer_score, GGML_SORT_ORDER_DESC); + } + cb(sorted, "dsa_sorted", il); + + return sorted; +} + +// Build an additive sparse causal mask {n_kv, n_tok} (F32): 0 for the top-k highest-scoring keys +// per query, a large negative value for the rest, then add the base causal KQ_mask so future/ +// padding keys stay masked. ggml_soft_max_ext only requires mask->ne[1] >= q n_tokens, and +// n_tok == n_tokens here, so no padding is needed. +// +// `sorted` is the FULL descending argsort of the indexer scores: {n_kv, n_tok} (I32), where +// sorted[rank, j] = key index with the rank-th highest score for query j. We scatter a rank-based +// penalty into EVERY key slot: pen(rank) = 0 if rank < n_top_k else -BIG. Because every key slot +// is written exactly once (sorted is a per-column permutation), the result does NOT depend on the +// scatter destination's initial contents — sidestepping the ggml in-place set_rows quirk where a +// partially-written CUDA destination keeps uninitialized (garbage) rows. +ggml_tensor * llm_build_context::build_deepseek2_dsa_sparse_mask( + ggml_tensor * sorted, + ggml_tensor * KQ_mask) { + const int64_t n_kv_local = KQ_mask->ne[0]; + const int64_t n_tok = sorted->ne[1]; + + int64_t n_top_k = (int64_t) hparams.indexer_top_k; + // Tuning knob: --dsa-top-k (cparams.dsa_top_k) lets us vary the kept-key count to characterize + // selection quality. <0 means use the model's configured top_k. With the model's configured + // top_k (2048) on heavily-quantized (IQ2_M) weights the indexer currently under-ranks some + // critical keys; a near-n_kv value stays coherent. + if (lctx.cparams.dsa_top_k >= 0) n_top_k = lctx.cparams.dsa_top_k; + if (n_top_k > n_kv_local) n_top_k = n_kv_local; + + // Penalty magnitude for non-top-k keys. On the soft_max (-fa 0) path this F32 -BIG is added to + // the score and softmaxed -> effectively -inf, while staying finite avoids -inf*0 = NaN. On the + // FA (-fa 1) path this mask is cast to F16 (build_deepseek2_dsa_fa_mask): 1e30 saturates to the + // F16 max (~6.5e4), which is still a large-enough negative bias to zero the key in the FA softmax + // (the dense FA mask uses -INFINITY/F16 -inf there; our finite-but-huge value is equivalent in + // effect and cannot produce NaN). So -BIG masks the key on BOTH paths. + const float BIG = 1e30f; + + // rank-based penalty vector: pen[rank] = 0 for rank < n_top_k, else -BIG. {n_kv} + // sel = step(n_top_k - 0.5 - rank) = 1 for rank <= n_top_k-1, else 0 + ggml_tensor * rank = ggml_arange(ctx0, 0.0f, (float) n_kv_local, 1.0f); // {n_kv} F32 + ggml_tensor * sel = ggml_step(ctx0, ggml_scale_bias(ctx0, rank, -1.0f, (float) n_top_k - 0.5f)); + ggml_tensor * pen = ggml_scale_bias(ctx0, sel, BIG, -BIG); // 0 or -BIG + + // shape penalty to {1, n_kv, n_tok} (broadcast the per-rank value across all query columns) + pen = ggml_reshape_3d(ctx0, pen, 1, n_kv_local, 1); + ggml_tensor * pen_b = ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, 1, n_kv_local, n_tok); + pen_b = ggml_repeat(ctx0, pen, pen_b); // {1, n_kv, n_tok} + + // destination base {1, n_kv, n_tok} (contents irrelevant — fully overwritten by set_rows) + ggml_tensor * base = ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, 1, n_kv_local, n_tok); + base = ggml_fill(ctx0, base, -BIG); + + // indices: {n_kv, n_tok, 1}. scatter pen_b[:, rank, j] into base[:, sorted[rank,j], j]. + ggml_tensor * idx = ggml_reshape_3d(ctx0, sorted, n_kv_local, n_tok, 1); + ggml_tensor * scattered = ggml_set_rows(ctx0, base, pen_b, idx); + + // {n_kv, n_tok} + ggml_tensor * sparse = ggml_reshape_2d(ctx0, ggml_cont(ctx0, scattered), n_kv_local, n_tok); + + // add base causal mask (first n_tok query columns) so future/padding keys stay masked + ggml_tensor * causal = ggml_view_2d(ctx0, KQ_mask, n_kv_local, n_tok, KQ_mask->nb[1], 0); + // see note in build_deepseek2_dsa_indexer: cast F16 (-fa 1) mask to F32 for the CPU add. + if (causal->type != GGML_TYPE_F32) causal = ggml_cast(ctx0, ggml_cont(ctx0, causal), GGML_TYPE_F32); + sparse = ggml_add(ctx0, sparse, causal); + cb(sparse, "dsa_sparse_mask", -1); + + 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: +// 1) cast the sparse mask to F16, +// 2) concat the dense FA mask's padding rows [n_tok, n_pad) (already F16, causal -inf for the +// non-existent padded queries) onto the bottom so ne[1] matches the dense mask, +// 3) ggml_cont so the result is contiguous (the FA assert requires it). +// The padded rows feed only the discarded outputs of padded query slots, so reusing the dense mask's +// padding region is both correct and the cheapest way to get the exact dense shape. +ggml_tensor * llm_build_context::build_deepseek2_dsa_fa_mask( + ggml_tensor * sparse, + ggml_tensor * KQ_mask) { + const int64_t n_kv_local = KQ_mask->ne[0]; + const int64_t n_tok = sparse->ne[1]; + const int64_t n_pad = KQ_mask->ne[1]; // GGML_PAD(n_tokens, GGML_KQ_MASK_PAD) + + GGML_ASSERT(KQ_mask->type == GGML_TYPE_F16 && "FA dense KQ_mask expected F16 on -fa 1"); + + ggml_tensor * fa_mask; + if (n_pad > n_tok) { + // dense padding rows: KQ_mask columns [n_tok, n_pad) -> {n_kv, n_pad - n_tok} (F16 view) + ggml_tensor * pad = ggml_view_2d(ctx0, KQ_mask, n_kv_local, n_pad - n_tok, + KQ_mask->nb[1], KQ_mask->nb[1] * n_tok); + // CPU ggml_concat only supports F16 along dim 0 (concat_any); the dim-1 row concat must be + // done in F32 (concat_f32 handles all dims), then cast the padded result to F16. On CUDA the + // F16 dim-1 concat is supported, so this path only needed adapting for the CPU build. + ggml_tensor * pad_f32 = ggml_cast(ctx0, ggml_cont(ctx0, pad), GGML_TYPE_F32); + ggml_tensor * fa_f32 = ggml_concat(ctx0, sparse, pad_f32, 1); // {n_kv, n_pad} F32 + fa_mask = ggml_cast(ctx0, fa_f32, GGML_TYPE_F16); // {n_kv, n_pad} F16 + } else { + fa_mask = ggml_cast(ctx0, sparse, GGML_TYPE_F16); + } + fa_mask = ggml_cont(ctx0, fa_mask); + cb(fa_mask, "dsa_fa_mask", -1); + return fa_mask; +} + // Layer-mode attention path (non-TP). Mirrors build_deepseek2_tp_attention's interface. ggml_tensor * llm_build_context::build_deepseek2_layer_attention( ggml_cgraph * gf, int il, @@ -343,6 +707,16 @@ ggml_tensor * llm_build_context::build_deepseek2_layer_attention( cur = llm_build_norm(ctx0, inpL, hparams, model.layers[il].attn_norm, NULL, LLM_NORM_RMS, cb, il); cb(cur, "attn_norm", il); + // DSA lightning indexer (GLM-5.2 / DeepSeek-V3.2). Built below from the q_lora latent + // and used to construct a sparse causal mask. Defaults to the dense KQ_mask. + // - sparse_mask : F32 additive sparse mask for the soft_max (-fa 0) path. + // - sparse_mask_fa : F16, padded variant for the ggml_flash_attn_ext (-fa 1) path. + // Both default to the dense KQ_mask so non-DSA / disabled builds are unchanged. + ggml_tensor * sparse_mask = KQ_mask; + ggml_tensor * sparse_mask_fa = KQ_mask; + //ggml_tensor * top_k = nullptr; + //(void) top_k; // captured for potential reuse/debug; only the masks are consumed downstream + // self_attention { ggml_tensor * q = nullptr; @@ -391,6 +765,38 @@ ggml_tensor * llm_build_context::build_deepseek2_layer_attention( q = llm_build_norm(ctx0, q, hparams, model.layers[il].attn_q_a_norm, NULL, LLM_NORM_RMS, cb, il); cb(q, "q", il); + // DSA lightning indexer (cache-backed): score the q_lora latent against the persistent + // indexer-key cache over the full n_kv, then build a sparse top-k causal mask. Correct + // for prefill AND decode (single sequence). Gate: --dsa opt-in (off by default) + + // GLM_DSA arch + indexer tensors + cache. When off, the model runs the dense MLA path, + // byte-identical to a build without this feature. + if (lctx.cparams.dsa && model.arch == LLM_ARCH_GLM_DSA && model.layers[il].indexer_attn_q_b + && kv_self.kr_l.size() > (size_t) il && kv_self.kr_l[il]) { + // GLM-5.2 IndexShare: "full" layers compute their own lightning-indexer top-k; + // "shared" layers reuse the previous full layer's top-k (transformers reference: + // shared layer indexer=None, topk_indices=prev_topk_indices). The full/shared map is + // hparams.indexer_is_full (GGUF metadata or derived config rule). At a given step all + // layers share the same n_kv/n_tokens, so a full layer's argsort is valid to reuse. + ggml_tensor * sorted; + if (hparams.indexer_is_full[il] || dsa_last_full_sorted == nullptr) { + ggml_tensor * qr = q; // q_lora latent (after attn_q_a_norm, before wq_b) + sorted = build_deepseek2_dsa_indexer(gf, il, qr, cur, KQ_mask, inp_pos); + dsa_last_full_sorted = sorted; + } else { + sorted = dsa_last_full_sorted; + } + if (lctx.cparams.flash_attn) { + sparse_mask_fa = ::build_deepseek2_dsa_fa_mask(lctx, ctx0, KQ_mask, sorted); + } else { + sparse_mask = build_deepseek2_dsa_sparse_mask(sorted, KQ_mask); + } + //// For the FA path the mask must be F16 + padded; build it from the F32 sparse mask. + //if (lctx.cparams.flash_attn) { + // sparse_mask_fa = build_deepseek2_dsa_fa_mask(sparse_mask, KQ_mask); + //} + //top_k = sorted; + } + q = ggml_mul_mat(ctx0, model.layers[il].wq_b, q); cb(q, "q", il); } else { @@ -566,7 +972,7 @@ ggml_tensor * llm_build_context::build_deepseek2_layer_attention( auto q_iter = ggml_view_3d(ctx0, q, q->ne[0], q->ne[1], n_max_head, q->nb[1], q->nb[2], q->nb[2]*n_max_head*iter); - kqv = ggml_flash_attn_ext(ctx0, q_iter, k, v, KQ_mask, kq_scale, hparams.f_max_alibi_bias, 0.f); + kqv = ggml_flash_attn_ext(ctx0, q_iter, k, v, sparse_mask_fa, kq_scale, hparams.f_max_alibi_bias, 0.f); if (use_f32_attn_precision || q->ne[1] <= 8) { ggml_flash_attn_ext_set_prec(kqv, GGML_PREC_F32); } @@ -607,7 +1013,7 @@ ggml_tensor * llm_build_context::build_deepseek2_layer_attention( ggml_row_size(kv_self.k_l[il]->type, n_embd_head_qk_rope)); cb(kv_cache_lora, "kv_cache_lora", il); - kqv_compressed = ggml_flash_attn_ext(ctx0, q, kv_cache, kv_cache_lora, KQ_mask, kq_scale, hparams.f_max_alibi_bias, 0.f); + kqv_compressed = ggml_flash_attn_ext(ctx0, q, kv_cache, kv_cache_lora, sparse_mask_fa, kq_scale, hparams.f_max_alibi_bias, 0.f); cb(kqv_compressed, "kqv_compressed", il); if (use_f32_attn_precision) { @@ -647,7 +1053,7 @@ ggml_tensor * llm_build_context::build_deepseek2_layer_attention( cb(kq, "kq_perm", il); } - kq = ggml_soft_max_ext(ctx0, kq, KQ_mask, kq_scale, hparams.f_max_alibi_bias); + kq = ggml_soft_max_ext(ctx0, kq, sparse_mask, kq_scale, hparams.f_max_alibi_bias); cb(kq, "kq_soft_max_ext", il); if (!pp_opt) { @@ -673,7 +1079,7 @@ ggml_tensor * llm_build_context::build_deepseek2_layer_attention( int this_ne12 = i_head + n_per_step <= q->ne[2] ? n_per_step : q->ne[2] - i_head; ggml_tensor * q_i = ggml_view_3d(ctx0, q, q->ne[0], q->ne[1], this_ne12, q->nb[1], q->nb[2], q->nb[2]*i_head); ggml_tensor * kq_i = ggml_mul_mat(ctx0, kv_cache, q_i); - kq_i = ggml_soft_max_ext(ctx0, kq_i, KQ_mask, kq_scale, hparams.f_max_alibi_bias); + kq_i = ggml_soft_max_ext(ctx0, kq_i, sparse_mask, kq_scale, hparams.f_max_alibi_bias); ggml_tensor * kqv_i = ggml_mul_mat(ctx0, kv_cache_trans, kq_i); if (i_head == 0) { kqv_compressed = kqv_i; @@ -762,6 +1168,7 @@ ggml_tensor * llm_build_context::build_deepseek2_layer_attention( } ggml_cgraph * llm_build_context::build_deepseek2() { + dsa_last_full_sorted = nullptr; // GLM-5.2 IndexShare: reset shared-layer top-k reuse state (before any layer, incl. the MTP early-return path) const bool tp_mode = (model.split_mode == LLAMA_SPLIT_MODE_GRAPH || model.split_mode == LLAMA_SPLIT_MODE_ATTN); #ifdef GGML_USE_VULKAN @@ -794,6 +1201,19 @@ ggml_cgraph * llm_build_context::build_deepseek2() { // KQ_mask (mask for 1 head, it will be broadcasted to all heads) struct ggml_tensor * KQ_mask = build_inp_KQ_mask(); + if (lctx.cparams.dsa && model.arch == LLM_ARCH_GLM_DSA) { + static const int n_sink = []{ const char * e = getenv("DSA_SINK"); return e ? atoi(e) : 1; }(); + if (n_sink > 0 && n_sink < (int) n_kv) { + lctx.inp_dsa_sink = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_kv, n_tokens); + cb(lctx.inp_dsa_sink, "dsa_sink", -1); + ggml_set_input(lctx.inp_dsa_sink); + } + 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; + } + // whether to use n_tokens as the matrix dimension during multiplication or n_head // n_tokens is higher during prompt processing, this allows to optimize for this case bool pp_opt = n_tokens >= 128 && lctx.cparams.mla_attn > 1; diff --git a/src/llama-build-context.cpp b/src/llama-build-context.cpp index 36f5f2743..27132606b 100644 --- a/src/llama-build-context.cpp +++ b/src/llama-build-context.cpp @@ -116,6 +116,7 @@ void llm_build_context::init() { lctx.inp_pos_bucket = nullptr; lctx.inp_embd_enc = nullptr; lctx.inp_KQ_mask_cross = nullptr; + lctx.inp_dsa_sink = nullptr; lctx.dflash.inputs.target_features = nullptr; lctx.dflash.inputs.pos_ctx = nullptr; lctx.dflash.inputs.kq_mask = nullptr; @@ -194,6 +195,103 @@ ggml_cgraph * llm_build_context::build_k_shift() { ggml_build_forward_expand(gf, tmp); } + // DSA lightning-indexer key cache (GLM-5.2 / DeepSeek-V3.2): the persistent indexer keys in + // kv_self.kr_l[il] are RoPE-position-encoded at write time (build_deepseek2_dsa_indexer rotates + // the first rope_dim dims with the cell's absolute pos), so a context-shift that re-RoPEs the + // main K cache MUST apply the identical delta-rotation to the indexer keys, or query/key + // rotations desync and the indexer top-k silently degrades. We mirror the main K-shift here. + // + // Subtlety: the cached key is H * concat(RoPE(k_pe,pos), k_nope), where H is the symmetric + // orthonormal Walsh-Hadamard rotation applied AFTER RoPE (head-wide, mixing pe+nope dims). So we + // cannot RoPE the cached key directly. We un-rotate by H (H==H^T, H*H==I), RoPE-delta the pe + // sub-block, then re-rotate by H. This is exact because (a) for this model ext_factor==0 and + // attn_factor==1 (no YaRN: GLM-DSA carries no rope-scaling metadata), so NEOX RoPE is a pure + // rotation and composable: RoPE(x,pos+delta)==RoPE(RoPE(x,pos),delta); and (b) when the Hadamard + // is disabled (DSA_HADAMARD_DISABLE) the un/re-rotate by H is simply absent (kr_had==false). + // + // NOTE (engine gate): build_k_shift only runs after get_can_shift() passes, which currently + // returns FALSE for ALL MLA models (llama.cpp llama_kv_cache_update_internal -> get_can_shift: + // is_mla_model() includes GLM_DSA). So on GLM-5.2 a RoPE context-shift is refused by the engine + // (decode returns 1 == "failed to eval") BEFORE any kr_l graph is built, identically with or + // without the indexer. This block is therefore correct-and-dormant: it never executes on the + // current MLA path, but keeps the indexer keys consistent the instant MLA K-shift is ever enabled. + // It is NOT a way for the indexer to desync — the shift simply does not happen for MLA. + if (model.arch == LLM_ARCH_GLM_DSA && hparams.indexer_head_size > 0 && !kv_self.kr_l.empty()) { + const int64_t head_size = hparams.indexer_head_size; + const int64_t rope_dim = n_rot; // n_embd_head_qk_rope (indexer pe dims) + const int64_t nope_dim = head_size - rope_dim; + + // Hadamard size used by the forward indexer: largest power-of-2 divisor of head_size that + // spans the full head row (else the forward path skips it). Rebuild the same matrix here so + // we can un/re-rotate the cached key. Must match build_deepseek2_dsa_indexer exactly. + bool do_hadamard = false; + static const bool dsa_had_disable = getenv("DSA_HADAMARD_DISABLE") != nullptr; + if (lctx.cparams.dsa_indexer_hadamard && !dsa_had_disable) { + GGML_ASSERT((head_size & ~(head_size-1)) == head_size); + do_hadamard = true; + } + int64_t nrot = 1; + while ((nrot * 2) <= head_size && head_size % (nrot * 2) == 0) nrot *= 2; + + for (int il = 0; il < n_layer; ++il) { + if ((size_t) il >= kv_self.kr_l.size() || kv_self.kr_l[il] == nullptr) { + continue; + } + ggml_tensor * kr = kv_self.kr_l[il]; // {head_size, kv_size} F16 + // The indexer forward RoPE (build_deepseek2_dsa_indexer) passes rope_factors==nullptr; + // mirror that exactly here so the delta-rotation matches the encoding. + struct ggml_tensor * rope_factors = nullptr; + + // {head_size, kv_size} -> work in F32 for the rotation, write back F16. + ggml_tensor * kr_f32 = ggml_cast(ctx0, kr, GGML_TYPE_F32); + for (auto * backend : lctx.backends) { + if (ggml_backend_supports_buft(backend, lctx.model.buft_layer[il].buft)) { + ggml_backend_sched_set_tensor_backend(lctx.sched, kr_f32, backend); + break; + } + } + cb(kr_f32, "kr_f32", il); + + // un-Hadamard: H * kr == concat(RoPE(k_pe,pos), k_nope) (H symmetric/orthonormal). + if (do_hadamard) { + kr_f32 = ggml_hadamard(ctx0, kr_f32, head_size); + cb(kr_f32, "kr_unhad", il); + } + + // view the pe sub-block {rope_dim, 1, kv_size} and the nope sub-block {nope_dim,1,kv_size}. + ggml_tensor * kr_pe = ggml_view_3d(ctx0, kr_f32, rope_dim, 1, n_ctx, + ggml_row_size(kr_f32->type, head_size), + ggml_row_size(kr_f32->type, head_size), 0); + ggml_tensor * kr_nope = ggml_view_3d(ctx0, kr_f32, nope_dim, 1, n_ctx, + ggml_row_size(kr_f32->type, head_size), + ggml_row_size(kr_f32->type, head_size), + ggml_row_size(kr_f32->type, rope_dim)); + // RoPE the pe block by the per-cell delta (inp_K_shift holds cells[i].delta) into a NEW + // tensor (non-in-place; ggml_cont copies the view first to avoid aliasing the source). + // NEOX, identical params to the indexer forward RoPE. + ggml_tensor * kr_pe_rot = ggml_rope_ext(ctx0, ggml_cont(ctx0, kr_pe), + lctx.inp_K_shift, rope_factors, n_rot, LLAMA_ROPE_TYPE_NEOX, n_ctx_orig, + freq_base, freq_scale, ext_factor, attn_factor, beta_fast, beta_slow); + cb(kr_pe_rot, "kr_pe_shifted", il); + + // reassemble {head_size, 1, kv_size} = concat(rotated pe, untouched nope) along dim0. + ggml_tensor * kr_cat = ggml_concat(ctx0, kr_pe_rot, ggml_cont(ctx0, kr_nope), 0); + kr_cat = ggml_reshape_2d(ctx0, kr_cat, head_size, n_ctx); + cb(kr_cat, "kr_cat_shifted", il); + + // re-Hadamard: H * concat(RoPE(k_pe,pos+delta), k_nope) == the new cached key. + ggml_tensor * kr_new = kr_cat; + if (do_hadamard) { + kr_new = ggml_hadamard(ctx0, kr_cat, head_size); + cb(kr_new, "kr_rehad", il); + } + // write back into the F16 cache. + ggml_tensor * kr_back = ggml_cpy(ctx0, kr_new, kr); + cb(kr_back, "kr_shifted", il); + ggml_build_forward_expand(gf, kr_back); + } + } + return gf; } @@ -302,6 +400,24 @@ ggml_cgraph * llm_build_context::build_defrag(const std::vector & ids) if (view_v_src && view_v_dst) { ggml_build_forward_expand(gf, ggml_cpy(ctx0, view_v_src, view_v_dst)); } + + // DSA lightning-indexer key cache: move the indexer keys alongside k_l. Each cell's + // indexer key (kr_l row, one per cache cell) must follow its cell, or after defrag the + // kr_l rows no longer match the cells the indexer scores by cell index. The indexer key + // is RoPE-encoded at its (unchanged) absolute pos, and defrag does NOT change pos, so a + // plain row move is correct (no re-RoPE needed; only seq_add/K-shift change pos). + if ((size_t) il < kv_self.kr_l.size() && kv_self.kr_l[il] != nullptr) { + const int64_t head_size = hparams.indexer_head_size; + ggml_tensor * view_kr_src = ggml_view_2d(ctx0, kv_self.kr_l[il], + head_size, nm, + ggml_row_size(kv_self.kr_l[il]->type, head_size), + ggml_row_size(kv_self.kr_l[il]->type, head_size*i)); + ggml_tensor * view_kr_dst = ggml_view_2d(ctx0, kv_self.kr_l[il], + head_size, nm, + ggml_row_size(kv_self.kr_l[il]->type, head_size), + ggml_row_size(kv_self.kr_l[il]->type, head_size*id)); + ggml_build_forward_expand(gf, ggml_cpy(ctx0, view_kr_src, view_kr_dst)); + } } i += nm - 1; diff --git a/src/llama-build-context.h b/src/llama-build-context.h index 28cb74c1e..8b557dd81 100644 --- a/src/llama-build-context.h +++ b/src/llama-build-context.h @@ -97,6 +97,10 @@ struct llm_build_context { struct ggml_context * ctx0 = nullptr; + // GLM-5.2 IndexShare: the most-recent "full" indexer layer's top-k selection (argsort indices). + // "shared" layers reuse this instead of computing their own. Reset to nullptr at each graph build. + struct ggml_tensor * dsa_last_full_sorted = nullptr; + // TODO: consider making the entire interface noexcept llm_build_context( llama_context & lctx, @@ -287,6 +291,27 @@ struct llm_build_context { bool is_lite, bool pp_opt); + // DSA lightning indexer (GLM-5.2 / DeepSeek-V3.2). Cache-backed (persistent indexer-key cache). + // Returns the FULL descending argsort of the per-query indexer scores [n_kv, n_tokens] (I32). + ggml_tensor * build_deepseek2_dsa_indexer( + ggml_cgraph * gf, + int il, + ggml_tensor * qr, // q_lora latent [q_lora_rank, n_tokens] (after attn_q_a_norm) + ggml_tensor * cur, // attn_norm output [n_embd, n_tokens] + ggml_tensor * KQ_mask, // F32 causal mask [n_kv, n_tokens_pad] + ggml_tensor * inp_pos); + + // Build the additive sparse causal mask from the full score ranking + the base causal KQ_mask. + ggml_tensor * build_deepseek2_dsa_sparse_mask( + ggml_tensor * sorted, // [n_kv, n_tokens] (I32) full descending argsort of scores + ggml_tensor * KQ_mask); // F32 causal mask [n_kv, n_tokens_pad] + + // Adapt the (F32, unpadded) sparse mask to the shape/dtype ggml_flash_attn_ext requires on this + // fork: F16, contiguous, ne[1] padded to GGML_KQ_MASK_PAD (== the dense KQ_mask's padded shape). + ggml_tensor * build_deepseek2_dsa_fa_mask( + ggml_tensor * sparse, // [n_kv, n_tokens] (F32) additive sparse causal mask + ggml_tensor * KQ_mask); // FA dense mask [n_kv, n_tokens_pad] (F16 when -fa 1) + ggml_cgraph * build_glm4_moe(); ggml_cgraph * build_bitnet(); diff --git a/src/llama-context.h b/src/llama-context.h index f7aaf4039..cd744c969 100644 --- a/src/llama-context.h +++ b/src/llama-context.h @@ -61,6 +61,12 @@ struct llama_kv_cache { std::vector v_l; std::vector s_l; // per layer recurrent state storage (Qwen3Next) + // DSA lightning-indexer key cache (GLM-5.2 / DeepSeek-V3.2). One per layer, MQA single + // head: [indexer_head_size, kv_size]. Mirrors k_l but stores the (Hadamard-rotated) + // indexer keys so a decoded token scores against ALL past indexer keys, not just the + // current batch. Empty unless the model has the DSA indexer. + std::vector kr_l; + // When true, the delta_net graph builder will enable per-step SSM state saves bool save_per_step_ssm = false; @@ -378,6 +384,8 @@ struct llama_context { struct ggml_tensor * inp_KQ_mask_cross; // F32 [n_outputs_enc, n_batch] struct ggml_tensor * inp_scale = nullptr; // F32 [n_tokens] struct ggml_tensor * inp_mtp_states = nullptr; + 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; ggml_backend_t ggml_backend_by_name(const char * name); @@ -393,6 +401,14 @@ struct llama_context { size_t step = 0; }; std::vector cache_copies; + // GLM-DSA lightning indexer: the indexer-key cache (kr_l) write is a separate ggml_cpy that + // the K/V cache_copies fixup does NOT cover. Under graph reuse (FA pads KV to 256, so n_kv + // stays constant across consecutive decode ubatches and the graph IS reused) its view_offs + // would stay baked at the first ubatch's kv_head, scattering this ubatch's indexer keys to a + // stale slot. Later ubatches never populate their own recent index-key cells (those cells read + // uninitialized -> wrong block-max-pool/top-k -> degraded/NaN sparse-FA decode). Register the + // kr_l cpy per layer here and patch its offset in update_cache_copies(), exactly like K/V. + std::vector dsa_cache_copies; bool update_cache_copies(); diff --git a/src/llama-cparams.h b/src/llama-cparams.h index c93739a53..ac24ba824 100644 --- a/src/llama-cparams.h +++ b/src/llama-cparams.h @@ -41,6 +41,9 @@ struct llama_cparams { bool graph_reuse; bool k_cache_hadamard; bool v_cache_hadamard; + bool dsa_indexer_hadamard = true; // apply Walsh-Hadamard rotation to DSA indexer q/k (precision) + bool dsa = false; // enable GLM DSA sparse attention (off by default; opt-in via --dsa) + int dsa_top_k = -1; // DSA top-k override (<0 => use the model's configured indexer_top_k) bool split_mode_graph_scheduling; //bool split_mode_f16; bool scheduler_async; diff --git a/src/llama-hparams.cpp b/src/llama-hparams.cpp index f3fcadf88..de2a2c874 100644 --- a/src/llama-hparams.cpp +++ b/src/llama-hparams.cpp @@ -1582,6 +1582,14 @@ void llm_load_hparams( { ml.get_key(LLM_KV_EXPERT_FEED_FORWARD_LENGTH, hparams.n_ff_exp); ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps); + // GLM-DSA lightning-indexer k_norm is a (non-RMS) LayerNorm built via LLM_NORM, + // which uses hparams.f_norm_eps in ggml_norm(). The GGUF only carries the RMS eps, + // so f_norm_eps stays 0 and CPU ggml_norm aborts (GGML_ASSERT(eps > 0)). On CUDA + // the kernel does not assert (eps=0 is numerically tolerable), which is why the + // CPU attention path was never exercised. Mirror the RMS eps so the indexer + // LayerNorm gets a valid epsilon on all backends. (HF hardcodes 1e-6 for this + // LayerNorm, but 1e-6 vs 1e-5 is within 4-chunk PPL noise here, so keep the mirror.) + if (hparams.f_norm_eps <= 0.0f) hparams.f_norm_eps = hparams.f_norm_rms_eps; ml.get_key_or_arr(LLM_KV_ROPE_DIMENSION_SECTIONS, hparams.rope_sections, 4, false); // MoE parameters @@ -1605,6 +1613,17 @@ void llm_load_hparams( ml.get_key(LLM_KV_ATTENTION_INDEXER_KEY_LENGTH, hparams.indexer_head_size); ml.get_key(LLM_KV_ATTENTION_INDEXER_TOP_K, hparams.indexer_top_k); + // GLM-5.2 IndexShare: per-layer full/shared indexer map. "full" layers compute their own + // top-k; "shared" layers reuse the previous full layer's selection (transformers + // modeling_glm_moe_dsa.py: shared layer indexer=None, topk_indices=prev_topk_indices). + // Derived from GLM-5.2's config indexer_types rule (full iff il<=1 or il%4==2), verified + // to reproduce the config's full set {0,1,2,6,10,...} exactly. Existing GGUFs carry no + // per-layer metadata, so the derivation is the source of truth; a future metadata key can + // override this for GLM-DSA variants with a different pattern. + for (uint32_t il = 0; il < hparams.n_layer; ++il) { + hparams.indexer_is_full[il] = (il <= 1) || (il % 4 == 2); + } + // Expert gating function (GLM-4.5 uses sigmoid) ml.get_key(LLM_KV_EXPERT_GATING_FUNC, hparams.expert_gating_func, false); if (hparams.expert_gating_func == LLM_EXPERT_GATING_FUNC_TYPE_NONE) { diff --git a/src/llama-hparams.h b/src/llama-hparams.h index 941c75f94..10c910698 100644 --- a/src/llama-hparams.h +++ b/src/llama-hparams.h @@ -126,6 +126,10 @@ struct llama_hparams { uint32_t indexer_n_head = 0; uint32_t indexer_head_size = 0; uint32_t indexer_top_k = 0; + // GLM-5.2 IndexShare: per-layer full/shared indexer map. "full" layers compute their own lightning- + // indexer top-k; "shared" layers reuse the previous full layer's top-k. Populated from GGUF + // indexer_types metadata if present, else derived from the GLM-5.2 config rule at load time. + std::array indexer_is_full = {}; // qwen3vl deepstack uint32_t n_deepstack_layers = 0; diff --git a/src/llama-load-tensors.cpp b/src/llama-load-tensors.cpp index 77b70d166..9981b6473 100644 --- a/src/llama-load-tensors.cpp +++ b/src/llama-load-tensors.cpp @@ -2814,16 +2814,24 @@ bool create_tensors_helper::create_glm_dsa_tensors(const LLM_TN & tn) { layer.wo = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_OUT, "weight", i), {n_head * n_embd_head_v, n_embd}, flags); // DSA indexer - // GLM-5.2 only ships the DSA indexer on a subset of layers; the rest omit it. - // The DSA indexer runtime is not implemented (graph is plain MLA), so these - // tensors are loaded-but-unused. Mark them optional so layers without an - // indexer load as nullptr (ported from ggml-org/llama.cpp#24770). + // DSA indexer tensors. GLM-5.2's reference marks "shared" layers indexer=None, but the GGUFs + // in circulation (e.g. unsloth) ship indexer weights on every layer, and the DSA runtime IS + // implemented (see build_deepseek2_dsa_indexer + IndexShare reuse). These are TENSOR_NOT_REQUIRED + // so a faithful conversion that omits shared-layer indexer weights still loads: full layers must + // have them (they compute the top-k), shared layers reuse the previous full layer's top-k. A + // "full" layer (per hparams.indexer_is_full) that loads a null indexer will warn below. layer.indexer_k_norm = create_tensor(ctx_split, tn(LLM_TENSOR_INDEXER_K_NORM, "weight", i), {hparams.indexer_head_size}, flags | llama_model_loader::TENSOR_NOT_REQUIRED); layer.indexer_k_norm_b = create_tensor(ctx_split, tn(LLM_TENSOR_INDEXER_K_NORM, "bias", i), {hparams.indexer_head_size}, flags | llama_model_loader::TENSOR_NOT_REQUIRED); layer.indexer_proj = create_tensor(ctx_split, tn(LLM_TENSOR_INDEXER_PROJ, "weight", i), {n_embd, hparams.indexer_n_head}, flags | llama_model_loader::TENSOR_NOT_REQUIRED); layer.indexer_attn_k = create_tensor(ctx_split, tn(LLM_TENSOR_INDEXER_ATTN_K, "weight", i), {n_embd, hparams.indexer_head_size}, flags | llama_model_loader::TENSOR_NOT_REQUIRED); layer.indexer_attn_q_b = create_tensor(ctx_split, tn(LLM_TENSOR_INDEXER_ATTN_Q_B, "weight", i), {q_lora_rank, hparams.indexer_n_head * hparams.indexer_head_size}, flags | llama_model_loader::TENSOR_NOT_REQUIRED); + // IndexShare: a "full" layer must carry indexer weights (it computes the top-k that shared layers + // reuse). If a conversion omits them on a full layer, DSA silently falls back to dense there. + if (i < (int) hparams.n_layer && hparams.indexer_is_full[i] && !layer.indexer_attn_q_b) { + LLAMA_LOG_WARN("%s: GLM-DSA layer %d is 'full' per indexer_types but has no indexer weights; DSA will run dense on it\n", __func__, i); + } + layer.ffn_norm = create_tensor(norm_ctx, tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, flags); if (i < (int) hparams.n_layer_dense_lead) { diff --git a/src/llama.cpp b/src/llama.cpp index f8c13eb06..de6a5f2ea 100644 --- a/src/llama.cpp +++ b/src/llama.cpp @@ -679,6 +679,28 @@ bool llama_context::update_cache_copies() { } } } + // GLM-DSA lightning indexer: patch the indexer-key (kr_l) cache write offset for the reused + // graph. Each registered cpy writes this ubatch's index keys into kr_l at the kv_head slot; + // like the K/V copies above, its baked view_offs must be re-pointed to the CURRENT kv_head, + // else a reused graph (FA pad-256, constant n_kv) keeps writing to the first ubatch's slot and + // the recent indexer-key cells read uninitialized (scattered selection -> degraded/NaN decode). + // step = kr_l->nb[1] (one index-key row = head_size * F16). No-op when DSA is off (entries null). + for (size_t il = 0; il < dsa_cache_copies.size(); ++il) { + auto & c = dsa_cache_copies[il]; + if (!c.cpy) continue; + // Sanity guard (the MSA fix omitted this): the registered cpy must still be a CPY whose + // destination view is rooted at this layer's kr_l cache (mirrors the K/V guard above, + // which checks c.cpy->view_src). If the graph was rebuilt with a different shape these no + // longer match, so refuse reuse and force a rebuild. + if (c.cpy->op != GGML_OP_CPY || + il >= kv_self.kr_l.size() || kv_self.kr_l[il] == nullptr || + c.cpy->view_src != kv_self.kr_l[il]) { + return false; + } + c.cpy->view_offs = kv_self.head * c.step; + c.cpy->src[1]->data = (char *) kv_self.kr_l[il]->data + c.cpy->view_offs; + c.cpy->data = c.cpy->src[1]->data; + } return true; } @@ -695,6 +717,9 @@ llama_context::llama_context(const llama_model & model) } else { cache_copies.resize(2*hparams.n_layer); } + // GLM-DSA lightning indexer: one indexer-key (kr_l) cache copy per layer. Entries stay null + // for non-DSA models / non-indexer layers, so update_cache_copies() is a no-op when DSA is off. + dsa_cache_copies.resize(hparams.n_layer); llama_all_contexts().push_back(this); } @@ -950,6 +975,13 @@ static bool llama_kv_cache_init( if (needs_v_cache) cache.v_l.reserve(n_layer); cache.s_l.resize(n_layer, nullptr); + // DSA lightning-indexer key cache: one [indexer_head_size, kv_size] tensor per layer. + // Allocated below (in the MLA branch) only when the model carries the indexer tensors. + const bool has_dsa_indexer = model.arch == LLM_ARCH_GLM_DSA && hparams.indexer_head_size > 0; + if (has_dsa_indexer) { + cache.kr_l.resize(n_layer, nullptr); + } + std::vector mem_split(model.splits.size(), 0); const uint32_t qnext_state_slots = llama_qwen3next_state_slots(cparams, kv_size); @@ -995,6 +1027,13 @@ static bool llama_kv_cache_init( ggml_tensor * kv = ggml_new_tensor_2d(ctx, primary_kv_type, kv_lora_rank + n_embd_head_qk_rope, kv_size); ggml_format_name(kv, "cache_k_l%d", i); cache.k_l.push_back(kv); + // DSA lightning-indexer key cache (MQA, single head). Store the Hadamard-rotated + // indexer keys in F16 so a decoded token can score against ALL past keys. + if (has_dsa_indexer && model.layers[i].indexer_attn_k && !is_mtp_tail_layer) { + ggml_tensor * kr = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, hparams.indexer_head_size, kv_size); + ggml_format_name(kr, "cache_kr_l%d", i); + cache.kr_l[i] = kr; + } if (!cparams.flash_attn && cparams.mla_attn == 1) { ggml_tensor * kvt = ggml_new_tensor_1d(ctx, cache.type_v, kv_lora_rank*kv_size); ggml_format_name(kvt, "cache_v_l%d", i); @@ -4284,6 +4323,58 @@ static void llama_set_inputs(llama_context & lctx, const llama_batch & batch) { const auto & cparams = lctx.cparams; const auto & kv_self = lctx.kv_self; + if (lctx.inp_dsa_sink) { + // Per-sequence attention-sink boost for the DSA lightning indexer top-k selection. + // inp_dsa_sink {n_kv, n_tokens}: 1e20 iff key cell i is one of query j's sequence's FIRST + // n_sink present tokens, else 0. Filled from kv_self.cells like the KQ_mask so it is correct + // for multi-sequence ubatches (each sequence's OWN sink is force-included). + // + // SERVING FIX: anchor the sink on each sequence's FIRST PRESENT position, not absolute + // pos < n_sink. After multi-turn seq_rm drops a sequence's early tokens, its earliest + // surviving token has pos >= n_sink; an absolute "pos < n_sink" test would then protect + // NOTHING for that sequence and let the (now-)sink token be masked out of top-k, collapsing + // it. Anchoring on per-sequence min(pos) keeps the sink protection following the sequence's + // actual first present cell. For a fresh sequence starting at pos 0, min(pos)==0 so the + // boosted set is identical to the old behaviour (n_seq==1 byte-identical). + GGML_ASSERT(ggml_backend_buffer_is_host(lctx.inp_dsa_sink->buffer)); + static const int n_sink = []{ const char * e = getenv("DSA_SINK"); return e ? atoi(e) : 1; }(); + const int64_t n_kv = lctx.inp_dsa_sink->ne[0]; + const int64_t n_tok_idx = lctx.inp_dsa_sink->ne[1]; + float * data = (float *) lctx.inp_dsa_sink->data; + std::memset(data, 0, ggml_nbytes(lctx.inp_dsa_sink)); + + // Per-sequence first present pos: min over present cells of that seq, restricted to the + // n_kv key span the indexer scores. Cache across query rows in this ubatch. + // ASSUMPTION (latent today): min(pos) over the whole [0,n_kv) span equals the sequence's + // genuine first-present token only for a NON-SLIDING cache. GLM_DSA has no SWA, so the + // oldest cell of a seq is its true sink; if SWA were ever added, the window could evict the + // real sink and min(pos) would anchor on the wrong (window-floor) cell — revisit then. + std::unordered_map seq_min_pos; + for (int64_t i = 0; i < n_kv; ++i) { + const auto & cell = kv_self.cells[i]; + if (cell.pos < 0) continue; + for (const llama_seq_id sid : cell.seq_id) { + auto it = seq_min_pos.find(sid); + if (it == seq_min_pos.end() || cell.pos < it->second) { + seq_min_pos[sid] = cell.pos; + } + } + } + + for (int64_t j = 0; j < n_tok_idx && j < (int64_t) batch.n_tokens; ++j) { + const llama_seq_id seq_id = batch.seq_id[j][0]; + auto it = seq_min_pos.find(seq_id); + if (it == seq_min_pos.end()) continue; // no present cell for this seq in the key span + const llama_pos first_pos = it->second; + for (int64_t i = 0; i < n_kv; ++i) { + const auto & cell = kv_self.cells[i]; + if (cell.pos >= first_pos && cell.pos < first_pos + n_sink && cell.has_seq_id(seq_id)) { + data[j*n_kv + i] = 1e20f; + } + } + } + } + if (batch.token && lctx.inp_tokens) { #if IK_PRINT_TIMING == 2 auto tim1 = ggml_time_us(); @@ -5932,7 +6023,12 @@ static void llama_kv_cache_defrag_internal(struct llama_context & lctx) { // - x2 for keys and values //const uint32_t max_moves = model.max_nodes()/(6*n_layer); // TODO: tmp fix https://github.com/ggerganov/llama.cpp/issues/6685#issuecomment-2057579516 - const uint32_t max_moves = (lctx.model.max_nodes(1) - 2*n_layer)/(6*n_layer); + // DSA: build_defrag additionally moves the indexer-key cache (kr_l), +3 tensors/layer/move + // (src view, dst view, copy), so budget 9*n_layer per move when the indexer cache is present. + const bool has_dsa_indexer_defrag = + lctx.model.arch == LLM_ARCH_GLM_DSA && !kv_self.kr_l.empty(); + const uint32_t tensors_per_move = has_dsa_indexer_defrag ? 9 : 6; + const uint32_t max_moves = (lctx.model.max_nodes(1) - 2*n_layer)/(tensors_per_move*n_layer); // determine which KV cells to move where // @@ -6541,6 +6637,8 @@ struct llama_context_params llama_context_default_params() { /*.fused_mmad =*/ true, /*.rope_cache =*/ false, /*.graph_reuse =*/ true, + /*.dsa =*/ false, + /*.dsa_top_k =*/ -1, /*.min_experts =*/ -1, /*.thtesh_experts =*/ 0.0f, /*.only_active_experts =*/ false, @@ -6958,6 +7056,14 @@ struct llama_context * llama_init_from_model( cparams.fused_mmad = params.fused_mmad; cparams.rope_cache = params.rope_cache; cparams.graph_reuse = params.graph_reuse; + cparams.dsa = params.dsa; + cparams.dsa_top_k = params.dsa_top_k; + // The DSA lightning indexer is built only in the layer-mode (non-TP) attention path. Under + // -sm graph / -sm attn the model runs the tensor-parallel attention path, which has no indexer, + // so --dsa would silently run dense MLA. Warn instead of degrading silently. + if (cparams.dsa && (model->split_mode == LLAMA_SPLIT_MODE_GRAPH || model->split_mode == LLAMA_SPLIT_MODE_ATTN)) { + LLAMA_LOG_WARN("%s: --dsa is not active under -sm graph/attn (tensor-parallel attention has no indexer); running dense MLA\n", __func__); + } cparams.k_cache_hadamard = params.k_cache_hadamard; cparams.v_cache_hadamard = params.v_cache_hadamard; // Folding H into wv_b/wk_b_pp permanently mutates the model; a later context @@ -7316,12 +7422,18 @@ struct llama_context * llama_init_from_model( { size_t memory_size_k = 0; size_t memory_size_v = 0; + size_t memory_size_k_indexer = 0; for (auto & k : ctx->kv_self.k_l) { if (k) { memory_size_k += ggml_nbytes(k); } } + for (auto & k : ctx->kv_self.kr_l) { + if (k) { + memory_size_k_indexer += ggml_nbytes(k); + } + } for (auto & v : ctx->kv_self.v_l) { if (v) { @@ -7329,8 +7441,8 @@ struct llama_context * llama_init_from_model( } } - if (memory_size_k + memory_size_v > 0) { - if (cparams.mla_attn != 0 && !cparams.flash_attn) { + if (memory_size_k + memory_size_v) { + if (cparams.mla_attn != 0 && !cparams.flash_attn) { LLAMA_LOG_INFO("%s: KV self size = %7.2f MiB, c^KV (%s): %7.2f MiB, kv^T (%s): %7.2f MiB\n", __func__, (float)(memory_size_k + memory_size_v) / (1024.0f * 1024.0f), ggml_type_name(type_k), (float)memory_size_k / (1024.0f * 1024.0f), @@ -7346,6 +7458,9 @@ struct llama_context * llama_init_from_model( ggml_type_name(type_v), (float)memory_size_v / (1024.0f * 1024.0f)); } } + if (memory_size_k_indexer > 0) { + LLAMA_LOG_INFO("%s: KV self indexer size = %7.2f MiB\n", __func__, memory_size_k_indexer/1024./1024.); + } } // graph outputs buffer