diff --git a/common/chat.cpp b/common/chat.cpp index 6da59f4dbd2c..74f99a575e2d 100644 --- a/common/chat.cpp +++ b/common/chat.cpp @@ -70,6 +70,11 @@ static bool has_content_or_tool_calls(const common_chat_msg & msg) { return !msg.content.empty() || !msg.tool_calls.empty(); } +static bool common_chat_template_uses_deepseek_dsml(const std::string & src) { + return src.find("dsml_token") != std::string::npos && + src.find("DSML") != std::string::npos; +} + std::string common_chat_msg::render_content(const std::string & delimiter) const { if (!content.empty() && !content_parts.empty()) { throw std::runtime_error("Cannot specify both content and content_parts"); @@ -1856,7 +1861,8 @@ static common_chat_params common_chat_params_init_gigachat_v3( } static common_chat_params common_chat_params_init_deepseek_v3_2(const common_chat_template & tmpl, - const autoparser::generation_params & inputs) { + const autoparser::generation_params & inputs, + const std::string & tool_calls_tag = "function_calls") { common_chat_params data; data.prompt = common_chat_template_direct_apply_impl(tmpl, inputs); @@ -1879,8 +1885,8 @@ static common_chat_params common_chat_params_init_deepseek_v3_2(const common_cha const std::string DSML = "｜DSML｜"; const std::string THINK_START = ""; const std::string THINK_END = ""; - const std::string FC_START = "<" + DSML + "function_calls>"; - const std::string FC_END = ""; + const std::string FC_START = "<" + DSML + tool_calls_tag + ">"; + const std::string FC_END = ""; const std::string INVOKE_START = "<" + DSML + "invoke"; const std::string INVOKE_END = ""; const std::string PARAM_START = "<" + DSML + "parameter"; @@ -1902,112 +1908,149 @@ static common_chat_params common_chat_params_init_deepseek_v3_2(const common_cha auto generation_prompt = p.literal(GEN_PROMPT); auto end = p.end(); - auto reasoning = p.eps(); - if (extract_reasoning && inputs.enable_thinking) { - reasoning = p.optional(THINK_START + p.reasoning(p.until(THINK_END)) + THINK_END); - } else if (extract_reasoning) { - // Thinking disabled but reasoning extraction requested: the generation prompt - // contains an empty pair that must still be consumed. - reasoning = p.optional(p.literal(THINK_START) + p.until(THINK_END) + p.literal(THINK_END)); - } - + auto after_reasoning = p.eps(); if (has_response_format) { auto response_format = p.rule("response-format", p.literal("```json") + p.space() + p.content(p.schema(p.json(), "response-format-schema", inputs.json_schema)) + p.space() + p.literal("```")); - return generation_prompt + reasoning + response_format + end; - } + after_reasoning = response_format; + } else if (!has_tools || inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_NONE) { + after_reasoning = p.content(p.rest()); + } else { + auto tool_choice = p.choice(); + foreach_function(inputs.tools, [&](const json & tool) { + const auto & function = tool.at("function"); + std::string name = function.at("name"); + auto params = function.contains("parameters") ? function.at("parameters") : json::object(); + const auto & props = params.contains("properties") ? params.at("properties") : json::object(); - if (!has_tools || inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_NONE) { - return generation_prompt + reasoning + p.content(p.rest()) + end; - } + std::set required; + if (params.contains("required")) { + params.at("required").get_to(required); + } - auto tool_choice = p.choice(); - foreach_function(inputs.tools, [&](const json & tool) { - const auto & function = tool.at("function"); - std::string name = function.at("name"); - auto params = function.contains("parameters") ? function.at("parameters") : json::object(); - const auto & props = params.contains("properties") ? params.at("properties") : json::object(); + auto schema_info = common_schema_info(); + schema_info.resolve_refs(params); + + std::vector required_indices; + std::vector arg_parsers; + for (const auto & [param_name, param_schema] : props.items()) { + bool is_required = required.find(param_name) != required.end(); + bool is_string = schema_info.resolves_to_string(param_schema); + + auto arg = p.tool_arg( + p.tool_arg_open( + p.literal(PARAM_START + " name=\"") + + p.tool_arg_name(p.literal(param_name)) + + p.literal("\" string=\"" + std::string(is_string ? "true" : "false") + "\">")) + + (is_string + ? p.tool_arg_string_value(p.until(PARAM_END)) + : p.tool_arg_json_value(p.schema(p.json(), + "tool-" + name + "-arg-" + param_name + "-schema", + param_schema, false))) + + p.tool_arg_close(p.literal(PARAM_END))); + + auto named_arg = p.rule("tool-" + name + "-arg-" + param_name, arg); + const auto arg_idx = arg_parsers.size(); + arg_parsers.push_back(named_arg); + if (is_required) { + required_indices.push_back(arg_idx); + } + } - std::set required; - if (params.contains("required")) { - params.at("required").get_to(required); - } + auto contains_index = [](const std::vector & values, size_t needle) { + for (const auto value : values) { + if (value == needle) { + return true; + } + } + return false; + }; - auto schema_info = common_schema_info(); - schema_info.resolve_refs(params); - - std::vector required_parsers; - std::vector optional_parsers; - for (const auto & [param_name, param_schema] : props.items()) { - bool is_required = required.find(param_name) != required.end(); - bool is_string = schema_info.resolves_to_string(param_schema); - - auto arg = p.tool_arg( - p.tool_arg_open( - p.literal(PARAM_START + " name=\"") + - p.tool_arg_name(p.literal(param_name)) + - p.literal("\" string=\"" + std::string(is_string ? "true" : "false") + "\">")) + - (is_string - ? p.tool_arg_string_value(p.until(PARAM_END)) - : p.tool_arg_json_value(p.schema(p.json(), - "tool-" + name + "-arg-" + param_name + "-schema", - param_schema, false))) + - p.tool_arg_close(p.literal(PARAM_END))); - - auto named_arg = p.rule("tool-" + name + "-arg-" + param_name, arg); - if (is_required) { - required_parsers.push_back(named_arg); - } else { - optional_parsers.push_back(named_arg); - } - } + std::function &)> build_args_seq = + [&](const std::vector & remaining_required) { + if (arg_parsers.empty()) { + return p.eps(); + } - common_peg_parser args_seq = p.eps(); - for (size_t i = 0; i < required_parsers.size(); i++) { - if (i > 0) { - args_seq = args_seq + p.space(); - } - args_seq = args_seq + required_parsers[i]; - } + common_peg_parser skippable_args = p.choice(); + bool has_skippable_args = false; + for (size_t i = 0; i < arg_parsers.size(); i++) { + if (contains_index(remaining_required, i)) { + continue; + } + skippable_args |= arg_parsers[i]; + has_skippable_args = true; + } - if (!optional_parsers.empty()) { - common_peg_parser any_opt = p.choice(); - for (const auto & opt : optional_parsers) { - any_opt |= opt; - } - args_seq = args_seq + p.repeat(p.space() + any_opt, 0, -1); - } + auto skipped = has_skippable_args ? p.repeat(skippable_args + p.space(), 0, -1) : p.eps(); + if (remaining_required.empty()) { + return skipped; + } + + common_peg_parser args_choice = p.choice(); + for (size_t i = 0; i < remaining_required.size(); i++) { + std::vector next_required; + next_required.reserve(remaining_required.size() - 1); + for (size_t j = 0; j < remaining_required.size(); j++) { + if (i != j) { + next_required.push_back(remaining_required[j]); + } + } + + auto required_arg = arg_parsers[remaining_required[i]] + p.space(); + args_choice |= skipped + required_arg + build_args_seq(next_required); + } + return args_choice; + }; - common_peg_parser invoke_body = args_seq; - auto func_parser = p.tool( - p.tool_open(p.literal(INVOKE_START + " name=\"") + - p.tool_name(p.literal(name)) + p.literal("\">\n")) + - invoke_body + p.space() + - p.tool_close(p.literal(INVOKE_END))); + common_peg_parser args_seq = build_args_seq(required_indices); - tool_choice |= p.rule("tool-" + name, func_parser); - }); + common_peg_parser invoke_body = args_seq; + auto func_parser = p.tool( + p.tool_open(p.literal(INVOKE_START + " name=\"") + + p.tool_name(p.literal(name)) + p.literal("\">\n")) + + invoke_body + p.space() + + p.tool_close(p.literal(INVOKE_END))); - auto require_tools = inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_REQUIRED; + tool_choice |= p.rule("tool-" + name, func_parser); + }); - common_peg_parser tool_calls = p.eps(); - if (inputs.parallel_tool_calls) { - tool_calls = p.trigger_rule("tool-call", - p.literal(FC_START) + p.space() + tool_choice + - p.zero_or_more(p.space() + tool_choice) + p.space() + p.literal(FC_END)); - } else { - tool_calls = p.trigger_rule("tool-call", - p.literal(FC_START) + p.space() + tool_choice + p.space() + p.literal(FC_END)); + auto require_tools = inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_REQUIRED; + + common_peg_parser tool_calls = p.eps(); + if (inputs.parallel_tool_calls) { + tool_calls = p.trigger_rule("tool-call", + p.literal(FC_START) + p.space() + tool_choice + + p.zero_or_more(p.space() + tool_choice) + p.space() + p.literal(FC_END)); + } else { + tool_calls = p.trigger_rule("tool-call", + p.literal(FC_START) + p.space() + tool_choice + p.space() + p.literal(FC_END)); + } + + if (!require_tools) { + tool_calls = p.optional(tool_calls); + } + + auto content_before_tools = p.content(p.until(FC_START)); + after_reasoning = content_before_tools + tool_calls; } - if (!require_tools) { - tool_calls = p.optional(tool_calls); + if (extract_reasoning && inputs.enable_thinking) { + auto closed_reasoning = p.literal(THINK_START) + p.reasoning(p.until(THINK_END)) + p.literal(THINK_END) + after_reasoning; + auto open_reasoning = p.literal(THINK_START) + p.reasoning(p.rest()); + return generation_prompt + p.choice({ closed_reasoning, open_reasoning }) + end; + } else if (extract_reasoning) { + // V3.2 pre-fills , while V4 pre-fills directly. + auto reasoning = p.optional(p.choice({ + p.literal(THINK_START) + p.until(THINK_END) + p.literal(THINK_END), + p.literal(THINK_END), + })); + return generation_prompt + reasoning + after_reasoning + end; } - auto content_before_tools = p.content(p.until(FC_START)); - return generation_prompt + reasoning + content_before_tools + tool_calls + end; + return generation_prompt + after_reasoning + end; }); data.parser = parser.save(); @@ -2214,6 +2257,71 @@ static void requires_non_null_content(json & messages) { } } +static void sort_tool_results_by_previous_tool_calls(json & messages) { + GGML_ASSERT(messages.is_array()); + + for (size_t i = 0; i + 1 < messages.size(); i++) { + auto & assistant = messages[i]; + if (!assistant.is_object() || + assistant.value("role", "") != "assistant" || + !assistant.contains("tool_calls") || + !assistant.at("tool_calls").is_array() || + assistant.at("tool_calls").empty()) { + continue; + } + + size_t begin = i + 1; + size_t end = begin; + while (end < messages.size() && + messages[end].is_object() && + messages[end].value("role", "") == "tool") { + end++; + } + if (begin == end) { + continue; + } + + std::vector used(end - begin, false); + std::vector ordered; + ordered.reserve(end - begin); + + for (const auto & tool_call : assistant.at("tool_calls")) { + if (!tool_call.is_object() || + !tool_call.contains("id") || + !tool_call.at("id").is_string()) { + continue; + } + + const auto id = tool_call.at("id").get(); + for (size_t j = begin; j < end; j++) { + auto & tool_result = messages[j]; + if (used[j - begin] || + !tool_result.contains("tool_call_id") || + !tool_result.at("tool_call_id").is_string() || + tool_result.at("tool_call_id").get() != id) { + continue; + } + + ordered.push_back(tool_result); + used[j - begin] = true; + break; + } + } + + for (size_t j = begin; j < end; j++) { + if (!used[j - begin]) { + ordered.push_back(messages[j]); + } + } + + for (size_t j = 0; j < ordered.size(); j++) { + messages[begin + j] = ordered[j]; + } + + i = end - 1; + } +} + // Gemma4 uses a custom tool_responses field instead of role:tool messages. // // This will transform a sequence of messages: @@ -2597,13 +2705,20 @@ std::optional common_chat_try_specialized_template( // DeepSeek V3.2 format detection: template defines dsml_token and uses it for tool calls. // The template source contains the token as a variable assignment, not as a literal in markup. - if (src.find("dsml_token") != std::string::npos && - src.find("function_calls") != std::string::npos && - src.find("DSML") != std::string::npos) { + if (common_chat_template_uses_deepseek_dsml(src) && + src.find("function_calls") != std::string::npos) { LOG_DBG("Using specialized template: DeepSeek V3.2\n"); return common_chat_params_init_deepseek_v3_2(tmpl, params); } + // DeepSeek V4 format detection: same DSML invoke/parameter shape as V3.2, + // but the outer block is tool_calls instead of function_calls. + if (common_chat_template_uses_deepseek_dsml(src) && + src.find("tool_calls>") != std::string::npos) { + LOG_DBG("Using specialized template: DeepSeek V4\n"); + return common_chat_params_init_deepseek_v3_2(tmpl, params, "tool_calls"); + } + // Gemma4 format detection if (src.find("'<|tool_call>call:'") != std::string::npos) { if (src.find("{#- OpenAI Chat Completions:") == std::string::npos) { @@ -2685,6 +2800,10 @@ static common_chat_params common_chat_templates_apply_jinja(const struct common_ workaround::func_args_not_string(params.messages); } + if (common_chat_template_uses_deepseek_dsml(src)) { + workaround::sort_tool_results_by_previous_tool_calls(params.messages); + } + params.extra_context = common_chat_extra_context(); for (auto el : inputs.chat_template_kwargs) { params.extra_context[el.first] = json::parse(el.second); diff --git a/common/preset.cpp b/common/preset.cpp index f0cc1fa1a242..4362c0621b78 100644 --- a/common/preset.cpp +++ b/common/preset.cpp @@ -7,6 +7,7 @@ #include #include #include +#include static std::string rm_leading_dashes(const std::string & str) { size_t pos = 0; @@ -16,6 +17,23 @@ static std::string rm_leading_dashes(const std::string & str) { return str.substr(pos); } +static std::string canonical_tag(const std::string & tag) { + static const std::regex re_tag("[-.]([A-Z0-9_]+)$", std::regex::icase); + std::smatch m; + if (std::regex_search(tag, m, re_tag)) { + std::string canon = m[1].str(); + for (char & c : canon) { + c = (char) std::toupper((unsigned char) c); + } + return canon; + } + std::string upper = tag; + for (char & c : upper) { + c = (char) std::toupper((unsigned char) c); + } + return upper; +} + std::vector common_preset::to_args(const std::string & bin_path) const { std::vector args; @@ -270,11 +288,18 @@ common_presets common_preset_context::load_from_ini(const std::string & path, co for (auto section : ini_data) { common_preset preset; - if (section.first.empty()) { - preset.name = COMMON_PRESET_DEFAULT_NAME; - } else { - preset.name = section.first; + std::string section_name = section.first.empty() ? std::string(COMMON_PRESET_DEFAULT_NAME) : section.first; + if (section_name != "*" && section_name != COMMON_PRESET_DEFAULT_NAME) { + auto colon_idx = section_name.rfind(':'); + if (colon_idx != std::string::npos) { + std::string tag = section_name.substr(colon_idx + 1); + std::string canon_tag = canonical_tag(tag); + if (canon_tag != tag) { + section_name = section_name.substr(0, colon_idx + 1) + canon_tag; + } + } } + preset.name = section_name; LOG_DBG("loading preset: %s\n", preset.name.c_str()); for (const auto & [key, value] : section.second) { if (key == "version") { diff --git a/ggml/include/ggml.h b/ggml/include/ggml.h index 0e6091a731f4..aa238f6bc56f 100644 --- a/ggml/include/ggml.h +++ b/ggml/include/ggml.h @@ -568,6 +568,7 @@ extern "C" { GGML_OP_RWKV_WKV7, GGML_OP_SOLVE_TRI, GGML_OP_GATED_DELTA_NET, + GGML_OP_LIGHTNING_INDEXER, GGML_OP_DSV4_HC_COMB, GGML_OP_DSV4_HC_PRE, GGML_OP_DSV4_HC_POST, @@ -2576,6 +2577,14 @@ extern "C" { struct ggml_tensor * state, int64_t K); + GGML_API struct ggml_tensor * ggml_lightning_indexer( + struct ggml_context * ctx, + struct ggml_tensor * q, + struct ggml_tensor * k, + struct ggml_tensor * weights, + float scale_embd, + float scale_heads); + // DeepSeek V4 hyper-connection helpers. // hc_comb: mixes [(2 + hc)*hc, n_tokens], scale [3], base [(2 + hc)*hc] // -> [dst_hc, src_hc, n_tokens] diff --git a/ggml/src/ggml-cpu/ggml-cpu.c b/ggml/src/ggml-cpu/ggml-cpu.c index 737e449ed351..80e67d85c007 100644 --- a/ggml/src/ggml-cpu/ggml-cpu.c +++ b/ggml/src/ggml-cpu/ggml-cpu.c @@ -2051,6 +2051,10 @@ static void ggml_compute_forward(struct ggml_compute_params * params, struct ggm { ggml_compute_forward_gated_delta_net(params, tensor); } break; + case GGML_OP_LIGHTNING_INDEXER: + { + ggml_compute_forward_lightning_indexer(params, tensor); + } break; case GGML_OP_DSV4_HC_COMB: { ggml_compute_forward_dsv4_hc_comb(params, tensor); @@ -2386,6 +2390,7 @@ static int ggml_get_n_tasks(struct ggml_tensor * node, int n_threads) { case GGML_OP_FLASH_ATTN_BACK: case GGML_OP_SSM_CONV: case GGML_OP_SSM_SCAN: + case GGML_OP_LIGHTNING_INDEXER: { n_tasks = n_threads; } break; @@ -2971,6 +2976,12 @@ struct ggml_cplan ggml_graph_plan( { GGML_ABORT("fatal error"); } + case GGML_OP_LIGHTNING_INDEXER: + { + // temp buffer for dequantizing lightning indexer keys + const int64_t ne10 = node->src[1]->ne[0]; + cur += sizeof(float)*ne10*n_tasks; + } break; default: break; } diff --git a/ggml/src/ggml-cpu/ops.cpp b/ggml/src/ggml-cpu/ops.cpp index af86add1fb2d..2876c6dc4166 100644 --- a/ggml/src/ggml-cpu/ops.cpp +++ b/ggml/src/ggml-cpu/ops.cpp @@ -1913,7 +1913,11 @@ static void ggml_compute_forward_concat_any( GGML_ASSERT(dim >= 0 && dim < 4); int64_t o[4] = {0, 0, 0, 0}; - o[dim] = src0->ne[dim]; + if (dim == 0) { + o[dim] = src0->ne[dim]/ggml_blck_size(src0->type); + } else { + o[dim] = src0->ne[dim]; + } const char * x; @@ -1921,8 +1925,8 @@ static void ggml_compute_forward_concat_any( for (int i3 = 0; i3 < ne3; i3++) { for (int i2 = ith; i2 < ne2; i2 += nth) { for (int i1 = 0; i1 < ne1; i1++) { - for (int i0 = 0; i0 < ne0; i0++) { - if (i0 < ne00 && i1 < ne01 && i2 < ne02 && i3 < ne03) { + for (int i0 = 0; i0 < ne0/ggml_blck_size(dst->type); i0++) { + if (i0 < ne00/ggml_blck_size(src0->type) && 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; @@ -2071,6 +2075,14 @@ void ggml_compute_forward_concat( ggml_tensor * dst) { const ggml_tensor * src0 = dst->src[0]; + const ggml_tensor * src1 = dst->src[1]; + + if (ggml_is_quantized(src0->type)) { + GGML_ASSERT(ggml_is_contiguous(src0)); + GGML_ASSERT(ggml_is_contiguous(src1)); + GGML_ASSERT(src0->ne[0] % ggml_blck_size(src0->type) == 0); + GGML_ASSERT(src1->ne[0] % ggml_blck_size(src1->type) == 0); + } switch (src0->type) { case GGML_TYPE_F16: @@ -11797,3 +11809,76 @@ void ggml_compute_forward_fwht(const ggml_compute_params * params, ggml_tensor * } } } + +// ggml_compute_forward_lightning_indexer + +void ggml_compute_forward_lightning_indexer( + const ggml_compute_params * params, + ggml_tensor * dst) { + + const ggml_tensor * src0 = dst->src[0]; // q + const ggml_tensor * src1 = dst->src[1]; // k + const ggml_tensor * src2 = dst->src[2]; // weights + + const float scale_embd = ggml_get_op_params_f32(dst, 0); + const float scale_heads = ggml_get_op_params_f32(dst, 1); + + GGML_ASSERT(dst->type == GGML_TYPE_F32); + GGML_ASSERT(src0->type == GGML_TYPE_F32); + GGML_ASSERT(src2->type == GGML_TYPE_F32); + + GGML_TENSOR_TERNARY_OP_LOCALS + + GGML_ASSERT( nb0 == sizeof(float)); + GGML_ASSERT(nb00 == sizeof(float)); + + int n_embd = src0->ne[0]; + int n_head = src0->ne[1]; + int n_batch = src0->ne[2]; + int n_stream = src0->ne[3]; + int n_kv = src1->ne[2]; + + ggml_to_float_t const k_to_float = ggml_get_type_traits(src1->type)->to_float; + GGML_ASSERT((src1->type == GGML_TYPE_F32 || k_to_float) && "lightning indexer: unsupported K-type"); + + const int nr = n_kv; + const int ith = params->ith; + const int nth = params->nth; + + // (temporary) buffer for K converted to float + float * src1_row_f32 = (float *) params->wdata + ith*(1*n_embd + CACHE_LINE_SIZE_F32); + + // rows per thread + const int dr = (nr + nth - 1)/nth; + + // row range for this thread + const int ir0 = dr*ith; + const int ir1 = MIN(ir0 + dr, nr); + + for (int i_stream = 0; i_stream < n_stream; ++i_stream) { + for (int i_batch = 0; i_batch < n_batch; ++i_batch) { + for (int i_kv = ir0; i_kv < ir1; ++i_kv) { + char * src1_row = (char *) src1->data + i_kv*nb12 + i_stream*nb13; + if (k_to_float) { + k_to_float(src1_row, src1_row_f32, n_embd); + } else { + src1_row_f32 = (float *) src1_row; + } + float * src2_row = (float *) ((char *) src2->data + i_batch*nb21 + i_stream*nb23); + float * dst_row = (float *) ((char *) dst->data + i_batch*nb1 + i_stream*nb3); + float score = 0.0f; + for (int i_head = 0; i_head < n_head; ++i_head) { + // dot product of q and k for head i_head + float qk = 0.0f; + float * src0_row = (float *) ((char *) src0->data + i_head*nb01 + i_batch*nb02 + i_stream*nb03); + ggml_vec_dot_f32(n_embd, &qk, 0, src0_row, 0, src1_row_f32, 0, 1); + qk *= scale_embd; + // ReLU and weights + score += MAX(qk, 0.0f) * src2_row[i_head]; + } + score *= scale_heads; + dst_row[i_kv] = score; + } + } + } +} diff --git a/ggml/src/ggml-cpu/ops.h b/ggml/src/ggml-cpu/ops.h index 4c1cc3bb4f77..4c1642a67603 100644 --- a/ggml/src/ggml-cpu/ops.h +++ b/ggml/src/ggml-cpu/ops.h @@ -105,6 +105,7 @@ void ggml_compute_forward_rwkv_wkv7(const struct ggml_compute_params * params, s void ggml_compute_forward_solve_tri(const struct ggml_compute_params * params, struct ggml_tensor * dst); void ggml_compute_forward_gla(const struct ggml_compute_params * params, struct ggml_tensor * dst); void ggml_compute_forward_gated_delta_net(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_lightning_indexer(const struct ggml_compute_params * params, struct ggml_tensor * dst); void ggml_compute_forward_dsv4_hc_comb(const struct ggml_compute_params * params, struct ggml_tensor * dst); void ggml_compute_forward_dsv4_hc_pre(const struct ggml_compute_params * params, struct ggml_tensor * dst); void ggml_compute_forward_dsv4_hc_post(const struct ggml_compute_params * params, struct ggml_tensor * dst); diff --git a/ggml/src/ggml-cuda/argsort.cu b/ggml/src/ggml-cuda/argsort.cu index c4f08091e79a..33a38c23e87e 100644 --- a/ggml/src/ggml-cuda/argsort.cu +++ b/ggml/src/ggml-cuda/argsort.cu @@ -28,6 +28,20 @@ static __global__ void init_offsets(int * offsets, const int ncols, const int nr #endif // STRIDED_ITERATOR_AVAILABLE #ifdef GGML_CUDA_USE_CUB + +// returns the suggested maximum number of rows to process during one argsort_f32_i32_cuda_cub() call +int argsort_f32_i32_cuda_cub_chunk_nrows(const size_t nb01, const int64_t nrows) { + // perform argsort in chunks up to approximately this size (currently 64MB) + // to avoid excessive temporary buffers memory usage + const int chunk_bytes = 1 << 26; + + // calculate how many rows will fit in one chunk (must be at least one) + const int chunk_nrows = chunk_bytes > nb01 ? chunk_bytes / nb01 : 1; + + // limit the resulting amount to total nrows + return nrows < chunk_nrows ? nrows : chunk_nrows; +} + void argsort_f32_i32_cuda_cub(ggml_cuda_pool & pool, const float * x, int * dst, @@ -254,11 +268,22 @@ void ggml_cuda_op_argsort(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { const size_t shared_mem = ncols_pad * sizeof(int); const size_t max_shared_mem = ggml_cuda_info().devices[ggml_cuda_get_device()].smpb; - if (shared_mem > max_shared_mem || ncols > 1024) { - ggml_cuda_pool & pool = ctx.pool(); - argsort_f32_i32_cuda_cub(pool, src0_d, (int *) dst_d, ncols, nrows, order, stream); - } else { - argsort_f32_i32_cuda_bitonic(src0_d, (int *) dst_d, ncols, nrows, order, stream); + // early return if we can use bitonic argsort + if (shared_mem <= max_shared_mem && ncols <= 1024) { + return argsort_f32_i32_cuda_bitonic(src0_d, (int *) dst_d, ncols, nrows, order, stream); + } + + const int chunk_nrows = argsort_f32_i32_cuda_cub_chunk_nrows(src0->nb[1], nrows); + + ggml_cuda_pool & pool = ctx.pool(); + + for (int64_t i = 0; i < nrows; i += chunk_nrows) { + int iter_nrows = chunk_nrows < nrows - i ? chunk_nrows : nrows - i; + + argsort_f32_i32_cuda_cub(pool, src0_d, (int *) dst_d, ncols, iter_nrows, order, stream); + + src0_d += ncols * iter_nrows; + dst_d += ncols * iter_nrows; } #else argsort_f32_i32_cuda_bitonic(src0_d, (int *) dst_d, ncols, nrows, order, stream); diff --git a/ggml/src/ggml-cuda/argsort.cuh b/ggml/src/ggml-cuda/argsort.cuh index 22b7306f2020..3abb6448a057 100644 --- a/ggml/src/ggml-cuda/argsort.cuh +++ b/ggml/src/ggml-cuda/argsort.cuh @@ -3,6 +3,7 @@ void ggml_cuda_op_argsort(ggml_backend_cuda_context & ctx, ggml_tensor * dst); #ifdef GGML_CUDA_USE_CUB +int argsort_f32_i32_cuda_cub_chunk_nrows(const size_t nb01, const int64_t nrows); void argsort_f32_i32_cuda_cub(ggml_cuda_pool & pool, const float * x, int * dst, diff --git a/ggml/src/ggml-cuda/concat.cu b/ggml/src/ggml-cuda/concat.cu index 8d557092b2be..cb9f91d793e5 100644 --- a/ggml/src/ggml-cuda/concat.cu +++ b/ggml/src/ggml-cuda/concat.cu @@ -147,13 +147,25 @@ static void concat_cuda(const ggml_tensor * src0, const ggml_tensor * src1, ggml T * dst_d = (T *) dst->data; if (dim != 3) { - for (int64_t i3 = 0; i3 < dst->ne[3]; i3++) { - concat_cont_cuda( - src0_d + i3*(src0->nb[3] / sizeof(T)), - src1_d + i3*(src1->nb[3] / sizeof(T)), - dst_d + i3*( dst->nb[3] / sizeof(T)), - src0->ne[0], src0->ne[1], src0->ne[2], - dst->ne[0], dst->ne[1], dst->ne[2], dim, stream); + if (ggml_is_quantized(src0->type)) { + // treat both tensors as byte tensors with ne[0] equal to nb[1] + for (int64_t i3 = 0; i3 < dst->ne[3]; i3++) { + concat_cont_cuda( + src0_d + i3*(src0->nb[3] / sizeof(T)), + src1_d + i3*(src1->nb[3] / sizeof(T)), + dst_d + i3*( dst->nb[3] / sizeof(T)), + src0->nb[1], src0->ne[1], src0->ne[2], + dst->nb[1], dst->ne[1], dst->ne[2], dim, stream); + } + } else { + for (int64_t i3 = 0; i3 < dst->ne[3]; i3++) { + concat_cont_cuda( + src0_d + i3*(src0->nb[3] / sizeof(T)), + src1_d + i3*(src1->nb[3] / sizeof(T)), + dst_d + i3*( dst->nb[3] / sizeof(T)), + src0->ne[0], src0->ne[1], src0->ne[2], + dst->ne[0], dst->ne[1], dst->ne[2], dim, stream); + } } } else { const size_t size0 = ggml_nbytes(src0); @@ -163,6 +175,8 @@ static void concat_cuda(const ggml_tensor * src0, const ggml_tensor * src1, ggml CUDA_CHECK(cudaMemcpyAsync((char *) dst->data + size0, src1->data, size1, cudaMemcpyDeviceToDevice, stream)); } } else { + GGML_ASSERT(!ggml_is_quantized(src0->type)); + dim3 grid_dim(dst->ne[1], dst->ne[2], dst->ne[3]); auto launch_kernel = [&](auto dim) { concat_non_cont<<>>( @@ -204,24 +218,34 @@ void ggml_cuda_op_concat(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { GGML_ASSERT(src0->type == src1->type); GGML_ASSERT(dst->type == src0->type); - GGML_ASSERT(!ggml_is_quantized(src0->type)); - GGML_ASSERT(ggml_blck_size(src0->type) == 1); - - switch (ggml_type_size(src0->type)) { - case 1: - concat_cuda(src0, src1, dst, dim, stream); - break; - case 2: - concat_cuda(src0, src1, dst, dim, stream); - break; - case 4: - concat_cuda(src0, src1, dst, dim, stream); - break; - case 8: - concat_cuda(src0, src1, dst, dim, stream); - break; - default: - GGML_ABORT("Unsupported type size: %zu", ggml_type_size(src0->type)); - break; + + if (ggml_is_quantized(src0->type)) { + GGML_ASSERT(ggml_is_contiguous(src0)); + GGML_ASSERT(ggml_is_contiguous(src1)); + GGML_ASSERT(src0->ne[0] % ggml_blck_size(src0->type) == 0); + GGML_ASSERT(src1->ne[0] % ggml_blck_size(src1->type) == 0); + + // if tensors are contiguous and ne[0] is multiple of the block size we can concat both tensors as byte tensors + concat_cuda(src0, src1, dst, dim, stream); + } else { + GGML_ASSERT(ggml_blck_size(src0->type) == 1); + + switch (ggml_type_size(src0->type)) { + case 1: + concat_cuda(src0, src1, dst, dim, stream); + break; + case 2: + concat_cuda(src0, src1, dst, dim, stream); + break; + case 4: + concat_cuda(src0, src1, dst, dim, stream); + break; + case 8: + concat_cuda(src0, src1, dst, dim, stream); + break; + default: + GGML_ABORT("Unsupported type size: %zu", ggml_type_size(src0->type)); + break; + } } } diff --git a/ggml/src/ggml-cuda/dsv4-hc.cu b/ggml/src/ggml-cuda/dsv4-hc.cu new file mode 100644 index 000000000000..4dc9055f519e --- /dev/null +++ b/ggml/src/ggml-cuda/dsv4-hc.cu @@ -0,0 +1,297 @@ +#include "common.cuh" +#include "dsv4-hc.cuh" + + +static __device__ void dsv4_hc_comb_norm_cols(float * comb, float eps) { + constexpr int64_t hc = 4; + + for (int64_t idst = 0; idst < hc; ++idst) { + float sum = eps; + for (int64_t isrc = 0; isrc < hc; ++isrc) { + sum += comb[idst + hc*isrc]; + } + + const float inv_sum = 1.0f / sum; + for (int64_t isrc = 0; isrc < hc; ++isrc) { + comb[idst + hc*isrc] *= inv_sum; + } + } +} + +static __device__ void dsv4_hc_comb_norm_rows(float * comb, float eps) { + constexpr int64_t hc = 4; + + for (int64_t isrc = 0; isrc < hc; ++isrc) { + float sum = eps; + for (int64_t idst = 0; idst < hc; ++idst) { + sum += comb[idst + hc*isrc]; + } + + const float inv_sum = 1.0f / sum; + for (int64_t idst = 0; idst < hc; ++idst) { + comb[idst + hc*isrc] *= inv_sum; + } + } +} + +static __global__ void dsv4_hc_comb_f32( + const float * mixes, + const float * scale, + const float * base, + float * dst, + int64_t n_tokens, + int64_t sm0, + int64_t sm1, + int64_t ss0, + int64_t sb0, + int64_t sd0, + int64_t sd1, + int64_t sd2, + float eps, + int32_t n_iter) { + constexpr int64_t hc = 4; + constexpr int64_t comb_offset = 2*hc; + + ggml_cuda_pdl_lc(); + const int64_t it = (int64_t) blockIdx.x * blockDim.x + threadIdx.x; + + if (it >= n_tokens) { + return; + } + + ggml_cuda_pdl_sync(); + + const float scale_comb = scale[2*ss0]; + float comb[hc*hc]; + + for (int64_t isrc = 0; isrc < hc; ++isrc) { + float max = -INFINITY; + for (int64_t idst = 0; idst < hc; ++idst) { + const int64_t idx = idst + hc*isrc; + const float v = mixes[(comb_offset + idx)*sm0 + it*sm1] * scale_comb + base[(comb_offset + idx)*sb0]; + comb[idx] = v; + max = fmaxf(max, v); + } + + float sum = 0.0f; + for (int64_t idst = 0; idst < hc; ++idst) { + const int64_t idx = idst + hc*isrc; + const float v = expf(comb[idx] - max); + comb[idx] = v; + sum += v; + } + + const float inv_sum = 1.0f / sum; + for (int64_t idst = 0; idst < hc; ++idst) { + const int64_t idx = idst + hc*isrc; + comb[idx] = comb[idx] * inv_sum + eps; + } + } + + dsv4_hc_comb_norm_cols(comb, eps); + for (int32_t i = 1; i < n_iter; ++i) { + dsv4_hc_comb_norm_rows(comb, eps); + dsv4_hc_comb_norm_cols(comb, eps); + } + + for (int64_t isrc = 0; isrc < hc; ++isrc) { + for (int64_t idst = 0; idst < hc; ++idst) { + const int64_t idx = idst + hc*isrc; + dst[idst*sd0 + isrc*sd1 + it*sd2] = comb[idx]; + } + } +} + +static __global__ void dsv4_hc_pre_f32( + const float * x, + const float * weights, + float * dst, + int64_t n_embd, + int64_t hc, + int64_t n_tokens, + int64_t sx0, + int64_t sx1, + int64_t sx2, + int64_t sw0, + int64_t sw1, + int64_t sd0, + int64_t sd1) { + ggml_cuda_pdl_lc(); + const int64_t ir = (int64_t) blockIdx.x * blockDim.x + threadIdx.x; + const int64_t nr = n_embd * n_tokens; + + if (ir >= nr) { + return; + } + + ggml_cuda_pdl_sync(); + + const int64_t i0 = ir % n_embd; + const int64_t it = ir / n_embd; + + float sum = __fmul_rn(x[i0*sx0 + it*sx2], weights[it*sw1]); + for (int64_t ih = 1; ih < hc; ++ih) { + const float xv = x[i0*sx0 + ih*sx1 + it*sx2]; + const float wv = weights[ih*sw0 + it*sw1]; + sum = __fadd_rn(sum, __fmul_rn(xv, wv)); + } + + dst[i0*sd0 + it*sd1] = sum; +} + +static __global__ void dsv4_hc_post_f32( + const float * x, + const float * residual, + const float * post, + const float * comb, + float * dst, + int64_t n_embd, + int64_t hc, + int64_t n_tokens, + int64_t sx0, + int64_t sx1, + int64_t sr0, + int64_t sr1, + int64_t sr2, + int64_t sp0, + int64_t sp1, + int64_t sc0, + int64_t sc1, + int64_t sc2, + int64_t sd0, + int64_t sd1, + int64_t sd2) { + ggml_cuda_pdl_lc(); + const int64_t ir = (int64_t) blockIdx.x * blockDim.x + threadIdx.x; + const int64_t nr = n_embd * hc * n_tokens; + + if (ir >= nr) { + return; + } + + ggml_cuda_pdl_sync(); + + const int64_t i0 = ir % n_embd; + const int64_t idst = (ir / n_embd) % hc; + const int64_t it = ir / (n_embd * hc); + + float sum = x[i0*sx0 + it*sx1] * post[idst*sp0 + it*sp1]; + for (int64_t isrc = 0; isrc < hc; ++isrc) { + sum += residual[i0*sr0 + isrc*sr1 + it*sr2] * comb[idst*sc0 + isrc*sc1 + it*sc2]; + } + + dst[i0*sd0 + idst*sd1 + it*sd2] = sum; +} + +void ggml_cuda_op_dsv4_hc_comb(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { + const ggml_tensor * mixes = dst->src[0]; + const ggml_tensor * scale = dst->src[1]; + const ggml_tensor * base = dst->src[2]; + + GGML_ASSERT(mixes->type == GGML_TYPE_F32); + GGML_ASSERT(scale->type == GGML_TYPE_F32); + GGML_ASSERT(base->type == GGML_TYPE_F32); + GGML_ASSERT(dst->type == GGML_TYPE_F32); + + constexpr int64_t hc = 4; + constexpr int64_t hc_mix_dim = (2 + hc)*hc; + + GGML_ASSERT(mixes->ne[0] == hc_mix_dim); + GGML_ASSERT(dst->ne[0] == hc); + GGML_ASSERT(dst->ne[1] == hc); + GGML_ASSERT(dst->ne[2] == mixes->ne[1]); + GGML_ASSERT(scale->ne[0] >= 3); + GGML_ASSERT(base->ne[0] == hc_mix_dim); + + GGML_TENSOR_LOCALS(size_t, nbm, mixes, nb); + GGML_TENSOR_LOCALS(size_t, nbs, scale, nb); + GGML_TENSOR_LOCALS(size_t, nbb, base, nb); + GGML_TENSOR_LOCALS(size_t, nbd, dst, nb); + + const int64_t n_tokens = mixes->ne[1]; + const float eps = ggml_get_op_params_f32(dst, 0); + const int32_t n_iter = ggml_get_op_params_i32(dst, 1); + + const int block_size = 256; + const dim3 block_dims(block_size, 1, 1); + const dim3 grid_dims((n_tokens + block_size - 1) / block_size, 1, 1); + const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params(grid_dims, block_dims, 0, ctx.stream()); + + ggml_cuda_kernel_launch(dsv4_hc_comb_f32, launch_params, + (const float *) mixes->data, (const float *) scale->data, (const float *) base->data, (float *) dst->data, + n_tokens, + nbm0 / sizeof(float), nbm1 / sizeof(float), + nbs0 / sizeof(float), + nbb0 / sizeof(float), + nbd0 / sizeof(float), nbd1 / sizeof(float), nbd2 / sizeof(float), + eps, n_iter); +} + +void ggml_cuda_op_dsv4_hc_pre(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { + const ggml_tensor * x = dst->src[0]; + const ggml_tensor * weights = dst->src[1]; + + GGML_ASSERT(x->type == GGML_TYPE_F32); + GGML_ASSERT(weights->type == GGML_TYPE_F32); + GGML_ASSERT(dst->type == GGML_TYPE_F32); + + GGML_TENSOR_LOCALS(size_t, nbx, x, nb); + GGML_TENSOR_LOCALS(size_t, nbw, weights, nb); + GGML_TENSOR_LOCALS(size_t, nbd, dst, nb); + + const int64_t n_embd = x->ne[0]; + const int64_t hc = x->ne[1]; + const int64_t n_tokens = x->ne[2]; + + const int block_size = 256; + const int64_t nr = n_embd * n_tokens; + const dim3 block_dims(block_size, 1, 1); + const dim3 grid_dims((nr + block_size - 1) / block_size, 1, 1); + const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params(grid_dims, block_dims, 0, ctx.stream()); + + ggml_cuda_kernel_launch(dsv4_hc_pre_f32, launch_params, + (const float *) x->data, (const float *) weights->data, (float *) dst->data, + n_embd, hc, n_tokens, + nbx0 / sizeof(float), nbx1 / sizeof(float), nbx2 / sizeof(float), + nbw0 / sizeof(float), nbw1 / sizeof(float), + nbd0 / sizeof(float), nbd1 / sizeof(float)); +} + +void ggml_cuda_op_dsv4_hc_post(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { + const ggml_tensor * x = dst->src[0]; + const ggml_tensor * residual = dst->src[1]; + const ggml_tensor * post = dst->src[2]; + const ggml_tensor * comb = dst->src[3]; + + GGML_ASSERT(x->type == GGML_TYPE_F32); + GGML_ASSERT(residual->type == GGML_TYPE_F32); + GGML_ASSERT(post->type == GGML_TYPE_F32); + GGML_ASSERT(comb->type == GGML_TYPE_F32); + GGML_ASSERT(dst->type == GGML_TYPE_F32); + + GGML_TENSOR_LOCALS(size_t, nbx, x, nb); + GGML_TENSOR_LOCALS(size_t, nbr, residual, nb); + GGML_TENSOR_LOCALS(size_t, nbp, post, nb); + GGML_TENSOR_LOCALS(size_t, nbc, comb, nb); + GGML_TENSOR_LOCALS(size_t, nbd, dst, nb); + + const int64_t n_embd = x->ne[0]; + const int64_t n_tokens = x->ne[1]; + const int64_t hc = residual->ne[1]; + + const int block_size = 256; + const int64_t nr = n_embd * hc * n_tokens; + const dim3 block_dims(block_size, 1, 1); + const dim3 grid_dims((nr + block_size - 1) / block_size, 1, 1); + const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params(grid_dims, block_dims, 0, ctx.stream()); + + ggml_cuda_kernel_launch(dsv4_hc_post_f32, launch_params, + (const float *) x->data, (const float *) residual->data, + (const float *) post->data, (const float *) comb->data, (float *) dst->data, + n_embd, hc, n_tokens, + nbx0 / sizeof(float), nbx1 / sizeof(float), + nbr0 / sizeof(float), nbr1 / sizeof(float), nbr2 / sizeof(float), + nbp0 / sizeof(float), nbp1 / sizeof(float), + nbc0 / sizeof(float), nbc1 / sizeof(float), nbc2 / sizeof(float), + nbd0 / sizeof(float), nbd1 / sizeof(float), nbd2 / sizeof(float)); +} diff --git a/ggml/src/ggml-cuda/dsv4-hc.cuh b/ggml/src/ggml-cuda/dsv4-hc.cuh new file mode 100644 index 000000000000..2379aaefb41b --- /dev/null +++ b/ggml/src/ggml-cuda/dsv4-hc.cuh @@ -0,0 +1,6 @@ +#include "common.cuh" +#include "ggml.h" + +void ggml_cuda_op_dsv4_hc_comb(ggml_backend_cuda_context & ctx, ggml_tensor * dst); +void ggml_cuda_op_dsv4_hc_pre(ggml_backend_cuda_context & ctx, ggml_tensor * dst); +void ggml_cuda_op_dsv4_hc_post(ggml_backend_cuda_context & ctx, ggml_tensor * dst); diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index 378d539220c6..db0dd3a3f5c1 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -66,6 +66,7 @@ #include "ggml-cuda/tri.cuh" #include "ggml-cuda/cumsum.cuh" #include "ggml-cuda/fill.cuh" +#include "ggml-cuda/lightning-indexer.cuh" #include "ggml.h" #include @@ -3128,6 +3129,9 @@ static bool ggml_cuda_compute_forward(ggml_backend_cuda_context & ctx, struct gg case GGML_OP_FILL: ggml_cuda_op_fill(ctx, dst); break; + case GGML_OP_LIGHTNING_INDEXER: + ggml_cuda_op_lightning_indexer(ctx, dst); + break; default: return false; } @@ -5330,12 +5334,24 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g ggml_type src1_type = op->src[1]->type; return src0_type == src1_type && src0_type == op->type && - !ggml_is_quantized(src0_type) && - ggml_blck_size(src0_type) == 1 && - (ggml_type_size(src0_type) == 1 || - ggml_type_size(src0_type) == 2 || - ggml_type_size(src0_type) == 4 || - ggml_type_size(src0_type) == 8); + ( + ( + ggml_is_quantized(src0_type) && + ggml_is_contiguous(op->src[0]) && + ggml_is_contiguous(op->src[1]) && + op->src[0]->ne[0] % ggml_blck_size(src0_type) == 0 && + op->src[1]->ne[0] % ggml_blck_size(src0_type) == 0 + ) || ( + !ggml_is_quantized(src0_type) && + ggml_blck_size(src0_type) == 1 && + ( + ggml_type_size(src0_type) == 1 || + ggml_type_size(src0_type) == 2 || + ggml_type_size(src0_type) == 4 || + ggml_type_size(src0_type) == 8 + ) + ) + ); } break; case GGML_OP_CONV_TRANSPOSE_1D: { @@ -5484,6 +5500,7 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g case GGML_OP_TRI: case GGML_OP_DIAG: case GGML_OP_SOLVE_TRI: + case GGML_OP_LIGHTNING_INDEXER: return true; default: diff --git a/ggml/src/ggml-cuda/lightning-indexer.cu b/ggml/src/ggml-cuda/lightning-indexer.cu new file mode 100644 index 000000000000..6cda36efd157 --- /dev/null +++ b/ggml/src/ggml-cuda/lightning-indexer.cu @@ -0,0 +1,507 @@ +#include "common.cuh" +#include "lightning-indexer.cuh" +#include "fattn-common.cuh" +#include "convert.cuh" + +#if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= GGML_CUDA_CC_AMPERE + +typedef union { + int2 i2; + half2 h2[2]; +} half4; + +#include +namespace wmma = nvcuda::wmma; + +template +static __global__ void lightning_indexer_kernel_wmma( + const float * src0, const char * src1, const float * src2, float * dst, + const float scale_embd, const float scale_heads, + int64_t n_stream, int64_t n_batch, int64_t n_kv, + size_t nb1, size_t nb2, size_t nb3, + size_t nb01, size_t nb02, size_t nb03, + size_t nb11, size_t nb12, size_t nb13, + size_t nb21, size_t nb22, size_t nb23 + ) { + + constexpr int K_VECS_PER_BLOCK = 32; + constexpr int WARPS_PER_BLOCK = 8; + constexpr int THREADS_PER_BLOCK = WARPS_PER_BLOCK * WARP_SIZE; + constexpr int HEADS_PER_INNER_LOOP = 8; + constexpr int K_EMBD_PER_INNER_LOOP = 16; + constexpr int n_embd_padded = n_embd + 8; + + const int i_batch = blockIdx.y; + const int i_stream = blockIdx.z; + const int i_warp = threadIdx.y; + const int i_lane = threadIdx.x; + const int tid = i_warp * WARP_SIZE + i_lane; + + // each block processes K_VECS_PER_BLOCK K vectors + const int start_kv = blockIdx.x * K_VECS_PER_BLOCK; + + const char * q_base = (const char *) src0 + i_batch*nb02 + i_stream*nb03; + const float * w_base = (const float *) ((const char *) src2 + i_batch*nb21 + i_stream*nb23); + + // phase 1 - load weights and first Q tile to shared memory + + __shared__ float w_shared[n_head]; + __shared__ int2 q_shared_h[HEADS_PER_INNER_LOOP][n_embd_padded / 4]; + + if (tid < n_head) { + w_shared[tid] = w_base[tid]; + } + + // total number of half4 elements in HEADS_PER_INNER_LOOP x n_embd Q tile + constexpr int n_q_tile = HEADS_PER_INNER_LOOP * (n_embd / 4); + // number of registers needed in each thread to store Q tile in thread block + constexpr int n_q_next = (n_q_tile + THREADS_PER_BLOCK - 1) / THREADS_PER_BLOCK; + + #pragma unroll + for (int i_q = tid; i_q < n_q_tile; i_q += THREADS_PER_BLOCK) { + const int i_head = i_q / (n_embd / 4); + const int i_embd = i_q % (n_embd / 4); + const float4 q = *(const float4 *) (q_base + i_head*nb01 + i_embd*sizeof(float4)); + half4 q_packed; + q_packed.h2[0] = __float22half2_rn(make_float2(q.x, q.y)); + q_packed.h2[1] = __float22half2_rn(make_float2(q.z, q.w)); + q_shared_h[i_head][i_embd] = q_packed.i2; + } + + // phase 2 - load (and dequantize if needed) K to shared mem + + __shared__ half2 k_shared_h[K_VECS_PER_BLOCK][n_embd_padded / 4][2]; + + constexpr int n_k = K_VECS_PER_BLOCK * (n_embd / 4); + + if constexpr (type_K == GGML_TYPE_F16) { + #pragma unroll + for (int i_k = tid; i_k < n_k; i_k += THREADS_PER_BLOCK) { + const int i_k_vec = i_k / (n_embd / 4); + const int i_embd = i_k % (n_embd / 4); + const int i_kv = start_kv + i_k_vec; + if (i_kv < n_kv) { + const int2 * k_base = (const int2 *) ((const char *) src1 + i_kv*nb12 + i_stream*nb13); + *(int2*) &k_shared_h[i_k_vec][i_embd] = k_base[i_embd]; + } else { + *(int2*) &k_shared_h[i_k_vec][i_embd] = make_int2(0, 0); + } + } + } else { + constexpr dequantize_V_t dequantize_k = get_dequantize_V(); + #pragma unroll + for (int i_k = tid; i_k < n_k; i_k += THREADS_PER_BLOCK) { + const int i_k_vec = i_k / (n_embd / 4); + const int i_embd = i_k % (n_embd / 4); + const int i_kv = start_kv + i_k_vec; + if (i_kv < n_kv) { + const void * k_base = (const void *) ((const char *) src1 + i_kv*nb12 + i_stream*nb13); + dequantize_k(k_base, &k_shared_h[i_k_vec][i_embd][0], i_embd * 4); + } else { + *(int2*) &k_shared_h[i_k_vec][i_embd] = make_int2(0, 0); + } + } + } + + __syncthreads(); + + // phase 3 - calculate lightning indexer scores + + __shared__ float qk_shared[WARPS_PER_BLOCK][HEADS_PER_INNER_LOOP][K_VECS_PER_BLOCK]; + + // load K fragment + wmma::fragment frag_k; + wmma::load_matrix_sync(frag_k, (half*) &k_shared_h[0][i_warp * K_EMBD_PER_INNER_LOOP / 4], n_embd_padded); + + float score_k = 0.0f; + + for (int i_head_0 = 0; i_head_0 < n_head; i_head_0 += HEADS_PER_INNER_LOOP) { + const int i_head_next = i_head_0 + HEADS_PER_INNER_LOOP; + + // we don't use accumulator for anything, fill it with zeros + wmma::fragment frag_acc; + wmma::fill_fragment(frag_acc, 0.0f); + + // load Q fragment + wmma::fragment frag_q; + wmma::load_matrix_sync(frag_q, (half*) &q_shared_h[0][i_warp * K_EMBD_PER_INNER_LOOP / 4], n_embd_padded); + + // preload next Q tile to registers during matrix multiplication + float4 q_next[n_q_next]; + + if (i_head_next < n_head) { + #pragma unroll + for (int i_q = tid, i_q_next = 0; i_q < n_q_tile; i_q += THREADS_PER_BLOCK) { + const int i_head = i_head_next + i_q / (n_embd / 4); + const int i_embd = i_q % (n_embd / 4); + q_next[i_q_next++] = *(const float4 *) (q_base + i_head*nb01 + i_embd*sizeof(float4)); + } + } + + // perform matrix multiplication + wmma::mma_sync(frag_acc, frag_q, frag_k, frag_acc); + wmma::store_matrix_sync((float*) &qk_shared[i_warp][0][0], frag_acc, K_VECS_PER_BLOCK, wmma::mem_row_major); + + // make sure all threads finished using q_shared_h so we can store next tile + __syncthreads(); + + // write preloaded Q tile to shared memory + if (i_head_next < n_head) { + #pragma unroll + for (int i_q = tid, i_q_next = 0; i_q < n_q_tile; i_q += THREADS_PER_BLOCK) { + const int i_head = i_q / (n_embd / 4); + const int i_embd = i_q % (n_embd / 4); + half4 q_packed; + q_packed.h2[0] = __float22half2_rn(make_float2(q_next[i_q_next].x, q_next[i_q_next].y)); + q_packed.h2[1] = __float22half2_rn(make_float2(q_next[i_q_next].z, q_next[i_q_next].w)); + q_shared_h[i_head][i_embd] = q_packed.i2; + ++i_q_next; + } + } + + // accumulate QK multiplication results from all block warps + // (there are 256 threads in block and 256 matmul outputs) + // TODO it will break if WARP_SIZE is not 32 + const int h = tid / K_VECS_PER_BLOCK; + const int k = tid % K_VECS_PER_BLOCK; + const float w_val = w_shared[i_head_0 + h]; + + float sum = 0.0f; + #pragma unroll + for (int w = 0; w < WARPS_PER_BLOCK; ++w) { + sum += qk_shared[w][h][k]; + } + + // scale_embd, ReLU, weight + sum *= scale_embd; + sum = sum > 0.0f ? sum : 0.0f; + sum *= w_val; + + // wait until qk_shared[0] is no longer used + __syncthreads(); + + // reuse qk_shared[0] for storing partial results + qk_shared[0][h][k] = sum; + + // wait until all threads write their results + __syncthreads(); + + // accumulate result over heads + if (tid < K_VECS_PER_BLOCK) { + #pragma unroll + for (int i_head = 0; i_head < HEADS_PER_INNER_LOOP; ++i_head) { + score_k += qk_shared[0][i_head][tid]; + } + } + + // make sure all threads finished using qk_shared + __syncthreads(); + } + + // phase 4 - store output to VRAM + + if (tid < K_VECS_PER_BLOCK) { + const int i_kv = start_kv + tid; + if (i_kv < n_kv) { + float * dst_base = (float *) ((char *) dst + i_batch*nb1 + i_stream*nb3); + dst_base[i_kv] = score_k * scale_heads; + } + } +} + +#else // defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= GGML_CUDA_CC_AMPERE + +template +static __global__ void lightning_indexer_kernel_wmma( + const float * src0, const char * src1, const float * src2, float * dst, + const float scale_embd, const float scale_heads, + int64_t n_stream, int64_t n_batch, int64_t n_kv, + size_t nb1, size_t nb2, size_t nb3, + size_t nb01, size_t nb02, size_t nb03, + size_t nb11, size_t nb12, size_t nb13, + size_t nb21, size_t nb22, size_t nb23 + ) { + GGML_UNUSED_VARS(src0, src1, src2, dst, + scale_embd, scale_heads, + n_stream, n_batch, n_kv, + nb1, nb2, nb3, + nb01, nb02, nb03, + nb11, nb12, nb13, + nb21, nb22, nb23); + NO_DEVICE_CODE; +} + +#endif // defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= GGML_CUDA_CC_AMPERE +#endif // !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) && defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= GGML_CUDA_CC_AMPERE + +// TODO there is one ugly assumption used in this kernel - that WARP_SIZE is equal to 32 +// thanks to that one warp operating on float4 processes whole indexer K/Q vectors +// 32 * 4 = 128 (n_embd) + +template +static __global__ void lightning_indexer_kernel_vec( + const float * src0, const char * src1, const float * src2, float * dst, + const float scale_embd, const float scale_heads, + int64_t n_stream, int64_t n_batch, int64_t n_kv, + size_t nb1, size_t nb2, size_t nb3, + size_t nb01, size_t nb02, size_t nb03, + size_t nb11, size_t nb12, size_t nb13, + size_t nb21, size_t nb22, size_t nb23 + ) { + + constexpr int K_VECS_PER_WARP = 8; + constexpr int WARPS_PER_BLOCK = 8; + constexpr int THREADS_PER_BLOCK = WARPS_PER_BLOCK * WARP_SIZE; + + const int i_batch = blockIdx.y; + const int i_stream = blockIdx.z; + const int i_warp = threadIdx.y; + const int i_lane = threadIdx.x; + const int tid = i_warp * WARP_SIZE + i_lane; + + // each warp processes K_VECS_PER_WARP K vectors + const int start_kv_block = blockIdx.x * (WARPS_PER_BLOCK * K_VECS_PER_WARP); + const int start_kv = start_kv_block + i_warp * K_VECS_PER_WARP; + + const char * q_base = (const char *) src0 + i_batch*nb02 + i_stream*nb03; + const float * w_base = (const float *) ((const char *) src2 + i_batch*nb21 + i_stream*nb23); + + // phase 1 - load (and dequantize if needed) K to registers + + float4 k_reg_f[K_VECS_PER_WARP]; + + if constexpr (type_K == GGML_TYPE_F32) { + // direct copy of float4 + #pragma unroll + for (int k = 0; k < K_VECS_PER_WARP; ++k) { + int i_kv = start_kv + k; + if (i_kv < n_kv) { + const float4 * k_base = (const float4 *) ((const char *) src1 + i_kv*nb12 + i_stream*nb13); + k_reg_f[k] = k_base[i_lane]; + } else { + k_reg_f[k] = make_float4(0, 0, 0, 0); + } + } + } else { + // dequantize remaining types to float + constexpr dequantize_V_t dequantize_k = get_dequantize_V(); + #pragma unroll + for (int k = 0; k < K_VECS_PER_WARP; ++k) { + int i_kv = start_kv + k; + if (i_kv < n_kv) { + const void * k_base = (const void *) ((const char *) src1 + i_kv*nb12 + i_stream*nb13); + dequantize_k(k_base, &k_reg_f[k], i_lane * 4); + } else { + k_reg_f[k] = make_float4(0, 0, 0, 0); + } + } + } + + float score_k[K_VECS_PER_WARP] = { 0.0f }; + + // load weights and Q only for n_head_inner heads at once to reduce shared memory usage + constexpr int n_head_inner = n_head / 4; + + for (int i_head_0 = 0; i_head_0 < n_head; i_head_0 += n_head_inner) { + // phase 2 - load weights and Q to shared memory + + __shared__ float w_shared[n_head_inner]; + __shared__ float4 q_shared_f[n_head_inner][n_embd / 4]; + + if (tid < n_head_inner) { + w_shared[tid] = w_base[i_head_0 + tid]; + } + + constexpr int n_q = n_head_inner * (n_embd / 4); + #pragma unroll + for (int i_q = tid; i_q < n_q; i_q += THREADS_PER_BLOCK) { + const int i_head_inner = i_q / (n_embd / 4); + const int i_head = i_head_0 + i_head_inner; + const int i_embd = i_q % (n_embd / 4); + q_shared_f[i_head_inner][i_embd] = *(const float4 *) (q_base + i_head*nb01 + i_embd*sizeof(float4)); + } + + __syncthreads(); + + // phase 3 - calculate lightning indexer scores + + for (int i_head_inner = 0; i_head_inner < n_head_inner; ++i_head_inner) { + const float w_val = w_shared[i_head_inner]; + float qk[K_VECS_PER_WARP] = { 0.0f }; + + // dot product of floats + const float4 q_vec = q_shared_f[i_head_inner][i_lane]; + + #pragma unroll + for (int k = 0; k < K_VECS_PER_WARP; ++k) { + ggml_cuda_mad(qk[k], q_vec.x, k_reg_f[k].x); + ggml_cuda_mad(qk[k], q_vec.y, k_reg_f[k].y); + ggml_cuda_mad(qk[k], q_vec.z, k_reg_f[k].z); + ggml_cuda_mad(qk[k], q_vec.w, k_reg_f[k].w); + } + + #pragma unroll + for (int k = 0; k < K_VECS_PER_WARP; ++k) { + float sum = warp_reduce_sum(qk[k]); + + // scale_embd, ReLU, weight + if (i_lane == 0) { + sum *= scale_embd; + sum = (sum > 0.0f) ? sum : 0.0f; + score_k[k] += sum * w_val; + } + } + } + + __syncthreads(); + } + + // phase 4 - store outputs to shared memory + + __shared__ float dst_shared[WARPS_PER_BLOCK * K_VECS_PER_WARP]; + + if (i_lane == 0) { + #pragma unroll + for (int k = 0; k < K_VECS_PER_WARP; ++k) { + dst_shared[i_warp * K_VECS_PER_WARP + k] = score_k[k] * scale_heads; + } + } + + __syncthreads(); + + // phase 5 - write from shared memory to VRAM in coalesced manner + + if (tid < WARPS_PER_BLOCK * K_VECS_PER_WARP) { + int i_kv = start_kv_block + tid; + if (i_kv < n_kv) { + float * dst_base = (float *) ((char *) dst + i_batch*nb1 + i_stream*nb3); + dst_base[i_kv] = dst_shared[tid]; + } + } +} + +#define DECL_LIGHTNING_INDEXER_CASE(lightning_indexer_kernel, n_embd, n_head, type_K) \ + template __global__ void lightning_indexer_kernel( \ + const float * src0, const char * src1, const float * src2, float * dst, \ + const float scale_embd, const float scale_heads, \ + int64_t n_stream, int64_t n_batch, int64_t n_kv, \ + size_t nb1, size_t nb2, size_t nb3, \ + size_t nb01, size_t nb02, size_t nb03, \ + size_t nb11, size_t nb12, size_t nb13, \ + size_t nb21, size_t nb22, size_t nb23); + +#if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) +DECL_LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_wmma, 128, 64, GGML_TYPE_F16) +DECL_LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_wmma, 128, 64, GGML_TYPE_Q4_0) +DECL_LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_wmma, 128, 64, GGML_TYPE_Q4_1) +DECL_LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_wmma, 128, 64, GGML_TYPE_Q5_0) +DECL_LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_wmma, 128, 64, GGML_TYPE_Q5_1) +DECL_LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_wmma, 128, 64, GGML_TYPE_Q8_0) +#endif // !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) + +DECL_LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_vec, 128, 64, GGML_TYPE_F16) +DECL_LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_vec, 128, 64, GGML_TYPE_Q4_0) +DECL_LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_vec, 128, 64, GGML_TYPE_Q4_1) +DECL_LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_vec, 128, 64, GGML_TYPE_Q5_0) +DECL_LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_vec, 128, 64, GGML_TYPE_Q5_1) +DECL_LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_vec, 128, 64, GGML_TYPE_Q8_0) +DECL_LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_vec, 128, 64, GGML_TYPE_BF16) +DECL_LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_vec, 128, 64, GGML_TYPE_F32) + +#define LIGHTNING_INDEXER_CASE(lightning_indexer_kernel, n_embd, n_head, K, type_K) \ + if (K->type == (type_K)) { \ + lightning_indexer_kernel<<>>( \ + src0_d, src1_d, src2_d, dst_d, scale_embd, scale_heads, \ + n_stream, n_batch, n_kv, \ + nb1, nb2, nb3, \ + nb01, nb02, nb03, \ + nb11, nb12, nb13, \ + nb21, nb22, nb23 \ + ); \ + } else + +void ggml_cuda_op_lightning_indexer(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { + const ggml_tensor * src0 = dst->src[0]; + const ggml_tensor * src1 = dst->src[1]; + const ggml_tensor * src2 = dst->src[2]; + + const float scale_embd = ggml_get_op_params_f32(dst, 0); + const float scale_heads = ggml_get_op_params_f32(dst, 1); + + GGML_ASSERT(dst->type == GGML_TYPE_F32); + GGML_ASSERT(src0->type == GGML_TYPE_F32); + GGML_ASSERT(src2->type == GGML_TYPE_F32); + + GGML_TENSOR_TERNARY_OP_LOCALS + + // input tensor rows must be contiguous + GGML_ASSERT(nb00 == ggml_type_size(src0->type)); + GGML_ASSERT(nb10 == ggml_type_size(src1->type)); + GGML_ASSERT(nb20 == ggml_type_size(src2->type)); + + // dst cannot be transposed or permuted + GGML_ASSERT(nb0 == sizeof(float)); + GGML_ASSERT(nb0 <= nb1); + GGML_ASSERT(nb1 <= nb2); + GGML_ASSERT(nb2 <= nb3); + + const int n_embd = src0->ne[0]; + const int n_head = src0->ne[1]; + const int n_batch = src0->ne[2]; + const int n_stream = src0->ne[3]; + const int n_kv = src1->ne[2]; + + const float * src0_d = (const float *) src0->data; + const char * src1_d = (const char *) src1->data; + const float * src2_d = (const float *) src2->data; + float * dst_d = (float *) dst->data; + + const int device = ggml_cuda_get_device(); + const int cc = ggml_cuda_info().devices[device].cc; + + if (n_embd == 128 && n_head == 64) { +#if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) + if (GGML_CUDA_CC_IS_NVIDIA(cc) && ampere_mma_available(cc) && src1->type != GGML_TYPE_F32 && src1->type != GGML_TYPE_BF16) { + // use wmma kernel + constexpr int K_VECS_PER_BLOCK = 32; + constexpr int WARPS_PER_BLOCK = 8; + + dim3 block(32, WARPS_PER_BLOCK); + int num_kv_blocks = (n_kv + (K_VECS_PER_BLOCK) - 1) / (K_VECS_PER_BLOCK); + dim3 grid(num_kv_blocks, n_batch, n_stream); + + LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_wmma, 128, 64, src1, GGML_TYPE_F16) + LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_wmma, 128, 64, src1, GGML_TYPE_Q4_0) + LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_wmma, 128, 64, src1, GGML_TYPE_Q4_1) + LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_wmma, 128, 64, src1, GGML_TYPE_Q5_0) + LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_wmma, 128, 64, src1, GGML_TYPE_Q5_1) + LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_wmma, 128, 64, src1, GGML_TYPE_Q8_0) + GGML_ABORT("fatal error"); + } else { +#else // !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) + { +#endif // !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) + // use vector kernel + constexpr int K_VECS_PER_WARP = 8; + constexpr int WARPS_PER_BLOCK = 8; + constexpr int K_VECS_PER_BLOCK = K_VECS_PER_WARP * WARPS_PER_BLOCK; + + dim3 block(32, WARPS_PER_BLOCK); + int num_kv_blocks = (n_kv + (K_VECS_PER_BLOCK) - 1) / (K_VECS_PER_BLOCK); + dim3 grid(num_kv_blocks, n_batch, n_stream); + + LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_vec, 128, 64, src1, GGML_TYPE_F16) + LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_vec, 128, 64, src1, GGML_TYPE_Q4_0) + LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_vec, 128, 64, src1, GGML_TYPE_Q4_1) + LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_vec, 128, 64, src1, GGML_TYPE_Q5_0) + LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_vec, 128, 64, src1, GGML_TYPE_Q5_1) + LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_vec, 128, 64, src1, GGML_TYPE_Q8_0) + LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_vec, 128, 64, src1, GGML_TYPE_BF16) + LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_vec, 128, 64, src1, GGML_TYPE_F32) + GGML_ABORT("fatal error"); + } + } else { + GGML_ABORT("fatal error"); + } +} diff --git a/ggml/src/ggml-cuda/lightning-indexer.cuh b/ggml/src/ggml-cuda/lightning-indexer.cuh new file mode 100644 index 000000000000..31fcc7d5ae0a --- /dev/null +++ b/ggml/src/ggml-cuda/lightning-indexer.cuh @@ -0,0 +1,3 @@ +#include "common.cuh" + +void ggml_cuda_op_lightning_indexer(ggml_backend_cuda_context & ctx, ggml_tensor * dst); diff --git a/ggml/src/ggml-cuda/top-k.cu b/ggml/src/ggml-cuda/top-k.cu index db1d39e2dc71..5e708e6c5ed4 100644 --- a/ggml/src/ggml-cuda/top-k.cu +++ b/ggml/src/ggml-cuda/top-k.cu @@ -75,17 +75,26 @@ void ggml_cuda_op_top_k(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { const int ncols_pad = next_power_of_2(ncols); const size_t shared_mem = ncols_pad * sizeof(int); const size_t max_shared_mem = ggml_cuda_info().devices[ggml_cuda_get_device()].smpb; + const bool use_bitonic = shared_mem <= max_shared_mem && ncols <= 1024; + const int chunk_nrows = argsort_f32_i32_cuda_cub_chunk_nrows(src0->nb[1], nrows); - ggml_cuda_pool_alloc temp_dst_alloc(pool, ncols * nrows); + ggml_cuda_pool_alloc temp_dst_alloc(pool, ncols * chunk_nrows); int * tmp_dst = temp_dst_alloc.get(); - if (shared_mem > max_shared_mem || ncols > 1024) { - argsort_f32_i32_cuda_cub(pool, src0_d, tmp_dst, ncols, nrows, GGML_SORT_ORDER_DESC, stream); - } else { - argsort_f32_i32_cuda_bitonic(src0_d, tmp_dst, ncols, nrows, GGML_SORT_ORDER_DESC, stream); + for (int64_t i = 0; i < nrows; i += chunk_nrows) { + int iter_nrows = chunk_nrows < nrows - i ? chunk_nrows : nrows - i; + + if (use_bitonic) { + argsort_f32_i32_cuda_bitonic(src0_d, tmp_dst, ncols, iter_nrows, GGML_SORT_ORDER_DESC, stream); + } else { + argsort_f32_i32_cuda_cub(pool, src0_d, tmp_dst, ncols, iter_nrows, GGML_SORT_ORDER_DESC, stream); + } + CUDA_CHECK(cudaMemcpy2DAsync(dst_d, k * sizeof(int), tmp_dst, ncols * sizeof(int), k * sizeof(int), iter_nrows, + cudaMemcpyDeviceToDevice, stream)); + + src0_d += ncols * iter_nrows; + dst_d += k * iter_nrows; } - CUDA_CHECK(cudaMemcpy2DAsync(dst_d, k * sizeof(int), tmp_dst, ncols * sizeof(int), k * sizeof(int), nrows, - cudaMemcpyDeviceToDevice, stream)); #else // GGML_CUDA_USE_CUB ggml_cuda_pool_alloc temp_dst_alloc(pool, ncols * nrows); int * tmp_dst = temp_dst_alloc.get(); diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index ea9191873282..ac021c7a6952 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -1907,6 +1907,38 @@ static bool vk_enable_sync_logger = false; static uint32_t vk_perf_logger_frequency = 1; static std::string vk_pipeline_stats_filter; +static uint64_t ggml_vk_get_node_flops(const ggml_tensor * node) { + if (node->op == GGML_OP_MUL_MAT || node->op == GGML_OP_MUL_MAT_ID) { + const uint64_t m = node->ne[0]; + const uint64_t n = node->ne[1]; + const uint64_t k = node->src[1]->ne[0]; + const uint64_t batch = node->ne[2] * node->ne[3]; + return m * n * (k + (k - 1)) * batch; + } + if (node->op == GGML_OP_CONV_2D || node->op == GGML_OP_CONV_TRANSPOSE_2D) { + const ggml_tensor * knl = node->src[0]; + const uint64_t Cout = node->ne[2]; + const uint64_t size_K = node->src[1]->ne[2] * knl->ne[0] * knl->ne[1]; + const uint64_t size_N = node->ne[3] * node->ne[0] * node->ne[1]; + return Cout * size_N * (size_K + (size_K - 1)); + } + if (node->op == GGML_OP_CONV_3D) { + const ggml_tensor * knl = node->src[0]; + const uint64_t OC = ggml_get_op_params_i32(node, 11); + const uint64_t IC = ggml_get_op_params_i32(node, 9); + const uint64_t size_K = IC * knl->ne[0] * knl->ne[1] * knl->ne[2]; + const uint64_t size_N = node->ne[3] / OC * node->ne[0] * node->ne[1] * node->ne[2]; + return OC * size_N * (size_K + (size_K - 1)); + } + if (node->op == GGML_OP_FLASH_ATTN_EXT) { + const ggml_tensor * q = node->src[0]; + const ggml_tensor * k = node->src[1]; + const ggml_tensor * v = node->src[2]; + return 2ull * q->ne[1] * q->ne[2] * (k->ne[0] + v->ne[0]) * k->ne[1] * q->ne[3]; + } + return 0; +} + class vk_perf_logger { public: void print_timings(bool force = false) { @@ -1955,7 +1987,7 @@ class vk_perf_logger { } std::string get_node_fusion_name(const ggml_tensor * node, const char *fusion_name, uint64_t *n_flops) { - *n_flops = 0; + *n_flops = ggml_vk_get_node_flops(node); std::string fusion_str; if (fusion_name) { fusion_str = fusion_name + std::string(" "); @@ -1982,35 +2014,22 @@ class vk_perf_logger { if (batch > 1) { name += " batch=" + std::to_string(batch); } - name = fusion_str + name; - *n_flops = m * n * (k + (k - 1)) * batch; - return name; + return fusion_str + name; } if (node->op == GGML_OP_CONV_2D || node->op == GGML_OP_CONV_TRANSPOSE_2D) { std::string name = ggml_op_name(node->op); - ggml_tensor * knl = node->src[0]; - uint64_t OW = node->ne[0]; - uint64_t OH = node->ne[1]; - uint64_t N = node->ne[3]; + const ggml_tensor * knl = node->src[0]; uint64_t Cout = node->ne[2]; - uint64_t KW = knl->ne[0]; - uint64_t KH = knl->ne[1]; - uint64_t Cin = node->src[1]->ne[2]; - // KxCRS @ CRSxNPQ = KxNPQ -> M=K, K=CRS, N=NPQ - uint64_t size_M = Cout; - uint64_t size_K = Cin * KW * KH; - uint64_t size_N = N * OW * OH; - *n_flops = size_M * size_N * (size_K + (size_K - 1)); - name += " M=Cout=" + std::to_string(size_M) + ", K=Cin*KW*KH=" + std::to_string(size_K) + + uint64_t size_K = node->src[1]->ne[2] * knl->ne[0] * knl->ne[1]; + uint64_t size_N = node->ne[3] * node->ne[0] * node->ne[1]; + name += " M=Cout=" + std::to_string(Cout) + ", K=Cin*KW*KH=" + std::to_string(size_K) + ", N=N*OW*OH=" + std::to_string(size_N); - name = fusion_str + name; - return name; + return fusion_str + name; } if (node->op == GGML_OP_RMS_NORM) { std::string name = ggml_op_name(node->op); name += "(" + std::to_string(node->ne[0]) + "," + std::to_string(node->ne[1]) + "," + std::to_string(node->ne[2]) + "," + std::to_string(node->ne[3]) + ")"; - name = fusion_str + name; - return name; + return fusion_str + name; } if (node->op == GGML_OP_FLASH_ATTN_EXT) { const ggml_tensor * dst = node; @@ -2026,7 +2045,6 @@ class vk_perf_logger { " k(" << k->ne[0] << "," << k->ne[1] << "," << k->ne[2] << "," << k->ne[3] << "), " << " v(" << v->ne[0] << "," << v->ne[1] << "," << v->ne[2] << "," << v->ne[3] << "), " << " m(" << (m?m->ne[0]:0) << "," << (m?m->ne[1]:0) << "," << (m?m->ne[2]:0) << "," << (m?m->ne[3]:0) << ")"; - *n_flops = 2ull * q->ne[1] * q->ne[2] * (k->ne[0] + v->ne[0]) * k->ne[1] * q->ne[3]; return name.str(); } if (node->op == GGML_OP_TOP_K) { @@ -2090,7 +2108,7 @@ struct ggml_backend_vk_context { bool do_add_rms_partials_offset_calculation; bool do_add_rms_partials; - uint64_t last_total_mul_mat_bytes {}; + uint64_t last_total_flops {UINT64_MAX}; // Cache most recent tensor that was converted into prealloc_y, and what pipeline it used to convert. vk_pipeline_struct * prealloc_y_last_pipeline_used {}; @@ -16188,22 +16206,23 @@ static ggml_status ggml_backend_vk_graph_compute(ggml_backend_t backend, ggml_cg } // Submit after enough work has accumulated, to overlap CPU cmdbuffer generation with GPU execution. - // Estimate the amount of matmul work by looking at the weight matrix size, and submit every 100MB - // (and scaled down based on model size, so smaller models submit earlier). - int submitted_nodes = 0; - int submit_count = 0; - uint64_t mul_mat_bytes = 0; - uint64_t total_mul_mat_bytes = 0; - uint64_t mul_mat_bytes_per_submit = std::min(uint64_t(100*1000*1000), ctx->last_total_mul_mat_bytes / 40u); + // Estimate the amount of compute work using flops, and submit every 200 GFLOP + // (and scaled down based on total graph flops, so smaller models submit earlier). + // Also submit at least every 100 nodes, in case there are workloads without heavy compute. + uint32_t submitted_nodes = 0; + uint32_t submit_count = 0; + uint64_t batch_flops = 0; + uint64_t total_flops = 0; + uint64_t flops_per_submit = std::min(uint64_t(200'000'000'000), ctx->last_total_flops / 40u); for (int i = 0; i < cgraph->n_nodes; i++) { if (first_node_in_batch) { submit_node_idx = i; } - if (cgraph->nodes[i]->op == GGML_OP_MUL_MAT || cgraph->nodes[i]->op == GGML_OP_MUL_MAT_ID) { - auto bytes = ggml_nbytes(cgraph->nodes[i]->src[0]); - mul_mat_bytes += bytes; - total_mul_mat_bytes += bytes; + { + auto node_flops = ggml_vk_get_node_flops(cgraph->nodes[i]); + batch_flops += node_flops; + total_flops += node_flops; } // op_srcs_fused_elementwise indicates whether an op's srcs all contribute to @@ -16415,8 +16434,8 @@ static ggml_status ggml_backend_vk_graph_compute(ggml_backend_t backend, ggml_cg // Signal the almost_ready fence when the graph is mostly complete (< 20% remaining) bool almost_ready = (cgraph->n_nodes - i) < cgraph->n_nodes / 5; - bool submit = ((uint32_t)submitted_nodes >= ctx->device->max_nodes_per_submit) || - (mul_mat_bytes_per_submit != 0 && mul_mat_bytes >= mul_mat_bytes_per_submit) || + bool submit = (submitted_nodes >= ctx->device->max_nodes_per_submit) || + (flops_per_submit != 0 && batch_flops >= flops_per_submit) || (i + ctx->num_additional_fused_ops >= last_node) || (almost_ready && !ctx->almost_ready_fence_pending); @@ -16450,9 +16469,9 @@ static ggml_status ggml_backend_vk_graph_compute(ggml_backend_t backend, ggml_cg if (submit && enqueued) { first_node_in_batch = true; submitted_nodes = 0; - mul_mat_bytes = 0; + batch_flops = 0; if (submit_count < 3) { - mul_mat_bytes_per_submit *= 2; + flops_per_submit *= 2; } submit_count++; } @@ -16461,7 +16480,7 @@ static ggml_status ggml_backend_vk_graph_compute(ggml_backend_t backend, ggml_cg ctx->fused_ops_write_mask = 0; } - ctx->last_total_mul_mat_bytes = total_mul_mat_bytes; + ctx->last_total_flops = total_flops; if (vk_perf_logger_enabled) { // End the command buffer and submit/wait diff --git a/ggml/src/ggml.c b/ggml/src/ggml.c index b66f9116cac4..c0a03c3b530a 100644 --- a/ggml/src/ggml.c +++ b/ggml/src/ggml.c @@ -1061,6 +1061,7 @@ static const char * GGML_OP_NAME[GGML_OP_COUNT] = { "RWKV_WKV7", "SOLVE_TRI", "GATED_DELTA_NET", + "LIGHTNING_INDEXER", "DSV4_HC_COMB", "DSV4_HC_PRE", "DSV4_HC_POST", @@ -1081,7 +1082,7 @@ static const char * GGML_OP_NAME[GGML_OP_COUNT] = { "GLU", }; -static_assert(GGML_OP_COUNT == 100, "GGML_OP_COUNT != 100"); +static_assert(GGML_OP_COUNT == 101, "GGML_OP_COUNT != 101"); static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = { "none", @@ -1175,6 +1176,7 @@ static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = { "rwkv_wkv7(r, w, k, v, a, b, s)", "A X = B, A triangular, solve X", "gated_delta_net(q, k, v, g, beta, s)", + "lightning_indexer(q, k, weights, scale_embd, scale_heads)", "dsv4_hc_comb(mixes, scale, base)", "dsv4_hc_pre(x, weights)", "dsv4_hc_post(x, residual, post, comb)", @@ -1195,7 +1197,7 @@ static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = { "glu(x)", }; -static_assert(GGML_OP_COUNT == 100, "GGML_OP_COUNT != 100"); +static_assert(GGML_OP_COUNT == 101, "GGML_OP_COUNT != 101"); static_assert(GGML_OP_POOL_COUNT == 2, "GGML_OP_POOL_COUNT != 2"); @@ -6274,6 +6276,40 @@ struct ggml_tensor * ggml_gated_delta_net( return result; } +// ggml_lightning_indexer + +struct ggml_tensor * ggml_lightning_indexer( + struct ggml_context * ctx, + struct ggml_tensor * q, + struct ggml_tensor * k, + struct ggml_tensor * weights, + float scale_embd, + float scale_heads) { + + GGML_ASSERT(q->type == GGML_TYPE_F32); + GGML_ASSERT(weights->type == GGML_TYPE_F32); + GGML_ASSERT(q->ne[0] == k->ne[0]); + GGML_ASSERT(q->ne[1] == weights->ne[0]); + GGML_ASSERT(k->ne[1] == 1); + GGML_ASSERT(q->ne[2] == weights->ne[1]); + GGML_ASSERT(weights->ne[2] == 1); + GGML_ASSERT(q->ne[3] == k->ne[3]); + GGML_ASSERT(k->ne[3] == weights->ne[3]); + + int64_t ne[4] = { k->ne[2], q->ne[2], 1, q->ne[3] }; + struct ggml_tensor * result = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, ne); + + ggml_set_op_params_f32(result, 0, scale_embd); + ggml_set_op_params_f32(result, 1, scale_heads); + + result->op = GGML_OP_LIGHTNING_INDEXER; + result->src[0] = q; + result->src[1] = k; + result->src[2] = weights; + + return result; +} + // ggml_dsv4_hc_comb struct ggml_tensor * ggml_dsv4_hc_comb( diff --git a/models/templates/deepseek-ai-DeepSeek-V4.jinja b/models/templates/deepseek-ai-DeepSeek-V4.jinja index f19f787b1b7e..b2ef0cc50004 100644 --- a/models/templates/deepseek-ai-DeepSeek-V4.jinja +++ b/models/templates/deepseek-ai-DeepSeek-V4.jinja @@ -11,6 +11,7 @@ {%- set dsml_token = '｜DSML｜' -%} {%- set thinking_start_token = '' -%} {%- set thinking_end_token = '' -%} +{%- set has_tools = false -%} {%- 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) -%} @@ -25,6 +26,7 @@ {%- endif -%} {%- endfor -%} {%- if tools is defined and tools -%} + {%- set has_tools = true -%} {%- set ts = namespace(schemas='') -%} {%- for tool in tools -%} {%- if tool['type'] == 'function' -%} @@ -67,7 +69,8 @@ {%- set state.in_user = false -%} {{- '<｜Assistant｜>' -}} {%- set is_after_last_user = loop.index0 > last_user_idx.value -%} - {%- if is_after_last_user and thinking -%} + {%- set preserve_reasoning = thinking and (has_tools or is_after_last_user) -%} + {%- if preserve_reasoning -%} {{- thinking_start_token -}} {%- if message['reasoning_content'] is defined and message['reasoning_content'] -%} {{- message['reasoning_content'] -}} @@ -109,4 +112,4 @@ {%- else -%} {{- thinking_end_token -}} {%- endif -%} -{%- endif -%} \ No newline at end of file +{%- endif -%} diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index 4c86e43c1f74..2e869a76996f 100644 --- a/src/llama-graph.cpp +++ b/src/llama-graph.cpp @@ -63,26 +63,6 @@ static bool can_reuse_kq_mask( // impl -static ggml_tensor * ggml_mul_mat_aux( - ggml_context * ctx, - ggml_tensor * cur, - ggml_tensor * rot) { - const auto n = rot->ne[0]; - - ggml_tensor * res; - - if (!ggml_is_contiguous(cur)) { - res = ggml_cont_2d (ctx, cur, n, ggml_nelements(cur)/n); - } else { - res = ggml_reshape_2d(ctx, cur, n, ggml_nelements(cur)/n); - } - res = ggml_mul_mat (ctx, rot, res); - ggml_mul_mat_set_hint(res, GGML_HINT_SRC0_IS_HADAMARD); - res = ggml_reshape_4d(ctx, res, cur->ne[0], cur->ne[1], cur->ne[2], cur->ne[3]); - - return res; -} - void llm_graph_input_embd::set_input(const llama_ubatch * ubatch) { if (ubatch->token) { const int64_t n_tokens = ubatch->n_tokens; @@ -881,6 +861,14 @@ void llm_graph_input_dsv4::set_input(const llama_ubatch * ubatch) { dsv4_set_comp_inputs(inp_hca, plan_hca, "hca", debug > 0, ubatch->n_tokens, n_stream); dsv4_set_comp_inputs(inp_lid, plan_lid, "lid", debug > 0, ubatch->n_tokens, n_stream); + if (inp_csa.k_rot && inp_csa.k_rot->buffer) { + mctx->get_csa()->set_input_k_rot(inp_csa.k_rot); + } + + if (inp_hca.k_rot && inp_hca.k_rot->buffer) { + mctx->get_hca()->set_input_k_rot(inp_hca.k_rot); + } + if (inp_lid.k_rot && inp_lid.k_rot->buffer) { mctx->get_lid()->set_input_k_rot(inp_lid.k_rot); } @@ -2633,12 +2621,12 @@ ggml_tensor * llm_graph_context::build_attn( GGML_ASSERT(v_mla == nullptr); if (inp->self_k_rot) { - q_cur = ggml_mul_mat_aux(ctx0, q_cur, inp->self_k_rot); - k_cur = ggml_mul_mat_aux(ctx0, k_cur, inp->self_k_rot); + q_cur = llama_mul_mat_hadamard(ctx0, q_cur, inp->self_k_rot); + k_cur = llama_mul_mat_hadamard(ctx0, k_cur, inp->self_k_rot); } if (inp->self_v_rot) { - v_cur = ggml_mul_mat_aux(ctx0, v_cur, inp->self_v_rot); + v_cur = llama_mul_mat_hadamard(ctx0, v_cur, inp->self_v_rot); } // these nodes are added to the graph together so that they are not reordered @@ -2669,7 +2657,7 @@ ggml_tensor * llm_graph_context::build_attn( cb(cur, "kqv_out", il); if (inp->self_v_rot) { - cur = ggml_mul_mat_aux(ctx0, cur, inp->self_v_rot); + cur = llama_mul_mat_hadamard(ctx0, cur, inp->self_v_rot); } if (wo) { @@ -2874,14 +2862,14 @@ ggml_tensor * llm_graph_context::build_attn( auto * v_rot = is_swa ? inp->self_v_rot_swa : inp->self_v_rot; if (k_rot) { - q_cur = ggml_mul_mat_aux(ctx0, q_cur, k_rot); + q_cur = llama_mul_mat_hadamard(ctx0, q_cur, k_rot); if (k_cur) { - k_cur = ggml_mul_mat_aux(ctx0, k_cur, k_rot); + k_cur = llama_mul_mat_hadamard(ctx0, k_cur, k_rot); } } if (v_rot) { if (v_cur) { - v_cur = ggml_mul_mat_aux(ctx0, v_cur, v_rot); + v_cur = llama_mul_mat_hadamard(ctx0, v_cur, v_rot); } } @@ -2924,7 +2912,7 @@ ggml_tensor * llm_graph_context::build_attn( cb(cur, "kqv_out", il); if (v_rot) { - cur = ggml_mul_mat_aux(ctx0, cur, v_rot); + cur = llama_mul_mat_hadamard(ctx0, cur, v_rot); } if (wo) { @@ -3084,6 +3072,8 @@ llm_graph_input_dsv4 * llm_graph_context::build_inp_dsv4() const { dsv4_build_comp_inputs(ctx0, inp->inp_csa, mctx_cur->get_csa_plan(ubatch), "csa", n_stream); dsv4_build_comp_inputs(ctx0, inp->inp_hca, mctx_cur->get_hca_plan(ubatch), "hca", n_stream); dsv4_build_comp_inputs(ctx0, inp->inp_lid, mctx_cur->get_lid_plan(ubatch), "lid", n_stream); + inp->inp_csa.k_rot = mctx_cur->get_csa()->build_input_k_rot(ctx0); + inp->inp_hca.k_rot = mctx_cur->get_hca()->build_input_k_rot(ctx0); inp->inp_lid.k_rot = mctx_cur->get_lid()->build_input_k_rot(ctx0); return (llm_graph_input_dsv4 *) res->add_input(std::move(inp)); diff --git a/src/llama-impl.h b/src/llama-impl.h index 7923c3f7ed55..2d06752c94c9 100644 --- a/src/llama-impl.h +++ b/src/llama-impl.h @@ -54,6 +54,26 @@ static inline dst_t llama_cast(src_t v) { } } +static inline ggml_tensor * llama_mul_mat_hadamard( + ggml_context * ctx, + ggml_tensor * cur, + ggml_tensor * rot) { + const auto n = rot->ne[0]; + + ggml_tensor * res; + + if (!ggml_is_contiguous(cur)) { + res = ggml_cont_2d(ctx, cur, n, ggml_nelements(cur)/n); + } else { + res = ggml_reshape_2d(ctx, cur, n, ggml_nelements(cur)/n); + } + res = ggml_mul_mat(ctx, rot, res); + ggml_mul_mat_set_hint(res, GGML_HINT_SRC0_IS_HADAMARD); + res = ggml_reshape_4d(ctx, res, cur->ne[0], cur->ne[1], cur->ne[2], cur->ne[3]); + + return res; +} + struct time_meas { time_meas(int64_t & t_acc, bool disable = false); ~time_meas(); diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp index 12bf5c37914d..680de5144a86 100644 --- a/src/llama-kv-cache.cpp +++ b/src/llama-kv-cache.cpp @@ -57,22 +57,6 @@ static void ggml_gen_hadamard(ggml_tensor * tensor) { } } -static ggml_tensor * ggml_mul_mat_aux( - ggml_context * ctx, - ggml_tensor * cur, - ggml_tensor * rot) { - const auto n = rot->ne[0]; - - ggml_tensor * res; - - res = ggml_reshape_2d(ctx, cur, n, ggml_nelements(cur)/n); - res = ggml_mul_mat (ctx, rot, res); - ggml_mul_mat_set_hint(res, GGML_HINT_SRC0_IS_HADAMARD); - res = ggml_reshape_4d(ctx, res, cur->ne[0], cur->ne[1], cur->ne[2], cur->ne[3]); - - return res; -} - // // llama_kv_cache // @@ -1875,14 +1859,14 @@ ggml_tensor * llama_kv_cache::build_rope_shift( tmp = ggml_cast(ctx, cur, GGML_TYPE_F32); // rotate back - tmp = ggml_mul_mat_aux(ctx, tmp, rot); + tmp = llama_mul_mat_hadamard(ctx, tmp, rot); tmp = ggml_rope_ext(ctx, tmp, shift, factors, n_rot, rope_type, n_ctx_orig, freq_base, freq_scale, yarn_ext_factor, yarn_attn_factor, yarn_beta_fast, yarn_beta_slow); // rotate fwd - tmp = ggml_mul_mat_aux(ctx, tmp, rot); + tmp = llama_mul_mat_hadamard(ctx, tmp, rot); tmp = ggml_cpy(ctx, tmp, cur); } else { diff --git a/src/models/deepseek4.cpp b/src/models/deepseek4.cpp index 770935d7d64c..40bba8760f7b 100644 --- a/src/models/deepseek4.cpp +++ b/src/models/deepseek4.cpp @@ -618,11 +618,10 @@ ggml_tensor * llama_model_deepseek4::graph::build_lid_top_k( cb(indexer_q_pe, "lid_q_pe", il); indexer_q = ggml_concat(ctx0, indexer_q_nope, indexer_q_pe, 0); - indexer_q = ggml_mul_mat(ctx0, inp_lid.k_rot, indexer_q); + indexer_q = llama_mul_mat_hadamard(ctx0, indexer_q, inp_lid.k_rot); cb(indexer_q, "lid_q_rot", il); ggml_tensor * indexer_weights = build_lora_mm(layer.indexer_proj, cur); - indexer_weights = ggml_scale(ctx0, indexer_weights, 1.0f/sqrtf(float(n_embd_indexer_head*n_indexer_head))); cb(indexer_weights, "lid_weights", il); ggml_tensor * indexer_k = inp_dsv4->mctx->get_lid()->get_k(ctx0, il); @@ -643,6 +642,10 @@ ggml_tensor * llama_model_deepseek4::graph::build_lid_top_k( 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); +#if 1 + ggml_tensor * indexer_score = ggml_lightning_indexer(ctx0, indexer_q, indexer_k, indexer_weights, 1.0f / sqrtf(float(n_embd_indexer_head)), 1.0f / sqrtf(float(n_indexer_head))); + cb(indexer_score, "indexer_score", il); +#else indexer_q = ggml_permute(ctx0, indexer_q, 0, 2, 1, 3); cb(indexer_q, "lid_q", il); indexer_k = ggml_permute(ctx0, indexer_k, 0, 2, 1, 3); @@ -654,12 +657,15 @@ ggml_tensor * llama_model_deepseek4::graph::build_lid_top_k( indexer_kq = ggml_cont(ctx0, ggml_permute(ctx0, indexer_kq, 2, 1, 0, 3)); cb(indexer_kq, "lid_kq", il); + indexer_weights = ggml_scale(ctx0, indexer_weights, 1.0f/sqrtf(float(n_embd_indexer_head*n_indexer_head))); + cb(indexer_weights, "lid_weights", il); + ggml_tensor * indexer_score = ggml_relu(ctx0, indexer_kq); 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)); cb(indexer_score, "lid_score", il); - +#endif indexer_score = ggml_add(ctx0, indexer_score, inp_lid.kq_mask); cb(indexer_score, "lid_score_masked", il); @@ -713,10 +719,15 @@ ggml_tensor * llama_model_deepseek4::graph::build_csa_lid_attention( int il) const { const auto & inp_csa = inp_dsv4->get_csa(); GGML_ASSERT(inp_csa.kq_mask); - GGML_ASSERT(inp_attn->self_k_rot == nullptr); ggml_tensor * top_k = build_lid_top_k(model, inp_dsv4, qr, cur, inp_pos, il); + ggml_tensor * k_rot = inp_attn->self_k_rot; + if (k_rot) { + q = llama_mul_mat_hadamard(ctx0, q, k_rot); + kv = llama_mul_mat_hadamard(ctx0, kv, k_rot); + } + ggml_build_forward_expand(gf, q); ggml_build_forward_expand(gf, kv); @@ -757,6 +768,9 @@ ggml_tensor * llama_model_deepseek4::graph::build_csa_lid_attention( ggml_tensor * kq_b = dsv4_build_kq_zero_bias(ctx0, cparams, kq_mask, q->ne[1]); ggml_tensor * out = build_attn_mha(q, k_all, k_all, kq_b, kq_mask, sinks, nullptr, kq_scale, il); + if (k_rot) { + out = llama_mul_mat_hadamard(ctx0, out, k_rot); + } cb(out, "attn_csa_lid", il); return out; @@ -772,7 +786,12 @@ ggml_tensor * llama_model_deepseek4::graph::build_hca_attention( int il) const { const auto & inp_hca = inp_dsv4->get_hca(); GGML_ASSERT(inp_hca.kq_mask); - GGML_ASSERT(inp_attn->self_k_rot == nullptr); + + ggml_tensor * k_rot = inp_attn->self_k_rot; + if (k_rot) { + q = llama_mul_mat_hadamard(ctx0, q, k_rot); + kv = llama_mul_mat_hadamard(ctx0, kv, k_rot); + } ggml_build_forward_expand(gf, q); ggml_build_forward_expand(gf, kv); @@ -814,6 +833,9 @@ ggml_tensor * llama_model_deepseek4::graph::build_hca_attention( ggml_tensor * kq_b = dsv4_build_kq_zero_bias(ctx0, cparams, kq_mask, q->ne[1]); ggml_tensor * out = build_attn_mha(q, k_all, k_all, kq_b, kq_mask, sinks, nullptr, kq_scale, il); + if (k_rot) { + out = llama_mul_mat_hadamard(ctx0, out, k_rot); + } cb(out, "attn_hca", il); return out; @@ -831,8 +853,8 @@ ggml_tensor * llama_model_deepseek4::graph::build_raw_attention( ggml_tensor * k_rot = inp_attn->self_k_rot; if (k_rot) { - q = ggml_mul_mat(ctx0, k_rot, q); - kv = ggml_mul_mat(ctx0, k_rot, kv); + q = llama_mul_mat_hadamard(ctx0, q, k_rot); + kv = llama_mul_mat_hadamard(ctx0, kv, k_rot); } ggml_build_forward_expand(gf, q); @@ -849,6 +871,9 @@ ggml_tensor * llama_model_deepseek4::graph::build_raw_attention( ggml_tensor * kq_b = dsv4_build_kq_zero_bias(ctx0, cparams, kq_mask, q->ne[1]); ggml_tensor * out = build_attn_mha(q, k, k, kq_b, kq_mask, sinks, nullptr, kq_scale, il); + if (k_rot) { + out = llama_mul_mat_hadamard(ctx0, out, k_rot); + } cb(out, "attn_raw", il); return out; @@ -978,6 +1003,11 @@ ggml_tensor * llama_model_deepseek4::graph::build_attention( "csa_state_compress", il); + if (inp_dsv4->get_csa().k_rot) { + kv_comp_csa_state = llama_mul_mat_hadamard(ctx0, kv_comp_csa_state, inp_dsv4->get_csa().k_rot); + cb(kv_comp_csa_state, "csa_state_compress_rot", il); + } + ggml_build_forward_expand(gf, inp_dsv4->mctx->get_csa()->cpy_k(ctx0, kv_comp_csa_state, inp_dsv4->get_csa().state_write_idxs, il)); @@ -1026,7 +1056,7 @@ ggml_tensor * llama_model_deepseek4::graph::build_attention( il); if (inp_dsv4->get_lid().k_rot) { - kv_comp_lid_state = ggml_mul_mat(ctx0, inp_dsv4->get_lid().k_rot, kv_comp_lid_state); + kv_comp_lid_state = llama_mul_mat_hadamard(ctx0, kv_comp_lid_state, inp_dsv4->get_lid().k_rot); cb(kv_comp_lid_state, "lid_state_compress_rot", il); } @@ -1068,6 +1098,11 @@ ggml_tensor * llama_model_deepseek4::graph::build_attention( "hca_state_compress", il); + if (inp_dsv4->get_hca().k_rot) { + kv_comp_hca = llama_mul_mat_hadamard(ctx0, kv_comp_hca, inp_dsv4->get_hca().k_rot); + cb(kv_comp_hca, "hca_state_compress_rot", il); + } + ggml_build_forward_expand(gf, inp_dsv4->mctx->get_hca()->cpy_k(ctx0, kv_comp_hca, inp_dsv4->get_hca().state_write_idxs, il)); hca_state_dep = kv_comp_hca; @@ -1096,13 +1131,11 @@ ggml_tensor * llama_model_deepseek4::graph::build_attention( if (ratio == DSV4_CSA_RATIO && inp_dsv4->get_csa().kq_mask && inp_dsv4->get_lid().kq_mask && - inp_dsv4->get_lid().k_rot && - inp_attn->self_k_rot == nullptr) { + inp_dsv4->get_lid().k_rot) { out = build_csa_lid_attention(model, inp_dsv4, inp_attn, q, kv, qr, cur, inp_pos, layer.attn_sinks, 1.0f/sqrtf(float(n_embd_head)), il); } else if (ratio == DSV4_HCA_RATIO && - inp_dsv4->get_hca().kq_mask && - inp_attn->self_k_rot == nullptr) { + inp_dsv4->get_hca().kq_mask) { out = build_hca_attention(inp_dsv4, inp_attn, q, kv, layer.attn_sinks, 1.0f/sqrtf(float(n_embd_head)), il); } else { diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index 42123c6fecf1..5eaaa7f64309 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -9427,6 +9427,12 @@ static std::vector> make_test_cases_eval() { } } + for (ggml_type type_a : { GGML_TYPE_Q4_0, GGML_TYPE_Q4_1, GGML_TYPE_Q5_0, GGML_TYPE_Q5_1, GGML_TYPE_Q8_0 }) { + for (int dim : { 0, 1, 2, 3, }) { + test_cases.emplace_back(new test_concat(type_a, {128, 12, 13, 14}, dim == 0 ? 256 : 7, dim, 0)); + } + } + for (ggml_sort_order order : {GGML_SORT_ORDER_ASC, GGML_SORT_ORDER_DESC}) { for (uint32_t i = 4; i <= 1024*1024; i *= 2) { test_cases.emplace_back(new test_argsort(GGML_TYPE_F32, {i-1, 1, 1, 1})); diff --git a/tests/test-chat.cpp b/tests/test-chat.cpp index 5f71e5da6e39..d3ed71c6c93b 100644 --- a/tests/test-chat.cpp +++ b/tests/test-chat.cpp @@ -3963,6 +3963,122 @@ static void test_template_output_peg_parsers(bool detailed_debug) { .run(); } + // DeepSeek V4 tests - format uses DSML markup: + // <｜DSML｜tool_calls> + // <｜DSML｜invoke name="foo"> + // <｜DSML｜parameter name="bar" string="true|false">value + // + // + { + auto tst = peg_tester("models/templates/deepseek-ai-DeepSeek-V4.jinja", detailed_debug); + + tst.test("Hello, world!\nWhat's up?") + .enable_thinking(false) + .reasoning_format(COMMON_REASONING_FORMAT_DEEPSEEK) + .expect(message_assist) + .run(); + + tst.test( + "Let me check the time\n\n" + "<｜DSML｜tool_calls>\n" + "<｜DSML｜invoke name=\"get_time\">\n" + "<｜DSML｜parameter name=\"city\" string=\"true\">Tokyo\n" + "\n" + "") + .enable_thinking(true) + .reasoning_format(COMMON_REASONING_FORMAT_DEEPSEEK) + .tools({ get_time_tool }) + .expect(message_with_tool_calls_and_reasoning("get_time", R"({"city": "Tokyo"})", "Let me check the time")) + .run(); + + tst.test( + "Calling both\n\n" + "<｜DSML｜tool_calls>\n" + "<｜DSML｜invoke name=\"get_time\">\n" + "<｜DSML｜parameter name=\"city\" string=\"true\">Paris\n" + "\n" + "<｜DSML｜invoke name=\"get_weather\">\n" + "<｜DSML｜parameter name=\"city\" string=\"true\">Paris\n" + "\n" + "") + .enable_thinking(true) + .reasoning_format(COMMON_REASONING_FORMAT_DEEPSEEK) + .parallel_tool_calls(true) + .tools({ get_time_tool, get_weather_tool }) + .expect(message_with_reasoning_content_and_multiple_tool_calls( + "Calling both", "", + { { "get_time", R"({"city": "Paris"})" }, { "get_weather", R"({"city": "Paris"})" } })) + .run(); + + tst.test( + "Optional first\n\n" + "<｜DSML｜tool_calls>\n" + "<｜DSML｜invoke name=\"magic_int\">\n" + "<｜DSML｜parameter name=\"name\" string=\"true\">foo bar\n" + "<｜DSML｜parameter name=\"ref\" string=\"false\">42\n" + "\n" + "") + .enable_thinking(true) + .reasoning_format(COMMON_REASONING_FORMAT_DEEPSEEK) + .tools({ magic_int_tool }) + .expect_reasoning("Optional first") + .expect_tool_calls({ + { "magic_int", R"({"name": "foo bar", "ref": 42})", {} }, + }) + .run(); + + tst.test( + "Still thinking\n\n" + "<｜DSML｜tool_calls>\n" + "<｜DSML｜invoke name=\"get_time\">\n" + "<｜DSML｜parameter name=\"city\" string=\"true\">Tokyo\n" + "\n" + "") + .enable_thinking(true) + .reasoning_format(COMMON_REASONING_FORMAT_DEEPSEEK) + .tools({ get_time_tool }) + .expect_reasoning( + "Still thinking\n\n" + "<｜DSML｜tool_calls>\n" + "<｜DSML｜invoke name=\"get_time\">\n" + "<｜DSML｜parameter name=\"city\" string=\"true\">Tokyo\n" + "\n" + "") + .expect_content("") + .expect_tool_calls({}) + .run(); + + { + auto tmpls = read_templates("models/templates/deepseek-ai-DeepSeek-V4.jinja"); + + common_chat_templates_inputs inputs; + inputs.messages = { message_user }; + inputs.reasoning_format = COMMON_REASONING_FORMAT_DEEPSEEK; + inputs.enable_thinking = true; + inputs.tools = { magic_int_tool }; + + make_peg_parser parser(tmpls.get(), inputs, detailed_debug); + + bool got_error = false; + try { + parser.parse( + "Missing required\n\n" + "<｜DSML｜tool_calls>\n" + "<｜DSML｜invoke name=\"magic_int\">\n" + "<｜DSML｜parameter name=\"name\" string=\"true\">foo bar\n" + "\n" + "", + false); + } catch (const std::runtime_error &) { + got_error = true; + } + + if (!got_error) { + throw std::runtime_error("Expected DeepSeek V4 parser to reject missing required parameter"); + } + } + } + // GLM-4.6 tests - format: function_name\n...\n...\n { auto tst = peg_tester("models/templates/GLM-4.6.jinja", detailed_debug); @@ -5812,6 +5928,54 @@ static void test_template_generation_prompt() { check(tmpls, continuation_reasoning(), "<｜Assistant｜>I'm"); } + { + auto tmpls = read_templates("models/templates/deepseek-ai-DeepSeek-V4.jinja"); + check(tmpls, basic(), "<｜Assistant｜>"); + check(tmpls, continuation_content(), "<｜Assistant｜>I'm thinkingHello, "); + check(tmpls, continuation_reasoning(), "<｜Assistant｜>I'm"); + + common_chat_msg user_start; + user_start.role = "user"; + user_start.content = "Check time and weather."; + + common_chat_msg assistant_tools; + assistant_tools.role = "assistant"; + assistant_tools.reasoning_content = "Need both"; + assistant_tools.tool_calls = { + { "get_time", R"({"city":"Paris"})", "call_time" }, + { "get_weather", R"({"city":"Paris"})", "call_weather" }, + }; + + common_chat_msg weather_result; + weather_result.role = "tool"; + weather_result.content = "weather result"; + weather_result.tool_call_id = "call_weather"; + + common_chat_msg time_result; + time_result.role = "tool"; + time_result.content = "time result"; + time_result.tool_call_id = "call_time"; + + common_chat_msg user_continue; + user_continue.role = "user"; + user_continue.content = "Continue."; + + common_chat_templates_inputs inputs; + inputs.messages = { user_start, assistant_tools, weather_result, time_result, user_continue }; + inputs.tools = { get_time_tool, get_weather_tool }; + + auto params = common_chat_templates_apply(tmpls.get(), inputs); + assert_contains(params.prompt, "<｜Assistant｜>Need both\n\n<｜DSML｜tool_calls>"); + + const auto time_pos = params.prompt.find("time result"); + const auto weather_pos = params.prompt.find("weather result"); + if (time_pos == std::string::npos || weather_pos == std::string::npos || time_pos > weather_pos) { + LOG_ERR("Expected tool results in tool-call order\nActual: %s\n", params.prompt.c_str()); + common_log_flush(common_log_main()); + throw std::runtime_error("Test failed"); + } + } + { auto tmpls = read_templates("models/templates/openbmb-MiniCPM5-1B.jinja"); check(tmpls, basic(), "<|im_start|>assistant\n\n");