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  = "<think>";
     const std::string THINK_END    = "</think>";
-    const std::string FC_START     = "<" + DSML + "function_calls>";
-    const std::string FC_END       = "</" + DSML + "function_calls>";
+    const std::string FC_START     = "<" + DSML + tool_calls_tag + ">";
+    const std::string FC_END       = "</" + DSML + tool_calls_tag + ">";
     const std::string INVOKE_START = "<" + DSML + "invoke";
     const std::string INVOKE_END   = "</" + DSML + "invoke>";
     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 <think></think> 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<std::string> 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<size_t>            required_indices;
+                std::vector<common_peg_parser> 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<std::string> required;
-            if (params.contains("required")) {
-                params.at("required").get_to(required);
-            }
+                auto contains_index = [](const std::vector<size_t> & 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<common_peg_parser> required_parsers;
-            std::vector<common_peg_parser> 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<common_peg_parser(const std::vector<size_t> &)> build_args_seq =
+                    [&](const std::vector<size_t> & 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<size_t> 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 <think></think>, while V4 pre-fills </think> 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<bool> used(end - begin, false);
+        std::vector<json> 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<std::string>();
+            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<std::string>() != 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_params> 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 <fstream>
 #include <sstream>
 #include <filesystem>
+#include <regex>
 
 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<std::string> common_preset::to_args(const std::string & bin_path) const {
     std::vector<std::string> 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<T, dim><<<grid_dim, CUDA_CONCAT_BLOCK_SIZE, 0, stream>>>(
@@ -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<uint8_t>(src0, src1, dst, dim, stream);
-            break;
-        case 2:
-            concat_cuda<uint16_t>(src0, src1, dst, dim, stream);
-            break;
-        case 4:
-            concat_cuda<uint32_t>(src0, src1, dst, dim, stream);
-            break;
-        case 8:
-            concat_cuda<uint64_t>(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<uint8_t>(src0, src1, dst, dim, stream);
+    } else {
+        GGML_ASSERT(ggml_blck_size(src0->type) == 1);
+
+        switch (ggml_type_size(src0->type)) {
+            case 1:
+                concat_cuda<uint8_t>(src0, src1, dst, dim, stream);
+                break;
+            case 2:
+                concat_cuda<uint16_t>(src0, src1, dst, dim, stream);
+                break;
+            case 4:
+                concat_cuda<uint32_t>(src0, src1, dst, dim, stream);
+                break;
+            case 8:
+                concat_cuda<uint64_t>(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 <algorithm>
@@ -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 <mma.h>
+namespace wmma = nvcuda::wmma;
+
+template <int64_t n_embd, int64_t n_head, ggml_type type_K>
+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<type_K, half, 4>();
+        #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<wmma::matrix_b, HEADS_PER_INNER_LOOP, K_VECS_PER_BLOCK, K_EMBD_PER_INNER_LOOP, half, wmma::col_major> 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<wmma::accumulator, HEADS_PER_INNER_LOOP, K_VECS_PER_BLOCK, K_EMBD_PER_INNER_LOOP, float> frag_acc;
+        wmma::fill_fragment(frag_acc, 0.0f);
+
+        // load Q fragment
+        wmma::fragment<wmma::matrix_a, HEADS_PER_INNER_LOOP, K_VECS_PER_BLOCK, K_EMBD_PER_INNER_LOOP, half, wmma::row_major> 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 <int64_t n_embd, int64_t n_head, ggml_type type_K>
+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 <int64_t n_embd, int64_t n_head, ggml_type type_K>
+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<type_K, float, 4>();
+        #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<n_embd, n_head, type_K>(        \
+        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<n_embd, n_head, type_K><<<grid, block, 0, ctx.stream()>>>( \
+            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<int> temp_dst_alloc(pool, ncols * nrows);
+    ggml_cuda_pool_alloc<int> 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<int> 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 = '<think>' -%}
 {%- set thinking_end_token = '</think>' -%}
+{%- 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</' + dsml_token + 'parameter>\n...\n</' + dsml_token + 'invoke>\n<' + dsml_token + 'invoke name="$TOOL_NAME2">\n...\n</' + dsml_token + 'invoke>\n</' + dsml_token + 'tool_calls>\n\nString parameters should be specified as is and set `string="true"`. For all other types (numbers, booleans, arrays, objects), pass the value in JSON format and set `string="false"`.\n\nIf thinking_mode is enabled (triggered by ' + thinking_start_token + '), you MUST output your complete reasoning inside ' + thinking_start_token + '...' + thinking_end_token + ' BEFORE any tool calls or final response.\n\nOtherwise, output directly after ' + thinking_end_token + ' with tool calls or final response.\n\n### Available Tool Schemas\n\n' -%}
 {%- set tools_footer = '\nYou MUST strictly follow the above defined tool name and parameter schemas to invoke tool calls.\n' -%}
 {%- set ns = namespace(system_prompt='', is_first_sp=true) -%}
@@ -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<std::unique_ptr<test_case>> 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</｜DSML｜parameter>
+    //   </｜DSML｜invoke>
+    //   </｜DSML｜tool_calls>
+    {
+        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</think>\n\n"
+               "<｜DSML｜tool_calls>\n"
+               "<｜DSML｜invoke name=\"get_time\">\n"
+               "<｜DSML｜parameter name=\"city\" string=\"true\">Tokyo</｜DSML｜parameter>\n"
+               "</｜DSML｜invoke>\n"
+               "</｜DSML｜tool_calls>")
+            .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</think>\n\n"
+               "<｜DSML｜tool_calls>\n"
+               "<｜DSML｜invoke name=\"get_time\">\n"
+               "<｜DSML｜parameter name=\"city\" string=\"true\">Paris</｜DSML｜parameter>\n"
+               "</｜DSML｜invoke>\n"
+               "<｜DSML｜invoke name=\"get_weather\">\n"
+               "<｜DSML｜parameter name=\"city\" string=\"true\">Paris</｜DSML｜parameter>\n"
+               "</｜DSML｜invoke>\n"
+               "</｜DSML｜tool_calls>")
+            .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</think>\n\n"
+               "<｜DSML｜tool_calls>\n"
+               "<｜DSML｜invoke name=\"magic_int\">\n"
+               "<｜DSML｜parameter name=\"name\" string=\"true\">foo bar</｜DSML｜parameter>\n"
+               "<｜DSML｜parameter name=\"ref\" string=\"false\">42</｜DSML｜parameter>\n"
+               "</｜DSML｜invoke>\n"
+               "</｜DSML｜tool_calls>")
+            .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</｜DSML｜parameter>\n"
+               "</｜DSML｜invoke>\n"
+               "</｜DSML｜tool_calls>")
+            .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</｜DSML｜parameter>\n"
+               "</｜DSML｜invoke>\n"
+               "</｜DSML｜tool_calls>")
+            .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</think>\n\n"
+                    "<｜DSML｜tool_calls>\n"
+                    "<｜DSML｜invoke name=\"magic_int\">\n"
+                    "<｜DSML｜parameter name=\"name\" string=\"true\">foo bar</｜DSML｜parameter>\n"
+                    "</｜DSML｜invoke>\n"
+                    "</｜DSML｜tool_calls>",
+                    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: <tool_call>function_name\n<arg_key>...</arg_key>\n<arg_value>...</arg_value>\n</tool_call>
     {
         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｜><think>I'm");
     }
 
+    {
+        auto tmpls = read_templates("models/templates/deepseek-ai-DeepSeek-V4.jinja");
+        check(tmpls, basic(),                  "<｜Assistant｜><think>");
+        check(tmpls, continuation_content(),   "<｜Assistant｜><think>I'm thinking</think>Hello, ");
+        check(tmpls, continuation_reasoning(), "<｜Assistant｜><think>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｜><think>Need both</think>\n\n<｜DSML｜tool_calls>");
+
+        const auto time_pos    = params.prompt.find("<tool_result>time result</tool_result>");
+        const auto weather_pos = params.prompt.find("<tool_result>weather result</tool_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<think>\n");
