diff --git a/ggml/include/ggml.h b/ggml/include/ggml.h index e3b212cc1..2446676f6 100644 --- a/ggml/include/ggml.h +++ b/ggml/include/ggml.h @@ -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, diff --git a/ggml/src/ggml.c b/ggml/src/ggml.c index 19ba0436c..4f62b1c5b 100644 --- a/ggml/src/ggml.c +++ b/ggml/src/ggml.c @@ -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( diff --git a/models/templates/README.md b/models/templates/README.md index 3a649b8f4..53792da45 100644 --- a/models/templates/README.md +++ b/models/templates/README.md @@ -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 ``` diff --git a/models/templates/deepseek-ai-DeepSeek-V4.jinja b/models/templates/deepseek-ai-DeepSeek-V4.jinja new file mode 100644 index 000000000..d4b0165de --- /dev/null +++ b/models/templates/deepseek-ai-DeepSeek-V4.jinja @@ -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 = '' -%} +{%- set thinking_end_token = '' -%} +{%- 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\n...\n\n<' + dsml_token + 'invoke name="$TOOL_NAME2">\n...\n\n\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 -%} + {{- '' + (message['content'] or '') + '' -}} + {%- 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 + '\n' -}} + {%- else -%} + {{- '<' + dsml_token + 'parameter name="' + key + '" string="false">' + (val | tojson) + '\n' -}} + {%- endif -%} + {%- endfor -%} + {{- '\n' -}} + {%- endfor -%} + {{- '' -}} + {%- endif -%} + {{- '<｜end▁of▁sentence｜>' -}} + {%- endif -%} +{%- endfor -%} +{%- if add_generation_prompt -%} + {{- '<｜Assistant｜>' -}} + {%- if thinking -%} + {{- thinking_start_token -}} + {%- else -%} + {{- thinking_end_token -}} + {%- endif -%} +{%- endif -%} diff --git a/src/graphs/build_deepseek4.cpp b/src/graphs/build_deepseek4.cpp index 74390ea8f..6bc990df2 100644 --- a/src/graphs/build_deepseek4.cpp +++ b/src/graphs/build_deepseek4.cpp @@ -5,9 +5,18 @@ #include #include +#include #include #include +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(1, n), std::max(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(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(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(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, diff --git a/src/llama-arch.h b/src/llama-arch.h index 85ab3388e..b40c11e42 100644 --- a/src/llama-arch.h +++ b/src/llama-arch.h @@ -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, diff --git a/src/llama-build-context.cpp b/src/llama-build-context.cpp index b9e669324..97fd29259 100644 --- a/src/llama-build-context.cpp +++ b/src/llama-build-context.cpp @@ -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, diff --git a/src/llama-build-context.h b/src/llama-build-context.h index 4fa9cb2c6..b24cb010b 100644 --- a/src/llama-build-context.h +++ b/src/llama-build-context.h @@ -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, diff --git a/src/llama-context.h b/src/llama-context.h index 0b891cb77..5d10bab14 100644 --- a/src/llama-context.h +++ b/src/llama-context.h @@ -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 csa_mask_data; std::vector hca_mask_data; std::vector lid_mask_data; + std::vector lid_k_rot_data; }; dsv4_runtime dsv4; diff --git a/src/llama-dsv4.cpp b/src/llama-dsv4.cpp index f3bc01e7e..8f886acd5 100644 --- a/src/llama-dsv4.cpp +++ b/src/llama-dsv4.cpp @@ -206,6 +206,55 @@ static uint32_t dsv4_comp_size(uint32_t kv_size, uint32_t ratio) { return std::max(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 & 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; } diff --git a/src/llama-hparams.cpp b/src/llama-hparams.cpp index 5849dfa4b..66e898454 100644 --- a/src/llama-hparams.cpp +++ b/src/llama-hparams.cpp @@ -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); diff --git a/src/llama-hparams.h b/src/llama-hparams.h index 5f5f771f0..57c3d9276 100644 --- a/src/llama-hparams.h +++ b/src/llama-hparams.h @@ -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 { diff --git a/src/llama-load-tensors.cpp b/src/llama-load-tensors.cpp index 3db64e96f..e72c59570 100644 --- a/src/llama-load-tensors.cpp +++ b/src/llama-load-tensors.cpp @@ -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); } } diff --git a/src/llama-model.cpp b/src/llama-model.cpp index 49e1610bc..e683dbd80 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -1071,19 +1071,19 @@ static const std::map> 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_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" }, diff --git a/src/llama.cpp b/src/llama.cpp index aa9ea66b4..28fc57ca6 100644 --- a/src/llama.cpp +++ b/src/llama.cpp @@ -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(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 top_ids = { -1, -1, -1, -1, -1 }; + std::array 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