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

wip: Enhance DeepSeek-V4 architecture with new tensor types and sqrtsoftplus gating function

Browse Source
This commit is contained in:
samuel
2026-07-05 12:54:17 -03:00
parent 40eb9f1515
commit 3de2b0a0cf
15 changed files with 702 additions and 114 deletions
Download Patch File Download Diff File Expand all files Collapse all files
ggml/include/ggml.h
+15
View File
@@ -2137,6 +2137,21 @@ extern "C" {
float beta_fast,
float beta_slow);
GGML_API struct ggml_tensor * ggml_rope_ext_back(
struct ggml_context * ctx,
struct ggml_tensor * a,
struct ggml_tensor * b,
struct ggml_tensor * c,
int n_dims,
int mode,
int n_ctx_orig,
float freq_base,
float freq_scale,
float ext_factor,
float attn_factor,
float beta_fast,
float beta_slow);
GGML_API struct ggml_tensor * ggml_rope_cache(
struct ggml_context * ctx,
struct ggml_tensor * b,
ggml/src/ggml.c
+59 -22
View File
@@ -9470,6 +9470,26 @@ struct ggml_tensor * ggml_rope_back(
return result;
}
struct ggml_tensor * ggml_rope_ext_back(
struct ggml_context * ctx,
struct ggml_tensor * a,
struct ggml_tensor * b,
struct ggml_tensor * c,
int n_dims,
int mode,
int n_ctx_orig,
float freq_base,
float freq_scale,
float ext_factor,
float attn_factor,
float beta_fast,
float beta_slow) {
struct ggml_tensor * result = ggml_rope_ext(
ctx, a, b, c, n_dims, mode, n_ctx_orig, freq_base, freq_scale, ext_factor, attn_factor, beta_fast, beta_slow);
result->op = GGML_OP_ROPE_BACK;
return result;
}
// ggml_clamp
struct ggml_tensor * ggml_clamp(
@@ -14972,34 +14992,51 @@ static void ggml_compute_forward_concat_any(
const struct ggml_tensor * src1 = dst->src[1];
GGML_ASSERT(src0->type == src1->type && src0->type == dst->type);
GGML_ASSERT(!ggml_is_quantized(src0->type));
const size_t len = ggml_type_size(src0->type);
const int ith = params->ith;
const int nth = params->nth;
GGML_TENSOR_BINARY_OP_LOCALS
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;
GGML_ASSERT(dim >= 0 && dim < 4);
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);
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);
for (int d = 0; d < 4; ++d) {
if (d == dim) {
GGML_ASSERT(dst->ne[d] == src0->ne[d] + src1->ne[d]);
} else {
GGML_ASSERT(src0->ne[d] == src1->ne[d]);
GGML_ASSERT(dst->ne[d] == src0->ne[d]);
}
}
int64_t o[4] = { 0, 0, 0, 0 };
o[dim] = src0->ne[dim];
const char * x;
// Keep the reference stride-aware behavior here. DSV4 naturally concatenates
// along dim 1/2, and assuming contiguous slices causes backend-only failures.
for (int64_t i3 = 0; i3 < ne3; ++i3) {
for (int64_t i2 = ith; i2 < ne2; i2 += nth) {
for (int64_t i1 = 0; i1 < ne1; ++i1) {
for (int64_t i0 = 0; i0 < ne0; ++i0) {
if (i0 < ne00 && i1 < ne01 && i2 < ne02 && i3 < ne03) {
x = (const char *) src0->data + (i0 )*nb00 + (i1 )*nb01 + (i2 )*nb02 + (i3 )*nb03;
} else {
x = (const char *) src1->data + (i0 - o[0])*nb10 + (i1 - o[1])*nb11 + (i2 - o[2])*nb12 + (i3 - o[3])*nb13;
}
char * y = (char *) dst->data + i0*nb0 + i1*nb1 + i2*nb2 + i3*nb3;
memcpy(y, x, len);
}
}
}
}
}
static void ggml_compute_forward_concat(
models/templates/README.md
+1
View File
@@ -23,4 +23,5 @@ These templates can be updated with the following commands:
./scripts/get_chat_template.py Qwen/Qwen3-0.6B > models/templates/Qwen-Qwen3-0.6B.jinja
./scripts/get_chat_template.py zai-org/GLM-4.5 > models/templates/zai-org-GLM-4.5.jinja
./scripts/get_chat_template.py deepseek-ai/DeepSeek-V3.1 > models/templates/deepseek-ai-DeepSeek-V3.1.jinja
./scripts/get_chat_template.py deepseek-ai/DeepSeek-V4 > models/templates/deepseek-ai-DeepSeek-V4.jinja
```
models/templates/deepseek-ai-DeepSeek-V4.jinja
+112
View File
@@ -0,0 +1,112 @@
{%- if not add_generation_prompt is defined -%}
{%- set add_generation_prompt = false -%}
{%- endif -%}
{%- if not thinking is defined -%}
{%- if enable_thinking is defined -%}
{%- set thinking = enable_thinking -%}
{%- else -%}
{%- set thinking = false -%}
{%- endif -%}
{%- endif -%}
{%- set dsml_token = '｜DSML｜' -%}
{%- set thinking_start_token = '<think>' -%}
{%- set thinking_end_token = '</think>' -%}
{%- set tools_header = '## Tools\n\nYou have access to a set of tools to help answer the user\'s question. You can invoke tools by writing a "<' + dsml_token + 'tool_calls>" block like the following:\n\n<' + dsml_token + 'tool_calls>\n<' + dsml_token + 'invoke name="$TOOL_NAME">\n<' + dsml_token + 'parameter name="$PARAMETER_NAME" string="true|false">$PARAMETER_VALUE</' + dsml_token + 'parameter>\n...\n</' + dsml_token + 'invoke>\n<' + dsml_token + 'invoke name="$TOOL_NAME2">\n...\n</' + dsml_token + 'invoke>\n</' + dsml_token + 'tool_calls>\n\nString parameters should be specified as is and set `string="true"`. For all other types (numbers, booleans, arrays, objects), pass the value in JSON format and set `string="false"`.\n\nIf thinking_mode is enabled (triggered by ' + thinking_start_token + '), you MUST output your complete reasoning inside ' + thinking_start_token + '...' + thinking_end_token + ' BEFORE any tool calls or final response.\n\nOtherwise, output directly after ' + thinking_end_token + ' with tool calls or final response.\n\n### Available Tool Schemas\n\n' -%}
{%- set tools_footer = '\nYou MUST strictly follow the above defined tool name and parameter schemas to invoke tool calls.\n' -%}
{%- set ns = namespace(system_prompt = '', is_first_sp = true) -%}
{%- for message in messages -%}
{%- if message['role'] == 'system' -%}
{%- if ns.is_first_sp -%}
{%- set ns.system_prompt = ns.system_prompt + (message['content'] or '') -%}
{%- set ns.is_first_sp = false -%}
{%- else -%}
{%- set ns.system_prompt = ns.system_prompt + '\n\n' + (message['content'] or '') -%}
{%- endif -%}
{%- endif -%}
{%- endfor -%}
{%- if tools is defined and tools -%}
{%- set ts = namespace(schemas = '') -%}
{%- for tool in tools -%}
{%- if tool['type'] == 'function' -%}
{%- set ts.schemas = ts.schemas + (tool['function'] | tojson) + '\n' -%}
{%- endif -%}
{%- endfor -%}
{%- if ns.system_prompt -%}
{%- set ns.system_prompt = ns.system_prompt + '\n\n' + tools_header + ts.schemas + tools_footer -%}
{%- else -%}
{%- set ns.system_prompt = tools_header + ts.schemas + tools_footer -%}
{%- endif -%}
{%- endif -%}
{{- bos_token -}}
{{- ns.system_prompt -}}
{%- set last_user_idx = namespace(value = -1) -%}
{%- for message in messages -%}
{%- if message['role'] == 'user' or message['role'] == 'developer' or message['role'] == 'tool' -%}
{%- set last_user_idx.value = loop.index0 -%}
{%- endif -%}
{%- endfor -%}
{%- set state = namespace(in_user = false) -%}
{%- for message in messages -%}
{%- if message['role'] == 'user' or message['role'] == 'developer' -%}
{%- if state.in_user -%}
{{- '\n\n' -}}
{%- else -%}
{{- '<｜User｜>' -}}
{%- set state.in_user = true -%}
{%- endif -%}
{{- message['content'] or '' -}}
{%- elif message['role'] == 'tool' -%}
{%- if state.in_user -%}
{{- '\n\n' -}}
{%- else -%}
{{- '<｜User｜>' -}}
{%- set state.in_user = true -%}
{%- endif -%}
{{- '<tool_result>' + (message['content'] or '') + '</tool_result>' -}}
{%- elif message['role'] == 'assistant' -%}
{%- set state.in_user = false -%}
{{- '<｜Assistant｜>' -}}
{%- set is_after_last_user = loop.index0 > last_user_idx.value -%}
{%- if is_after_last_user and thinking -%}
{{- thinking_start_token -}}
{%- if message['reasoning_content'] is defined and message['reasoning_content'] -%}
{{- message['reasoning_content'] -}}
{%- endif -%}
{{- thinking_end_token -}}
{%- else -%}
{{- thinking_end_token -}}
{%- endif -%}
{%- if message['content'] is defined and message['content'] -%}
{{- message['content'] -}}
{%- endif -%}
{%- if message['tool_calls'] -%}
{{- '\n\n<' + dsml_token + 'tool_calls>\n' -}}
{%- for tool in message['tool_calls'] -%}
{%- set func = tool['function'] -%}
{{- '<' + dsml_token + 'invoke name="' + func['name'] + '">\n' -}}
{%- set args = func['arguments'] -%}
{%- if args is string -%}
{%- set args = args | from_json -%}
{%- endif -%}
{%- for key, val in args.items() -%}
{%- if val is string -%}
{{- '<' + dsml_token + 'parameter name="' + key + '" string="true">' + val + '</' + dsml_token + 'parameter>\n' -}}
{%- else -%}
{{- '<' + dsml_token + 'parameter name="' + key + '" string="false">' + (val | tojson) + '</' + dsml_token + 'parameter>\n' -}}
{%- endif -%}
{%- endfor -%}
{{- '</' + dsml_token + 'invoke>\n' -}}
{%- endfor -%}
{{- '</' + dsml_token + 'tool_calls>' -}}
{%- endif -%}
{{- '<｜end▁of▁sentence｜>' -}}
{%- endif -%}
{%- endfor -%}
{%- if add_generation_prompt -%}
{{- '<｜Assistant｜>' -}}
{%- if thinking -%}
{{- thinking_start_token -}}
{%- else -%}
{{- thinking_end_token -}}
{%- endif -%}
{%- endif -%}
src/graphs/build_deepseek4.cpp
+197 -37
View File
@@ -5,9 +5,18 @@
#include <algorithm>
#include <cmath>
#include <cstring>
#include <sstream>
#include <string>
static float dsv4_rope_attn_factor(float freq_scale, float ext_factor) {
if (ext_factor == 0.0f) {
return 1.0f;
}
return 1.0f / (1.0f + 0.1f*logf(1.0f/freq_scale));
}
static size_t dsv4_elem_offset(const ggml_tensor * t, int64_t i) {
return ggml_row_size(t->type, i);
}
@@ -39,6 +48,32 @@ static void dsv4_trace_tensor_shape(std::ostringstream & out, const char * tag,
<< t->ne[3] << "]";
}
static ggml_tensor * dsv4_concat_named(
ggml_context * ctx,
ggml_tensor * a,
ggml_tensor * b,
int dim,
const char * name,
int il) {
if (llama_dsv4_trace_enabled()) {
std::ostringstream out;
out << "{\"event\":\"dsv4_concat\""
<< ",\"name\":\"" << name << "\""
<< ",\"layer\":" << il
<< ",\"dim\":" << dim;
dsv4_trace_tensor_shape(out, "a", a);
dsv4_trace_tensor_shape(out, "b", b);
out << ",\"atype\":" << int(a ? a->type : -1)
<< ",\"btype\":" << int(b ? b->type : -1)
<< "}";
llama_dsv4_trace_jsonl(out.str());
}
ggml_tensor * r = ggml_concat(ctx, a, b, dim);
ggml_set_name(r, name);
return r;
}
static ggml_tensor * dsv4_hc_affine(
ggml_context * ctx,
ggml_tensor * x,
@@ -70,6 +105,26 @@ static ggml_tensor * dsv4_new_mask_input(ggml_context * ctx, ggml_tensor ** dst,
return *dst;
}
static ggml_tensor * dsv4_new_f32_mat_input(ggml_context * ctx, ggml_tensor ** dst, int64_t n, const char * name) {
*dst = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, std::max<int64_t>(1, n), std::max<int64_t>(1, n));
ggml_set_input(*dst);
ggml_set_name(*dst, name);
return *dst;
}
static int64_t dsv4_k_rot_size(int64_t n_embd) {
if (n_embd < 64 || n_embd % 64 != 0) {
return 0;
}
int64_t n_rot = 64;
do {
n_rot *= 2;
} while (n_embd % n_rot == 0);
return n_rot / 2;
}
static void dsv4_build_plan_inputs(
ggml_context * ctx,
llama_context & lctx,
@@ -84,6 +139,12 @@ static void dsv4_build_plan_inputs(
dsv4_new_i64_input(ctx, &inputs.state_write_idxs, (int64_t) plan.state_write_idxs.size(), (std::string(tag) + "_state_write").c_str());
dsv4_new_i32_input(ctx, &inputs.state_write_pos, (int64_t) plan.state_write_pos.size(), (std::string(tag) + "_write_pos").c_str());
dsv4_new_mask_input(ctx, &inputs.kq_mask, std::max<int64_t>(1, plan.n_kv), n_tokens, (std::string(tag) + "_kq_mask").c_str());
if (std::strcmp(tag, "dsv4_lid") == 0) {
const int64_t n_rot = dsv4_k_rot_size(lctx.model.hparams.indexer_head_size);
if (n_rot > 0) {
dsv4_new_f32_mat_input(ctx, &inputs.k_rot, n_rot, (std::string(tag) + "_k_rot").c_str());
}
}
GGML_UNUSED(lctx);
}
@@ -91,7 +152,7 @@ static ggml_tensor * dsv4_append_zero_row(ggml_context * ctx, ggml_tensor * t, b
ggml_tensor * row = ggml_view_1d(ctx, t, t->ne[0], 0);
row = neg_inf ? ggml_scale_bias(ctx, row, 0.0f, -INFINITY) : ggml_scale(ctx, row, 0.0f);
row = ggml_reshape_2d(ctx, row, t->ne[0], 1);
return ggml_concat(ctx, t, row, 1);
return dsv4_concat_named(ctx, t, row, 1, "dsv4_append_zero_row", -1);
}
static ggml_tensor * dsv4_with_zero_dep(ggml_context * ctx, ggml_tensor * t, ggml_tensor * dep) {
@@ -114,8 +175,9 @@ static ggml_tensor * dsv4_cache_view_2d(
static ggml_tensor * dsv4_build_raw_mask_view(
ggml_context * ctx,
ggml_tensor * mask,
int64_t n_kv) {
return ggml_view_2d(ctx, mask, n_kv, mask->ne[1], mask->nb[1], 0);
int64_t n_kv,
int64_t n_tokens) {
return ggml_view_2d(ctx, mask, n_kv, n_tokens, mask->nb[1], 0);
}
static ggml_tensor * dsv4_cache_view_3d(
@@ -130,6 +192,33 @@ static ggml_tensor * dsv4_cache_view_3d(
0);
}
static ggml_tensor * dsv4_repeat_streams(ggml_context * ctx, ggml_tensor * t, int64_t n_stream) {
if (t->ne[3] == n_stream) {
return t;
}
GGML_ASSERT(t->ne[3] == 1);
return ggml_repeat_4d(ctx, t, t->ne[0], t->ne[1], t->ne[2], n_stream);
}
static ggml_tensor * dsv4_build_kq_zero_bias(
ggml_context * ctx,
const llama_cparams & cparams,
ggml_tensor * kq_mask,
int64_t n_head) {
GGML_UNUSED(ctx);
GGML_UNUSED(n_head);
if (!cparams.flash_attn || kq_mask->ne[3] == 1) {
return nullptr;
}
// Mainline uses this to keep unified multi-stream DSV4 on the explicit
// attention path. This fork keeps DSV4 stores non-unified, so there is no
// direct kv_unified equivalent to trigger here yet.
return nullptr;
}
static ggml_tensor * dsv4_build_attn(
ggml_context * ctx,
const llama_hparams & hparams,
@@ -137,6 +226,7 @@ static ggml_tensor * dsv4_build_attn(
ggml_tensor * q,
ggml_tensor * k,
ggml_tensor * v,
ggml_tensor * kq_b,
ggml_tensor * kq_mask,
ggml_tensor * sinks,
float kq_scale) {
@@ -160,7 +250,8 @@ static ggml_tensor * dsv4_build_attn(
llama_dsv4_trace_jsonl(out.str());
}
if (cparams.flash_attn) {
const bool use_flash_attn = cparams.flash_attn && kq_b == nullptr;
if (use_flash_attn) {
if (v_trans) {
v = ggml_transpose(ctx, v);
}
@@ -172,6 +263,9 @@ static ggml_tensor * dsv4_build_attn(
}
ggml_tensor * kq = ggml_mul_mat(ctx, k, q);
if (kq_b != nullptr) {
kq = ggml_add(ctx, kq, kq_b);
}
if (hparams.attn_soft_cap) {
kq = ggml_softcap_max(ctx, kq, kq_mask, kq_scale, hparams.f_max_alibi_bias,
1.0f / hparams.f_attn_logit_softcapping, hparams.f_attn_logit_softcapping);
@@ -312,7 +406,7 @@ static ggml_tensor * build_hc_post(
}
cur = ggml_reshape_3d(ctx0, cur, n_embd, 1, nt);
out = out ? ggml_concat(ctx0, out, cur, 1) : cur;
out = out ? dsv4_concat_named(ctx0, out, cur, 1, "dsv4_hc_weighted_sum", -1) : cur;
}
return out;
@@ -386,7 +480,8 @@ static ggml_tensor * build_hca_compressed_kv_from_state(
ggml_row_size(comp->type, n_embd_head),
ggml_row_size(comp->type, n_embd_head_nope));
comp_pe = ggml_rope_ext(ctx0, comp_pe, comp_pos, nullptr, n_embd_head_rope, llm.rope_type, llm.n_ctx_orig,
llm.hparams.dsv4_compress_rope_base, llm.freq_scale, llm.ext_factor, llm.attn_factor, llm.beta_fast, llm.beta_slow);
llm.hparams.dsv4_compress_rope_base, llm.freq_scale, llm.ext_factor,
dsv4_rope_attn_factor(llm.freq_scale, llm.ext_factor), llm.beta_fast, llm.beta_slow);
comp = ggml_concat(ctx0, comp_nope, comp_pe, 0);
llm.cb(comp, "hca_comp_out", il);
@@ -436,8 +531,8 @@ static ggml_tensor * build_overlap_compressed_kv_from_state(
ggml_row_size(score_cur->type, n_embd_head)));
score_cur = ggml_reshape_3d(ctx0, score_cur, n_embd_head, ratio, n_blocks);
ggml_tensor * values = ggml_concat(ctx0, kv_prev, kv_cur, 1);
ggml_tensor * scores = ggml_concat(ctx0, score_prev, score_cur, 1);
ggml_tensor * values = dsv4_concat_named(ctx0, kv_prev, kv_cur, 1, "dsv4_comp_values", il);
ggml_tensor * scores = dsv4_concat_named(ctx0, score_prev, score_cur, 1, "dsv4_comp_scores", il);
values = ggml_cont(ctx0, ggml_permute(ctx0, values, 1, 0, 2, 3));
scores = ggml_cont(ctx0, ggml_permute(ctx0, scores, 1, 0, 2, 3));
@@ -459,7 +554,8 @@ static ggml_tensor * build_overlap_compressed_kv_from_state(
ggml_row_size(comp->type, n_embd_head),
ggml_row_size(comp->type, n_embd_head_nope));
comp_pe = ggml_rope_ext(ctx0, comp_pe, comp_pos, nullptr, n_embd_head_rope, llm.rope_type, llm.n_ctx_orig,
llm.hparams.dsv4_compress_rope_base, llm.freq_scale, llm.ext_factor, llm.attn_factor, llm.beta_fast, llm.beta_slow);
llm.hparams.dsv4_compress_rope_base, llm.freq_scale, llm.ext_factor,
dsv4_rope_attn_factor(llm.freq_scale, llm.ext_factor), llm.beta_fast, llm.beta_slow);
comp = ggml_concat(ctx0, comp_nope, comp_pe, 0);
llm.cb(comp, tag, il);
@@ -470,6 +566,15 @@ static ggml_tensor * build_top_k_mask(
ggml_context * ctx0,
ggml_tensor * kq_mask,
ggml_tensor * top_k) {
if (llama_dsv4_trace_enabled()) {
std::ostringstream out;
out << "{\"event\":\"top_k_mask_inputs\"";
dsv4_trace_tensor_shape(out, "kq_mask", kq_mask);
dsv4_trace_tensor_shape(out, "top_k", top_k);
out << "}";
llama_dsv4_trace_jsonl(out.str());
}
ggml_tensor * kq_mask_all = ggml_fill(ctx0, kq_mask, -INFINITY);
kq_mask_all = ggml_view_4d(ctx0, kq_mask_all, 1, kq_mask_all->ne[0], kq_mask_all->ne[1], 1,
kq_mask_all->nb[0], kq_mask_all->nb[1], kq_mask_all->nb[2], 0);
@@ -500,9 +605,11 @@ static ggml_tensor * dsv4_build_lid_top_k(
const int64_t n_indexer_head = hparams.indexer_n_head;
const int64_t n_tokens = cur->ne[1];
const int64_t n_lid = llm.lctx.dsv4.lid_plan.n_kv;
ggml_tensor * k_rot = llm.lctx.dsv4.inputs.lid.k_rot;
GGML_ASSERT(n_embd_indexer_head >= n_embd_indexer_head_rope);
GGML_ASSERT(n_lid > 0);
GGML_ASSERT(k_rot != nullptr);
ggml_tensor * indexer_q = llm.llm_build_lora_mm(llm.lctx, ctx0, layer.indexer_attn_q_b, qr);
indexer_q = ggml_reshape_3d(ctx0, indexer_q, n_embd_indexer_head, n_indexer_head, n_tokens);
@@ -519,30 +626,64 @@ static ggml_tensor * dsv4_build_lid_top_k(
indexer_q_pe = ggml_rope_ext(ctx0, indexer_q_pe, inp_pos, nullptr, n_embd_indexer_head_rope,
llm.rope_type, llm.n_ctx_orig,
hparams.dsv4_compress_rope_base, llm.freq_scale,
llm.ext_factor, llm.attn_factor, llm.beta_fast, llm.beta_slow);
llm.ext_factor, dsv4_rope_attn_factor(llm.freq_scale, llm.ext_factor), llm.beta_fast, llm.beta_slow);
indexer_q = ggml_concat(ctx0, indexer_q_nope, indexer_q_pe, 0);
indexer_q = ggml_hadamard(ctx0, indexer_q, (int) n_embd_indexer_head);
llm.cb(indexer_q, "lid_q_rope", il);
indexer_q = ggml_mul_mat(ctx0, k_rot, indexer_q);
llm.cb(indexer_q, "lid_q_rot", il);
ggml_tensor * indexer_weights = llm.llm_build_lora_mm(llm.lctx, ctx0, layer.indexer_proj, cur);
indexer_weights = ggml_scale(ctx0, indexer_weights, 1.0f / std::sqrt(float(n_embd_indexer_head * n_indexer_head)));
llm.cb(indexer_weights, "lid_weights", il);
ggml_tensor * indexer_k = dsv4_cache_view_3d(ctx0, llm.lctx.dsv4.cache.lid_k[(size_t) il], n_embd_indexer_head, n_lid);
indexer_k = ggml_view_4d(ctx0, indexer_k,
indexer_k->ne[0], indexer_k->ne[1], n_lid, indexer_k->ne[3],
indexer_k->nb[1], indexer_k->nb[2], indexer_k->nb[3], 0);
llm.cb(indexer_k, "lid_k", il);
const int64_t n_stream = std::max<int64_t>(1, indexer_k->ne[3]);
indexer_q = ggml_view_4d(ctx0, indexer_q,
indexer_q->ne[0], indexer_q->ne[1], indexer_q->ne[2] / n_stream, n_stream,
indexer_q->nb[1], indexer_q->nb[2], indexer_q->nb[3] / n_stream, 0);
indexer_weights = ggml_view_4d(ctx0, indexer_weights,
indexer_weights->ne[0], indexer_weights->ne[1] / n_stream, indexer_weights->ne[2], n_stream,
indexer_weights->nb[1], indexer_weights->nb[2] / n_stream, indexer_weights->nb[3] / n_stream, 0);
indexer_q = ggml_permute(ctx0, indexer_q, 0, 2, 1, 3);
llm.cb(indexer_q, "lid_q_stream", il);
indexer_k = ggml_permute(ctx0, indexer_k, 0, 2, 1, 3);
llm.cb(indexer_k, "lid_k_stream", il);
ggml_tensor * indexer_kq = ggml_mul_mat(ctx0, indexer_k, indexer_q);
llm.cb(indexer_kq, "lid_kq", il);
indexer_kq = ggml_relu(ctx0, indexer_kq);
ggml_tensor * indexer_weights_3d = ggml_reshape_3d(ctx0, indexer_weights, 1, n_indexer_head, n_tokens);
ggml_tensor * indexer_score = ggml_mul(ctx0, indexer_kq, indexer_weights_3d);
indexer_score = ggml_cont(ctx0, ggml_permute(ctx0, indexer_score, 1, 0, 2, 3));
indexer_kq = ggml_cont(ctx0, ggml_permute(ctx0, indexer_kq, 2, 1, 0, 3));
llm.cb(indexer_kq, "lid_kq_perm", il);
ggml_tensor * indexer_score = ggml_relu(ctx0, indexer_kq);
if (llama_dsv4_trace_enabled()) {
std::ostringstream out;
out << "{\"event\":\"lid_mul_inputs\",\"layer\":" << il;
dsv4_trace_tensor_shape(out, "score_pre_mul", indexer_score);
dsv4_trace_tensor_shape(out, "weights", indexer_weights);
out << "}";
llama_dsv4_trace_jsonl(out.str());
}
indexer_score = ggml_mul(ctx0, indexer_score, indexer_weights);
indexer_score = ggml_sum_rows(ctx0, indexer_score);
indexer_score = ggml_cont(ctx0, ggml_permute(ctx0, indexer_score, 2, 1, 0, 3));
indexer_score = ggml_view_2d(ctx0, indexer_score, n_lid, n_tokens, indexer_score->nb[2], 0);
llm.cb(indexer_score, "lid_score", il);
indexer_score = ggml_add(ctx0, indexer_score, dsv4_build_raw_mask_view(ctx0, llm.lctx.dsv4.inputs.lid.kq_mask, n_lid));
if (llama_dsv4_trace_enabled()) {
std::ostringstream out;
out << "{\"event\":\"lid_mask_inputs\",\"layer\":" << il;
dsv4_trace_tensor_shape(out, "score", indexer_score);
dsv4_trace_tensor_shape(out, "mask", llm.lctx.dsv4.inputs.lid.kq_mask);
out << "}";
llama_dsv4_trace_jsonl(out.str());
}
indexer_score = ggml_add(ctx0, indexer_score, dsv4_build_raw_mask_view(ctx0, llm.lctx.dsv4.inputs.lid.kq_mask, n_lid, n_tokens));
llm.cb(indexer_score, "lid_score_masked", il);
const uint32_t n_top_k = (uint32_t) std::min<int64_t>(indexer_score->ne[0], hparams.indexer_top_k);
@@ -602,7 +743,7 @@ ggml_cgraph * llm_build_context::build_deepseek4() {
const float freq_base_l = use_compress_rope ? hparams.dsv4_compress_rope_base : freq_base;
const float freq_scale_l = use_compress_rope ? freq_scale : 1.0f;
const float ext_factor_l = use_compress_rope ? ext_factor : 0.0f;
const float attn_factor_l = use_compress_rope ? attn_factor : 1.0f;
const float attn_factor_l = dsv4_rope_attn_factor(freq_scale_l, ext_factor_l);
const float beta_fast_l = use_compress_rope ? beta_fast : 0.0f;
const float beta_slow_l = use_compress_rope ? beta_slow : 0.0f;
const int32_t n_ctx_orig_l = use_compress_rope ? n_ctx_orig : 0;
@@ -664,8 +805,8 @@ ggml_cgraph * llm_build_context::build_deepseek4() {
csa_state_score = ggml_add(ctx0, csa_state_score, csa_ape_rows);
if (lctx.dsv4.inputs.csa.state_write_idxs != nullptr && lctx.dsv4.csa_plan.state_write_idxs.size() > 0) {
ggml_tensor * csa_source_kv = ggml_concat(ctx0, lctx.dsv4.cache.csa_state_kv[(size_t) il], csa_state_kv, 1);
ggml_tensor * csa_source_score = ggml_concat(ctx0, lctx.dsv4.cache.csa_state_score[(size_t) il], csa_state_score, 1);
ggml_tensor * csa_source_kv = dsv4_concat_named(ctx0, lctx.dsv4.cache.csa_state_kv[(size_t) il], csa_state_kv, 1, "dsv4_csa_source_kv", il);
ggml_tensor * csa_source_score = dsv4_concat_named(ctx0, lctx.dsv4.cache.csa_state_score[(size_t) il], csa_state_score, 1, "dsv4_csa_source_score", il);
ggml_tensor * csa_comp = build_overlap_compressed_kv_from_state(
ctx0, *this,
csa_source_kv, csa_source_score,
@@ -697,8 +838,8 @@ ggml_cgraph * llm_build_context::build_deepseek4() {
lid_state_score = ggml_add(ctx0, lid_state_score, lid_ape_rows);
if (lctx.dsv4.inputs.lid.state_write_idxs != nullptr && lctx.dsv4.lid_plan.state_write_idxs.size() > 0) {
ggml_tensor * lid_source_kv = ggml_concat(ctx0, lctx.dsv4.cache.lid_state_kv[(size_t) il], lid_state_kv, 1);
ggml_tensor * lid_source_score = ggml_concat(ctx0, lctx.dsv4.cache.lid_state_score[(size_t) il], lid_state_score, 1);
ggml_tensor * lid_source_kv = dsv4_concat_named(ctx0, lctx.dsv4.cache.lid_state_kv[(size_t) il], lid_state_kv, 1, "dsv4_lid_source_kv", il);
ggml_tensor * lid_source_score = dsv4_concat_named(ctx0, lctx.dsv4.cache.lid_state_score[(size_t) il], lid_state_score, 1, "dsv4_lid_source_score", il);
ggml_tensor * lid_comp = build_overlap_compressed_kv_from_state(
ctx0, *this,
lid_source_kv, lid_source_score,
@@ -709,8 +850,10 @@ ggml_cgraph * llm_build_context::build_deepseek4() {
hparams.indexer_head_size,
il,
"lid_state_compress");
lid_comp = ggml_hadamard(ctx0, lid_comp, (int) hparams.indexer_head_size);
cb(lid_comp, "lid_state_compress_rot", il);
if (lctx.dsv4.inputs.lid.k_rot != nullptr) {
lid_comp = ggml_mul_mat(ctx0, lctx.dsv4.inputs.lid.k_rot, lid_comp);
cb(lid_comp, "lid_state_compress_rot", il);
}
ggml_tensor * lid_comp_2d = ggml_reshape_2d(ctx0, lid_comp, hparams.indexer_head_size, lctx.dsv4.inputs.lid.state_write_idxs->ne[0]);
ggml_tensor * lid_k_full = dsv4_cache_view_2d(ctx0, lctx.dsv4.cache.lid_k[(size_t) il], hparams.indexer_head_size, lctx.dsv4.cache.lid_k[(size_t) il]->ne[1]);
ggml_tensor * lid_write = ggml_set_rows(ctx0, lid_k_full, lid_comp_2d, lctx.dsv4.inputs.lid.state_write_idxs);
@@ -730,8 +873,8 @@ ggml_cgraph * llm_build_context::build_deepseek4() {
if (ratio == llama_context::dsv4_runtime::HCA_RATIO && hca_state_kv != nullptr && hca_state_score != nullptr) {
ggml_tensor * hca_dep = nullptr;
if (lctx.dsv4.inputs.hca.state_write_idxs != nullptr && lctx.dsv4.hca_plan.state_write_idxs.size() > 0) {
ggml_tensor * hca_source_kv = ggml_concat(ctx0, lctx.dsv4.cache.hca_state_kv[(size_t) il], hca_state_kv, 1);
ggml_tensor * hca_source_score = ggml_concat(ctx0, lctx.dsv4.cache.hca_state_score[(size_t) il], hca_state_score, 1);
ggml_tensor * hca_source_kv = dsv4_concat_named(ctx0, lctx.dsv4.cache.hca_state_kv[(size_t) il], hca_state_kv, 1, "dsv4_hca_source_kv", il);
ggml_tensor * hca_source_score = dsv4_concat_named(ctx0, lctx.dsv4.cache.hca_state_score[(size_t) il], hca_state_score, 1, "dsv4_hca_source_score", il);
ggml_tensor * hca_comp = build_hca_compressed_kv_from_state(
ctx0, *this,
hca_source_kv, hca_source_score,
@@ -765,9 +908,12 @@ ggml_cgraph * llm_build_context::build_deepseek4() {
ggml_row_size(kv_self.k_l[il]->type, n_embd_head),
ggml_row_size(kv_self.k_l[il]->type, n_embd_head) * hparams.n_head_kv(il),
0);
if (2*il < (int64_t) lctx.cache_copies.size()) {
raw_k = dsv4_with_zero_dep(ctx0, raw_k, lctx.cache_copies[2*il + 0].cpy);
}
cb(raw_k, "raw_k", il);
ggml_tensor * raw_mask = dsv4_build_raw_mask_view(ctx0, KQ_mask, n_kv);
ggml_tensor * raw_mask = dsv4_build_raw_mask_view(ctx0, KQ_mask, n_kv, n_tokens);
ggml_tensor * attn = nullptr;
const char * attn_path = "raw";
@@ -783,14 +929,15 @@ ggml_cgraph * llm_build_context::build_deepseek4() {
lctx.dsv4.csa_plan.n_kv);
ggml_tensor * top_k = dsv4_build_lid_top_k(ctx0, *this, qr, cur, inp_pos, il);
ggml_tensor * csa_mask = build_top_k_mask(ctx0,
dsv4_build_raw_mask_view(ctx0, lctx.dsv4.inputs.csa.kq_mask, lctx.dsv4.csa_plan.n_kv),
dsv4_build_raw_mask_view(ctx0, lctx.dsv4.inputs.csa.kq_mask, lctx.dsv4.csa_plan.n_kv, n_tokens),
top_k);
ggml_tensor * k_all = ggml_concat(ctx0, raw_k, csa_k, 2);
ggml_tensor * k_all = dsv4_concat_named(ctx0, raw_k, csa_k, 2, "dsv4_raw_plus_csa_k", il);
ggml_tensor * kq_mask = ggml_concat(ctx0, raw_mask, csa_mask, 0);
ggml_tensor * kq_b = dsv4_build_kq_zero_bias(ctx0, cparams, kq_mask, q->ne[1]);
cb(csa_k, "csa_k", il);
cb(k_all, "csa_k_all", il);
cb(kq_mask, "csa_kq_mask", il);
attn = dsv4_build_attn(ctx0, hparams, cparams, q, k_all, k_all, kq_mask, model.layers[il].attn_sinks, kq_scale);
attn = dsv4_build_attn(ctx0, hparams, cparams, q, k_all, k_all, kq_b, kq_mask, model.layers[il].attn_sinks, kq_scale);
cb(attn, "attn_csa", il);
} else if (ratio == llama_context::dsv4_runtime::HCA_RATIO &&
lctx.dsv4.inputs.hca.kq_mask != nullptr &&
@@ -800,16 +947,18 @@ ggml_cgraph * llm_build_context::build_deepseek4() {
lctx.dsv4.cache.hca_k[(size_t) il],
n_embd_head,
lctx.dsv4.hca_plan.n_kv);
ggml_tensor * hca_mask = dsv4_build_raw_mask_view(ctx0, lctx.dsv4.inputs.hca.kq_mask, lctx.dsv4.hca_plan.n_kv);
ggml_tensor * k_all = ggml_concat(ctx0, raw_k, hca_k, 2);
ggml_tensor * hca_mask = dsv4_build_raw_mask_view(ctx0, lctx.dsv4.inputs.hca.kq_mask, lctx.dsv4.hca_plan.n_kv, n_tokens);
ggml_tensor * k_all = dsv4_concat_named(ctx0, raw_k, hca_k, 2, "dsv4_raw_plus_hca_k", il);
ggml_tensor * kq_mask = ggml_concat(ctx0, raw_mask, hca_mask, 0);
ggml_tensor * kq_b = dsv4_build_kq_zero_bias(ctx0, cparams, kq_mask, q->ne[1]);
cb(hca_k, "hca_k", il);
cb(k_all, "hca_k_all", il);
cb(kq_mask, "hca_kq_mask", il);
attn = dsv4_build_attn(ctx0, hparams, cparams, q, k_all, k_all, kq_mask, model.layers[il].attn_sinks, kq_scale);
attn = dsv4_build_attn(ctx0, hparams, cparams, q, k_all, k_all, kq_b, kq_mask, model.layers[il].attn_sinks, kq_scale);
cb(attn, "attn_hca", il);
} else {
attn = dsv4_build_attn(ctx0, hparams, cparams, q, raw_k, raw_k, raw_mask, model.layers[il].attn_sinks, kq_scale);
ggml_tensor * kq_b = dsv4_build_kq_zero_bias(ctx0, cparams, raw_mask, q->ne[1]);
attn = dsv4_build_attn(ctx0, hparams, cparams, q, raw_k, raw_k, kq_b, raw_mask, model.layers[il].attn_sinks, kq_scale);
cb(attn, "attn_raw", il);
}
@@ -858,7 +1007,7 @@ ggml_cgraph * llm_build_context::build_deepseek4() {
ggml_row_size(attn->type, n_embd_head),
ggml_row_size(attn->type, n_embd_head) * n_head,
ggml_row_size(attn->type, n_embd_head_nope));
attn_pe = ggml_rope_back(ctx0, attn_pe, inp_pos, nullptr, n_embd_head_rope, rope_type, n_ctx_orig_l,
attn_pe = ggml_rope_ext_back(ctx0, attn_pe, inp_pos, nullptr, n_embd_head_rope, rope_type, n_ctx_orig_l,
freq_base_l, freq_scale_l, ext_factor_l, attn_factor_l, beta_fast_l, beta_slow_l);
cb(attn_pe, "attn_derope", il);
attn = ggml_concat(ctx0, attn_nope, attn_pe, 0);
@@ -906,17 +1055,28 @@ ggml_cgraph * llm_build_context::build_deepseek4() {
nullptr,
LLM_FFN_SILU, LLM_FFN_PAR, cb, il);
} else {
ggml_tensor * selected_experts = nullptr;
ggml_tensor * exp_probs_b = model.layers[il].ffn_exp_probs_b;
if ((uint32_t) il < hparams.dsv4_hash_layer_count) {
selected_experts = ggml_get_rows(ctx0, model.layers[il].ffn_gate_tid2eid, lctx.inp_tokens);
exp_probs_b = nullptr;
}
ggml_tensor * moe_out = llm_build_moe_ffn(ctx0, lctx, cur,
model.layers[il].ffn_gate_inp,
nullptr,
model.layers[il].ffn_up_exps,
nullptr,
model.layers[il].ffn_gate_exps,
nullptr,
model.layers[il].ffn_down_exps,
model.layers[il].ffn_exp_probs_b,
nullptr,
exp_probs_b,
n_expert, n_expert_used,
LLM_FFN_SILU, hparams.expert_weights_norm,
true, hparams.expert_weights_scale,
(enum llm_expert_gating_func_type) hparams.expert_gating_func,
cb, il, gf, false, model.layers[il].ffn_up_gate_exps);
cb, il, gf, false, model.layers[il].ffn_up_gate_exps, nullptr, nullptr, nullptr, selected_experts);
cb(moe_out, "ffn_moe_out", il);
ggml_tensor * ffn_shexp = llm_build_ffn(ctx0, lctx, nullptr, cur,
src/llama-arch.h
+8
View File
@@ -377,6 +377,14 @@ enum llm_tensor {
LLM_TENSOR_ATTN_KV_LATENT,
LLM_TENSOR_ATTN_OUT_A,
LLM_TENSOR_ATTN_OUT_B,
LLM_TENSOR_ATTN_COMP_KV,
LLM_TENSOR_ATTN_COMP_GATE,
LLM_TENSOR_ATTN_COMP_APE,
LLM_TENSOR_ATTN_COMP_NORM,
LLM_TENSOR_INDEXER_COMP_KV,
LLM_TENSOR_INDEXER_COMP_GATE,
LLM_TENSOR_INDEXER_COMP_APE,
LLM_TENSOR_INDEXER_COMP_NORM,
LLM_TENSOR_FFN_GATE_TID2EID,
LLM_TENSOR_HC_HEAD_BASE,
LLM_TENSOR_HC_HEAD_FN,
src/llama-build-context.cpp
+18 -9
View File
@@ -1050,7 +1050,8 @@ ggml_tensor * llm_build_context::llm_build_moe_ffn(
llm_expert_gating_func_type gating_op,
const llm_build_cb & cb, int il, ggml_cgraph * graph, bool add_input,
ggml_tensor * up_gate_exps, ggml_tensor * up_gate_exps_b,
ggml_tensor * input_logits, ggml_tensor * down_exps_s) {
ggml_tensor * input_logits, ggml_tensor * down_exps_s,
ggml_tensor * selected_experts) {
GGML_ASSERT(gate_inp || input_logits);
@@ -1068,6 +1069,9 @@ llm_expert_gating_func_type gating_op,
cb(logits, "ffn_moe_logits_biased", il);
}
if (gating_op == LLM_EXPERT_GATING_FUNC_TYPE_SQRT_SOFTPLUS) {
ggml_mul_mat_set_prec(logits, GGML_PREC_F32);
}
//ggml_tensor * probs = ggml_soft_max(ctx, logits); // [n_expert, n_tokens]
ggml_tensor * probs = nullptr;
@@ -1084,6 +1088,10 @@ llm_expert_gating_func_type gating_op,
{
probs = logits; // [n_expert, n_tokens]
} break;
case LLM_EXPERT_GATING_FUNC_TYPE_SQRT_SOFTPLUS:
{
probs = ggml_sqrt(ctx, ggml_softplus(ctx, logits)); // [n_expert, n_tokens]
} break;
default:
GGML_ABORT("fatal error");
}
@@ -1104,14 +1112,15 @@ llm_expert_gating_func_type gating_op,
}
// select experts
ggml_tensor * selected_experts;
if (lctx.cparams.grouped_expert_routing && lctx.model.arch == LLM_ARCH_BAILINGMOE2 && n_tokens > 0) {
auto& hparams = lctx.model.hparams;
selected_experts = ggml_grouped_topk(ctx, selection_probs, hparams.n_expert_groups, hparams.n_group_used, 2, n_expert_used);
} else {
//selected_experts = ggml_top_k_thresh(ctx, selection_probs, n_expert_used,
// lctx.cparams.min_experts, lctx.cparams.thresh_experts); // [n_expert_used, n_tokens]
selected_experts = ggml_top_k(ctx, selection_probs, n_expert_used); // [n_expert_used, n_tokens]
if (selected_experts == nullptr) {
if (lctx.cparams.grouped_expert_routing && lctx.model.arch == LLM_ARCH_BAILINGMOE2 && n_tokens > 0) {
auto& hparams = lctx.model.hparams;
selected_experts = ggml_grouped_topk(ctx, selection_probs, hparams.n_expert_groups, hparams.n_group_used, 2, n_expert_used);
} else {
//selected_experts = ggml_top_k_thresh(ctx, selection_probs, n_expert_used,
// lctx.cparams.min_experts, lctx.cparams.thresh_experts); // [n_expert_used, n_tokens]
selected_experts = ggml_top_k(ctx, selection_probs, n_expert_used); // [n_expert_used, n_tokens]
}
}
cb(selected_experts, "ffn_moe_topk", il);
ggml_tensor * weights = ggml_get_rows(ctx,
src/llama-build-context.h
+3 -2
View File
@@ -411,7 +411,8 @@ struct llm_build_context {
llm_expert_gating_func_type gating_op,
const llm_build_cb & cb, int il, ggml_cgraph * graph = nullptr, bool add_input = false,
ggml_tensor * up_gate_exps = nullptr, ggml_tensor * up_gate_exps_b = nullptr,
ggml_tensor * input_logits = nullptr, ggml_tensor * down_exps_s = nullptr);
ggml_tensor * input_logits = nullptr, ggml_tensor * down_exps_s = nullptr,
ggml_tensor * selected_experts = nullptr);
static ggml_tensor * llm_build_moe_ffn(ggml_context * ctx, llama_context & lctx,
ggml_tensor * cur,
@@ -439,7 +440,7 @@ llm_expert_gating_func_type gating_op,
n_expert, n_expert_used,
type_op, norm_w, scale_w, w_scale,
gating_op, cb, il, graph, add_input, up_gate_exps, up_gate_exps_b,
input_logits, down_exps_s);
input_logits, down_exps_s, nullptr);
}
static ggml_tensor * llm_build_std_moe_ffn(ggml_context * ctx, llama_context & lctx,
src/llama-context.h
+2
View File
@@ -383,6 +383,7 @@ struct llama_context {
struct ggml_tensor * state_write_idxs = nullptr;
struct ggml_tensor * state_write_pos = nullptr;
struct ggml_tensor * kq_mask = nullptr;
struct ggml_tensor * k_rot = nullptr;
};
struct storage {
@@ -417,6 +418,7 @@ struct llama_context {
std::vector<float> csa_mask_data;
std::vector<float> hca_mask_data;
std::vector<float> lid_mask_data;
std::vector<float> lid_k_rot_data;
};
dsv4_runtime dsv4;
src/llama-dsv4.cpp
+57
View File
@@ -206,6 +206,55 @@ static uint32_t dsv4_comp_size(uint32_t kv_size, uint32_t ratio) {
return std::max<uint32_t>(1, (kv_size + ratio - 1)/ratio);
}
static int64_t dsv4_k_rot_size(int64_t n_embd) {
if (n_embd < 64 || n_embd % 64 != 0) {
return 0;
}
int64_t n_rot = 64;
do {
n_rot *= 2;
} while (n_embd % n_rot == 0);
return n_rot / 2;
}
static void dsv4_build_hadamard_matrix(std::vector<float> & storage, int64_t n_rot) {
if (n_rot <= 0) {
storage.clear();
return;
}
const size_t n_elem = (size_t) n_rot*(size_t) n_rot;
if (storage.size() == n_elem) {
return;
}
storage.assign(n_elem, 0.0f);
storage[0] = 1.0f;
for (int64_t size = 1; size < n_rot; size *= 2) {
for (int64_t row = 0; row < size; ++row) {
const float * src = storage.data() + row*n_rot;
float * dst_top = storage.data() + row*n_rot + size;
float * dst_bottom = storage.data() + (row + size)*n_rot;
float * dst_bottom_right = dst_bottom + size;
for (int64_t col = 0; col < size; ++col) {
const float v = src[col];
dst_top[col] = v;
dst_bottom[col] = v;
dst_bottom_right[col] = -v;
}
}
}
const float scale = 1.0f / std::sqrt((float) n_rot);
for (float & v : storage) {
v *= scale;
}
}
static void dsv4_batch_shape(
const llama_batch & batch,
uint32_t & n_seqs,
@@ -652,5 +701,13 @@ bool llama_prepare_dsv4_graph_inputs(llama_context & lctx, const llama_batch & b
set_comp(lctx.dsv4.inputs.hca, lctx.dsv4.hca_plan, lctx.dsv4.hca_mask_data);
set_comp(lctx.dsv4.inputs.lid, lctx.dsv4.lid_plan, lctx.dsv4.lid_mask_data);
if (lctx.dsv4.inputs.lid.k_rot != nullptr && lctx.dsv4.inputs.lid.k_rot->buffer != nullptr) {
const int64_t n_rot = dsv4_k_rot_size(lctx.model.hparams.indexer_head_size);
dsv4_build_hadamard_matrix(lctx.dsv4.lid_k_rot_data, n_rot);
if (!lctx.dsv4.lid_k_rot_data.empty()) {
ggml_backend_tensor_set(lctx.dsv4.inputs.lid.k_rot, lctx.dsv4.lid_k_rot_data.data(), 0, lctx.dsv4.lid_k_rot_data.size()*sizeof(float));
}
}
return true;
}
src/llama-hparams.cpp
+32 -19
View File
@@ -34,6 +34,7 @@ static inline const char * llm_expert_gating_func_name(llm_expert_gating_func_ty
case LLM_EXPERT_GATING_FUNC_SOFTMAX: return "softmax";
case LLM_EXPERT_GATING_FUNC_SIGMOID: return "sigmoid";
case LLM_EXPERT_GATING_FUNC_TYPE_SOFTMAX_WEIGHT: return "weight";
case LLM_EXPERT_GATING_FUNC_TYPE_SQRT_SOFTPLUS: return "sqrtsoftplus";
default: return "none";
}
}
@@ -1598,26 +1599,33 @@ void llm_load_hparams(
ml.get_key(LLM_KV_ATTENTION_Q_LORA_RANK, hparams.n_lora_q);
ml.get_key(LLM_KV_ATTENTION_KV_LORA_RANK, hparams.n_lora_kv, false);
if (model.arch == LLM_ARCH_DEEPSEEK4 && hparams.n_lora_kv == 0) {
auto * kv_a = ml.get_tensor_meta("blk.0.attn_kv_latent.weight");
bool subtract_rope = false;
if (kv_a == nullptr) {
kv_a = ml.require_tensor_meta("blk.0.attn_kv_a_mqa.weight");
subtract_rope = true;
}
const int64_t kv_a_inner = kv_a->ne[0] == hparams.n_embd ? kv_a->ne[1] : kv_a->ne[0];
if (subtract_rope) {
if (kv_a_inner <= hparams.n_rot) {
throw std::runtime_error(format(
"%s: unable to infer %s from blk.0.attn_kv_a_mqa.weight shape [%lld, %lld]",
__func__,
ml.llm_kv(LLM_KV_ATTENTION_KV_LORA_RANK).c_str(),
(long long) kv_a->ne[0],
(long long) kv_a->ne[1]));
}
hparams.n_lora_kv = (uint32_t) (kv_a_inner - hparams.n_rot);
if (auto * kv_norm = ml.get_tensor_meta("blk.0.attn_kv_a_norm.weight")) {
hparams.n_lora_kv = (uint32_t) kv_norm->ne[0];
} else if (auto * kv = ml.get_tensor_meta("blk.0.attn_kv.weight")) {
const int64_t kv_inner = kv->ne[0] == hparams.n_embd ? kv->ne[1] : kv->ne[0];
hparams.n_lora_kv = (uint32_t) kv_inner;
} else {
hparams.n_lora_kv = (uint32_t) kv_a_inner;
auto * kv_a = ml.get_tensor_meta("blk.0.attn_kv_latent.weight");
bool subtract_rope = false;
if (kv_a == nullptr) {
kv_a = ml.require_tensor_meta("blk.0.attn_kv_a_mqa.weight");
subtract_rope = true;
}
const int64_t kv_a_inner = kv_a->ne[0] == hparams.n_embd ? kv_a->ne[1] : kv_a->ne[0];
if (subtract_rope) {
if (kv_a_inner <= hparams.n_rot) {
throw std::runtime_error(format(
"%s: unable to infer %s from blk.0.attn_kv_a_mqa.weight shape [%lld, %lld]",
__func__,
ml.llm_kv(LLM_KV_ATTENTION_KV_LORA_RANK).c_str(),
(long long) kv_a->ne[0],
(long long) kv_a->ne[1]));
}
hparams.n_lora_kv = (uint32_t) (kv_a_inner - hparams.n_rot);
} else {
hparams.n_lora_kv = (uint32_t) kv_a_inner;
}
}
}
//ml.get_key(LLM_KV_ATTENTION_KEY_LENGTH_MLA, hparams.n_embd_head_k_mla_impl, false);
@@ -1716,6 +1724,11 @@ void llm_load_hparams(
hparams.expert_gating_func = LLM_EXPERT_GATING_FUNC_SIGMOID;
}
if (model.arch == LLM_ARCH_DEEPSEEK4 &&
hparams.expert_gating_func != LLM_EXPERT_GATING_FUNC_TYPE_SQRT_SOFTPLUS) {
throw std::runtime_error("DeepSeek-V4 loader currently expects sqrtsoftplus MoE scoring");
}
// NextN/MTP parameters
ml.get_key(LLM_KV_NEXTN_PREDICT_LAYERS, hparams.nextn_predict_layers, false);
src/llama-hparams.h
+1
View File
@@ -13,6 +13,7 @@ enum llm_expert_gating_func_type {
LLM_EXPERT_GATING_FUNC_SOFTMAX = 1,
LLM_EXPERT_GATING_FUNC_SIGMOID = 2,
LLM_EXPERT_GATING_FUNC_TYPE_SOFTMAX_WEIGHT = 3,
LLM_EXPERT_GATING_FUNC_TYPE_SQRT_SOFTPLUS = 4,
};
struct llama_hparams {
src/llama-load-tensors.cpp
+21 -16
View File
@@ -2771,16 +2771,16 @@ bool create_tensors_helper::create_deepseek4_tensors(const LLM_TN & tn) {
model.output_norm = create_tensor_from_meta(ctx_output, "output_norm.weight");
model.output = create_tensor_from_meta(ctx_output, "output.weight");
model.hc_head_base = create_tensor_from_meta(ctx_output, pick_tensor_name({"hc_head_base", "output_hc_base"}), llama_model_loader::TENSOR_NOT_REQUIRED);
model.hc_head_fn = create_tensor_from_meta(ctx_output, pick_tensor_name({"hc_head_fn", "output_hc_fn"}), llama_model_loader::TENSOR_NOT_REQUIRED);
model.hc_head_scale = create_tensor_from_meta(ctx_output, pick_tensor_name({"hc_head_scale", "output_hc_scale"}), llama_model_loader::TENSOR_NOT_REQUIRED);
model.hc_head_base = create_tensor_from_meta(ctx_output, pick_tensor_name({"hc_head_base.weight", "output_hc_base.weight"}), llama_model_loader::TENSOR_NOT_REQUIRED);
model.hc_head_fn = create_tensor_from_meta(ctx_output, pick_tensor_name({"hc_head_fn.weight", "output_hc_fn.weight"}), llama_model_loader::TENSOR_NOT_REQUIRED);
model.hc_head_scale = create_tensor_from_meta(ctx_output, pick_tensor_name({"hc_head_scale.weight", "output_hc_scale.weight"}), llama_model_loader::TENSOR_NOT_REQUIRED);
for (int i = 0; i < n_layer; ++i) {
ggml_context * ctx_split = ctx_for_layer_split(i);
auto & layer = model.layers[i];
layer.attn_norm = create_tensor_from_meta(ctx_split, layer_weight_name(i, "attn_norm"));
layer.attn_sinks = create_tensor_from_meta(ctx_split, format("blk.%d.attn_sinks", i), llama_model_loader::TENSOR_NOT_REQUIRED);
layer.attn_sinks = create_tensor_from_meta(ctx_split, format("blk.%d.attn_sinks.weight", i), llama_model_loader::TENSOR_NOT_REQUIRED);
layer.wq_a = create_tensor_from_meta(ctx_split, layer_weight_name(i, "attn_q_a"));
layer.attn_q_a_norm = create_tensor_from_meta(ctx_split, layer_weight_name(i, "attn_q_a_norm"));
layer.wq_b = create_tensor_from_meta(ctx_split, layer_weight_name(i, "attn_q_b"));
@@ -2797,12 +2797,12 @@ bool create_tensors_helper::create_deepseek4_tensors(const LLM_TN & tn) {
layer.wo_b = create_tensor_from_meta(ctx_split, layer_weight_name(i, "attn_output_b"));
layer.wo = layer.wo_b;
layer.hc_attn_base = create_tensor_from_meta(ctx_split, format("blk.%d.hc_attn_base", i), llama_model_loader::TENSOR_NOT_REQUIRED);
layer.hc_attn_fn = create_tensor_from_meta(ctx_split, format("blk.%d.hc_attn_fn", i), llama_model_loader::TENSOR_NOT_REQUIRED);
layer.hc_attn_scale = create_tensor_from_meta(ctx_split, format("blk.%d.hc_attn_scale", i), llama_model_loader::TENSOR_NOT_REQUIRED);
layer.hc_ffn_base = create_tensor_from_meta(ctx_split, format("blk.%d.hc_ffn_base", i), llama_model_loader::TENSOR_NOT_REQUIRED);
layer.hc_ffn_fn = create_tensor_from_meta(ctx_split, format("blk.%d.hc_ffn_fn", i), llama_model_loader::TENSOR_NOT_REQUIRED);
layer.hc_ffn_scale = create_tensor_from_meta(ctx_split, format("blk.%d.hc_ffn_scale", i), llama_model_loader::TENSOR_NOT_REQUIRED);
layer.hc_attn_base = create_tensor_from_meta(ctx_split, format("blk.%d.hc_attn_base.weight", i), llama_model_loader::TENSOR_NOT_REQUIRED);
layer.hc_attn_fn = create_tensor_from_meta(ctx_split, format("blk.%d.hc_attn_fn.weight", i), llama_model_loader::TENSOR_NOT_REQUIRED);
layer.hc_attn_scale = create_tensor_from_meta(ctx_split, format("blk.%d.hc_attn_scale.weight", i), llama_model_loader::TENSOR_NOT_REQUIRED);
layer.hc_ffn_base = create_tensor_from_meta(ctx_split, format("blk.%d.hc_ffn_base.weight", i), llama_model_loader::TENSOR_NOT_REQUIRED);
layer.hc_ffn_fn = create_tensor_from_meta(ctx_split, format("blk.%d.hc_ffn_fn.weight", i), llama_model_loader::TENSOR_NOT_REQUIRED);
layer.hc_ffn_scale = create_tensor_from_meta(ctx_split, format("blk.%d.hc_ffn_scale.weight", i), llama_model_loader::TENSOR_NOT_REQUIRED);
layer.attn_comp_wkv = create_tensor_from_meta(ctx_split, pick_tensor_name({
layer_weight_name(i, "attn_compress_kv"),
@@ -2813,14 +2813,16 @@ bool create_tensors_helper::create_deepseek4_tensors(const LLM_TN & tn) {
layer_weight_name(i, "attn_compressor_gate"),
}), llama_model_loader::TENSOR_NOT_REQUIRED);
layer.attn_comp_ape = create_tensor_from_meta(ctx_split, pick_tensor_name({
format("blk.%d.attn_compress_ape", i),
format("blk.%d.attn_compressor_ape", i),
format("blk.%d.attn_compress_ape.weight", i),
format("blk.%d.attn_compressor_ape.weight", i),
}), llama_model_loader::TENSOR_NOT_REQUIRED);
layer.attn_comp_norm = create_tensor_from_meta(ctx_split, pick_tensor_name({
layer_weight_name(i, "attn_compress_norm"),
layer_weight_name(i, "attn_compressor_norm"),
}), llama_model_loader::TENSOR_NOT_REQUIRED);
layer.indexer_k_norm = create_tensor_from_meta(ctx_split, layer_weight_name(i, "indexer.k_norm"), llama_model_loader::TENSOR_NOT_REQUIRED);
layer.indexer_attn_k = create_tensor_from_meta(ctx_split, layer_weight_name(i, "indexer.attn_k"), llama_model_loader::TENSOR_NOT_REQUIRED);
layer.indexer_proj = create_tensor_from_meta(ctx_split, layer_weight_name(i, "indexer.proj"), llama_model_loader::TENSOR_NOT_REQUIRED);
layer.indexer_attn_q_b = create_tensor_from_meta(ctx_split, layer_weight_name(i, "indexer.attn_q_b"), llama_model_loader::TENSOR_NOT_REQUIRED);
layer.indexer_comp_wkv = create_tensor_from_meta(ctx_split, pick_tensor_name({
@@ -2832,8 +2834,8 @@ bool create_tensors_helper::create_deepseek4_tensors(const LLM_TN & tn) {
layer_weight_name(i, "indexer_compressor_gate"),
}), llama_model_loader::TENSOR_NOT_REQUIRED);
layer.indexer_comp_ape = create_tensor_from_meta(ctx_split, pick_tensor_name({
format("blk.%d.indexer.compress_ape", i),
format("blk.%d.indexer_compressor_ape", i),
format("blk.%d.indexer.compress_ape.weight", i),
format("blk.%d.indexer_compressor_ape.weight", i),
}), llama_model_loader::TENSOR_NOT_REQUIRED);
layer.indexer_comp_norm = create_tensor_from_meta(ctx_split, pick_tensor_name({
layer_weight_name(i, "indexer.compress_norm"),
@@ -2849,9 +2851,12 @@ bool create_tensors_helper::create_deepseek4_tensors(const LLM_TN & tn) {
layer.ffn_down_shexp = create_tensor_from_meta(ctx_split, layer_weight_name(i, "ffn_down_shexp"));
layer.ffn_up_shexp = create_tensor_from_meta(ctx_split, layer_weight_name(i, "ffn_up_shexp"));
layer.ffn_gate_tid2eid = create_tensor_from_meta(ctx_split, format("blk.%d.ffn_gate_tid2eid", i), llama_model_loader::TENSOR_NOT_REQUIRED);
layer.ffn_gate_tid2eid = create_tensor_from_meta(ctx_split, format("blk.%d.ffn_gate_tid2eid.weight", i), llama_model_loader::TENSOR_NOT_REQUIRED);
if (layer.ffn_gate_tid2eid == nullptr) {
layer.ffn_exp_probs_b = create_tensor_from_meta(ctx_split, format("blk.%d.exp_probs_b", i), llama_model_loader::TENSOR_NOT_REQUIRED);
layer.ffn_exp_probs_b = create_tensor_from_meta(ctx_split, pick_tensor_name({
format("blk.%d.exp_probs_b.bias", i),
format("blk.%d.exp_probs_b.weight", i),
}), llama_model_loader::TENSOR_NOT_REQUIRED);
}
}
src/llama-model.cpp
+13 -9
View File
@@ -1071,19 +1071,19 @@ static const std::map<llm_arch, std::map<llm_tensor, std::string>> LLM_TENSOR_NA
{ LLM_TENSOR_OUTPUT_NORM, "output_norm" },
{ LLM_TENSOR_OUTPUT, "output" },
{ LLM_TENSOR_ATTN_NORM, "blk.%d.attn_norm" },
{ LLM_TENSOR_ATTN_SINKS, "blk.%d.attn_sinks" },
{ LLM_TENSOR_ATTN_Q_A_NORM, "blk.%d.attn_q_a_norm" },
{ LLM_TENSOR_ATTN_KV_A_NORM, "blk.%d.attn_kv_a_norm" },
{ LLM_TENSOR_ATTN_Q, "blk.%d.attn_q" },
{ LLM_TENSOR_ATTN_Q_A, "blk.%d.attn_q_a" },
{ LLM_TENSOR_ATTN_Q_B, "blk.%d.attn_q_b" },
{ LLM_TENSOR_ATTN_KV_LATENT, "blk.%d.attn_kv_latent" },
{ LLM_TENSOR_ATTN_KV_A_MQA, "blk.%d.attn_kv_a_mqa" },
{ LLM_TENSOR_ATTN_KQ_A_MQA, "blk.%d.attn_kq_a_mqa" },
{ LLM_TENSOR_ATTN_KV_B, "blk.%d.attn_kv_b" },
{ LLM_TENSOR_ATTN_K_B, "blk.%d.attn_k_b" },
{ LLM_TENSOR_ATTN_V_B, "blk.%d.attn_v_b" },
{ LLM_TENSOR_ATTN_KV_LATENT, "blk.%d.attn_kv" },
{ LLM_TENSOR_ATTN_OUT_A, "blk.%d.attn_output_a" },
{ LLM_TENSOR_ATTN_OUT_B, "blk.%d.attn_output_b" },
{ LLM_TENSOR_ATTN_COMP_KV, "blk.%d.attn_compressor_kv" },
{ LLM_TENSOR_ATTN_COMP_GATE, "blk.%d.attn_compressor_gate" },
{ LLM_TENSOR_ATTN_COMP_APE, "blk.%d.attn_compressor_ape" },
{ LLM_TENSOR_ATTN_COMP_NORM, "blk.%d.attn_compressor_norm" },
{ LLM_TENSOR_ATTN_OUT, "blk.%d.attn_output" },
{ LLM_TENSOR_FFN_NORM, "blk.%d.ffn_norm" },
{ LLM_TENSOR_FFN_GATE, "blk.%d.ffn_gate" },
@@ -1104,9 +1104,13 @@ static const std::map<llm_arch, std::map<llm_tensor, std::string>> LLM_TENSOR_NA
{ LLM_TENSOR_INDEXER_PROJ, "blk.%d.indexer.proj" },
{ LLM_TENSOR_INDEXER_ATTN_K, "blk.%d.indexer.attn_k" },
{ LLM_TENSOR_INDEXER_ATTN_Q_B, "blk.%d.indexer.attn_q_b" },
{ LLM_TENSOR_HC_HEAD_BASE, "hc_head_base" },
{ LLM_TENSOR_HC_HEAD_FN, "hc_head_fn" },
{ LLM_TENSOR_HC_HEAD_SCALE, "hc_head_scale" },
{ LLM_TENSOR_INDEXER_COMP_KV, "blk.%d.indexer_compressor_kv" },
{ LLM_TENSOR_INDEXER_COMP_GATE, "blk.%d.indexer_compressor_gate" },
{ LLM_TENSOR_INDEXER_COMP_APE, "blk.%d.indexer_compressor_ape" },
{ LLM_TENSOR_INDEXER_COMP_NORM, "blk.%d.indexer_compressor_norm" },
{ LLM_TENSOR_HC_HEAD_BASE, "output_hc_base" },
{ LLM_TENSOR_HC_HEAD_FN, "output_hc_fn" },
{ LLM_TENSOR_HC_HEAD_SCALE, "output_hc_scale" },
{ LLM_TENSOR_HC_ATTN_BASE, "blk.%d.hc_attn_base" },
{ LLM_TENSOR_HC_ATTN_FN, "blk.%d.hc_attn_fn" },
{ LLM_TENSOR_HC_ATTN_SCALE, "blk.%d.hc_attn_scale" },
src/llama.cpp
+163
View File
@@ -410,6 +410,7 @@ static const char * llama_expert_gating_func_name(llm_expert_gating_func_type ty
case LLM_EXPERT_GATING_FUNC_SOFTMAX: return "softmax";
case LLM_EXPERT_GATING_FUNC_SIGMOID: return "sigmoid";
case LLM_EXPERT_GATING_FUNC_TYPE_SOFTMAX_WEIGHT: return "softmax_weight";
case LLM_EXPERT_GATING_FUNC_TYPE_SQRT_SOFTPLUS: return "sqrtsoftplus";
default: return "unknown";
}
}
@@ -5058,6 +5059,137 @@ static size_t llama_output_reserve(llama_context & lctx, size_t n_outputs) {
return n_outputs_max;
}
static void llama_dsv4_trace_decode_stage(
const llama_context & lctx,
const char * stage,
const llama_batch * ubatch = nullptr,
const ggml_cgraph * gf = nullptr,
const ggml_tensor * tensor = nullptr) {
if (lctx.model.arch != LLM_ARCH_DEEPSEEK4 || !llama_dsv4_trace_enabled()) {
return;
}
std::ostringstream out;
out << "{\"event\":\"decode_stage\",\"stage\":\"" << stage << "\"";
if (ubatch != nullptr) {
out << ",\"n_tokens\":" << ubatch->n_tokens
<< ",\"all_seq_id\":" << ubatch->all_seq_id;
}
out << ",\"kv_head\":" << lctx.kv_self.head
<< ",\"kv_n\":" << lctx.kv_self.n
<< ",\"n_outputs\":" << lctx.n_outputs;
if (gf != nullptr) {
out << ",\"graph_nodes\":" << gf->n_nodes
<< ",\"graph_leafs\":" << gf->n_leafs;
}
if (tensor != nullptr) {
out << ",\"tensor\":\"" << tensor->name << "\""
<< ",\"type\":\"" << ggml_type_name(tensor->type) << "\""
<< ",\"shape\":[" << tensor->ne[0] << ',' << tensor->ne[1] << ','
<< tensor->ne[2] << ',' << tensor->ne[3] << "]";
}
out << "}";
llama_dsv4_trace_jsonl(out.str());
}
static bool llama_dsv4_trace_eval_enabled() {
const char * env = std::getenv("LLAMA_DSV4_TRACE_EVAL");
return env != nullptr && *env != '\0' &&
std::strcmp(env, "0") != 0 &&
std::strcmp(env, "false") != 0 &&
std::strcmp(env, "off") != 0;
}
static int llama_dsv4_trace_eval_callback(struct ggml_tensor * tensor, bool ask, void * user_data) {
auto * lctx = static_cast<llama_context *>(user_data);
if (lctx == nullptr || lctx->model.arch != LLM_ARCH_DEEPSEEK4 || !llama_dsv4_trace_enabled()) {
return 0;
}
if (ask) {
return 2;
}
std::ostringstream out;
out << "{\"event\":\"eval_node\",\"name\":\"" << tensor->name
<< "\",\"op\":\"" << ggml_op_name(tensor->op)
<< "\",\"type\":\"" << ggml_type_name(tensor->type)
<< "\",\"shape\":[" << tensor->ne[0] << ',' << tensor->ne[1] << ','
<< tensor->ne[2] << ',' << tensor->ne[3] << "]}";
llama_dsv4_trace_jsonl(out.str());
return 2;
}
static void llama_dsv4_trace_logits_summary(
ggml_backend_t backend,
const float * logits,
int32_t n_vocab,
int32_t n_outputs) {
if (!backend || logits == nullptr || n_vocab <= 0 || n_outputs <= 0 || !llama_dsv4_trace_enabled()) {
return;
}
ggml_backend_synchronize(backend);
const float * row = logits + (size_t) (n_outputs - 1) * (size_t) n_vocab;
int32_t finite_count = 0;
int32_t nan_count = 0;
int32_t inf_count = 0;
float min_val = 0.0f;
float max_val = 0.0f;
bool have_finite = false;
std::array<int32_t, 5> top_ids = { -1, -1, -1, -1, -1 };
std::array<float, 5> top_vals = { -INFINITY, -INFINITY, -INFINITY, -INFINITY, -INFINITY };
for (int32_t i = 0; i < n_vocab; ++i) {
const float v = row[i];
if (std::isnan(v)) {
++nan_count;
continue;
}
if (!std::isfinite(v)) {
++inf_count;
continue;
}
++finite_count;
if (!have_finite) {
min_val = max_val = v;
have_finite = true;
} else {
min_val = std::min(min_val, v);
max_val = std::max(max_val, v);
}
for (int k = 0; k < 5; ++k) {
if (v > top_vals[(size_t) k]) {
for (int s = 4; s > k; --s) {
top_vals[(size_t) s] = top_vals[(size_t) (s - 1)];
top_ids[(size_t) s] = top_ids[(size_t) (s - 1)];
}
top_vals[(size_t) k] = v;
top_ids[(size_t) k] = i;
break;
}
}
}
std::ostringstream out;
out << "{\"event\":\"logits_summary\",\"finite_count\":" << finite_count
<< ",\"nan_count\":" << nan_count
<< ",\"inf_count\":" << inf_count
<< ",\"min\":" << (have_finite ? min_val : 0.0f)
<< ",\"max\":" << (have_finite ? max_val : 0.0f)
<< ",\"top_ids\":[" << top_ids[0] << ',' << top_ids[1] << ',' << top_ids[2] << ',' << top_ids[3] << ',' << top_ids[4] << "]"
<< ",\"top_vals\":[" << top_vals[0] << ',' << top_vals[1] << ',' << top_vals[2] << ',' << top_vals[3] << ',' << top_vals[4] << "]}";
llama_dsv4_trace_jsonl(out.str());
}
static void llama_graph_compute(
llama_context & lctx,
@@ -5443,6 +5575,7 @@ static int llama_decode_internal(
}
gf = llm_build_context::llama_build_graph(lctx, u_batch, false);
llama_dsv4_trace_decode_stage(lctx, "graph_built", &u_batch, gf);
#if IK_PRINT_TIMING
tim2 = ggml_time_us();
printf("build_graph(...): %d us\n", int(tim2-tim1));
@@ -5452,6 +5585,7 @@ static int llama_decode_internal(
tim1 = ggml_time_us();
#endif
ggml_backend_sched_alloc_graph(lctx.sched, gf);
llama_dsv4_trace_decode_stage(lctx, "graph_allocated", &u_batch, gf);
#if IK_PRINT_TIMING
tim2 = ggml_time_us();
printf("sched_alloc_graph(...): %d us\n", int(tim2-tim1));
@@ -5468,6 +5602,7 @@ static int llama_decode_internal(
} else {
//printf("Reusing graph with type = %d, n_kv = %d, n_tokens = %d\n", cparams.mtp_op_type, (int)prev->n_kv, (int)prev->n_tokens);
gf = prev->graph;
llama_dsv4_trace_decode_stage(lctx, "graph_reused", &u_batch, gf);
}
if (cparams.mtp_op_type != MTP_OP_NONE) {
@@ -5483,6 +5618,7 @@ static int llama_decode_internal(
if (lctx.model.arch == LLM_ARCH_DEEPSEEK4 && !llama_prepare_dsv4_graph_inputs(lctx, u_batch, true, false)) {
return GGML_STATUS_FAILED;
}
llama_dsv4_trace_decode_stage(lctx, "inputs_ready", &u_batch, gf);
// the output is always the last tensor in the graph
struct ggml_tensor * res = gf->nodes[gf->n_nodes - 1];
@@ -5536,11 +5672,13 @@ static int llama_decode_internal(
}
}
}
llama_dsv4_trace_decode_stage(lctx, "result_selected", &u_batch, gf, res);
// LLAMA_LOG_INFO("graph build time: %.3f ms (%d nodes, %d leafs)\n", (ggml_time_us() - t_start_us)/1000.0, gf->n_nodes, gf->n_leafs);
#if IK_PRINT_TIMING == 1
tim1 = ggml_time_us();
#endif
llama_set_inputs(lctx, u_batch);
llama_dsv4_trace_decode_stage(lctx, "inputs_set", &u_batch, gf);
#if IK_PRINT_TIMING == 1
tim2 = ggml_time_us();
printf("set_inputs(...): %d us\n", int(tim2-tim1));
@@ -5548,7 +5686,27 @@ static int llama_decode_internal(
#if IK_PRINT_TIMING
tim1 = ggml_time_us();
#endif
llama_dsv4_trace_decode_stage(lctx, "compute_submit", &u_batch, gf);
ggml_backend_sched_eval_callback prev_eval_cb = nullptr;
void * prev_eval_cb_user_data = nullptr;
const bool install_dsv4_eval_trace =
lctx.model.arch == LLM_ARCH_DEEPSEEK4 &&
llama_dsv4_trace_eval_enabled() &&
lctx.cparams.cb_eval == nullptr;
if (install_dsv4_eval_trace) {
prev_eval_cb = lctx.cparams.cb_eval;
prev_eval_cb_user_data = lctx.cparams.cb_eval_user_data;
lctx.cparams.cb_eval = llama_dsv4_trace_eval_callback;
lctx.cparams.cb_eval_user_data = &lctx;
ggml_backend_sched_set_eval_callback(lctx.sched, lctx.cparams.cb_eval, lctx.cparams.cb_eval_user_data);
}
llama_graph_compute(lctx, gf, n_threads);
if (install_dsv4_eval_trace) {
lctx.cparams.cb_eval = prev_eval_cb;
lctx.cparams.cb_eval_user_data = prev_eval_cb_user_data;
ggml_backend_sched_set_eval_callback(lctx.sched, prev_eval_cb, prev_eval_cb_user_data);
}
llama_dsv4_trace_decode_stage(lctx, "compute_returned", &u_batch, gf);
#if IK_PRINT_TIMING
llama_synchronize(&lctx);
tim2 = ggml_time_us();
@@ -5605,6 +5763,7 @@ static int llama_decode_internal(
ggml_backend_t backend_res = ggml_backend_sched_get_tensor_backend(lctx.sched, res);
GGML_ASSERT(backend_res != nullptr);
GGML_ASSERT(lctx.logits != nullptr);
llama_dsv4_trace_decode_stage(lctx, "logits_backend_ready", &u_batch, gf, res);
float * logits_out = lctx.logits + n_outputs_prev*n_vocab;
const int32_t n_outputs_new = lctx.n_outputs;
@@ -5628,6 +5787,10 @@ static int llama_decode_internal(
} else {
ggml_backend_tensor_get_async(backend_res, res, logits_out, 0, n_outputs_new*n_vocab*sizeof(float));
}
llama_dsv4_trace_decode_stage(lctx, "logits_copy_queued", &u_batch, gf, res);
if (lctx.model.arch == LLM_ARCH_DEEPSEEK4 && llama_dsv4_trace_enabled()) {
llama_dsv4_trace_logits_summary(backend_res, logits_out, n_vocab, n_outputs_new);
}
}
}
#if IK_PRINT_TIMING