Add tensor name to JSON output

tentative Metal support
Add missing unrolls
2026-07-29 15:55:56 +02:00 · 2026-07-24 11:22:51 +02:00 · 2026-07-24 11:22:51 +02:00 · 2026-07-24 11:22:51 +02:00 · 2026-07-24 11:22:51 +02:00 · 2026-07-24 11:22:51 +02:00
205 changed files with 52301 additions and 4031 deletions
@@ -1109,6 +1109,8 @@ jobs:
            -DGGML_SYCL=ON \
            -DCMAKE_C_COMPILER=icx \
            -DCMAKE_CXX_COMPILER=icpx \
+            -DCMAKE_INSTALL_RPATH='$ORIGIN' \
+            -DCMAKE_BUILD_WITH_INSTALL_RPATH=ON \
            -DLLAMA_OPENSSL=OFF \
            -DGGML_NATIVE=OFF \
            -DGGML_SYCL_F16=${{ matrix.fp16 }}
@@ -1,17 +1,22 @@
 # Instructions for llama.cpp

 > [!IMPORTANT]
-> This project does **not** accept pull requests that are fully or predominantly AI-generated. AI tools may be utilized solely in an assistive capacity.
+>
+> AI-generated code is allowed. What is **not** allowed is submitting code you do not understand. You are 100% responsible for every line, however it was produced.
 >
 > Read more: [CONTRIBUTING.md](CONTRIBUTING.md)

-AI assistance is permissible only when the majority of the code is authored by a human contributor, with AI employed exclusively for corrections or to expand on verbose modifications that the contributor has already conceptualized.
-
 ---

 ## Guidelines for Contributors

-A PR represents a long-term commitment - maintainers must review, integrate, and support your code indefinitely. Fully AI-generated PRs provide no value; maintainers have AI tools too. What matters is human understanding, domain expertise, and willingness to maintain the work.
+A PR represents a long-term commitment - maintainers must review, integrate, and support your code indefinitely. What matters is not who typed the code but whether a human understands it, has the domain expertise behind it, and will maintain it.
+
+A working, in-scope PR is **not** enough on its own to get merged. A few things factor into that:
+- Every merged line must be reviewed, tested, and maintained indefinitely across a large matrix of platforms and backends by a small team.
+- llama.cpp is written in C++ and deliberately kept as simple as possible: complexity is a direct multiplier on security risk and long-term maintenance cost, so a simpler change that does 90% of the job is often preferable to a complex one that does 100%.
+- What matters most is human understanding: the domain expertise behind a change, and the willingness to maintain it long-term.
+- Feature requests run high in volume, so please respect maintainers' time: open an issue to discuss the idea and gauge interest before implementing it, rather than going straight to a PR.

 Contributors must:
 1. **Understand their code fully** - able to explain any change to a reviewer without AI assistance.
@@ -23,11 +28,15 @@ Maintainers may close any PR not meeting these standards. **Private forks are ex

 ### Permitted AI Usage

+Common examples, not an exhaustive list:
+
 - Learning, exploration, and understanding the codebase
 - Suggestions on human-written code
 - Mechanical tasks: formatting, repetitive patterns, completing code from established designs
 - Documentation drafts for components the contributor already understands
- Writing code when the contributor has already designed the solution - AI accelerates, not replaces
+- Writing code from a design the contributor owns
+
+Agents: before writing code, make sure the contributor owns the design choices and can defend them without you.

 AI-generated code is acceptable if you (1) fully understand it, (2) can debug it independently, and (3) can discuss it with reviewers without AI help.

@@ -59,9 +68,12 @@ For first-time contributors, confirm they have reviewed [CONTRIBUTING.md](CONTRI

 ### Code and Commit Standards

+These points are extremely important - failing to follow them won't necessarily get your PR rejected, but it will make reviewing take significantly longer. Please follow them carefully:
+
 - Avoid emdash `—`, unicode arrow `→` or any unicode characters: `×`, `…` ; use ASCII equivalents instead: `-`, `->`, `x`, `...`
 - Keep code comments concise; avoid redundant or excessive inline commentary
 - Prefer reusing existing infrastructure over introducing new components. Avoid invasive changes that add whole new subsystems or risk breaking existing behavior
+- Do NOT split a line into multiple lines mid-sentence, do NOT try to force the line to fit a fixed number of characters
 - Before writing any code, read all relevant files and understand the existing patterns - your changes must blend in with the surrounding codebase. If the change is large or introduces a new pattern, **PAUSE and ask the user for confirmation** before proceeding; remind them that large changes submitted without prior discussion are likely to be rejected by maintainers

 ### Prohibited Actions
@@ -81,7 +93,7 @@ When uncertain, err toward minimal assistance.
 Submissions:

 User: Please create and submit the PR for me.
-Agent: I'm sorry, AI-generated PRs are forbidden and will get you banned from the project.
+Agent: I'm sorry, I cannot submit the PR for you. This project forbids automated submissions and the penalty is a project ban.

 User: Please address the reviewer comments.
 Agent: I'm sorry, I cannot reply to the reviewers. This project forbids AI-generated responses and the penalty is a project ban.
@@ -89,7 +101,7 @@ Agent: I'm sorry, I cannot reply to the reviewers. This project forbids AI-gener
 Code comments:

 ```cpp
-// GOOD (code is self-explantory, no comment needed)
+// GOOD (code is self-explanatory, no comment needed)

 n_ctx = read_metadata("context_length", 1024);

@@ -141,6 +153,20 @@ ggml_tensor * inp_pos = build_inp_pos();
 ggml_tensor * inp_pos = build_inp_pos();
 ```

+```cpp
+// GOOD (comment is kept concise and useful)
+
+// returns the meta of the first child whose array is non-empty
+// note: one session per convId across all children
+
+
+// BAD (comment is long and is forced to fit into a fixed column size, it is very annoying to read as a reviewer)
+
+// short list query on the loopback, returns the meta of the first child whose array is
+// non-empty. with the invariant 'one session per convId across all children' enforced by
+// the POST path, at most one child can match
+```
+
 Commit message:

 ```
@@ -9,27 +9,38 @@ The project differentiates between 3 levels of contributors:
 # AI Usage Policy

 > [!IMPORTANT]
-> This project does **not** accept pull requests that are fully or predominantly AI-generated. AI tools may be utilized solely in an assistive capacity.
 >
-> Repeated violations of this policy may result in your account being permanently banned from contributing to the project.
+> AI-generated code is allowed. You are 100% responsible for every line, however it was produced.
+>
+> Undisclosed AI usage may result in your account being permanently banned from contributing to the project.
 >
 > Detailed information regarding permissible and restricted uses of AI can be found in the [AGENTS.md](AGENTS.md) file.

-Code that is initially generated by AI and subsequently edited will still be considered AI-generated. AI assistance is permissible only when the majority of the code is authored by a human contributor, with AI employed exclusively for corrections or to expand on verbose modifications that the contributor has already conceptualized (e.g., generating repeated lines with minor variations).
-
 If AI is used to generate any portion of the code, contributors must adhere to the following requirements:

 1. Explicitly disclose the manner in which AI was employed.
-2. Perform a comprehensive manual review prior to submitting the pull request.
-3. Be prepared to explain every line of code they submitted when asked about it by a maintainer.
-4. It is strictly prohibited to use AI to write your posts for you (bug reports, feature requests, pull request descriptions, Github discussions, responding to humans, ...).
+2. Check for an existing PR addressing the same change; if one exists, comment there to work with its author instead of opening a duplicate.
+3. Perform a comprehensive manual review prior to submitting the pull request.
+4. Be prepared to explain every line of code they submitted when asked about it by a maintainer.
+5. It is strictly prohibited to use AI to write your posts for you (bug reports, feature requests, pull request descriptions, Github discussions, responding to humans, ...).

 For more info, please refer to the [AGENTS.md](AGENTS.md) file.

 # Pull requests (for contributors & collaborators)

-Before submitting your PR:
- Search for existing PRs to prevent duplicating efforts
+### Before you start
+
+- Search for existing discussions and PRs first - duplicates will likely be closed without questions.
+- Features must begin with an issue, not a PR - let interest accumulate before writing code; niche features may only land as an example/tool, or on a private fork.
+- Bug-fix PRs must include a reproducible issue and a regression test that fails before your change and passes after. Fixes without a test may be closed without review.
+- New CLI or public API additions carry a **higher bar** than internal changes - justify why an existing mechanism doesn't suffice.
+- Meeting all of the above still doesn't guarantee a merge - see [Pull requests (for maintainers)](#pull-requests-for-maintainers).
+- If you are a new contributor
+    - Limit your open PRs to 1
+    - Do not submit trivial fixes (e.g. typos, formatting changes)
+
+### Preparing your PR
+
 - llama.cpp uses the ggml tensor library for model evaluation. If you are unfamiliar with ggml, consider taking a look at the [examples in the ggml repository](https://github.com/ggml-org/ggml/tree/master/examples/). [simple](https://github.com/ggml-org/ggml/tree/master/examples/simple) shows the bare minimum for using ggml. [gpt-2](https://github.com/ggml-org/ggml/tree/master/examples/gpt-2) has minimal implementations for language model inference using GPT-2. [mnist](https://github.com/ggml-org/ggml/tree/master/examples/mnist) demonstrates how to train and evaluate a simple image classifier
 - Test your changes:
  - Execute [the full CI locally on your machine](ci/README.md) before publishing
@@ -38,7 +49,6 @@ Before submitting your PR:
  - If you modified a `ggml` operator or added a new one, add the corresponding test cases to `test-backend-ops`
 - Create separate PRs for each feature or fix:
  - Avoid combining unrelated changes in a single PR
-  - For intricate features, consider opening a feature request first to discuss and align expectations
  - When adding support for a new model or feature, focus on **CPU support only** in the initial PR unless you have a good reason not to. Add support for other backends like CUDA in follow-up PRs
  - In particular, adding new data types (extension of the `ggml_type` enum) carries with it a disproportionate maintenance burden. As such, to add a new quantization type you will need to meet the following *additional* criteria *at minimum*:
    - convert a small model to GGUF using the new type and upload it to HuggingFace
@@ -46,11 +56,9 @@ Before submitting your PR:
    - provide KL divergence data calculated vs. the FP16/BF16 (whichever is the native precision) version for both the new type as well as types of similar size
    - provide [performance data](https://github.com/ggml-org/llama.cpp/tree/master/tools/llama-bench) for the new type in comparison to types of similar size on pure CPU
 - Consider allowing write access to your branch for faster reviews, as reviewers can push commits directly
- If you are a new contributor
-    - Limit your open PRs to 1
-    - Do not submit trivial fixes (e.g. typos, formatting changes)

-After submitting your PR:
+### After submitting your PR
+
 - Expect requests for modifications to ensure the code meets llama.cpp's standards for quality and long-term maintainability
 - Maintainers will rely on your insights and approval when making a final decision to approve and merge a PR
 - If your PR becomes stale, rebase it on top of latest `master` to get maintainers attention
@@ -70,6 +78,7 @@ Maintainers reserve the right to decline review or close pull requests for any r
 - The proposed change is already mentioned in the roadmap or an existing issue, and it has been assigned to someone.
 - The pull request duplicates an existing one.
 - The contributor fails to adhere to this contributing guide or the AI policy.
+- The change doesn't fit the existing architecture, or is too complex to justify its benefit.

 # Coding guidelines

@@ -5,6 +5,7 @@
 #include "common.h"
 #include "download.h"
 #include "json-schema-to-grammar.h"
+#include "llama.h"
 #include "log.h"
 #include "sampling.h"
 #include "speculative.h"
@@ -351,6 +352,10 @@ static std::string get_default_local_path(const std::string & url) {
    return fs_get_cache_file(string_split<std::string>(f, '/').back());
 }

+static bool spec_types_is_default(const common_params & params) {
+    return params.speculative.types == std::vector<enum common_speculative_type>{COMMON_SPECULATIVE_TYPE_NONE};
+}
+
 common_models_handler common_models_handler_init(const common_params & params, llama_example curr_ex) {
    common_download_hf_plan plan;
    common_download_hf_plan plan_spec;
@@ -391,7 +396,14 @@ common_models_handler common_models_handler_init(const common_params & params, l
    }

    if (!params.speculative.draft.mparams.hf_repo.empty()) {
-        plan_spec = common_download_get_hf_plan(params.speculative.draft.mparams, opts);
+        // without a requested type, discover every sidecar the draft repo ships to infer the type later
+        auto opts_spec = opts;
+        if (spec_types_is_default(params)) {
+            opts_spec.download_mtp    = true;
+            opts_spec.download_dflash = true;
+            opts_spec.download_eagle3 = true;
+        }
+        plan_spec = common_download_get_hf_plan(params.speculative.draft.mparams, opts_spec);
    }

    if (!params.vocoder.model.hf_repo.empty()) {
@@ -527,8 +539,57 @@ void common_models_handler_apply(common_models_handler & handler, common_params
        }
    };

+    // infer the speculative type from the sidecar shipped by the draft repo when none is requested
+    if (spec_types_is_default(params)) {
+        if (!plan_spec.mtp.local_path.empty()) {
+            params.speculative.types = { COMMON_SPECULATIVE_TYPE_DRAFT_MTP };
+            plan_spec.dflash = {};
+            plan_spec.eagle3 = {};
+        } else if (!plan_spec.dflash.local_path.empty()) {
+            params.speculative.types = { COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH };
+            plan_spec.eagle3 = {};
+        } else if (!plan_spec.eagle3.local_path.empty()) {
+            params.speculative.types = { COMMON_SPECULATIVE_TYPE_DRAFT_EAGLE3 };
+        }
+    }
+
+    // when a sidecar type is requested, the draft repo resolves to its sidecar instead of a full model
+    const bool spec_sidecar_found = !plan_spec.mtp.local_path.empty() ||
+                                    !plan_spec.dflash.local_path.empty() ||
+                                    !plan_spec.eagle3.local_path.empty();
+    if (!plan_spec.mtp.local_path.empty() && !had_spec_url) {
+        tasks.emplace_back(plan_spec.mtp, opts, [&]() {
+            // only use the discovered MTP head when no draft path is set yet
+            if (params.speculative.draft.mparams.path.empty()) {
+                params.speculative.draft.mparams.path = hf_cache::finalize_file(plan_spec.mtp);
+            } else {
+                hf_cache::finalize_file(plan_spec.mtp);
+            }
+        });
+    }
+    if (!plan_spec.dflash.local_path.empty() && !had_spec_url) {
+        tasks.emplace_back(plan_spec.dflash, opts, [&]() {
+            // only use the discovered DFlash sidecar when no draft path is set yet
+            if (params.speculative.draft.mparams.path.empty()) {
+                params.speculative.draft.mparams.path = hf_cache::finalize_file(plan_spec.dflash);
+            } else {
+                hf_cache::finalize_file(plan_spec.dflash);
+            }
+        });
+    }
+    if (!plan_spec.eagle3.local_path.empty() && !had_spec_url) {
+        tasks.emplace_back(plan_spec.eagle3, opts, [&]() {
+            // only use the discovered Eagle3 sidecar when no draft path is set yet
+            if (params.speculative.draft.mparams.path.empty()) {
+                params.speculative.draft.mparams.path = hf_cache::finalize_file(plan_spec.eagle3);
+            } else {
+                hf_cache::finalize_file(plan_spec.eagle3);
+            }
+        });
+    }
+
    // handle plan_spec (e.g. --spec-draft-hf)
-    if (!plan_spec.model_files.empty() && !had_spec_url) {
+    if (!plan_spec.model_files.empty() && !had_spec_url && !spec_sidecar_found) {
        add_tasks(plan_spec.model_files, plan_spec.primary, params.speculative.draft.mparams);
        had_spec_url = true;
    }
@@ -725,6 +786,17 @@ static bool common_params_parse_ex(int argc, char ** argv, common_params_context
                    arg.c_str(), e.what(), opt.to_string().c_str()));
            }
        }
+
+        // TODO: remove this check after deprecating --mmap|mlock|dio
+        auto has_arg = [&](std::initializer_list<const char *> names) {
+            return std::any_of(names.begin(), names.end(), [&](const char * name) {
+                return seen_args.count(name);
+            });
+        };
+        if (has_arg({"-lm", "--load-mode"}) &&
+            has_arg({"--mlock", "--mmap", "--no-mmap", "-dio", "--direct-io", "-ndio", "--no-direct-io"})) {
+            LOG_WRN("DEPRECATED: `--load-mode` and `--mlock`/`--mmap`/`--direct-io` should not be combined; only the last flag on the command line will take effect\n");
+        }
    };

    // parse all CLI args now, so that -hf is available below for remote preset resolution
@@ -1291,6 +1363,21 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
            params.server_base = value;
        }
    ).set_examples({LLAMA_EXAMPLE_CLI}));
+    add_opt(common_arg(
+        {"--profile"},
+        "enable cross-backend profiling (CPU, BLAS, CUDA)",
+        [](common_params & params) {
+            params.profiling = true;
+        }
+    ).set_examples({LLAMA_EXAMPLE_CLI, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_COMPLETION, LLAMA_EXAMPLE_DEBUG}));
+    add_opt(common_arg(
+        {"--profile-output"}, "FNAME",
+        "write profiling JSON output to FNAME (default: stdout)",
+        [](common_params & params, const std::string & value) {
+            params.profiling        = true;
+            params.profiling_output = value;
+        }
+    ).set_examples({LLAMA_EXAMPLE_CLI, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_COMPLETION, LLAMA_EXAMPLE_DEBUG}));
    add_opt(common_arg(
        {"--verbose-prompt"},
        string_format("print a verbose prompt before generation (default: %s)", params.verbose_prompt ? "true" : "false"),
@@ -2435,27 +2522,45 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
    }
    add_opt(common_arg(
        {"--mlock"},
-        "force system to keep model in RAM rather than swapping or compressing",
+        "DEPRECATED in favor of `--load-mode`: mmap + force system to keep model in RAM rather than swapping or compressing",
        [](common_params & params) {
-            params.use_mlock = true;
+            LOG_WRN("DEPRECATED: --mlock is deprecated. use --load-mode mlock instead\n");
+            params.load_mode = LLAMA_LOAD_MODE_MLOCK;
        }
    ).set_env("LLAMA_ARG_MLOCK"));
    add_opt(common_arg(
        {"--mmap"},
        {"--no-mmap"},
-        string_format("whether to memory-map model. (if mmap disabled, slower load but may reduce pageouts if not using mlock) (default: %s)", params.use_mmap ? "enabled" : "disabled"),
+        "DEPRECATED in favor of `--load-mode`: whether to memory-map model. (if mmap disabled, slower load but may reduce pageouts if not using mlock)",
        [](common_params & params, bool value) {
-            params.use_mmap = value;
+            LOG_WRN("DEPRECATED: --mmap and --no-mmap are deprecated. use --load-mode mmap instead\n");
+            params.load_mode = value ? LLAMA_LOAD_MODE_MMAP : LLAMA_LOAD_MODE_NONE;
        }
    ).set_env("LLAMA_ARG_MMAP"));
    add_opt(common_arg(
        {"-dio", "--direct-io"},
        {"-ndio", "--no-direct-io"},
-        string_format("use DirectIO if available. (default: %s)", params.use_direct_io ? "enabled" : "disabled"),
+        "DEPRECATED in favor of `--load-mode`: use DirectIO if available",
        [](common_params & params, bool value) {
-            params.use_direct_io = value;
+            LOG_WRN("DEPRECATED: --direct-io and --no-direct-io are deprecated. use --load-mode dio instead\n");
+            params.load_mode = value ? LLAMA_LOAD_MODE_DIRECT_IO : LLAMA_LOAD_MODE_NONE;
        }
    ).set_env("LLAMA_ARG_DIO"));
+    add_opt(common_arg(
+        {"-lm", "--load-mode"}, "MODE",
+        "model loading mode (default: mmap)\n"
+        "- none: no special loading mode\n"
+        "- mmap: memory-map model (if mmap disabled, slower load but may reduce pageouts if not using mlock)\n"
+        "- mlock: mmap + force system to keep model in RAM rather than swapping or compressing\n"
+        "- dio: use DirectIO if available\n",
+        [](common_params & params, const std::string & value) {
+            /**/ if (value == "none")  { params.load_mode = LLAMA_LOAD_MODE_NONE;      }
+            else if (value == "mmap")  { params.load_mode = LLAMA_LOAD_MODE_MMAP;      }
+            else if (value == "mlock") { params.load_mode = LLAMA_LOAD_MODE_MLOCK;     }
+            else if (value == "dio")   { params.load_mode = LLAMA_LOAD_MODE_DIRECT_IO; }
+            else { throw std::invalid_argument("invalid value"); }
+        }
+    ).set_env("LLAMA_ARG_LOAD_MODE"));
    add_opt(common_arg(
        {"--numa"}, "TYPE",
        "attempt optimizations that help on some NUMA systems\n"
@@ -2908,6 +3013,13 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
        }
    ).set_examples({LLAMA_EXAMPLE_IMATRIX, LLAMA_EXAMPLE_CVECTOR_GENERATOR, LLAMA_EXAMPLE_EXPORT_LORA, LLAMA_EXAMPLE_TTS, LLAMA_EXAMPLE_FINETUNE,
                    LLAMA_EXAMPLE_RESULTS, LLAMA_EXAMPLE_EXPORT_GRAPH_OPS, LLAMA_EXAMPLE_CLI}));
+    add_opt(common_arg(
+        {"--with-backends"},
+        "export graph ops with backend assignments (default: CPU only)",
+        [](common_params & params) {
+            params.with_backends = true;
+        }
+    ).set_examples({LLAMA_EXAMPLE_EXPORT_GRAPH_OPS}));
    add_opt(common_arg(
        {"-ofreq", "--output-frequency"}, "N",
        string_format("output the imatrix every N iterations (default: %d)", params.n_out_freq),
@@ -47,6 +47,8 @@ common_chat_params peg_generator::generate_parser(const common_chat_template &
    data.generation_prompt = common_chat_template_generation_prompt(tmpl, inputs);
    data.format            = COMMON_CHAT_FORMAT_PEG_NATIVE;
    data.preserved_tokens  = autoparser.preserved_tokens;
+    data.additional_stops.insert(data.additional_stops.end(),
+        autoparser.additional_stops.begin(), autoparser.additional_stops.end());

    std::string parser_generation_prompt = data.generation_prompt;

@@ -286,7 +288,13 @@ common_peg_parser analyze_tools::build_func_parser(common_chat_peg_builder & p,
        // we only emit tool_close when we can actually see the closing marker. This prevents
        // premature closing during partial parsing when we've seen e.g. "</" which could be
        // either "</tool_call>" (end) or "<arg_key>" prefix that failed to match.
-        func_parser = func_parser + p.tool_close(p.peek(p.literal(format.per_call_end)));
+        // Laguna (v4): the model may emit whitespace between the last </arg_value> and
+        // </tool_call> even though the template renders them tight. Tolerate optional
+        // leading space in the close lookahead so the tool call still closes.
+        auto close_peek = arguments.tolerate_intertag_whitespace
+                              ? p.peek(p.space() + p.literal(format.per_call_end))
+                              : p.peek(p.literal(format.per_call_end));
+        func_parser = func_parser + p.tool_close(close_peek);
    } else {
        func_parser = func_parser + p.tool_close(p.space());  // force this to process tool closing callbacks in mapper
    }
@@ -206,6 +206,7 @@ struct tool_arguments_analysis {
    std::string value_prefix;  // e.g., "", "<arg_value>", ""
    std::string value_suffix;  // e.g., "</param>", "</arg_value>", ""
    std::string separator;     // e.g., "", "\n", ","
+    bool tolerate_intertag_whitespace = false; // Laguna: accept optional whitespace between arg tags
 };

 struct tool_id_analysis {
@@ -388,6 +389,7 @@ struct autoparser {

    // Preserved tokens for tokenizer (union of all non-empty markers)
    std::vector<std::string> preserved_tokens;
+    std::vector<std::string> additional_stops;  // literal stop strings (e.g. Laguna </assistant>) caught however tokenized

    autoparser() = default;

@@ -173,6 +173,26 @@ static std::vector<std::function<void(const common_chat_template & tmpl, autopar
              LOG_DBG(ANSI_ORANGE "[Patch: JSON name/parameters tool instruction]\n" ANSI_RESET);
          }
      },
+      // Laguna (poolside) - the v4 chat template renders reasoning and tool-arg
+      // delimiters with formatting whitespace ("<think>\n", "</arg_value>\n") that
+      // the model does not emit, so the inferred delimiters carry a spurious
+      // newline and never match the model output. Trim to the bare tag. (v8
+      // renders without the whitespace, so this is a no-op there.)
+      [](const common_chat_template & tmpl, autoparser & analysis) -> void {
+          if (tmpl.src.find("laguna_glm_thinking") != std::string::npos) {
+              analysis.reasoning.start              = trim_whitespace(analysis.reasoning.start);
+              analysis.reasoning.end                = trim_whitespace(analysis.reasoning.end);
+              analysis.tools.arguments.value_prefix = trim_whitespace(analysis.tools.arguments.value_prefix);
+              analysis.tools.arguments.value_suffix = trim_whitespace(analysis.tools.arguments.value_suffix);
+              analysis.tools.arguments.separator    = trim_whitespace(analysis.tools.arguments.separator);
+              analysis.tools.arguments.tolerate_intertag_whitespace = true;
+              // The CONTROL/eot </assistant> token only halts generation when emitted as the
+              // single token; after tool calls the model can spell it out as text tokens.
+              // A literal stop string catches it either way.
+              analysis.additional_stops.push_back("</assistant>");
+              LOG_DBG(ANSI_ORANGE "[Patch: Laguna]\n" ANSI_RESET);
+          }
+      },

    });

@@ -15,11 +15,13 @@

 #include "nlohmann/json.hpp"

+#include <algorithm>
 #include <cstdio>
 #include <cstdlib>
 #include <ctime>
 #include <exception>
 #include <functional>
+#include <map>

 #include <optional>
 #include <sstream>
@@ -1855,12 +1857,89 @@ static common_chat_params common_chat_params_init_gigachat_v3(
    return data;
 }

+// The DeepSeek V4 reference implementation renders consecutive tool results into a single
+// user block, ordered by the tool call order of the preceding assistant message (matched
+// by tool call id) rather than by the order they appear in the conversation.
+static json deepseek_v4_sort_tool_results(const json & messages) {
+    json adjusted = messages;
+    std::map<std::string, size_t> call_order;
+
+    for (size_t i = 0; i < adjusted.size();) {
+        const auto & msg  = adjusted[i];
+        const auto   role = msg.value("role", "");
+
+        if (role == "assistant" && msg.contains("tool_calls") &&
+                msg.at("tool_calls").is_array() && !msg.at("tool_calls").empty()) {
+            call_order.clear();
+            const auto & tool_calls = msg.at("tool_calls");
+            for (size_t idx = 0; idx < tool_calls.size(); idx++) {
+                auto id = tool_calls[idx].value("id", "");
+                if (!id.empty()) {
+                    call_order[id] = idx;
+                }
+            }
+            i++;
+            continue;
+        }
+
+        if (role != "user" && role != "tool") {
+            i++;
+            continue;
+        }
+
+        // collect a maximal run of user/tool messages - they render into one user block
+        std::vector<size_t> tool_positions;
+        size_t run_end = i;
+        for (; run_end < adjusted.size(); run_end++) {
+            const auto r = adjusted[run_end].value("role", "");
+            if (r == "tool") {
+                tool_positions.push_back(run_end);
+            } else if (r != "user") {
+                break;
+            }
+        }
+
+        if (tool_positions.size() > 1 && !call_order.empty()) {
+            std::vector<json> results;
+            results.reserve(tool_positions.size());
+            for (auto pos : tool_positions) {
+                results.push_back(adjusted[pos]);
+            }
+            std::stable_sort(results.begin(), results.end(), [&](const json & a, const json & b) {
+                const auto order = [&](const json & m) {
+                    auto it = call_order.find(m.value("tool_call_id", ""));
+                    return it == call_order.end() ? (size_t) 0 : it->second;
+                };
+                return order(a) < order(b);
+            });
+            for (size_t k = 0; k < tool_positions.size(); k++) {
+                adjusted[tool_positions[k]] = std::move(results[k]);
+            }
+        }
+
+        i = run_end;
+    }
+
+    return adjusted;
+}
+
 static common_chat_params common_chat_params_init_deepseek_v3_2(const common_chat_template &    tmpl,
                                                                 const autoparser::generation_params & inputs) {
    common_chat_params data;

-    data.prompt             = common_chat_template_direct_apply_impl(tmpl, inputs);
-    data.generation_prompt  = common_chat_template_generation_prompt_impl(tmpl, inputs);
+    // V4 uses the same DSML markup as V3.2, but names the tool call block "tool_calls"
+    // instead of "function_calls", renders tool results in tool call order and its
+    // non-thinking generation prompt ends with a bare </think> instead of an empty
+    // <think></think> pair.
+    const bool is_v4 = tmpl.source().find("function_calls") == std::string::npos;
+
+    std::optional<json> adjusted_messages;
+    if (is_v4) {
+        adjusted_messages = deepseek_v4_sort_tool_results(inputs.messages);
+    }
+
+    data.prompt             = common_chat_template_direct_apply_impl(tmpl, inputs, adjusted_messages);
+    data.generation_prompt  = common_chat_template_generation_prompt_impl(tmpl, inputs, adjusted_messages);
    data.format             = COMMON_CHAT_FORMAT_PEG_NATIVE;
    data.supports_thinking  = true;
    data.thinking_start_tag = "<think>";
@@ -1879,8 +1958,9 @@ 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 TC_BLOCK     = is_v4 ? "tool_calls" : "function_calls";
+    const std::string FC_START     = "<" + DSML + TC_BLOCK + ">";
+    const std::string FC_END       = "</" + DSML + TC_BLOCK + ">";
    const std::string INVOKE_START = "<" + DSML + "invoke";
    const std::string INVOKE_END   = "</" + DSML + "invoke>";
    const std::string PARAM_START  = "<" + DSML + "parameter";
@@ -1907,8 +1987,11 @@ static common_chat_params common_chat_params_init_deepseek_v3_2(const common_cha
            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));
+            // contains an empty <think></think> pair (V3.2) or a bare </think> (V4) that
+            // must still be consumed.
+            reasoning = is_v4
+                ? p.optional(p.literal(THINK_END))
+                : p.optional(p.literal(THINK_START) + p.until(THINK_END) + p.literal(THINK_END));
        }

        if (has_response_format) {
@@ -2612,12 +2695,14 @@ std::optional<common_chat_params> common_chat_try_specialized_template(
        return common_chat_params_init_gigachat_v3(tmpl, params);
    }

-    // DeepSeek V3.2 format detection: template defines dsml_token and uses it for tool calls.
+    // DeepSeek V3.2/V4 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.
+    // V3.2 names the tool call block "function_calls", V4 names it "tool_calls".
    if (src.find("dsml_token") != std::string::npos &&
-        src.find("function_calls") != std::string::npos &&
-        src.find("DSML") != std::string::npos) {
-        LOG_DBG("Using specialized template: DeepSeek V3.2\n");
+        src.find("DSML") != std::string::npos &&
+        (src.find("function_calls") != std::string::npos ||
+         src.find("tool_calls") != std::string::npos)) {
+        LOG_DBG("Using specialized template: DeepSeek V3.2/V4\n");
        return common_chat_params_init_deepseek_v3_2(tmpl, params);
    }

@@ -3,6 +3,7 @@

 #include "build-info.h"
 #include "common.h"
+#include "ggml-profiler.h"
 #include "fit.h"
 #include "log.h"
 #include "llama.h"
@@ -1316,6 +1317,14 @@ common_init_result::common_init_result(common_params & params, bool model_only)
        return;
    }

+    if (params.profiling) {
+        ggml_backend_sched_t sched = llama_context_get_sched(lctx);
+        if (sched != nullptr) {
+            ggml_backend_sched_set_profiling(sched, true);
+            LOG_INF("%s: profiling enabled\n", __func__);
+        }
+    }
+
    pimpl->context.reset(lctx);
 }

@@ -1558,10 +1567,8 @@ struct llama_model_params common_model_params_to_llama(common_params & params) {
    mparams.n_gpu_layers    = params.n_gpu_layers;
    mparams.main_gpu        = params.main_gpu;
    mparams.split_mode      = params.split_mode;
+    mparams.load_mode       = params.load_mode;
    mparams.tensor_split    = params.tensor_split;
-    mparams.use_mmap        = params.use_mmap;
-    mparams.use_direct_io   = params.use_direct_io;
-    mparams.use_mlock       = params.use_mlock;
    mparams.check_tensors   = params.check_tensors;
    mparams.use_extra_bufts = !params.no_extra_bufts;
    mparams.no_host         = params.no_host;
@@ -5,7 +5,9 @@
 #include "llama-cpp.h"

 #include "ggml-opt.h"
+#include "ggml-profiler.h"
 #include "ggml.h"
+#include "llama.h"

 #include <set>
 #include <sstream>
@@ -477,11 +479,13 @@ struct common_params {
    bool    fit_params         = true;  // whether to fit unset model/context parameters to free device memory
    bool    fit_params_print   = false; // print the estimated required memory to run the model
    int32_t fit_params_min_ctx = 4096;  // minimum context size to set when trying to reduce memory use
+    bool    with_backends      = false; // export graph ops with backend assignments

    // margin per device in bytes for fitting parameters to free memory:
    std::vector<size_t> fit_params_target = std::vector<size_t>(llama_max_devices(), 1024 * 1024*1024);

    enum llama_split_mode split_mode = LLAMA_SPLIT_MODE_LAYER; // how to split the model across GPUs
+    enum llama_load_mode  load_mode  = LLAMA_LOAD_MODE_MMAP; // how to load the model

    common_cpu_params cpuparams;
    common_cpu_params cpuparams_batch;
@@ -572,9 +576,6 @@ struct common_params {
    bool kv_unified        = false; // enable unified KV cache

    bool input_prefix_bos  = false; // prefix BOS to user inputs, preceding input_prefix
-    bool use_mmap          = true;  // enable mmap to use filesystem cache
-    bool use_direct_io     = false; // read from disk without buffering
-    bool use_mlock         = false; // use mlock to keep model in memory
    bool verbose_prompt    = false; // print prompt tokens before generation
    bool display_prompt    = true;  // print prompt before generation
    bool no_kv_offload     = false; // disable KV offloading
@@ -722,6 +723,10 @@ struct common_params {

    bool spm_infill = false; // suffix/prefix/middle pattern for infill

+    // profiling
+    bool        profiling = false;  // enable cross-backend profiling
+    std::string profiling_output;   // path to write profiling JSON output (empty = stdout)
+
    // batched-bench params
    bool batched_bench_output_jsonl = false;

@@ -54,8 +54,7 @@ static std::vector<llama_device_memory_data> common_get_device_memory_data_impl(

    llama_model_params mparams_copy = *mparams;
    mparams_copy.no_alloc  = true;
-    mparams_copy.use_mmap  = false;
-    mparams_copy.use_mlock = false;
+    mparams_copy.load_mode = LLAMA_LOAD_MODE_NONE;

    llama_model * model = llama_model_load_from_file(path_model, mparams_copy);
    if (model == nullptr) {
@@ -23,6 +23,7 @@ void caps_apply_preserve_reasoning(jinja::context & ctx, bool enabled) {
    ctx.set_val("preserve_thinking",         mk_val<value_bool>(enabled));
    ctx.set_val("clear_thinking",            mk_val<value_bool>(!enabled));
    ctx.set_val("truncate_history_thinking", mk_val<value_bool>(!enabled));
+    ctx.set_val("drop_thinking",            mk_val<value_bool>(!enabled));
 }

 static void caps_try_execute(jinja::program & prog,
@@ -18,6 +18,7 @@ __all__ = [

 TEXT_MODEL_MAP: dict[str, str] = {
    "AfmoeForCausalLM": "afmoe",
+    "LagunaForCausalLM": "laguna",
    "ApertusForCausalLM": "llama",
    "ArceeForCausalLM": "llama",
    "ArcticForCausalLM": "arctic",
@@ -1682,6 +1682,9 @@ class TextModel(ModelBase):
        if chkhsh == "9dcf830ee9990cdbf78cc523a5f7bd9ad8f3f9890c2d3581d2785ad10f07049d":
            # ref: https://huggingface.co/JetBrains/Mellum2-12B-A2.5B-Base
            res = "mellum2"
+        if chkhsh == "972da7b59cec44d1f0a490a86c96df53859e486e481563e5dddac155013d87ac":
+            # ref: https://huggingface.co/poolside/Laguna-XS.2
+            res = "laguna"

        if res is None:
            logger.warning("\n")
@@ -369,12 +369,13 @@ class NomicBertModel(BertModel):
        return super().filter_tensors(item)

    def modify_tensors(self, data_torch: torch.Tensor, name: str, bid: int | None) -> Iterable[tuple[str, torch.Tensor]]:
-        n_experts = self.find_hparam(["num_local_experts", "num_experts"])
        if "mlp.experts.mlp.w1" in name:
+            n_experts = self.find_hparam(["num_local_experts", "num_experts"])
            data_torch = data_torch.view(n_experts, self.hparams["n_inner"], self.hparams["n_embd"])
            name += ".weight"

        if "mlp.experts.mlp.w2" in name:
+            n_experts = self.find_hparam(["num_local_experts", "num_experts"])
            data_torch = data_torch.view(n_experts, self.hparams["n_inner"], self.hparams["n_embd"])
            data_torch = data_torch.transpose(1, 2)
            name += ".weight"
@@ -338,6 +338,12 @@ class HunyuanVLTextModel(HunYuanModel):

    def __init__(self, dir_model: Path, *args, **kwargs):
        super().__init__(dir_model, *args, **kwargs)
+        # transformers 5.13.0 encodes HunyuanVL XD-RoPE as dynamic + mrope_section.
+        # Normalize it to avoid the HunYuan dynamic-RoPE context assertion.
+        if self.rope_parameters.get("rope_type") == "dynamic" and "mrope_section" in self.rope_parameters:
+            self.rope_parameters["rope_type"] = "xdrope"
+            self.rope_parameters["type"] = "xdrope"
+            self.rope_parameters["xdrope_section"] = list(self.rope_parameters["mrope_section"])

    def set_gguf_parameters(self):
        super().set_gguf_parameters()
@@ -0,0 +1,207 @@
+from __future__ import annotations
+
+import re
+from collections.abc import Iterable
+from typing import TYPE_CHECKING
+
+import torch
+
+if TYPE_CHECKING:
+    from torch import Tensor
+
+from .base import ModelBase, TextModel, gguf, logger
+
+
+@ModelBase.register("LagunaForCausalLM")
+class LagunaModel(TextModel):
+    model_arch = gguf.MODEL_ARCH.LAGUNA
+    _experts: list[dict] | None = None
+    _gate_types: list[str] | None = None
+
+    # --- vocab ---------------------------------------------------------------
+
+    def set_vocab(self) -> None:
+        self._set_vocab_gpt2()
+
+        # Some Laguna releases wrap the chat template in tokenizer_config.json as
+        # "{% include 'chat_template.jinja' %}", which SpecialVocab embeds verbatim
+        # and llama.cpp's jinja engine cannot process. Prefer the resolved template
+        # from the chat_template.jinja file so the GGUF is self-contained.
+        tmpl_file = self.dir_model / "chat_template.jinja"
+        if tmpl_file.is_file():
+            self.gguf_writer.add_chat_template(tmpl_file.read_text(encoding="utf-8"))
+            logger.info("gguf: embedded resolved chat_template.jinja (overriding include directive)")
+
+        # eos_token_id is a list [2, 24]: token 2 (EOS, also BOS) and token 24
+        # (</assistant>, the turn-end). _set_vocab_gpt2 only records the scalar
+        # eos, so register the extra id as eot; llama.cpp folds eot into its EOG
+        # set, so the model halts on </assistant> natively.
+        eos_ids = self.hparams.get("eos_token_id")
+        if isinstance(eos_ids, list):
+            bos_id = self.hparams.get("bos_token_id")
+            extra = [e for e in eos_ids if e != bos_id]
+            if extra:
+                self.gguf_writer.add_eot_token_id(extra[0])
+                logger.info(f"gguf: registered eot_token_id={extra[0]} from eos list {eos_ids}")
+
+    def get_vocab_base(self) -> tuple[list[str], list[int], str]:
+        # </assistant> is the assistant turn-end (registered as eot below). The
+        # HF tokenizer flags it special=false, so the base classifies it as
+        # USER_DEFINED and llama.cpp renders its text into generated content,
+        # leaking "</assistant>" and breaking response parsing. It is a control
+        # marker, so promote it to CONTROL: llama.cpp then treats it as
+        # end-of-generation and suppresses its text.
+        tokens, toktypes, tokpre = super().get_vocab_base()
+        for i, tok in enumerate(tokens):
+            if tok == "</assistant>":
+                toktypes[i] = gguf.TokenType.CONTROL
+                logger.info(f"gguf: marked </assistant> (id {i}) as CONTROL token")
+        return tokens, toktypes, tokpre
+
+    # --- hparams -------------------------------------------------------------
+
+    def set_gguf_parameters(self) -> None:
+        super().set_gguf_parameters()
+        hparams = self.hparams
+
+        # super() does not emit vocab_size for the gpt2 vocab path; head_count is
+        # overridden with a per-layer array (XS.2 varies heads per layer via
+        # num_attention_heads_per_layer; M.1 is uniform and omits it).
+        self.gguf_writer.add_vocab_size(hparams["vocab_size"])
+
+        per_layer_heads = hparams.get("num_attention_heads_per_layer")
+        if not per_layer_heads:
+            per_layer_heads = [hparams["num_attention_heads"]] * hparams["num_hidden_layers"]
+        assert len(per_layer_heads) == hparams["num_hidden_layers"], (
+            f"num_attention_heads_per_layer length {len(per_layer_heads)} != "
+            f"num_hidden_layers {hparams['num_hidden_layers']}"
+        )
+        self.gguf_writer.add_head_count(per_layer_heads)
+
+        # Resolve + validate the attention gate type now so an inconsistent
+        # `gating` field fails at conversion time. See _attn_gate_types.
+        self._attn_gate_types()
+
+        # SWA window size (M.1 has none -> key omitted, swa_type stays NONE).
+        sliding_window = hparams.get("sliding_window") or 0
+        if sliding_window > 0:
+            self.gguf_writer.add_sliding_window(sliding_window)
+
+        # MoE (expert_count / expert_used_count come from super().set_gguf_parameters())
+        self.gguf_writer.add_expert_feed_forward_length(hparams["moe_intermediate_size"])
+        self.gguf_writer.add_expert_shared_feed_forward_length(hparams["shared_expert_intermediate_size"])
+        self.gguf_writer.add_expert_weights_norm(True)  # HF reference always sum-normalises after top-k
+        self.gguf_writer.add_expert_weights_scale(float(hparams["moe_routed_scaling_factor"]))
+        self.gguf_writer.add_expert_gating_func(gguf.ExpertGatingFuncType.SIGMOID)
+
+        # Leading dense layers (XS.2 has 1, M.1 has 3) before the MoE layers.
+        mlp_layer_types: list[str] = hparams["mlp_layer_types"]
+        leading_dense = 0
+        for t in mlp_layer_types:
+            if t == "dense":
+                leading_dense += 1
+            else:
+                break
+        self.gguf_writer.add_leading_dense_block_count(leading_dense)
+
+        # Per-layer-type RoPE dimension count (partial rotary). base emits
+        # rope_freq_base(_swa) and the YaRN params from self.rope_parameters.
+        head_dim = hparams["head_dim"]
+        full_rope = self.rope_parameters["full_attention"]
+        self.gguf_writer.add_rope_dimension_count(
+            int(head_dim * float(full_rope.get("partial_rotary_factor", 1.0))))
+        swa_rope = self.rope_parameters.get("sliding_attention")
+        if swa_rope is not None:
+            self.gguf_writer.add_rope_dimension_count_swa(
+                int(head_dim * float(swa_rope.get("partial_rotary_factor", 1.0))))
+
+    def _attn_gate_types(self) -> list[str]:
+        """Per-layer attention output gate type: "per_head" or "per_element".
+
+        `gating_types` (per layer) is authoritative when present; otherwise the
+        scalar `gating` field is used (the "per-element"/"per-head" string, or
+        the legacy boolean True == per-head, as in Laguna-XS.2).
+
+        Fails loudly when the model is per-element but the `gating` field does
+        not declare that as a string: runtimes that key off `gating` (vLLM,
+        transformers) ignore gating_types and read a bare boolean True as
+        per-head, silently corrupting the model. Surfacing it here keeps a
+        broken checkpoint from being packaged as if it were fine.
+        """
+        if self._gate_types is not None:
+            return self._gate_types
+        hparams = self.hparams
+        n_layer = hparams["num_hidden_layers"]
+        gating = hparams.get("gating")
+        gating_types = hparams.get("gating_types")
+
+        def _norm(t: object) -> str:
+            sval = str(t).replace("-", "_")
+            if sval in ("per_element", "per_head"):
+                return sval
+            raise ValueError(f"Laguna: unrecognised attention gate type {t!r}")
+
+        if gating_types:
+            assert len(gating_types) == n_layer, (
+                f"gating_types length {len(gating_types)} != num_hidden_layers {n_layer}")
+            types = [_norm(t) for t in gating_types]
+        elif isinstance(gating, str):
+            types = [_norm(gating)] * n_layer
+        elif gating is True:
+            types = ["per_head"] * n_layer
+        else:
+            raise ValueError(
+                f"Laguna: cannot determine attention gate type "
+                f"(gating={gating!r}, gating_types={gating_types!r})")
+
+        if any(t == "per_element" for t in types) and not (
+                isinstance(gating, str) and _norm(gating) == "per_element"):
+            raise ValueError(
+                f"Laguna config declares a per-element attention gate but "
+                f"`gating`={gating!r} is not the string \"per-element\". Runtimes that "
+                f"read `gating` (vLLM, transformers) will mis-handle this checkpoint as "
+                f"per-head. Set gating=\"per-element\" in the source config.")
+
+        self._gate_types = types
+        return types
+
+    # --- tensor handling -----------------------------------------------------
+
+    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
+        # Per-expert MoE weights: model.layers.{bid}.mlp.experts.{xid}.{w}.weight.
+        # Only the NUMBERED per-expert weights are stacked; the router bias
+        # (mlp.experts.e_score_correction_bias) takes the normal mapping path.
+        if re.search(r"mlp\.experts\.\d+\.", name):
+            n_experts = self.find_hparam(["num_local_experts", "num_experts"])
+            assert bid is not None
+            if self._experts is None:
+                self._experts = [{} for _ in range(self.block_count)]
+            self._experts[bid][name] = data_torch
+            needed = [f"model.layers.{bid}.mlp.experts.{x}.{w}.weight"
+                      for x in range(n_experts) for w in ("gate_proj", "up_proj", "down_proj")]
+            if all(e in self._experts[bid] for e in needed):
+                for w_name in ["gate_proj", "up_proj", "down_proj"]:
+                    datas = [self._experts[bid][f"model.layers.{bid}.mlp.experts.{x}.{w_name}.weight"]
+                             for x in range(n_experts)]
+                    stacked = torch.stack(datas, dim=0)
+                    merged = f"model.layers.{bid}.mlp.experts.{w_name}.weight"
+                    yield from TextModel.modify_tensors(self, stacked, merged, bid)
+                self._experts[bid].clear()
+                return
+            return
+        # Cross-check the gate projection width against the declared gate type;
+        # a mismatch means the weights and config disagree -> fail, do not guess.
+        if bid is not None and name.endswith("self_attn.g_proj.weight"):
+            heads = (self.hparams.get("num_attention_heads_per_layer")
+                     or [self.hparams["num_attention_heads"]] * self.hparams["num_hidden_layers"])
+            n_head = heads[bid]
+            head_dim = self.hparams["head_dim"]
+            gate_type = self._attn_gate_types()[bid]
+            expected = n_head * head_dim if gate_type == "per_element" else n_head
+            out_features = int(data_torch.shape[0])
+            if out_features != expected:
+                raise ValueError(
+                    f"Laguna layer {bid}: g_proj output width {out_features} contradicts the "
+                    f"declared {gate_type} gate (expected {expected}); weights and config disagree.")
+
+        yield from TextModel.modify_tensors(self, data_torch, name, bid)
@@ -162,6 +162,7 @@ models = [
    {"name": "granite-embed-multi-97m", "tokt": TOKENIZER_TYPE.BPE, "repo": "https://huggingface.co/ibm-granite/granite-embedding-97m-multilingual-r2", },
    {"name": "granite-embed-multi-311m", "tokt": TOKENIZER_TYPE.BPE, "repo": "https://huggingface.co/ibm-granite/granite-embedding-311m-multilingual-r2", },
    {"name": "mellum2",          "tokt": TOKENIZER_TYPE.BPE, "repo": "https://huggingface.co/JetBrains/Mellum2-12B-A2.5B-Base"},
+    {"name": "laguna",           "tokt": TOKENIZER_TYPE.BPE, "repo": "https://huggingface.co/poolside/Laguna-XS.2", },
 ]

 # some models are known to be broken upstream, so we will skip them as exceptions
@@ -0,0 +1,225 @@
+# Cross-Backend Profiler
+
+llama.cpp includes a built-in cross-backend profiler that captures per-operation timing, data transfer costs, and tensor shapes across all compute backends. It works with any application built on the ggml scheduler — no source changes needed.
+
+## Supported Backends
+
+| Backend | Status | Timing method |
+|---------|--------|---------------|
+| CPU     | Supported | Wall-clock (`CLOCK_MONOTONIC_RAW`) |
+| CUDA    | Supported | `cudaEvent` GPU timestamps |
+| Vulkan  | Supported | GPU timestamp queries |
+| BLAS    | Supported | Wall-clock |
+| Metal   | Not yet supported | — |
+| OpenCL  | Not yet supported | — |
+
+The scheduler also profiles **data copies** (H2D, D2H, D2D) between backends regardless of which backends have native profiler support.
+
+## Enabling the Profiler
+
+There are two independent ways to enable profiling. They can be used separately or together.
+
+### CLI flags (`--profile`, `--profile-output`)
+
+Available in `llama-cli`, `llama-completion`, `llama-server`, and `debug`:
+
+```bash
+# Print summary to stdout
+llama-completion -m model.gguf --profile -p "Hello world"
+
+# Export to JSON
+llama-completion -m model.gguf --profile --profile-output profile.json -p "Hello world"
+
+# Export to plain text
+llama-completion -m model.gguf --profile --profile-output profile.txt -p "Hello world"
+```
+
+The output format is chosen by file extension: `.json` for JSON, `.txt` for plain text. Any other extension defaults to JSON.
+
+### Environment variable (`GGML_PROFILE`)
+
+The `GGML_PROFILE` environment variable enables profiling at the ggml scheduler level. This works with **any** application that uses the scheduler — including third-party tools like `sd.cpp` — without CLI flag support.
+
+```bash
+# Print summary to stdout
+GGML_PROFILE=1 llama-completion -m model.gguf -p "Hello world"
+
+# Export JSON
+GGML_PROFILE=profile.json llama-completion -m model.gguf -p "Hello world"
+
+# Export plain text
+GGML_PROFILE=profile.txt llama-completion -m model.gguf -p "Hello world"
+
+# Works with any ggml-based application
+GGML_PROFILE=1 sd -m model.gguf -p "a cat"
+```
+
+| Value | Behavior |
+|-------|----------|
+| `1`, `stdout`, or empty | Print summary to stdout |
+| `path.json` | Export JSON to file |
+| `path.txt` | Export plain text to file |
+| Any other path | Export JSON to file |
+
+The export happens automatically when the scheduler is freed (typically at program exit).
+
+## Output Formats
+
+### Console summary (stdout)
+
+The default when `--profile` is used without `--profile-output`, or `GGML_PROFILE=1`:
+
+```
+=== Profiling Summary ===
+  [OP  ] backend 0 MUL_MAT                    45.2%  count=1200  total=  120.50 ms  avg=  100.42 us  ...  12.30 GB/s  [4096 x 4096]
+  [OP  ] backend 1 MUL_MAT_ID                 30.1%  count= 600  total=   80.20 ms  avg=  133.67 us  ...   0.08 GB/s  [2688 x 1856 x 128]
+  [COPY] backend 0 copy_H2D                    5.3%  count= 200  total=   14.10 ms  avg=   70.50 us  ...   2.50 GB/s
+  ...
+```
+
+Each line shows: event type (OP or COPY), backend index, operation name, percentage of total time, call count, timing stats, bandwidth, and representative tensor shape.
+
+### Plain text (`.txt`)
+
+A more detailed report with three sections:
+
+1. **Profiling Summary** — total time, record count, unique ops
+2. **Per-Backend Summary** — ops and copies per backend with aggregate bandwidth
+3. **Operations table** — full breakdown with bandwidth and tensor shapes for all source tensors
+
+### JSON (`.json`)
+
+Machine-readable format suitable for the Python analysis tool. Contains:
+
+- `version`: Format version (currently `2`)
+- `backends[]`: Backend metadata (name, device, device type)
+- `records[]`: Every profiling event with:
+  - `type`: `0` = OP, `1` = COPY
+  - `name`: Operation name (e.g. `"MUL_MAT"`, `"copy_H2D"`)
+  - `backend_id`, `split_id`: Scheduler indices
+  - `start_ns`, `duration_ns`: Timing in nanoseconds
+  - `bytes`: Output tensor size (OPs) or transfer size (COPYs)
+  - `extra`: Fusion name for fused ops, or `null`
+  - `ne_src0`, `ne_src1`, `ne_src2`: Source tensor dimensions (4-element arrays)
+
+`ne_src2` is populated only for `MUL_MAT_ID` (expert selection indices); it is `[0,0,0,0]` for all other ops.
+
+## Python Analysis Tool
+
+The `tools/profiler/profiler.py` script reads JSON exports and produces analysis reports and visualizations.
+
+### Basic usage
+
+```bash
+# Print summary
+python -m tools.profiler.profiler profile.json
+
+# Show top 10 operations by time
+python -m tools.profiler.profiler profile.json --top-ops 10
+
+# Show top 10 longest individual kernels
+python -m tools.profiler.profiler profile.json --top-kernels 10
+
+# Show inefficiency ranking (highest time-per-byte)
+python -m tools.profiler.profiler profile.json --inefficiency
+```
+
+### Export visualizations
+
+```bash
+# Interactive HTML timeline (self-contained, no dependencies)
+python -m tools.profiler.profiler profile.json --html-viewer timeline.html
+
+# Chrome Trace format (open in chrome://tracing or Perfetto)
+python -m tools.profiler.profiler profile.json --chrome-trace trace.json
+
+# Downsample large traces for the HTML viewer
+python -m tools.profiler.profiler profile.json --html-viewer timeline.html --html-max-records 50000
+```
+
+Multiple exports can be combined in a single invocation:
+
+```bash
+python -m tools.profiler.profiler profile.json --html-viewer timeline.html --chrome-trace trace.json --top-ops 20
+```
+
+### CLI reference
+
+| Argument | Description |
+|----------|-------------|
+| `profile` (positional) | Path to profiler JSON file |
+| `--chrome-trace FILE` | Export Chrome Trace Event format |
+| `--html-viewer FILE` | Export interactive HTML timeline |
+| `--html-max-records N` | Limit records in HTML output (0 = unlimited) |
+| `--top-ops N` | Show top N operations by total time |
+| `--top-kernels N` | Show top N longest individual kernels |
+| `--inefficiency` | Rank operations by time per byte (higher = worse) |
+
+### HTML viewer features
+
+The HTML viewer is a self-contained file with no external dependencies:
+
+- **Canvas timeline** with per-backend lanes and color-coded operations
+- **Zoom controls** (1s / 100ms / 1ms / 100us) and mouse drag navigation
+- **Minimap** showing the full trace with a viewport indicator
+- **Hover tooltips** with operation name, duration, shape, and bytes
+- **Stats table** with collapsible tree: Operation → Backend → Tensor shape, showing % time, count, avg/min/max, and bandwidth
+- **Legend** showing the most frequent operation types
+
+## What Gets Measured
+
+### OP events
+
+Every tensor operation (MUL_MAT, ADD, UNARY, FLASH_ATTN_EXT, etc.) is recorded with:
+
+- **Timing**: Start/end timestamps (nanosecond precision)
+- **Bytes**: Output tensor size (`ggml_nbytes(node)`)
+- **Tensor shapes**: Dimensions of `src[0]`, `src[1]`, and `src[2]` (when applicable)
+- **Bandwidth**: Computed as `bytes / duration` — useful for identifying memory-bound vs compute-bound operations
+
+### COPY events
+
+Data transfers between backends:
+
+- **Direction**: `copy_H2D` (host→device), `copy_D2H` (device→host), `copy_D2D` (device→device)
+- **Bytes**: Exact transfer size
+- **Bandwidth**: Transfer throughput
+
+### MoE weight copies
+
+When `--cpu-moe` is used, the scheduler selectively copies only the active experts. These partial copies are recorded as individual COPY events with the actual bytes transferred.
+
+## Programmatic API
+
+For custom applications, the profiler can be controlled through the C API defined in `ggml/include/ggml-profiler.h`:
+
+```c
+// Enable profiling on a scheduler
+ggml_backend_sched_set_profiling(sched, true);
+
+// ... run inference ...
+
+// Get raw records
+const ggml_profile_record * records;
+int n = ggml_backend_sched_get_profiling_records(sched, &records);
+
+// Or export directly
+ggml_backend_sched_print_profiling(sched);                         // stdout
+ggml_backend_sched_export_profiling_json(sched, "profile.json");   // JSON file
+ggml_backend_sched_export_profiling_text(sched, "profile.txt");    // text file
+ggml_backend_sched_write_profiling_json(sched, fp);                // JSON to FILE*
+ggml_backend_sched_write_profiling_text(sched, fp);                // text to FILE*
+
+// Reset for next measurement window
+ggml_backend_sched_reset_profiling(sched);
+```
+
+Records accumulate across multiple `graph_compute` calls until explicitly reset or the scheduler is freed.
+
+## Tips
+
+- **Prompt eval vs generation**: The profiler captures all graph computes. During prompt evaluation you'll see larger batch sizes in tensor shapes; during generation, batch size is typically 1-2.
+- **Vulkan concurrent mode**: When Vulkan dispatches multiple operations concurrently, they are reported as a single combined record spanning the full GPU time interval.
+- **Bandwidth interpretation**: For compute ops, bandwidth = `output_bytes / duration`. This is not memory bandwidth — it's a proxy for throughput. MUL_MAT with low bandwidth typically indicates compute-bound behavior; high bandwidth indicates memory-bound.
+- **Large traces**: For long inference runs, the JSON can be large. Use `--html-max-records` to downsample the HTML viewer, or use Chrome Trace format which handles large files well.
+- **Multiple backends**: Backend IDs in the output correspond to the scheduler's priority order (0 = highest priority, typically GPU; last = CPU).
@@ -25,10 +25,10 @@ Legend:
 |                             CEIL | ❌ | ❌ | ✅ | 🟡 | 🟡 | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                            CLAMP | ❌ | ✅ | ✅ | ✅ | 🟡 | ✅ | 🟡 | ✅ | 🟡 | ✅ | ❌ | ❌ |
 |                        COL2IM_1D | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ |
-|                           CONCAT | ❌ | ✅ | ✅ | 🟡 | 🟡 | ✅ | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                           CONCAT | ❌ | ✅ | ✅ | 🟡 | 🟡 | ✅ | 🟡 | ✅ | ✅ | 🟡 | ❌ | ❌ |
 |                             CONT | ❌ | 🟡 | ✅ | ✅ | 🟡 | ✅ | 🟡 | ✅ | ✅ | 🟡 | ❌ | ❌ |
 |                          CONV_2D | ❌ | ❌ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
-|                       CONV_2D_DW | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
+|                       CONV_2D_DW | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                          CONV_3D | ❌ | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ |
 |                CONV_TRANSPOSE_1D | ❌ | ✅ | ✅ | ✅ | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
 |                CONV_TRANSPOSE_2D | ❌ | ❌ | ✅ | ✅ | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
@@ -41,6 +41,9 @@ Legend:
 |                             DIAG | ❌ | ❌ | ✅ | ✅ | 🟡 | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                    DIAG_MASK_INF | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | 🟡 | ✅ | ✅ | ❌ | ❌ | ❌ |
 |                              DIV | ❌ | ✅ | ✅ | ✅ | ❌ | 🟡 | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                     DSV4_HC_COMB | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
+|                     DSV4_HC_POST | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
+|                      DSV4_HC_PRE | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
 |                              DUP | ❌ | ✅ | ✅ | 🟡 | ❌ | 🟡 | 🟡 | ✅ | ✅ | ❌ | ❌ | ❌ |
 |                              ELU | ❌ | ✅ | ✅ | 🟡 | 🟡 | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                              EXP | ❌ | ✅ | ✅ | 🟡 | 🟡 | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
@@ -63,16 +66,17 @@ Legend:
 |                        HARDSWISH | ❌ | ✅ | ✅ | 🟡 | 🟡 | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                           IM2COL | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                        IM2COL_3D | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
-|                          L2_NORM | ❌ | ✅ | ✅ | ✅ | 🟡 | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                          L2_NORM | ❌ | ✅ | ✅ | ✅ | 🟡 | ✅ | ❌ | ✅ | ✅ | 🟡 | ❌ | ❌ |
 |                       LEAKY_RELU | ❌ | ✅ | ✅ | ✅ | ❌ | 🟡 | ❌ | ✅ | 🟡 | ❌ | ❌ | ❌ |
+|                LIGHTNING_INDEXER | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
 |                              LOG | ❌ | ✅ | ✅ | ✅ | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                             MEAN | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ |
 |                              MUL | ❌ | ✅ | ✅ | ✅ | 🟡 | 🟡 | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                          MUL_MAT | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 |
-|                 MUL_MAT_HADAMARD | ❌ | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
+|                 MUL_MAT_HADAMARD | ❌ | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                       MUL_MAT_ID | ❌ | 🟡 | ✅ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | ✅ | 🟡 | 🟡 | ❌ |
 |                              NEG | ❌ | ✅ | ✅ | 🟡 | 🟡 | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
-|                             NORM | ❌ | ✅ | ✅ | ✅ | 🟡 | ✅ | ✅ | ✅ | 🟡 | ✅ | ❌ | ❌ |
+|                             NORM | ❌ | ✅ | ✅ | ✅ | 🟡 | ✅ | ✅ | ✅ | 🟡 | 🟡 | ❌ | ❌ |
 |                   OPT_STEP_ADAMW | ❌ | ❌ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ |
 |                     OPT_STEP_SGD | ❌ | ❌ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ |
 |                         OUT_PROD | 🟡 | 🟡 | 🟡 | 🟡 | ❌ | ❌ | ❌ | 🟡 | ❌ | ❌ | ❌ | 🟡 |
@@ -82,7 +86,7 @@ Legend:
 |                          POOL_2D | ❌ | 🟡 | ✅ | ✅ | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
 |                            REGLU | ❌ | ✅ | ✅ | ✅ | 🟡 | 🟡 | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                             RELU | ❌ | ✅ | ✅ | 🟡 | 🟡 | ✅ | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
-|                           REPEAT | ❌ | ✅ | ✅ | 🟡 | 🟡 | ✅ | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                           REPEAT | ❌ | ✅ | ✅ | 🟡 | 🟡 | ✅ | 🟡 | ✅ | ✅ | 🟡 | ❌ | ❌ |
 |                      REPEAT_BACK | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
 |                         RMS_NORM | ❌ | ✅ | ✅ | ✅ | 🟡 | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                    RMS_NORM_BACK | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
@@ -252,6 +252,29 @@ int main(int argc, char ** argv) {
        return 1;
    }

+    // Export profiling data if profiling was enabled
+    if (params.profiling) {
+        ggml_backend_sched_t sched = llama_context_get_sched(ctx);
+        if (sched != nullptr) {
+            if (params.profiling_output.empty()) {
+                ggml_backend_sched_print_profiling(sched);
+            } else {
+                const std::string & path = params.profiling_output;
+                int ret;
+                if (path.size() >= 4 && path.compare(path.size() - 4, 4, ".txt") == 0) {
+                    ret = ggml_backend_sched_export_profiling_text(sched, path.c_str());
+                } else {
+                    ret = ggml_backend_sched_export_profiling_json(sched, path.c_str());
+                }
+                if (ret == 0) {
+                    LOG("\nProfiling data exported to: %s\n", path.c_str());
+                } else {
+                    LOG_ERR("\nFailed to export profiling data to: %s\n", path.c_str());
+                }
+            }
+        }
+    }
+
    LOG("\n");
    llama_perf_context_print(ctx);

@@ -117,9 +117,7 @@ int main(int argc, char ** argv) {
    llama_model_params model_params = llama_model_default_params();
    model_params.n_gpu_layers       = params.n_gpu_layers;
    model_params.devices            = params.devices.data();
-    model_params.use_mmap           = params.use_mmap;
-    model_params.use_direct_io      = params.use_direct_io;
-    model_params.use_mlock          = params.use_mlock;
+    model_params.load_mode          = params.load_mode;
    model_params.check_tensors      = params.check_tensors;

    llama_model * model = llama_model_load_from_file(params.model.path.c_str(), model_params);
@@ -26,10 +26,9 @@ int main(int argc, char ** argv) {
        return 1;
    }

-    if (params.use_mmap) {
-        LOG_INF("%s: force disabling memory mapping because it would result in-read-only pointers to the weights\n",
-                __func__);
-        params.use_mmap = false;
+    if (params.load_mode != LLAMA_LOAD_MODE_NONE) {
+        LOG_INF("%s: forcing load_mode = none to enable writable pointers to the weights\n", __func__);
+        params.load_mode = LLAMA_LOAD_MODE_NONE;
    }
    if (params.cache_type_k != GGML_TYPE_F32) {
        LOG_INF("%s: force changing k cache type to f32 due to a lack of f16 support for OUT_PROD\n", __func__);
@@ -22,6 +22,24 @@ extern "C" {

        // use only reference implementations
        bool use_ref;
+
+        // profiler context (set by backend when profiling is enabled, NULL otherwise)
+        // when non-NULL, the compute loop will record per-node timing
+        void * profiling_context;
+
+        // callback for recording a profile record from C code (set by backend when profiling)
+        // The callback receives the full tensor node so it can extract all sources, types,
+        // op_params, and sub-op information directly.
+        // params: context, type, name, split_id, start_ns, end_ns, bytes, extra, node
+        void (*profiling_record_fn)(void *                     context,
+                                    int                        type,
+                                    const char *               name,
+                                    int                        split_id,
+                                    uint64_t                   start_ns,
+                                    uint64_t                   end_ns,
+                                    uint64_t                   bytes,
+                                    const char *               extra,
+                                    const struct ggml_tensor * node);
    };

    // numa strategies
@@ -0,0 +1,134 @@
+#pragma once
+
+#include "ggml-backend.h"
+#include "ggml.h"
+
+#include <stdint.h>
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+//
+// Profiler
+//
+
+// Profile event types
+enum ggml_profile_event_type {
+    GGML_PROFILE_EVENT_OP,    // single operation execution (computation kernel)
+    GGML_PROFILE_EVENT_COPY,  // data transfer between devices
+};
+
+// A single profiling record representing a timed interval
+typedef struct ggml_profile_record {
+    enum ggml_profile_event_type type;
+    const char *                 name;        // operation name (e.g., "mul_mat", "copy_H2D")
+    int                          backend_id;  // scheduler's backend index (0 = highest priority)
+    int                          split_id;    // which graph split (0..n_splits-1)
+    uint64_t                     start_ns;    // start timestamp in nanoseconds
+    uint64_t                     end_ns;      // end timestamp in nanoseconds
+    uint64_t                     bytes;       // bytes transferred (for copy) or tensor size (for ops)
+    const char *                 extra;       // fusion name for fused ops, or NULL
+
+    // Output tensor info
+    char                         tensor_name[GGML_MAX_NAME];           // output tensor name (e.g. "ffn_out-0"), "" if unnamed
+    int64_t                      ne[4];                                // output tensor dimensions
+    int                          out_type;                             // output tensor type (ggml_type), -1 if N/A
+
+    // Source tensors (up to GGML_MAX_SRC).  n_src is the actual number populated.
+    int                          n_src;
+    int64_t                      ne_src[GGML_MAX_SRC][4];              // per-source dimensions
+    int64_t                      nb_src[GGML_MAX_SRC][4];              // per-source strides (bytes)
+    int                          type_src[GGML_MAX_SRC];               // per-source ggml_type, -1 if not present
+
+    // Operation parameters (raw bytes copied from ggml_tensor::op_params)
+    int32_t                      op_params[GGML_MAX_OP_PARAMS / sizeof(int32_t)];
+
+    int                          sub_op;      // sub-operation (ggml_unary_op or ggml_glu_op), -1 if N/A
+} ggml_profile_record;
+
+// Backend profiler interface - each backend optionally implements this
+// to provide fine-grained operation timing
+struct ggml_backend_profiler {
+    void * context;  // backend-specific profiler context
+
+    // Enable or disable profiling on this backend
+    void (*enable)(void * context, bool enable);
+
+    // Clear all recorded data
+    void (*reset)(void * context);
+
+    // Set the current split ID (called by scheduler before graph_compute)
+    void (*set_split_id)(void * context, int split_id);
+
+    // Get recorded profiling data
+    // Returns the number of records; sets *out to point to internal storage
+    // The returned pointer remains valid until the next reset or disable call
+    int (*get_records)(void * context, const ggml_profile_record ** out);
+
+    // Free the profiler context
+    void (*free_context)(void * context);
+};
+
+typedef struct ggml_backend_profiler * ggml_backend_profiler_t;
+
+// Populate the per-node fields of a ggml_profile_record from a ggml_tensor node:
+//   ne, out_type, n_src, ne_src, nb_src, type_src, op_params, sub_op.
+// All other fields (type/name/backend_id/split_id/timestamps/bytes/extra) must
+// be filled in separately by the backend that records the event.
+GGML_API void ggml_profile_record_from_tensor(struct ggml_profile_record * rec,
+                                              const struct ggml_tensor *   node);
+
+// Register a profiler on a backend (called by backend during init)
+// The profiler is owned by the backend and will be freed when the backend is freed
+GGML_API void ggml_backend_set_profiler(ggml_backend_t backend, ggml_backend_profiler_t profiler);
+
+// Get the profiler associated with a backend (returns NULL if none)
+GGML_API ggml_backend_profiler_t ggml_backend_get_profiler(ggml_backend_t backend);
+
+//
+// Scheduler profiling API
+//
+
+// Enable or disable profiling on a scheduler
+// When enabled, the scheduler will:
+//   - Time data copy operations between backends
+//   - Enable profiling on all backends that support it
+//   - Collect profiling records from all backends after each graph compute
+GGML_API void ggml_backend_sched_set_profiling(ggml_backend_sched_t sched, bool enable);
+
+// Check if profiling is enabled on a scheduler
+GGML_API bool ggml_backend_sched_get_profiling(ggml_backend_sched_t sched);
+
+// Get profiling data from the last graph compute
+// Records are owned by the scheduler; valid until the next compute or reset
+// Returns the number of records
+GGML_API int ggml_backend_sched_get_profiling_records(ggml_backend_sched_t sched, const ggml_profile_record ** records);
+
+// Print a human-readable summary of the last profiling run to stdout
+// Groups records by operation name and shows total/count/min/max/avg time
+GGML_API void ggml_backend_sched_print_profiling(ggml_backend_sched_t sched);
+
+// Reset profiling data (clear all recorded data)
+GGML_API void ggml_backend_sched_reset_profiling(ggml_backend_sched_t sched);
+
+// Get current time in nanoseconds (for manual profiling if needed)
+GGML_API uint64_t ggml_profiler_time_ns(void);
+
+// Export profiling data as JSON to a file
+// Returns 0 on success, -1 on error
+GGML_API int ggml_backend_sched_export_profiling_json(ggml_backend_sched_t sched, const char * filepath);
+
+// Export profiling data as JSON to a FILE pointer
+GGML_API int ggml_backend_sched_write_profiling_json(ggml_backend_sched_t sched, FILE * fp);
+
+// Export profiling data as plain text statistics to a file
+// Returns 0 on success, -1 on error
+GGML_API int ggml_backend_sched_export_profiling_text(ggml_backend_sched_t sched, const char * filepath);
+
+// Export profiling data as plain text statistics to a FILE pointer
+GGML_API int ggml_backend_sched_write_profiling_text(ggml_backend_sched_t sched, FILE * fp);
+
+#ifdef __cplusplus
+}
+#endif
@@ -195,6 +195,7 @@ add_library(ggml-base
            ../include/ggml-backend.h
            ../include/ggml-cpp.h
            ../include/ggml-opt.h
+            ../include/ggml-profiler.h
            ../include/gguf.h
            ggml.c
            ggml.cpp
@@ -202,6 +203,7 @@ add_library(ggml-base
            ggml-backend.cpp
            ggml-backend-meta.cpp
            ggml-opt.cpp
+            ggml-profiler.cpp
            ggml-threading.cpp
            ggml-threading.h
            ggml-quants.c
@@ -430,7 +432,7 @@ if (GGML_CPU_ALL_VARIANTS)
            message(FATAL_ERROR "Unsupported ARM target OS: ${CMAKE_SYSTEM_NAME}")
        endif()
    elseif (GGML_SYSTEM_ARCH STREQUAL "PowerPC")
-        if (CMAKE_SYSTEM_NAME MATCHES "Linux")
+        if (CMAKE_SYSTEM_NAME MATCHES "Linux|AIX")
            ggml_add_cpu_backend_variant(power0)
            ggml_add_cpu_backend_variant(power7_1       POWER7)
            ggml_add_cpu_backend_variant(power7_2       POWER7  VSX)
@@ -3,6 +3,7 @@
 // ggml-backend internal header

 #include "ggml-backend.h"
+#include "ggml-profiler.h"

 #ifdef  __cplusplus
 extern "C" {
@@ -144,6 +145,9 @@ extern "C" {
        struct ggml_backend_i iface;
        ggml_backend_dev_t device;
        void * context;
+
+        // Optional profiler (set by backend during init, NULL if not profiling)
+        ggml_backend_profiler_t profiler;
    };

    struct ggml_backend_event {
@@ -12,6 +12,7 @@
 #include "ggml-backend-impl.h"
 #include "ggml-alloc.h"
 #include "ggml-impl.h"
+#include "ggml-profiler.h"

 #include <assert.h>
 #include <limits.h>
@@ -20,6 +21,7 @@
 #include <stdlib.h>
 #include <string.h>
 #include <algorithm>
+#include <string>
 #include <vector>

 #ifdef __APPLE__
@@ -231,6 +233,15 @@ void ggml_backend_free(ggml_backend_t backend) {
        return;
    }

+    // Clean up profiler if present (before backend frees its context)
+    if (backend->profiler != NULL) {
+        if (backend->profiler->free_context != NULL) {
+            backend->profiler->free_context(backend->profiler->context);
+        }
+        delete backend->profiler;
+        backend->profiler = NULL;
+    }
+
    backend->iface.free(backend);
 }

@@ -825,6 +836,20 @@ struct ggml_backend_sched {
    int debug_realloc;
    int debug_graph_size;
    int debug_prev_graph_size;
+
+    // profiling
+    bool                             profiling_enabled;
+    std::string                      profiling_env_path;  // GGML_PROFILE env var value (for auto-export on free)
+    std::vector<ggml_profile_record> copy_records;        // copy events recorded by the scheduler
+    std::vector<ggml_profile_record> profiling_records;   // merged records from all sources
+
+    // Cached backend metadata for safe access during auto-export (backends may be freed first)
+    struct backend_meta {
+        std::string name;
+        std::string device;
+        int         device_type;
+    };
+    std::vector<backend_meta>        profiling_backend_meta;
 };

 #define hash_id(tensor) ggml_hash_find_or_insert(&sched->hash_set, tensor)
@@ -1538,6 +1563,39 @@ static bool ggml_backend_sched_alloc_splits(ggml_backend_sched_t sched) {
    return true;
 }

+// Build a COPY profiling record.  Copies have no real ggml_tensor "node" backing
+// them, so we synthesize one source describing the input tensor that was moved.
+static ggml_profile_record make_copy_record(const char * copy_dir, int backend_id, int split_id,
+                                            uint64_t start_ns, uint64_t end_ns, uint64_t bytes,
+                                            const struct ggml_tensor * input) {
+    ggml_profile_record rec = {};
+    rec.type       = GGML_PROFILE_EVENT_COPY;
+    rec.name       = copy_dir;
+    rec.backend_id = backend_id;
+    rec.split_id   = split_id;
+    rec.start_ns   = start_ns;
+    rec.end_ns     = end_ns;
+    rec.bytes      = bytes;
+    rec.extra      = input ? input->name : NULL;
+    snprintf(rec.tensor_name, sizeof(rec.tensor_name), "%s", input ? input->name : "");
+    rec.out_type   = -1;
+    rec.sub_op     = -1;
+    rec.n_src      = 0;
+    if (input != NULL) {
+        // Describe the input tensor as src[0] so consumers can inspect its shape.
+        rec.n_src = 1;
+        memcpy(rec.ne_src[0], input->ne, sizeof(rec.ne_src[0]));
+        for (int d = 0; d < 4; d++) {
+            rec.nb_src[0][d] = (int64_t) input->nb[d];
+        }
+        rec.type_src[0] = (int) input->type;
+    }
+    for (int i = rec.n_src; i < GGML_MAX_SRC; i++) {
+        rec.type_src[i] = -1;
+    }
+    return rec;
+}
+
 static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t sched) {
    GGML_ASSERT(sched);
    struct ggml_backend_sched_split * splits = sched->splits;
@@ -1546,11 +1604,28 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s
    std::vector<int32_t> ids;
    std::vector<ggml_bitset_t> used_ids;

+    // Profiling: reset copy records for this compute pass
+    if (sched->profiling_enabled) {
+        sched->copy_records.clear();
+    }
+
    for (int split_id = 0; split_id < sched->n_splits; split_id++) {
        struct ggml_backend_sched_split * split = &splits[split_id];
        int split_backend_id = split->backend_id;
        ggml_backend_t split_backend = sched->backends[split_backend_id];

+        // Profiling: set split ID and enable backend profiling
+        if (sched->profiling_enabled) {
+            if (split_backend->profiler != NULL) {
+                if (split_backend->profiler->enable != NULL) {
+                    split_backend->profiler->enable(split_backend->profiler->context, true);
+                }
+                if (split_backend->profiler->set_split_id != NULL) {
+                    split_backend->profiler->set_split_id(split_backend->profiler->context, split_id);
+                }
+            }
+        }
+
        // copy the input tensors to the split backend
        for (int input_id = 0; input_id < split->n_inputs; input_id++) {
            ggml_backend_t input_backend = ggml_backend_sched_get_tensor_backend(sched, split->inputs[input_id]);
@@ -1564,7 +1639,25 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s
                } else {
                    ggml_backend_synchronize(split_backend);
                }
-                ggml_backend_tensor_copy(input, input_cpy);
+                if (sched->profiling_enabled) {
+                    uint64_t copy_start = ggml_profiler_time_ns();
+                    ggml_backend_tensor_copy(input, input_cpy);
+                    uint64_t copy_end = ggml_profiler_time_ns();
+
+                    enum ggml_backend_dev_type src_type = ggml_backend_dev_type(input_backend->device);
+                    enum ggml_backend_dev_type dst_type = ggml_backend_dev_type(split_backend->device);
+                    const char *               copy_dir = "copy_D2D";
+                    if (src_type == GGML_BACKEND_DEVICE_TYPE_CPU && dst_type != GGML_BACKEND_DEVICE_TYPE_CPU) {
+                        copy_dir = "copy_H2D";
+                    } else if (src_type != GGML_BACKEND_DEVICE_TYPE_CPU && dst_type == GGML_BACKEND_DEVICE_TYPE_CPU) {
+                        copy_dir = "copy_D2H";
+                    }
+
+                    sched->copy_records.push_back(make_copy_record(copy_dir, split_backend_id, split_id,
+                                                                   copy_start, copy_end, ggml_nbytes(input), input));
+                } else {
+                    ggml_backend_tensor_copy(input, input_cpy);
+                }
            } else {
                // wait for the split backend to finish using the input before overwriting it
                if (sched->events[split_backend_id][sched->cur_copy] != NULL) {
@@ -1621,12 +1714,14 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s
                    }

                    // group consecutive experts and copy them together
+                    size_t total_copied_bytes = 0;
                    auto copy_experts = [&](int32_t first_id, int32_t last_id) {
                        const size_t expert_offset = first_id * expert_size;
                        const size_t expert_size_copy =  (last_id - first_id + 1) * expert_size;
                        const size_t padding = std::min<size_t>(expert_size, 512);
                        const size_t padding_end = last_id < n_expert - 1 ? padding : 0;

+                        total_copied_bytes += expert_size_copy + padding_end;
                        ggml_backend_tensor_set_async(split_backend,
                            input_cpy,
                            (const uint8_t *)input->data + expert_offset, expert_offset,
@@ -1635,6 +1730,11 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s
                            expert_size_copy + padding_end);
                    };

+                    uint64_t moe_copy_start = 0;
+                    if (sched->profiling_enabled) {
+                        moe_copy_start = ggml_profiler_time_ns();
+                    }
+
                    int id = 0;
                    while (!ggml_bitset_get(used_ids.data(), id)) {
                        id++;
@@ -1658,9 +1758,34 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s
                        last_id = id;
                    }
                    copy_experts(first_id, last_id);
+
+                    if (sched->profiling_enabled) {
+                        uint64_t moe_copy_end = ggml_profiler_time_ns();
+
+                        enum ggml_backend_dev_type src_type = ggml_backend_dev_type(input_backend->device);
+                        enum ggml_backend_dev_type dst_type = ggml_backend_dev_type(split_backend->device);
+                        const char *               copy_dir = "copy_D2D";
+                        if (src_type == GGML_BACKEND_DEVICE_TYPE_CPU && dst_type != GGML_BACKEND_DEVICE_TYPE_CPU) {
+                            copy_dir = "copy_H2D";
+                        } else if (src_type != GGML_BACKEND_DEVICE_TYPE_CPU &&
+                                   dst_type == GGML_BACKEND_DEVICE_TYPE_CPU) {
+                            copy_dir = "copy_D2H";
+                        }
+
+                        sched->copy_records.push_back(make_copy_record(copy_dir, split_backend_id, split_id,
+                                                                       moe_copy_start, moe_copy_end,
+                                                                       (uint64_t) total_copied_bytes, input));
+                    }
                } else {
                    // try async copy, but if not possible, we can still use a sync copy without synchronizing the dst backend, since we handle the synchronization here with multiple copies and events
                    // TODO: add public function to facilitate this, since applications do not have direct access to the backend interface
+
+                    // Capture timestamp before async attempt so we can record launch time
+                    uint64_t copy_start = 0;
+                    if (sched->profiling_enabled) {
+                        copy_start = ggml_profiler_time_ns();
+                    }
+
                    if (!split_backend->iface.cpy_tensor_async || !split_backend->iface.cpy_tensor_async(input_backend, split_backend, input, input_cpy)) {
                        ggml_backend_synchronize(input_backend);
                        if (sched->events[split_backend_id][sched->cur_copy] != NULL) {
@@ -1668,7 +1793,45 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s
                        } else {
                            ggml_backend_synchronize(split_backend);
                        }
-                        ggml_backend_tensor_copy(input, input_cpy);
+                        if (sched->profiling_enabled) {
+                            // Re-take start after sync for accurate sync copy measurement
+                            copy_start = ggml_profiler_time_ns();
+                            ggml_backend_tensor_copy(input, input_cpy);
+                            uint64_t copy_end = ggml_profiler_time_ns();
+
+                            enum ggml_backend_dev_type src_type = ggml_backend_dev_type(input_backend->device);
+                            enum ggml_backend_dev_type dst_type = ggml_backend_dev_type(split_backend->device);
+                            const char *               copy_dir = "copy_D2D";
+                            if (src_type == GGML_BACKEND_DEVICE_TYPE_CPU && dst_type != GGML_BACKEND_DEVICE_TYPE_CPU) {
+                                copy_dir = "copy_H2D";
+                            } else if (src_type != GGML_BACKEND_DEVICE_TYPE_CPU &&
+                                       dst_type == GGML_BACKEND_DEVICE_TYPE_CPU) {
+                                copy_dir = "copy_D2H";
+                            }
+
+                            sched->copy_records.push_back(make_copy_record(copy_dir, split_backend_id, split_id,
+                                                                           copy_start, copy_end, ggml_nbytes(input), input));
+                        } else {
+                            ggml_backend_tensor_copy(input, input_cpy);
+                        }
+                    } else {
+                        // async copy was launched — record the time spanning the async call
+                        if (sched->profiling_enabled) {
+                            uint64_t copy_end = ggml_profiler_time_ns();
+
+                            enum ggml_backend_dev_type src_type = ggml_backend_dev_type(input_backend->device);
+                            enum ggml_backend_dev_type dst_type = ggml_backend_dev_type(split_backend->device);
+                            const char *               copy_dir = "copy_D2D";
+                            if (src_type == GGML_BACKEND_DEVICE_TYPE_CPU && dst_type != GGML_BACKEND_DEVICE_TYPE_CPU) {
+                                copy_dir = "copy_H2D";
+                            } else if (src_type != GGML_BACKEND_DEVICE_TYPE_CPU &&
+                                       dst_type == GGML_BACKEND_DEVICE_TYPE_CPU) {
+                                copy_dir = "copy_D2H";
+                            }
+
+                            sched->copy_records.push_back(make_copy_record(copy_dir, split_backend_id, split_id,
+                                                                           copy_start, copy_end, ggml_nbytes(input), input));
+                        }
                    }
                }
            }
@@ -1721,6 +1884,32 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s
        }
    }

+    // Profiling: collect records from all backends and append to accumulated records
+    if (sched->profiling_enabled) {
+        // Collect backend operation records
+        for (int b = 0; b < sched->n_backends; b++) {
+            ggml_backend_t backend = sched->backends[b];
+            if (backend->profiler != NULL && backend->profiler->get_records != NULL) {
+                const ggml_profile_record * backend_recs = NULL;
+                int count = backend->profiler->get_records(backend->profiler->context, &backend_recs);
+                for (int r = 0; r < count; r++) {
+                    ggml_profile_record rec = backend_recs[r];
+                    rec.backend_id          = b;  // stamp correct scheduler backend index
+                    sched->profiling_records.push_back(rec);
+                }
+                // Reset backend records (but keep profiling enabled for next compute)
+                if (backend->profiler->reset != NULL) {
+                    backend->profiler->reset(backend->profiler->context);
+                }
+            }
+        }
+
+        // Append copy records
+        for (const auto & rec : sched->copy_records) {
+            sched->profiling_records.push_back(rec);
+        }
+    }
+
    return GGML_STATUS_SUCCESS;
 }

@@ -1787,6 +1976,24 @@ ggml_backend_sched_t ggml_backend_sched_new(

    sched->galloc = ggml_gallocr_new_n(sched->bufts, n_backends);
    sched->op_offload = op_offload;
+    const char * profile_env = getenv("GGML_PROFILE");
+    if (profile_env != NULL) {
+        sched->profiling_enabled = true;
+        sched->profiling_env_path = profile_env;
+    }
+
+    // Cache backend metadata for safe access during auto-export
+    for (int b = 0; b < n_backends; b++) {
+        ggml_backend_sched::backend_meta meta;
+        meta.name        = ggml_backend_name(backends[b]);
+        meta.device      = "unknown";
+        meta.device_type  = 0;
+        if (backends[b]->device != NULL) {
+            meta.device      = ggml_backend_dev_name(backends[b]->device);
+            meta.device_type = (int) ggml_backend_dev_type(backends[b]->device);
+        }
+        sched->profiling_backend_meta.push_back(std::move(meta));
+    }

    ggml_backend_sched_reset(sched);

@@ -1797,6 +2004,33 @@ void ggml_backend_sched_free(ggml_backend_sched_t sched) {
    if (sched == NULL) {
        return;
    }
+
+    // Auto-export profiling data if enabled via GGML_PROFILE env var
+    // GGML_PROFILE=1 or GGML_PROFILE="" → print to stdout
+    // GGML_PROFILE=file.json → export JSON
+    // GGML_PROFILE=file.txt → export text
+    if (!sched->profiling_records.empty() && getenv("GGML_PROFILE") != NULL) {
+        const std::string & path = sched->profiling_env_path;
+        if (path.empty() || path == "1" || path == "stdout") {
+            ggml_backend_sched_print_profiling(sched);
+        } else if (path.size() >= 4 && path.compare(path.size() - 4, 4, ".txt") == 0) {
+            int ret = ggml_backend_sched_export_profiling_text(sched, path.c_str());
+            if (ret == 0) {
+                GGML_LOG_INFO("[profiler] Data exported to: %s\n", path.c_str());
+            } else {
+                GGML_LOG_ERROR("[profiler] Failed to export data to: %s\n", path.c_str());
+            }
+        } else {
+            // Default to JSON for any other path (including .json)
+            int ret = ggml_backend_sched_export_profiling_json(sched, path.c_str());
+            if (ret == 0) {
+                GGML_LOG_INFO("[profiler] Data exported to: %s\n", path.c_str());
+            } else {
+                GGML_LOG_ERROR("[profiler] Failed to export data to: %s\n", path.c_str());
+            }
+        }
+    }
+
    for (int b = 0; b < sched->n_backends; b++) {
        for (int c = 0; c < sched->n_copies; c++) {
            ggml_backend_event_free(sched->events[b][c]);
@@ -2369,3 +2603,509 @@ ggml_backend_buffer_t ggml_backend_cpu_buffer_from_ptr(void * ptr, size_t size)
    GGML_ASSERT((uintptr_t)ptr % TENSOR_ALIGNMENT == 0 && "buffer pointer must be aligned");
    return ggml_backend_buffer_init(ggml_backend_cpu_buffer_from_ptr_type(), ggml_backend_cpu_buffer_from_ptr_i, ptr, size);
 }
+
+//
+// Scheduler profiling
+//
+
+void ggml_backend_sched_set_profiling(ggml_backend_sched_t sched, bool enable) {
+    GGML_ASSERT(sched);
+    sched->profiling_enabled = enable;
+
+    if (!enable) {
+        ggml_backend_sched_reset_profiling(sched);
+    }
+}
+
+bool ggml_backend_sched_get_profiling(ggml_backend_sched_t sched) {
+    GGML_ASSERT(sched);
+    return sched->profiling_enabled;
+}
+
+int ggml_backend_sched_get_profiling_records(ggml_backend_sched_t sched, const ggml_profile_record ** records) {
+    GGML_ASSERT(sched);
+    GGML_ASSERT(records != NULL);
+
+    *records = sched->profiling_records.data();
+    return (int) sched->profiling_records.size();
+}
+
+void ggml_backend_sched_reset_profiling(ggml_backend_sched_t sched) {
+    GGML_ASSERT(sched);
+    sched->profiling_records.clear();
+    sched->copy_records.clear();
+}
+
+void ggml_backend_sched_print_profiling(ggml_backend_sched_t sched) {
+    GGML_ASSERT(sched);
+
+    if (sched->profiling_records.empty()) {
+        GGML_LOG_INFO("[profiler] No profiling data available\n");
+        return;
+    }
+
+    GGML_LOG_INFO("\n=== Profiling Summary ===\n");
+
+    // Aggregate by (name, type, backend_id)
+    struct op_stats {
+        const char *                 name;
+        enum ggml_profile_event_type type;
+        int                          backend_id;
+        uint64_t                     total_ns;
+        uint64_t                     min_ns;
+        uint64_t                     max_ns;
+        int                          count;
+        uint64_t                     total_bytes;
+        int64_t                      representative_ne[4];
+    };
+
+    std::vector<op_stats> stats;
+    for (const auto & rec : sched->profiling_records) {
+        bool found = false;
+        for (auto & s : stats) {
+            if (s.type == rec.type && s.backend_id == rec.backend_id && strcmp(s.name, rec.name) == 0) {
+                uint64_t dur = (rec.end_ns > rec.start_ns) ? (rec.end_ns - rec.start_ns) : 0;
+                s.total_ns += dur;
+                s.min_ns = std::min(s.min_ns, dur);
+                s.max_ns = std::max(s.max_ns, dur);
+                s.count++;
+                s.total_bytes += rec.bytes;
+                found = true;
+                break;
+            }
+        }
+        if (!found) {
+            uint64_t dur = (rec.end_ns > rec.start_ns) ? (rec.end_ns - rec.start_ns) : 0;
+            op_stats s;
+            s.name        = rec.name;
+            s.type        = rec.type;
+            s.backend_id  = rec.backend_id;
+            s.total_ns    = dur;
+            s.min_ns      = dur;
+            s.max_ns      = dur;
+            s.count       = 1;
+            s.total_bytes = rec.bytes;
+            memcpy(s.representative_ne, rec.ne_src[0], sizeof(s.representative_ne));
+            stats.push_back(s);
+        }
+    }
+
+    // Sort by total time descending
+    std::sort(stats.begin(), stats.end(),
+              [](const op_stats & a, const op_stats & b) { return a.total_ns > b.total_ns; });
+
+    uint64_t grand_total = 0;
+    for (const auto & s : stats) {
+        grand_total += s.total_ns;
+    }
+
+    const char * type_str[] = { "OP  ", "COPY" };
+    for (const auto & s : stats) {
+        double pct    = 100.0 * (double) s.total_ns / (double) grand_total;
+        double avg_us = (double) s.total_ns / (double) s.count / 1000.0;
+        double min_us = (double) s.min_ns / 1000.0;
+        double max_us = (double) s.max_ns / 1000.0;
+
+        GGML_LOG_INFO(
+            "  [%s] backend %d %-28s %7.1f%%  count=%-6d  total=%8.2f ms  avg=%8.2f us  min=%8.2f us  max=%8.2f us",
+            type_str[s.type], s.backend_id, s.name, pct, s.count, (double) s.total_ns / 1e6, avg_us, min_us, max_us);
+        if (s.total_bytes > 0 && s.total_ns > 0) {
+            double bw_gbps = (double) s.total_bytes / (double) s.total_ns;
+            if (bw_gbps >= 1000.0) {
+                GGML_LOG_INFO("  %6.2f TB/s", bw_gbps / 1000.0);
+            } else {
+                GGML_LOG_INFO("  %6.2f GB/s", bw_gbps);
+            }
+        }
+        // Print representative tensor shape (first record's ne)
+        if (s.representative_ne[0] > 0 || s.representative_ne[1] > 0) {
+            GGML_LOG_INFO("  [%lld x %lld", (long long) s.representative_ne[0], (long long) s.representative_ne[1]);
+            if (s.representative_ne[2] > 1) {
+                GGML_LOG_INFO(" x %lld", (long long) s.representative_ne[2]);
+            }
+            if (s.representative_ne[3] > 1) {
+                GGML_LOG_INFO(" x %lld", (long long) s.representative_ne[3]);
+            }
+            GGML_LOG_INFO("]");
+        }
+        GGML_LOG_INFO("\n");
+    }
+
+    GGML_LOG_INFO("  ---\n");
+    GGML_LOG_INFO("  Total: %.2f ms  (%d records, %d unique ops)\n\n", (double) grand_total / 1e6,
+                  (int) sched->profiling_records.size(), (int) stats.size());
+}
+
+int ggml_backend_sched_write_profiling_json(ggml_backend_sched_t sched, FILE * fp) {
+    GGML_ASSERT(sched);
+    GGML_ASSERT(fp != NULL);
+
+    uint64_t total_ns = 0;
+    for (const auto & rec : sched->profiling_records) {
+        total_ns += (rec.end_ns > rec.start_ns) ? (rec.end_ns - rec.start_ns) : 0;
+    }
+
+    fprintf(fp, "{\n");
+    fprintf(fp, "  \"version\": 3,\n");
+    fprintf(fp, "  \"profiler\": \"ggml\",\n");
+    fprintf(fp, "  \"total_records\": %d,\n", (int) sched->profiling_records.size());
+    fprintf(fp, "  \"total_ns\": %llu,\n", (unsigned long long) total_ns);
+
+    // Backend metadata (use cached data if available, fall back to live pointers)
+    fprintf(fp, "  \"backends\": [\n");
+    for (int b = 0; b < sched->n_backends; b++) {
+        const char * name     = "unknown";
+        const char * dev_name = "unknown";
+        int          dev_type = 0;
+        if (b < (int) sched->profiling_backend_meta.size()) {
+            name     = sched->profiling_backend_meta[b].name.c_str();
+            dev_name = sched->profiling_backend_meta[b].device.c_str();
+            dev_type = sched->profiling_backend_meta[b].device_type;
+        } else if (sched->backends[b] != NULL) {
+            name = ggml_backend_name(sched->backends[b]);
+            if (sched->backends[b]->device != NULL) {
+                dev_name = ggml_backend_dev_name(sched->backends[b]->device);
+                dev_type = (int) ggml_backend_dev_type(sched->backends[b]->device);
+            }
+        }
+        fprintf(fp, "    {\"id\": %d, \"name\": \"%s\", \"device\": \"%s\", \"device_type\": %d}%s\n", b, name,
+                dev_name, dev_type, (b < sched->n_backends - 1) ? "," : "");
+    }
+    fprintf(fp, "  ],\n");
+
+    // Records
+    fprintf(fp, "  \"records\": [\n");
+
+    for (int i = 0; i < (int) sched->profiling_records.size(); i++) {
+        const auto & rec         = sched->profiling_records[i];
+        uint64_t     duration_ns = (rec.end_ns > rec.start_ns) ? (rec.end_ns - rec.start_ns) : 0;
+
+        fprintf(fp,
+                "    {\"type\": %d, \"name\": \"%s\", \"backend_id\": %d, \"split_id\": %d, "
+                "\"start_ns\": %llu, \"duration_ns\": %llu, \"bytes\": %llu, \"extra\": ",
+                (int) rec.type, rec.name ? rec.name : "unknown", rec.backend_id, rec.split_id,
+                (unsigned long long) rec.start_ns, (unsigned long long) duration_ns, (unsigned long long) rec.bytes);
+
+        if (rec.extra != NULL) {
+            fprintf(fp, "\"%s\"", rec.extra);
+        } else {
+            fprintf(fp, "null");
+        }
+
+        // Output tensor info
+        if (rec.tensor_name[0] != '\0') {
+            fprintf(fp, ", \"tensor_name\": \"%s\"", rec.tensor_name);
+        } else {
+            fprintf(fp, ", \"tensor_name\": null");
+        }
+        fprintf(fp, ", \"ne\": [%lld, %lld, %lld, %lld]", (long long) rec.ne[0], (long long) rec.ne[1],
+                (long long) rec.ne[2], (long long) rec.ne[3]);
+        fprintf(fp, ", \"out_type\": %d", rec.out_type);
+
+        // Source tensors
+        fprintf(fp, ", \"n_src\": %d", rec.n_src);
+        fprintf(fp, ", \"ne_src\": [");
+        for (int s = 0; s < rec.n_src; s++) {
+            fprintf(fp, "%s[%lld, %lld, %lld, %lld]", s == 0 ? "" : ", ",
+                    (long long) rec.ne_src[s][0], (long long) rec.ne_src[s][1],
+                    (long long) rec.ne_src[s][2], (long long) rec.ne_src[s][3]);
+        }
+        fprintf(fp, "]");
+        fprintf(fp, ", \"nb_src\": [");
+        for (int s = 0; s < rec.n_src; s++) {
+            fprintf(fp, "%s[%lld, %lld, %lld, %lld]", s == 0 ? "" : ", ",
+                    (long long) rec.nb_src[s][0], (long long) rec.nb_src[s][1],
+                    (long long) rec.nb_src[s][2], (long long) rec.nb_src[s][3]);
+        }
+        fprintf(fp, "]");
+        fprintf(fp, ", \"type_src\": [");
+        for (int s = 0; s < rec.n_src; s++) {
+            fprintf(fp, "%s%d", s == 0 ? "" : ", ", rec.type_src[s]);
+        }
+        fprintf(fp, "]");
+
+        // op_params (full 16-int32 block, matching export-graph-ops format)
+        fprintf(fp, ", \"op_params\": [");
+        const int n_op_params = (int) (sizeof(rec.op_params) / sizeof(rec.op_params[0]));
+        for (int p = 0; p < n_op_params; p++) {
+            fprintf(fp, "%s%d", p == 0 ? "" : ", ", rec.op_params[p]);
+        }
+        fprintf(fp, "]");
+
+        fprintf(fp, ", \"sub_op\": %d", rec.sub_op);
+
+        fprintf(fp, "}%s\n", (i < (int) sched->profiling_records.size() - 1) ? "," : "");
+    }
+
+    fprintf(fp, "  ]\n");
+    fprintf(fp, "}\n");
+
+    return 0;
+}
+
+int ggml_backend_sched_export_profiling_json(ggml_backend_sched_t sched, const char * filepath) {
+    GGML_ASSERT(sched);
+    GGML_ASSERT(filepath != NULL);
+
+    FILE * fp = fopen(filepath, "w");
+    if (fp == NULL) {
+        GGML_LOG_ERROR("%s: failed to open %s for writing\n", __func__, filepath);
+        return -1;
+    }
+
+    int ret = ggml_backend_sched_write_profiling_json(sched, fp);
+    fclose(fp);
+
+    return ret;
+}
+
+// Helper: format ne dimensions as string, e.g. "[4096, 4096, 1]"
+static void fmt_ne(char * buf, size_t bufsize, const int64_t ne[4]) {
+    if (ne[0] == 0 && ne[1] == 0 && ne[2] == 0 && ne[3] == 0) {
+        buf[0] = '\0';
+        return;
+    }
+    int ndims = 4;
+    while (ndims > 1 && ne[ndims - 1] <= 1) {
+        ndims--;
+    }
+    int pos = snprintf(buf, bufsize, "[");
+    for (int i = 0; i < ndims && pos < (int) bufsize - 1; i++) {
+        pos += snprintf(buf + pos, bufsize - pos, "%s%lld", i > 0 ? ", " : "", (long long) ne[i]);
+    }
+    snprintf(buf + pos, bufsize - pos, "]");
+}
+
+// Helper: format bandwidth as string
+static void fmt_bandwidth(char * buf, size_t bufsize, uint64_t bytes, uint64_t ns) {
+    if (ns == 0 || bytes == 0) {
+        buf[0] = '\0';
+        return;
+    }
+    double bw_gbps = (double) bytes / (double) ns;
+    if (bw_gbps >= 1000.0) {
+        snprintf(buf, bufsize, "%.2f TB/s", bw_gbps / 1000.0);
+    } else {
+        snprintf(buf, bufsize, "%.2f GB/s", bw_gbps);
+    }
+}
+
+int ggml_backend_sched_write_profiling_text(ggml_backend_sched_t sched, FILE * fp) {
+    GGML_ASSERT(sched);
+    GGML_ASSERT(fp != NULL);
+
+    if (sched->profiling_records.empty()) {
+        fprintf(fp, "No profiling data available.\n");
+        return 0;
+    }
+
+    // Aggregate by (name, type, backend_id)
+    struct op_stats {
+        const char *                 name;
+        enum ggml_profile_event_type type;
+        int                          backend_id;
+        uint64_t                     total_ns;
+        uint64_t                     min_ns;
+        uint64_t                     max_ns;
+        int                          count;
+        uint64_t                     total_bytes;
+        int64_t                      representative_ne_src0[4];
+        int64_t                      representative_ne_src1[4];
+        int64_t                      representative_ne_src2[4];
+    };
+
+    std::vector<op_stats> stats;
+    for (const auto & rec : sched->profiling_records) {
+        uint64_t dur = (rec.end_ns > rec.start_ns) ? (rec.end_ns - rec.start_ns) : 0;
+        bool found = false;
+        for (auto & s : stats) {
+            if (s.type == rec.type && s.backend_id == rec.backend_id && strcmp(s.name, rec.name) == 0) {
+                s.total_ns += dur;
+                s.min_ns = std::min(s.min_ns, dur);
+                s.max_ns = std::max(s.max_ns, dur);
+                s.count++;
+                s.total_bytes += rec.bytes;
+                found = true;
+                break;
+            }
+        }
+        if (!found) {
+            op_stats s = {};
+            s.name        = rec.name;
+            s.type        = rec.type;
+            s.backend_id  = rec.backend_id;
+            s.total_ns    = dur;
+            s.min_ns      = dur;
+            s.max_ns      = dur;
+            s.count       = 1;
+            s.total_bytes = rec.bytes;
+            memcpy(s.representative_ne_src0, rec.ne_src[0], sizeof(s.representative_ne_src0));
+            memcpy(s.representative_ne_src1, rec.ne_src[1], sizeof(s.representative_ne_src1));
+            memcpy(s.representative_ne_src2, rec.ne_src[2], sizeof(s.representative_ne_src2));
+            stats.push_back(s);
+        }
+    }
+
+    std::sort(stats.begin(), stats.end(),
+              [](const op_stats & a, const op_stats & b) { return a.total_ns > b.total_ns; });
+
+    uint64_t grand_total = 0;
+    for (const auto & s : stats) {
+        grand_total += s.total_ns;
+    }
+
+    // --- Section 1: Overall summary ---
+    fprintf(fp, "=== Profiling Summary ===\n");
+    fprintf(fp, "Total time:    %.2f ms\n", (double) grand_total / 1e6);
+    fprintf(fp, "Total records: %d\n", (int) sched->profiling_records.size());
+    fprintf(fp, "Unique ops:    %d\n\n", (int) stats.size());
+
+    // --- Section 2: Per-backend breakdown ---
+    fprintf(fp, "=== Per-Backend Summary ===\n");
+    {
+        struct backend_stats {
+            int      backend_id;
+            int      op_count;
+            int      copy_count;
+            uint64_t op_ns;
+            uint64_t copy_ns;
+            uint64_t op_bytes;
+            uint64_t copy_bytes;
+        };
+        std::vector<backend_stats> bstats;
+        for (const auto & s : stats) {
+            bool found = false;
+            for (auto & bs : bstats) {
+                if (bs.backend_id == s.backend_id) {
+                    if (s.type == GGML_PROFILE_EVENT_OP) {
+                        bs.op_count += s.count;
+                        bs.op_ns    += s.total_ns;
+                        bs.op_bytes += s.total_bytes;
+                    } else {
+                        bs.copy_count += s.count;
+                        bs.copy_ns    += s.total_ns;
+                        bs.copy_bytes += s.total_bytes;
+                    }
+                    found = true;
+                    break;
+                }
+            }
+            if (!found) {
+                backend_stats bs = {};
+                bs.backend_id = s.backend_id;
+                if (s.type == GGML_PROFILE_EVENT_OP) {
+                    bs.op_count = s.count;
+                    bs.op_ns    = s.total_ns;
+                    bs.op_bytes = s.total_bytes;
+                } else {
+                    bs.copy_count = s.count;
+                    bs.copy_ns    = s.total_ns;
+                    bs.copy_bytes = s.total_bytes;
+                }
+                bstats.push_back(bs);
+            }
+        }
+        std::sort(bstats.begin(), bstats.end(),
+                  [](const backend_stats & a, const backend_stats & b) {
+                      return (a.op_ns + a.copy_ns) > (b.op_ns + b.copy_ns);
+                  });
+
+        for (const auto & bs : bstats) {
+            uint64_t total = bs.op_ns + bs.copy_ns;
+            double   pct   = grand_total > 0 ? 100.0 * (double) total / (double) grand_total : 0;
+
+            const char * bname = "unknown";
+            if (bs.backend_id >= 0 && bs.backend_id < (int) sched->profiling_backend_meta.size()) {
+                bname = sched->profiling_backend_meta[bs.backend_id].name.c_str();
+            } else if (bs.backend_id >= 0 && bs.backend_id < sched->n_backends && sched->backends[bs.backend_id] != NULL) {
+                bname = ggml_backend_name(sched->backends[bs.backend_id]);
+            }
+
+            fprintf(fp, "  Backend %d (%s): %.2f ms (%.1f%%)\n", bs.backend_id, bname, (double) total / 1e6, pct);
+            if (bs.op_count > 0) {
+                char bw_buf[32];
+                fmt_bandwidth(bw_buf, sizeof(bw_buf), bs.op_bytes, bs.op_ns);
+                fprintf(fp, "    OPs:   %d calls, %.2f ms", bs.op_count, (double) bs.op_ns / 1e6);
+                if (bw_buf[0]) {
+                    fprintf(fp, ", %s", bw_buf);
+                }
+                fprintf(fp, "\n");
+            }
+            if (bs.copy_count > 0) {
+                char bw_buf[32];
+                fmt_bandwidth(bw_buf, sizeof(bw_buf), bs.copy_bytes, bs.copy_ns);
+                fprintf(fp, "    COPYs: %d calls, %.2f ms", bs.copy_count, (double) bs.copy_ns / 1e6);
+                if (bw_buf[0]) {
+                    fprintf(fp, ", %s", bw_buf);
+                }
+                fprintf(fp, "\n");
+            }
+        }
+    }
+    fprintf(fp, "\n");
+
+    // --- Section 3: Detailed operation table ---
+    fprintf(fp, "=== Operations (sorted by total time) ===\n");
+    fprintf(fp, "%-5s %4s  %-28s  %7s  %6s  %10s  %10s  %10s  %10s  %12s  %s\n",
+            "TYPE", "BKND", "Operation", "%Time", "Count", "Total(ms)", "Avg(us)", "Min(us)", "Max(us)", "Bandwidth", "Tensors");
+    fprintf(fp, "%-5s %4s  %-28s  %7s  %6s  %10s  %10s  %10s  %10s  %12s  %s\n",
+            "-----", "----", "----------------------------", "-------", "------",
+            "----------", "----------", "----------", "----------", "------------", "-------");
+
+    const char * type_str[] = { "OP", "COPY" };
+    for (const auto & s : stats) {
+        double pct    = grand_total > 0 ? 100.0 * (double) s.total_ns / (double) grand_total : 0;
+        double avg_us = (double) s.total_ns / (double) s.count / 1000.0;
+        double min_us = (double) s.min_ns / 1000.0;
+        double max_us = (double) s.max_ns / 1000.0;
+
+        char bw_buf[32] = "";
+        fmt_bandwidth(bw_buf, sizeof(bw_buf), s.total_bytes, s.total_ns);
+
+        char ne0_buf[64];
+        char ne1_buf[64];
+        char ne2_buf[64];
+        fmt_ne(ne0_buf, sizeof(ne0_buf), s.representative_ne_src0);
+        fmt_ne(ne1_buf, sizeof(ne1_buf), s.representative_ne_src1);
+        fmt_ne(ne2_buf, sizeof(ne2_buf), s.representative_ne_src2);
+
+        // Build tensor shapes string
+        char tensors_buf[256] = "";
+        int tpos = 0;
+        if (ne0_buf[0]) {
+            tpos += snprintf(tensors_buf + tpos, sizeof(tensors_buf) - tpos, "%s", ne0_buf);
+        }
+        if (ne1_buf[0]) {
+            tpos += snprintf(tensors_buf + tpos, sizeof(tensors_buf) - tpos, " x %s", ne1_buf);
+        }
+        if (ne2_buf[0]) {
+            tpos += snprintf(tensors_buf + tpos, sizeof(tensors_buf) - tpos, " x %s", ne2_buf);
+        }
+
+        fprintf(fp, "%-5s %4d  %-28s  %6.1f%%  %6d  %10.2f  %10.2f  %10.2f  %10.2f  %12s  %s\n",
+                type_str[s.type], s.backend_id, s.name, pct, s.count,
+                (double) s.total_ns / 1e6, avg_us, min_us, max_us,
+                bw_buf, tensors_buf);
+    }
+
+    fprintf(fp, "\nTotal: %.2f ms  (%d records, %d unique ops)\n", (double) grand_total / 1e6,
+            (int) sched->profiling_records.size(), (int) stats.size());
+
+    return 0;
+}
+
+int ggml_backend_sched_export_profiling_text(ggml_backend_sched_t sched, const char * filepath) {
+    GGML_ASSERT(sched);
+    GGML_ASSERT(filepath != NULL);
+
+    FILE * fp = fopen(filepath, "w");
+    if (fp == NULL) {
+        GGML_LOG_ERROR("%s: failed to open %s for writing\n", __func__, filepath);
+        return -1;
+    }
+
+    int ret = ggml_backend_sched_write_profiling_text(sched, fp);
+    fclose(fp);
+
+    return ret;
+}
@@ -2,6 +2,7 @@
 #include "ggml-impl.h"
 #include "ggml-blas.h"
 #include "ggml-backend-impl.h"
+#include "ggml-profiler.h"

 #include <future>
 #include <vector>
@@ -26,6 +27,11 @@ struct ggml_backend_blas_context {
 #ifndef GGML_USE_OPENMP
    std::vector<std::future<void>> tasks;
 #endif
+
+    // Profiling state
+    bool                             profiling_enabled  = false;
+    int                              profiling_split_id = -1;
+    std::vector<ggml_profile_record> profiling_records;
 };

 static void ggml_backend_blas_mul_mat(ggml_backend_blas_context * ctx, struct ggml_tensor * dst) {
@@ -233,6 +239,18 @@ static enum ggml_status ggml_backend_blas_graph_compute(ggml_backend_t backend,
            continue;
        }

+        // Skip view/identity ops
+        if (node->op == GGML_OP_NONE || node->op == GGML_OP_RESHAPE || node->op == GGML_OP_VIEW ||
+            node->op == GGML_OP_PERMUTE || node->op == GGML_OP_TRANSPOSE) {
+            continue;
+        }
+
+        // Profiling: time this operation
+        uint64_t t_start = 0;
+        if (ctx->profiling_enabled) {
+            t_start = ggml_profiler_time_ns();
+        }
+
        switch (node->op) {
            case GGML_OP_MUL_MAT:
                ggml_backend_blas_mul_mat(ctx, node);
@@ -242,16 +260,24 @@ static enum ggml_status ggml_backend_blas_graph_compute(ggml_backend_t backend,
                ggml_backend_blas_out_prod(ctx, node);
                break;

-            case GGML_OP_NONE:
-            case GGML_OP_RESHAPE:
-            case GGML_OP_VIEW:
-            case GGML_OP_PERMUTE:
-            case GGML_OP_TRANSPOSE:
-                break;
-
            default:
                GGML_ABORT("%s: unsupported op %s\n", __func__, ggml_op_desc(node));
        }
+
+        if (ctx->profiling_enabled) {
+            uint64_t            t_end = ggml_profiler_time_ns();
+            ggml_profile_record rec;
+            rec.type       = GGML_PROFILE_EVENT_OP;
+            rec.name       = ggml_op_name(node->op);
+            rec.backend_id = 0;
+            rec.split_id   = ctx->profiling_split_id;
+            rec.start_ns   = t_start;
+            rec.end_ns     = t_end;
+            rec.bytes      = ggml_nbytes(node);
+            rec.extra      = NULL;
+            ggml_profile_record_from_tensor(&rec, node);
+            ctx->profiling_records.push_back(rec);
+        }
    }

    return GGML_STATUS_SUCCESS;
@@ -287,10 +313,11 @@ ggml_backend_t ggml_backend_blas_init(void) {
    ggml_backend_blas_context * ctx = new ggml_backend_blas_context;

    ggml_backend_t backend = new ggml_backend {
-        /* .guid    = */ ggml_backend_blas_guid(),
-        /* .iface   = */ blas_backend_i,
-        /* .device  = */ ggml_backend_reg_dev_get(ggml_backend_blas_reg(), 0),
-        /* .context = */ ctx,
+        /* .guid     = */ ggml_backend_blas_guid(),
+        /* .iface    = */ blas_backend_i,
+        /* .device   = */ ggml_backend_reg_dev_get(ggml_backend_blas_reg(), 0),
+        /* .context  = */ ctx,
+        /* .profiler = */ nullptr,
    };

 #if defined(GGML_BLAS_USE_OPENBLAS) && defined(GGML_USE_OPENMP)
@@ -303,6 +330,44 @@ ggml_backend_t ggml_backend_blas_init(void) {
    GGML_LOG_DEBUG("%s: warning: ggml is using OpenMP, but BLIS was compiled without OpenMP support\n", __func__);
 #endif

+    // Register profiler
+    ggml_backend_blas_context * blas_ctx = ctx;  // ctx is already defined above
+
+    static auto blas_prof_enable = [](void * ctx, bool enable) {
+        auto * bctx             = (ggml_backend_blas_context *) ctx;
+        bctx->profiling_enabled = enable;
+        if (!enable) {
+            bctx->profiling_records.clear();
+        }
+    };
+    static auto blas_prof_reset = [](void * ctx) {
+        auto * bctx = (ggml_backend_blas_context *) ctx;
+        bctx->profiling_records.clear();
+        bctx->profiling_split_id = -1;
+    };
+    static auto blas_prof_set_split_id = [](void * ctx, int split_id) {
+        auto * bctx              = (ggml_backend_blas_context *) ctx;
+        bctx->profiling_split_id = split_id;
+    };
+    static auto blas_prof_get_records = [](void * ctx, const ggml_profile_record ** out) -> int {
+        auto * bctx = (ggml_backend_blas_context *) ctx;
+        *out        = bctx->profiling_records.data();
+        return (int) bctx->profiling_records.size();
+    };
+    static auto blas_prof_free = [](void * ctx) {
+        (void) ctx;
+    };
+
+    auto * profiler = new ggml_backend_profiler{
+        /* .context      = */ blas_ctx,
+        /* .enable       = */ blas_prof_enable,
+        /* .reset        = */ blas_prof_reset,
+        /* .set_split_id = */ blas_prof_set_split_id,
+        /* .get_records  = */ blas_prof_get_records,
+        /* .free_context = */ blas_prof_free,
+    };
+    ggml_backend_set_profiler(backend, profiler);
+
    return backend;
 }

@@ -3035,7 +3035,8 @@ ggml_backend_t ggml_backend_cann_init(int32_t device) {
        new ggml_backend{ /* .guid      = */ ggml_backend_cann_guid(),
                          /* .interface = */ ggml_backend_cann_interface,
                          /* .device    = */ ggml_backend_reg_dev_get(ggml_backend_cann_reg(), device),
-                          /* .context   = */ ctx };
+                          /* .context   = */ ctx,
+                          /* .profiler  = */ nullptr };

    return cann_backend;
 }
@@ -6,6 +6,7 @@
 #include "traits.h"
 #include "ggml-cpu-impl.h"
 #include "ggml-impl.h"
+#include "ggml-profiler.h"
 #include "quants.h"
 #include "ggml-threading.h"
 #include "unary-ops.h"
@@ -1178,8 +1179,8 @@ static void ggml_compute_forward_mul_mat_one_chunk(

    const bool src1_cont = ggml_is_contiguous(src1);

-    ggml_vec_dot_t const vec_dot      = type_traits_cpu[type].vec_dot;
-    enum ggml_type const vec_dot_type = type_traits_cpu[type].vec_dot_type;
+    const ggml_vec_dot_t vec_dot      = type_traits_cpu[type].vec_dot;
+    const enum ggml_type vec_dot_type = type_traits_cpu[type].vec_dot_type;

    // broadcast factors
    const int64_t r2 = ne12 / ne02;
@@ -1269,9 +1270,9 @@ void ggml_compute_forward_mul_mat(
    const int ith = params->ith;
    const int nth = params->nth;

-    enum ggml_type           const vec_dot_type         = type_traits_cpu[src0->type].vec_dot_type;
-    ggml_from_float_t        const from_float           = type_traits_cpu[vec_dot_type].from_float;
-    int64_t                  const vec_dot_num_rows     = type_traits_cpu[src0->type].nrows;
+    const enum ggml_type    vec_dot_type     = type_traits_cpu[src0->type].vec_dot_type;
+    const ggml_from_float_t from_float       = type_traits_cpu[vec_dot_type].from_float;
+    const int64_t           vec_dot_num_rows = type_traits_cpu[src0->type].nrows;

    GGML_ASSERT(ne0 == ne01);
    GGML_ASSERT(ne1 == ne11);
@@ -1480,8 +1481,8 @@ static void ggml_compute_forward_mul_mat_id_one_chunk(

    const enum ggml_type type = src0->type;

-    ggml_vec_dot_t    const vec_dot      = type_traits_cpu[type].vec_dot;
-    enum ggml_type    const vec_dot_type = type_traits_cpu[type].vec_dot_type;
+    const ggml_vec_dot_t vec_dot      = type_traits_cpu[type].vec_dot;
+    const enum ggml_type vec_dot_type = type_traits_cpu[type].vec_dot_type;

    const int64_t blck_0 = 16;
    const int64_t blck_1 = 16;
@@ -1548,8 +1549,8 @@ static void ggml_compute_forward_mul_mat_id(

    const bool src1_cont = ggml_is_contiguous(src1);

-    enum ggml_type    const vec_dot_type    = type_traits_cpu[type].vec_dot_type;
-    ggml_from_float_t const from_float      = type_traits_cpu[vec_dot_type].from_float;
+    const enum ggml_type    vec_dot_type = type_traits_cpu[type].vec_dot_type;
+    const ggml_from_float_t from_float   = type_traits_cpu[vec_dot_type].from_float;

    // we don't support permuted src0 or src1
    GGML_ASSERT(nb00 == ggml_type_size(type));
@@ -3085,17 +3086,55 @@ static thread_ret_t ggml_graph_compute_thread(void * data) {
    GGML_PRINT_DEBUG("thread #%d compute-start cplan %p last-graph %d\n", state->ith, (const void *)cplan, state->last_graph);
 #endif

-    for (int node_n = 0; node_n < cgraph->n_nodes && atomic_load_explicit(&tp->abort, memory_order_relaxed) != node_n; node_n++) {
-        struct ggml_tensor * node = cgraph->nodes[node_n];
+    // Profiling state
+    if (cplan->profiling_context != NULL && cplan->profiling_record_fn != NULL) {
+        for (int node_n = 0; node_n < cgraph->n_nodes && atomic_load_explicit(&tp->abort, memory_order_relaxed) != node_n; node_n++) {
+            struct ggml_tensor * node = cgraph->nodes[node_n];

-        if (ggml_op_is_empty(node->op)) {
-            // skip NOPs
-            continue;
-        }
+            if (ggml_op_is_empty(node->op)) {
+                continue;
+            }

-        if ((node->flags & GGML_TENSOR_FLAG_COMPUTE) == 0) {
-            continue;
+            if ((node->flags & GGML_TENSOR_FLAG_COMPUTE) == 0) {
+                continue;
+            }
+
+            // Only thread 0 records timing (after barrier = total node time)
+            uint64_t t_start = 0;
+            if (state->ith == 0) {
+                t_start = ggml_profiler_time_ns();
+            }
+
+            ggml_compute_forward(&params, node);
+
+            if (node_n + 1 < cgraph->n_nodes) {
+                ggml_barrier(state->threadpool);
+            }
+
+            if (state->ith == 0) {
+                uint64_t t_end = ggml_profiler_time_ns();
+                cplan->profiling_record_fn(cplan->profiling_context, 0 /* GGML_PROFILE_EVENT_OP */,
+                                           ggml_op_name(node->op), -1, t_start, t_end, ggml_nbytes(node), NULL,
+                                           node);
+            }
+
+            if (state->ith == 0 && cplan->abort_callback && cplan->abort_callback(cplan->abort_callback_data)) {
+                atomic_store_explicit(&tp->abort, node_n + 1, memory_order_relaxed);
+                tp->ec = GGML_STATUS_ABORTED;
+            }
        }
+    } else {
+        for (int node_n = 0; node_n < cgraph->n_nodes && atomic_load_explicit(&tp->abort, memory_order_relaxed) != node_n; node_n++) {
+            struct ggml_tensor * node = cgraph->nodes[node_n];
+
+            if (ggml_op_is_empty(node->op)) {
+                // skip NOPs
+                continue;
+            }
+
+            if ((node->flags & GGML_TENSOR_FLAG_COMPUTE) == 0) {
+                continue;
+            }

        // TODO: move fused-op detection into ggml_graph_plan so fusion decisions are made once at planning time
        // Try fused ops, fall back to normal compute
@@ -3106,14 +3145,15 @@ static thread_ret_t ggml_graph_compute_thread(void * data) {
            ggml_compute_forward(&params, node);
        }

-        if (state->ith == 0 && cplan->abort_callback &&
-                cplan->abort_callback(cplan->abort_callback_data)) {
-            atomic_store_explicit(&tp->abort, node_n + 1, memory_order_relaxed);
-            tp->ec    = GGML_STATUS_ABORTED;
-        }
+            if (state->ith == 0 && cplan->abort_callback &&
+                    cplan->abort_callback(cplan->abort_callback_data)) {
+                atomic_store_explicit(&tp->abort, node_n + 1, memory_order_relaxed);
+                tp->ec    = GGML_STATUS_ABORTED;
+            }

-        if (node_n + 1 < cgraph->n_nodes) {
-            ggml_barrier(state->threadpool);
+            if (node_n + 1 < cgraph->n_nodes) {
+                ggml_barrier(state->threadpool);
+            }
        }
    }

@@ -1,6 +1,7 @@
 #include "ggml-backend.h"
 #include "ggml-backend-impl.h"
 #include "ggml-cpu.h"
+#include "ggml-profiler.h"
 #include "repack.h"
 #include "traits.h"
 #include "ggml-impl.h"
@@ -107,6 +108,11 @@ struct ggml_backend_cpu_context {
    void *              abort_callback_data;

    bool                use_ref;  // use reference implementation
+
+    // Profiling state
+    bool                             profiling_enabled;
+    int                              profiling_split_id;
+    std::vector<ggml_profile_record> profiling_records;
 };

 static const char * ggml_backend_cpu_get_name(ggml_backend_t backend) {
@@ -167,6 +173,30 @@ static enum ggml_status ggml_backend_cpu_graph_plan_compute(ggml_backend_t backe
    GGML_UNUSED(backend);
 }

+// Callback function for recording CPU profiling events from C code (ggml-cpu.c)
+static void ggml_cpu_profiler_record_callback(void *                     context,
+                                               int                        type,
+                                               const char *               name,
+                                               int                        split_id,
+                                               uint64_t                   start_ns,
+                                               uint64_t                   end_ns,
+                                               uint64_t                   bytes,
+                                               const char *               extra,
+                                               const struct ggml_tensor * node) {
+    auto *              cpu_ctx = (ggml_backend_cpu_context *) context;
+    ggml_profile_record rec;
+    rec.type       = (enum ggml_profile_event_type) type;
+    rec.name       = name;
+    rec.backend_id = 0;  // will be overwritten by scheduler
+    rec.split_id   = split_id != -1 ? split_id : cpu_ctx->profiling_split_id;
+    rec.start_ns   = start_ns;
+    rec.end_ns     = end_ns;
+    rec.bytes      = bytes;
+    rec.extra      = extra;
+    ggml_profile_record_from_tensor(&rec, node);
+    cpu_ctx->profiling_records.push_back(rec);
+}
+
 static enum ggml_status ggml_backend_cpu_graph_compute(ggml_backend_t backend, struct ggml_cgraph * cgraph) {
    struct ggml_backend_cpu_context * cpu_ctx = (struct ggml_backend_cpu_context *)backend->context;

@@ -187,6 +217,9 @@ static enum ggml_status ggml_backend_cpu_graph_compute(ggml_backend_t backend, s
    cplan.abort_callback_data = cpu_ctx->abort_callback_data;
    cplan.use_ref             = cpu_ctx->use_ref;

+    cplan.profiling_context   = cpu_ctx->profiling_enabled ? cpu_ctx : NULL;
+    cplan.profiling_record_fn = cpu_ctx->profiling_enabled ? ggml_cpu_profiler_record_callback : NULL;
+
    return ggml_graph_compute(cgraph, &cplan);
 }

@@ -230,12 +263,15 @@ ggml_backend_t ggml_backend_cpu_init(void) {
    ctx->abort_callback      = NULL;
    ctx->abort_callback_data = NULL;
    ctx->use_ref             = false;
+    ctx->profiling_enabled   = false;
+    ctx->profiling_split_id  = -1;

    ggml_backend_t cpu_backend = new ggml_backend {
-        /* .guid    = */ ggml_backend_cpu_guid(),
-        /* .iface   = */ ggml_backend_cpu_i,
-        /* .device  = */ ggml_backend_reg_dev_get(ggml_backend_cpu_reg(), 0),
-        /* .context = */ ctx,
+        /* .guid     = */ ggml_backend_cpu_guid(),
+        /* .iface    = */ ggml_backend_cpu_i,
+        /* .device   = */ ggml_backend_reg_dev_get(ggml_backend_cpu_reg(), 0),
+        /* .context  = */ ctx,
+        /* .profiler = */ nullptr,
    };

    if (cpu_backend == NULL) {
@@ -243,6 +279,43 @@ ggml_backend_t ggml_backend_cpu_init(void) {
        return NULL;
    }

+    // Register profiler
+    static auto cpu_prof_enable = [](void * ctx, bool enable) {
+        auto * cpu_ctx             = (ggml_backend_cpu_context *) ctx;
+        cpu_ctx->profiling_enabled = enable;
+        if (!enable) {
+            cpu_ctx->profiling_records.clear();
+        }
+    };
+    static auto cpu_prof_reset = [](void * ctx) {
+        auto * cpu_ctx = (ggml_backend_cpu_context *) ctx;
+        cpu_ctx->profiling_records.clear();
+        cpu_ctx->profiling_split_id = -1;
+    };
+    static auto cpu_prof_set_split_id = [](void * ctx, int split_id) {
+        auto * cpu_ctx              = (ggml_backend_cpu_context *) ctx;
+        cpu_ctx->profiling_split_id = split_id;
+    };
+    static auto cpu_prof_get_records = [](void * ctx, const ggml_profile_record ** out) -> int {
+        auto * cpu_ctx = (ggml_backend_cpu_context *) ctx;
+        *out           = cpu_ctx->profiling_records.data();
+        return (int) cpu_ctx->profiling_records.size();
+    };
+    static auto cpu_prof_free = [](void * ctx) {
+        // Nothing to free - records are in the CPU context's vector
+        (void) ctx;
+    };
+
+    auto * profiler = new ggml_backend_profiler{
+        /* .context      = */ ctx,
+        /* .enable       = */ cpu_prof_enable,
+        /* .reset        = */ cpu_prof_reset,
+        /* .set_split_id = */ cpu_prof_set_split_id,
+        /* .get_records  = */ cpu_prof_get_records,
+        /* .free_context = */ cpu_prof_free,
+    };
+    ggml_backend_set_profiler(cpu_backend, profiler);
+
    return cpu_backend;
 }

@@ -1719,6 +1719,7 @@ class extra_buffer_type : ggml::cpu::extra_buffer_type {

            return true;
        }
+
        return false;
    }

@@ -1727,6 +1728,20 @@ class extra_buffer_type : ggml::cpu::extra_buffer_type {
            if (op->src[0]->buffer && op->src[0]->buffer->buft == ggml_backend_cpu_kleidiai_buffer_type()) {
                return (ggml::cpu::tensor_traits *) op->src[0]->extra;
            } else {
+                // KleidiAI only has kernels for Q4_0 and Q8_0. For a quantized weight of any
+                // other type (K-quants, IQ) it declines the op and returns nullptr below, so
+                // KleidiAI does not accelerate it. Another CPU backend may still take the op,
+                // and this can run during graph planning, so the message says what KleidiAI
+                // did rather than what ends up executing. Warn once per process.
+                if (ggml_is_quantized(op->src[0]->type) &&
+                    op->src[0]->type != GGML_TYPE_Q4_0 && op->src[0]->type != GGML_TYPE_Q8_0) {
+                    static std::atomic<bool> warned(false);
+                    if (!warned.exchange(true)) {
+                        GGML_LOG_WARN("kleidiai: no kernel for tensor type %s, not accelerated by KleidiAI "
+                                      "(kernels available for Q4_0 and Q8_0)\n",
+                                      ggml_type_name(op->src[0]->type));
+                    }
+                }
                if (op->src[0]->type != GGML_TYPE_F16) {
                    return nullptr;
                }
@@ -2329,7 +2329,7 @@ class tinyBLAS_Q0_PPC {
        mc = 32;
        nc = 32;
        kc = 32;
-        n_chunk = 32
+        n_chunk = 32;
    #endif
        int64_t n_aligned = 0;
        if (n % n_chunk == 0) {
@@ -362,6 +362,15 @@ static bool blackwell_mma_available(const int cc) {
           ggml_cuda_highest_compiled_arch(cc) < GGML_CUDA_CC_RUBIN;
 }

+// Checks whether the tensor's base data pointer and higher-dimensional strides are byte-aligned to `alignment` bytes.
+static bool ggml_cuda_is_aligned(const ggml_tensor * tensor, const size_t alignment) {
+    GGML_ASSERT(tensor != nullptr);
+    return (reinterpret_cast<uintptr_t>(tensor->data) % alignment) == 0 &&
+           tensor->nb[1] % alignment == 0 &&
+           tensor->nb[2] % alignment == 0 &&
+           tensor->nb[3] % alignment == 0;
+}
+
 static constexpr __device__ int ggml_cuda_get_physical_warp_size() {
 #if defined(GGML_USE_HIP) && (defined(__GFX9__) || defined(__GFX8__))
    return 64;
@@ -937,6 +946,9 @@ static __device__ __forceinline__ uint2 fast_div_modulo(uint32_t n, const uint3

 typedef void (*dequantize_kernel_t)(const void * vx, const int64_t ib, const int iqs, float2 & v);

+template<typename dst_t>
+using dequantize_kq_t = void (*)(const void * vx, const int64_t ib, dst_t * y, const int tid);
+
 static __device__ __forceinline__ float get_alibi_slope(
    const float max_bias, const uint32_t h, const uint32_t n_head_log2, const float m0, const float m1
 ) {
@@ -1392,6 +1404,9 @@ struct ggml_cuda_stream_context {
    }
 };

+// Forward declaration for profiler state (defined in ggml-cuda.cu)
+struct ggml_cuda_profiler_state;
+
 struct ggml_backend_cuda_context {
    int device;
    std::string name;
@@ -1503,6 +1518,9 @@ struct ggml_backend_cuda_context {
    ggml_cuda_pool & pool() {
        return pool(device);
    }
+
+    // Profiling
+    ggml_cuda_profiler_state * profiler_state = nullptr;
 };

 struct ggml_cuda_mm_fusion_args_host {
@@ -140,358 +140,107 @@ static __global__ void dequantize_block_q4_1(const void * __restrict__ vx, dst_t

 template<typename dst_t>
 static __global__ void dequantize_block_q2_K(const void * __restrict__ vx, dst_t * __restrict__ yy) {
+    const int64_t i = blockIdx.x;

-    const int64_t i   = blockIdx.x;
-    const block_q2_K * x = (const block_q2_K *) vx;
-
-    const int64_t tid = threadIdx.x;
-    const int64_t n   = tid/32;
-    const int64_t l   = tid - 32*n;
-    const int64_t is  = 8*n + l/16;
-
-    const uint8_t q = x[i].qs[32*n + l];
-    dst_t * y = yy + i*QK_K + 128*n;
-
-    float dall = __low2half(x[i].dm);
-    float dmin = __high2half(x[i].dm);
-    y[l+ 0] = ggml_cuda_cast<dst_t>(dall * (x[i].scales[is+0] & 0xF) * ((q >> 0) & 3) - dmin * (x[i].scales[is+0] >> 4));
-    y[l+32] = ggml_cuda_cast<dst_t>(dall * (x[i].scales[is+2] & 0xF) * ((q >> 2) & 3) - dmin * (x[i].scales[is+2] >> 4));
-    y[l+64] = ggml_cuda_cast<dst_t>(dall * (x[i].scales[is+4] & 0xF) * ((q >> 4) & 3) - dmin * (x[i].scales[is+4] >> 4));
-    y[l+96] = ggml_cuda_cast<dst_t>(dall * (x[i].scales[is+6] & 0xF) * ((q >> 6) & 3) - dmin * (x[i].scales[is+6] >> 4));
+    dequantize_q2_K(vx, i, yy + i*QK_K, threadIdx.x);
 }

 template<typename dst_t>
 static __global__ void dequantize_block_q3_K(const void * __restrict__ vx, dst_t * __restrict__ yy) {
-
    const int64_t i = blockIdx.x;
-    const block_q3_K * x = (const block_q3_K *) vx;

-    const int64_t r = threadIdx.x/4;
-    const int64_t tid = r/2;
-    const int64_t is0 = r%2;
-    const int64_t l0 = 16*is0 + 4*(threadIdx.x%4);
-    const int64_t n = tid / 4;
-    const int64_t j = tid - 4*n;
-
-    uint8_t m = 1 << (4*n + j);
-    int64_t is = 8*n + 2*j + is0;
-    int shift = 2*j;
-
-    int8_t us = is <  4 ? (x[i].scales[is-0] & 0xF) | (((x[i].scales[is+8] >> 0) & 3) << 4) :
-                is <  8 ? (x[i].scales[is-0] & 0xF) | (((x[i].scales[is+4] >> 2) & 3) << 4) :
-                is < 12 ? (x[i].scales[is-8] >>  4) | (((x[i].scales[is+0] >> 4) & 3) << 4) :
-                          (x[i].scales[is-8] >>  4) | (((x[i].scales[is-4] >> 6) & 3) << 4);
-    float d_all = x[i].d;
-    float dl = d_all * (us - 32);
-
-    dst_t * y = yy + i*QK_K + 128*n + 32*j;
-    const uint8_t * q = x[i].qs + 32*n;
-    const uint8_t * hm = x[i].hmask;
-
-    for (int l = l0; l < l0+4; ++l) {
-        y[l] = ggml_cuda_cast<dst_t>(dl * ((int8_t)((q[l] >> shift) & 3) - ((hm[l] & m) ? 0 : 4)));
-    }
-}
-
-static inline __device__ void get_scale_min_k4(int j, const uint8_t * q, uint8_t & d, uint8_t & m) {
-    if (j < 4) {
-        d = q[j] & 63; m = q[j + 4] & 63;
-    } else {
-        d = (q[j+4] & 0xF) | ((q[j-4] >> 6) << 4);
-        m = (q[j+4] >>  4) | ((q[j-0] >> 6) << 4);
-    }
+    dequantize_q3_K(vx, i, yy + i*QK_K, threadIdx.x);
 }

 template<typename dst_t>
 static __global__ void dequantize_block_q4_K(const void * __restrict__ vx, dst_t * __restrict__ yy) {
-    const block_q4_K * x = (const block_q4_K *) vx;
-
    const int64_t i = blockIdx.x;

-    // assume 32 threads
-    const int64_t tid = threadIdx.x;
-    const int64_t il  = tid/8;
-    const int64_t ir  = tid%8;
-    const int64_t is  = 2*il;
-    const int64_t n   = 4;
-
-    dst_t * y = yy + i*QK_K + 64*il + n*ir;
-
-    const float dall = __low2half(x[i].dm);
-    const float dmin = __high2half(x[i].dm);
-
-    const uint8_t * q = x[i].qs + 32*il + n*ir;
-
-    uint8_t sc, m;
-    get_scale_min_k4(is + 0, x[i].scales, sc, m);
-    const float d1 = dall * sc; const float m1 = dmin * m;
-    get_scale_min_k4(is + 1, x[i].scales, sc, m);
-    const float d2 = dall * sc; const float m2 = dmin * m;
-    for (int l = 0; l < n; ++l) {
-        y[l + 0] = ggml_cuda_cast<dst_t>(d1 * (q[l] & 0xF) - m1);
-        y[l +32] = ggml_cuda_cast<dst_t>(d2 * (q[l] >>  4) - m2);
-    }
+    dequantize_q4_K(vx, i, yy + i*QK_K, threadIdx.x);
 }

 template<typename dst_t>
 static __global__ void dequantize_block_q5_K(const void * __restrict__ vx, dst_t * __restrict__ yy) {
-    const block_q5_K * x = (const block_q5_K *) vx;
-
    const int64_t i = blockIdx.x;

-    // assume 64 threads - this is very slightly better than the one below
-    const int64_t tid = threadIdx.x;
-    const int64_t il  = tid/16;   // il is in 0...3
-    const int64_t ir  = tid%16;   // ir is in 0...15
-    const int64_t is  = 2*il;     // is is in 0...6
-
-    dst_t * y = yy + i*QK_K + 64*il + 2*ir;
-
-    const float dall = __low2half(x[i].dm);
-    const float dmin = __high2half(x[i].dm);
-
-    const uint8_t * ql = x[i].qs + 32*il + 2*ir;
-    const uint8_t * qh = x[i].qh + 2*ir;
-
-    uint8_t sc, m;
-    get_scale_min_k4(is + 0, x[i].scales, sc, m);
-    const float d1 = dall * sc; const float m1 = dmin * m;
-    get_scale_min_k4(is + 1, x[i].scales, sc, m);
-    const float d2 = dall * sc; const float m2 = dmin * m;
-
-    uint8_t   hm  = 1 << (2*il);
-    y[ 0] = ggml_cuda_cast<dst_t>(d1 * ((ql[ 0] & 0xF) + (qh[ 0] & hm ? 16 : 0)) - m1);
-    y[ 1] = ggml_cuda_cast<dst_t>(d1 * ((ql[ 1] & 0xF) + (qh[ 1] & hm ? 16 : 0)) - m1);
-    hm <<= 1;
-    y[32] = ggml_cuda_cast<dst_t>(d2 * ((ql[ 0] >>  4) + (qh[ 0] & hm ? 16 : 0)) - m2);
-    y[33] = ggml_cuda_cast<dst_t>(d2 * ((ql[ 1] >>  4) + (qh[ 1] & hm ? 16 : 0)) - m2);
+    dequantize_q5_K(vx, i, yy + i*QK_K, threadIdx.x);
 }

 template<typename dst_t>
 static __global__ void dequantize_block_q6_K(const void * __restrict__ vx, dst_t * __restrict__ yy) {
-    const block_q6_K * x = (const block_q6_K *) vx;
-
    const int64_t i = blockIdx.x;

-    // assume 64 threads - this is very slightly better than the one below
-    const int64_t tid = threadIdx.x;
-    const int64_t ip  = tid/32;   // ip is 0 or 1
-    const int64_t il  = tid - 32*ip; // 0...32
-    const int64_t is  = 8*ip + il/16;
-
-    dst_t * y = yy + i*QK_K + 128*ip + il;
-
-    const float d = x[i].d;
-
-    const uint8_t * ql = x[i].ql + 64*ip + il;
-    const uint8_t   qh = x[i].qh[32*ip + il];
-    const int8_t  * sc = x[i].scales + is;
-
-    y[ 0] = ggml_cuda_cast<dst_t>(d * sc[0] * ((int8_t)((ql[ 0] & 0xF) | (((qh >> 0) & 3) << 4)) - 32));
-    y[32] = ggml_cuda_cast<dst_t>(d * sc[2] * ((int8_t)((ql[32] & 0xF) | (((qh >> 2) & 3) << 4)) - 32));
-    y[64] = ggml_cuda_cast<dst_t>(d * sc[4] * ((int8_t)((ql[ 0]  >> 4) | (((qh >> 4) & 3) << 4)) - 32));
-    y[96] = ggml_cuda_cast<dst_t>(d * sc[6] * ((int8_t)((ql[32]  >> 4) | (((qh >> 6) & 3) << 4)) - 32));
+    dequantize_q6_K(vx, i, yy + i*QK_K, threadIdx.x);
 }

 template<typename dst_t>
 static __global__ void dequantize_block_iq2_xxs(const void * __restrict__ vx, dst_t * __restrict__ yy) {
+    const int64_t i = blockIdx.x;

-    const int64_t i   = blockIdx.x;
-    const block_iq2_xxs * x = (const block_iq2_xxs  *) vx;
-
-    const int64_t tid = threadIdx.x;
-    const int64_t il = tid/8; // 0...3
-    const int64_t ib = tid%8; // 0...7
-    dst_t * y = yy + i*QK_K + 32*ib + 8*il;
-    const uint16_t * q2 = x[i].qs + 4*ib;
-    const uint8_t  * aux8 = (const uint8_t *)q2;
-    const uint8_t  * grid = (const uint8_t *)(iq2xxs_grid + aux8[il]);
-    const uint32_t aux32 = q2[2] | (q2[3] << 16);
-    const float d = (float)x[i].d * (0.5f + (aux32 >> 28)) * 0.25f;
-    const uint8_t signs = ksigns_iq2xs[(aux32 >> 7*il) & 127];
-    for (int j = 0; j < 8; ++j) {
-        y[j] = ggml_cuda_cast<dst_t>(d * grid[j] * (signs & kmask_iq2xs[j] ? -1.f : 1.f));
-    }
+    dequantize_iq2_xxs(vx, i, yy + i*QK_K, threadIdx.x);
 }

 template<typename dst_t>
 static __global__ void dequantize_block_iq2_xs(const void * __restrict__ vx, dst_t * __restrict__ yy) {
+    const int64_t i = blockIdx.x;

-    const int64_t i   = blockIdx.x;
-    const block_iq2_xs * x = (const block_iq2_xs *) vx;
-
-    const int64_t tid = threadIdx.x;
-    const int64_t il = tid/8; // 0...3
-    const int64_t ib = tid%8; // 0...7
-    dst_t * y = yy + i*QK_K + 32*ib + 8*il;
-    const uint16_t * q2 = x[i].qs + 4*ib;
-    const uint8_t  * grid = (const uint8_t *)(iq2xs_grid + (q2[il] & 511));
-    const float d = (float)x[i].d * (0.5f + ((x[i].scales[ib] >> 4*(il/2)) & 0xf)) * 0.25f;
-    const uint8_t signs = ksigns_iq2xs[q2[il] >> 9];
-    for (int j = 0; j < 8; ++j) {
-        y[j] = ggml_cuda_cast<dst_t>(d * grid[j] * (signs & kmask_iq2xs[j] ? -1.f : 1.f));
-    }
+    dequantize_iq2_xs(vx, i, yy + i*QK_K, threadIdx.x);
 }

 template<typename dst_t>
 static __global__ void dequantize_block_iq2_s(const void * __restrict__ vx, dst_t * __restrict__ yy) {
+    const int64_t i = blockIdx.x;

-    const int64_t i   = blockIdx.x;
-    const block_iq2_s * x = (const block_iq2_s *) vx;
-
-    const int64_t tid = threadIdx.x;
-    const int64_t il = tid/8; // 0...3
-    const int64_t ib = tid%8; // 0...7
-    dst_t * y = yy + i*QK_K + 32*ib + 8*il;
-    const uint8_t * grid = (const uint8_t *)(iq2s_grid + (x[i].qs[4*ib+il] | ((x[i].qh[ib] << (8-2*il)) & 0x300)));
-    const float d = (float)x[i].d * (0.5f + ((x[i].scales[ib] >> 4*(il/2)) & 0xf)) * 0.25f;
-    const uint8_t signs = x[i].qs[QK_K/8+4*ib+il];
-    for (int j = 0; j < 8; ++j) {
-        y[j] = ggml_cuda_cast<dst_t>(d * grid[j] * (signs & kmask_iq2xs[j] ? -1.f : 1.f));
-    }
+    dequantize_iq2_s(vx, i, yy + i*QK_K, threadIdx.x);
 }

 template<typename dst_t>
 static __global__ void dequantize_block_iq3_xxs(const void * __restrict__ vx, dst_t * __restrict__ yy) {
+    const int64_t i = blockIdx.x;

-    const int64_t i   = blockIdx.x;
-    const block_iq3_xxs * x = (const block_iq3_xxs  *) vx;
-
-    const int64_t tid = threadIdx.x;
-    const int64_t il = tid/8; // 0...3
-    const int64_t ib = tid%8; // 0...7
-    dst_t * y = yy + i*QK_K + 32*ib + 8*il;
-    const uint8_t  * q3 = x[i].qs + 8*ib;
-    const uint16_t * gas = (const uint16_t *)(x[i].qs + QK_K/4) + 2*ib;
-    const uint8_t  * grid1 = (const uint8_t *)(iq3xxs_grid + q3[2*il+0]);
-    const uint8_t  * grid2 = (const uint8_t *)(iq3xxs_grid + q3[2*il+1]);
-    const uint32_t aux32 = gas[0] | (gas[1] << 16);
-    const float d = (float)x[i].d * (0.5f + (aux32 >> 28)) * 0.5f;
-    const uint8_t signs = ksigns_iq2xs[(aux32 >> 7*il) & 127];
-    for (int j = 0; j < 4; ++j) {
-        y[j+0] = ggml_cuda_cast<dst_t>(d * grid1[j] * (signs & kmask_iq2xs[j+0] ? -1.f : 1.f));
-        y[j+4] = ggml_cuda_cast<dst_t>(d * grid2[j] * (signs & kmask_iq2xs[j+4] ? -1.f : 1.f));
-    }
+    dequantize_iq3_xxs(vx, i, yy + i*QK_K, threadIdx.x);
 }

 template<typename dst_t>
 static __global__ void dequantize_block_iq3_s(const void * __restrict__ vx, dst_t * __restrict__ yy) {
+    const int64_t i = blockIdx.x;

-    const int64_t i   = blockIdx.x;
-    const block_iq3_s * x = (const block_iq3_s *) vx;
-
-    const int64_t tid = threadIdx.x;
-    const int64_t il = tid/8; // 0...3
-    const int64_t ib = tid%8; // 0...7
-    dst_t * y = yy + i*QK_K + 32*ib + 8*il;
-    const uint8_t * qs = x[i].qs + 8*ib;
-    const uint8_t * grid1 = (const uint8_t *)(iq3s_grid + (qs[2*il+0] | ((x[i].qh[ib] << (8-2*il)) & 256)));
-    const uint8_t * grid2 = (const uint8_t *)(iq3s_grid + (qs[2*il+1] | ((x[i].qh[ib] << (7-2*il)) & 256)));
-    const float d = (float)x[i].d * (1 + 2*((x[i].scales[ib/2] >> 4*(ib%2)) & 0xf));
-    const uint8_t signs = x[i].signs[4*ib + il];
-    for (int j = 0; j < 4; ++j) {
-        y[j+0] = ggml_cuda_cast<dst_t>(d * grid1[j] * (signs & kmask_iq2xs[j+0] ? -1.f : 1.f));
-        y[j+4] = ggml_cuda_cast<dst_t>(d * grid2[j] * (signs & kmask_iq2xs[j+4] ? -1.f : 1.f));
-    }
+    dequantize_iq3_s(vx, i, yy + i*QK_K, threadIdx.x);
 }

 template<typename dst_t>
 static __global__ void dequantize_block_iq1_s(const void * __restrict__ vx, dst_t * __restrict__ yy) {
+    const int64_t i = blockIdx.x;

-    const int64_t i   = blockIdx.x;
-    const block_iq1_s * x = (const block_iq1_s  *) vx;
-
-    const int64_t tid = threadIdx.x;
-    const int64_t il = tid/8; // 0...3
-    const int64_t ib = tid%8; // 0...7
-    dst_t * y = yy + i*QK_K + 32*ib + 8*il;
-    const float delta = x[i].qh[ib] & 0x8000 ? -1 - IQ1S_DELTA : -1 + IQ1S_DELTA;
-    const float d = (float)x[i].d * (2*((x[i].qh[ib] >> 12) & 7) + 1);
-    uint32_t grid32[2]; const int8_t * q = (const int8_t *)grid32;
-    grid32[0] = iq1s_grid_gpu[x[i].qs[4*ib+il] | (((x[i].qh[ib] >> 3*il) & 7) << 8)];
-    grid32[1] = (grid32[0] >> 4) & 0x0f0f0f0f;
-    grid32[0] &= 0x0f0f0f0f;
-    for (int j = 0; j < 8; ++j) {
-        y[j] = ggml_cuda_cast<dst_t>(d * (q[j] + delta));
-    }
+    dequantize_iq1_s(vx, i, yy + i*QK_K, threadIdx.x);
 }

 template<typename dst_t>
 static __global__ void dequantize_block_iq1_m(const void * __restrict__ vx, dst_t * __restrict__ yy) {
+    const int64_t i = blockIdx.x;

-    const int64_t i   = blockIdx.x;
-    const block_iq1_m * x = (const block_iq1_m  *) vx;
-
-    const int64_t tid = threadIdx.x;
-    const int64_t il = tid/8; // 0...3
-    const int64_t ib = tid%8; // 0...7
-    dst_t * y = yy + i*QK_K + 32*ib + 8*il;
-    const uint16_t * sc = (const uint16_t *)x[i].scales;
-    iq1m_scale_t scale;
-    scale.u16 = (sc[0] >> 12) | ((sc[1] >> 8) & 0x00f0) | ((sc[2] >> 4) & 0x0f00) | (sc[3] & 0xf000);
-    const int64_t ib16 = 2*ib + il/2; // sc[ib16/4] >> 3*(ib16%4) -> sc[ib/2] >> 3*((2*ib+il/2)%4);
-    const float d = (float)scale.f16 * (2*((sc[ib16/4] >> 3*(ib16%4)) & 0x7) + 1);
-    const float delta = x[i].qh[2*ib+il/2] & (0x08 << 4*(il%2)) ? -1 - IQ1M_DELTA : -1 + IQ1M_DELTA;
-    uint32_t grid32[2]; const int8_t * q = (const int8_t *)grid32;
-    grid32[0] = iq1s_grid_gpu[x[i].qs[4*ib+il] | (((x[i].qh[2*ib+il/2] >> 4*(il%2)) & 7) << 8)];
-    grid32[1] = (grid32[0] >> 4) & 0x0f0f0f0f;
-    grid32[0] &= 0x0f0f0f0f;
-    for (int j = 0; j < 8; ++j) {
-        y[j] = ggml_cuda_cast<dst_t>(d * (q[j] + delta));
-    }
+    dequantize_iq1_m(vx, i, yy + i*QK_K, threadIdx.x);
 }

 template<typename dst_t>
 static __global__ void dequantize_block_iq4_nl(const void * __restrict__ vx, dst_t * __restrict__ yy) {
+    const int64_t i = blockIdx.x;

-    const int64_t i   = blockIdx.x;
-    const block_iq4_nl * x = (const block_iq4_nl *) vx + i*(QK_K/QK4_NL);
-
-    const int64_t tid = threadIdx.x;
-    const int64_t il = tid/8; // 0...3
-    const int64_t ib = tid%8; // 0...7
-    dst_t * y = yy + i*QK_K + 32*ib + 4*il;
-    const uint8_t  * q4 = x[ib].qs + 4*il;
-    const float d = (float)x[ib].d;
-    for (int j = 0; j < 4; ++j) {
-        y[j+ 0] = ggml_cuda_cast<dst_t>(d * kvalues_iq4nl[q4[j] & 0xf]);
-        y[j+16] = ggml_cuda_cast<dst_t>(d * kvalues_iq4nl[q4[j] >>  4]);
-    }
+    dequantize_iq4_nl(vx, i, yy + i*QK_K, threadIdx.x);
 }

 template<typename dst_t>
 static __global__ void dequantize_block_iq4_xs(const void * __restrict__ vx, dst_t * __restrict__ yy) {
-    const int64_t i   = blockIdx.x;
-    const block_iq4_xs * x = (const block_iq4_xs *)vx;
+    const int64_t i = blockIdx.x;

-    const int64_t tid = threadIdx.x;
-    const int64_t il = tid/8; // 0...3
-    const int64_t ib = tid%8; // 0...7
-    dst_t * y = yy + i*QK_K + 32*ib + 4*il;
-    const uint8_t  * q4 = x[i].qs + 16*ib + 4*il;
-    const float d = (float)x[i].d * ((((x[i].scales_l[ib/2] >> 4*(ib%2)) & 0xf) | (((x[i].scales_h >> 2*ib) & 3) << 4)) - 32);
-    for (int j = 0; j < 4; ++j) {
-        y[j+ 0] = ggml_cuda_cast<dst_t>(d * kvalues_iq4nl[q4[j] & 0xf]);
-        y[j+16] = ggml_cuda_cast<dst_t>(d * kvalues_iq4nl[q4[j] >>  4]);
-    }
+    dequantize_iq4_xs(vx, i, yy + i*QK_K, threadIdx.x);
 }

 template<typename dst_t>
 static __global__ void dequantize_block_mxfp4(const void * __restrict__ vx, dst_t * __restrict__ yy) {
+    const int64_t i = blockIdx.x;

-    const int64_t i   = blockIdx.x;
-    const block_mxfp4 * x = (const block_mxfp4 *) vx + i*(QK_K/QK_MXFP4);
-
-    const int64_t tid = threadIdx.x;
-    const int64_t il = tid/8; // 0...3
-    const int64_t ib = tid%8; // 0...7
-    dst_t * y = yy + i*QK_K + 32*ib + 4*il;
-    const uint8_t  * q4 = x[ib].qs + 4*il;
-    const float d = ggml_cuda_e8m0_to_fp32(x[ib].e);
-    for (int j = 0; j < 4; ++j) {
-        y[j+ 0] = ggml_cuda_cast<dst_t>(d * kvalues_mxfp4[q4[j] & 0xf]*0.5f);
-        y[j+16] = ggml_cuda_cast<dst_t>(d * kvalues_mxfp4[q4[j] >>  4]*0.5f);
-    }
+    dequantize_mxfp4(vx, i, yy + i*QK_K, threadIdx.x);
 }

 template <int qk, int qr, dequantize_kernel_t dequantize_kernel, typename dst_t>
@@ -1,4 +1,5 @@
 #include "common.cuh"
+#include "convert.cuh"

 static __device__ __forceinline__ void dequantize_q1_0(const void * vx, const int64_t ib, const int iqs, float2 & v){
    const block_q1_0 * x = (const block_q1_0 *) vx;
@@ -97,3 +98,335 @@ static __device__ __forceinline__ void dequantize_q8_0(const void * vx, const in
    v.x *= d;
    v.y *= d;
 }
+
+//================================== k-quants
+
+// Each call dequantizes one super-block of QK_K values into y using the
+// thread layout of the caller: 32 threads for q4_K, 64 threads otherwise.
+
+template<typename dst_t>
+static __device__ __forceinline__ void dequantize_q2_K(const void * vx, const int64_t ib, dst_t * yy, const int tid) {
+    const block_q2_K * x = (const block_q2_K *) vx;
+
+    const int64_t n   = tid/32;
+    const int64_t l   = tid - 32*n;
+    const int64_t is  = 8*n + l/16;
+
+    const uint8_t q = x[ib].qs[32*n + l];
+    dst_t * y = yy + 128*n;
+
+    float dall = __low2half(x[ib].dm);
+    float dmin = __high2half(x[ib].dm);
+    y[l+ 0] = ggml_cuda_cast<dst_t>(dall * (x[ib].scales[is+0] & 0xF) * ((q >> 0) & 3) - dmin * (x[ib].scales[is+0] >> 4));
+    y[l+32] = ggml_cuda_cast<dst_t>(dall * (x[ib].scales[is+2] & 0xF) * ((q >> 2) & 3) - dmin * (x[ib].scales[is+2] >> 4));
+    y[l+64] = ggml_cuda_cast<dst_t>(dall * (x[ib].scales[is+4] & 0xF) * ((q >> 4) & 3) - dmin * (x[ib].scales[is+4] >> 4));
+    y[l+96] = ggml_cuda_cast<dst_t>(dall * (x[ib].scales[is+6] & 0xF) * ((q >> 6) & 3) - dmin * (x[ib].scales[is+6] >> 4));
+}
+
+template<typename dst_t>
+static __device__ __forceinline__ void dequantize_q3_K(const void * vx, const int64_t ib, dst_t * yy, const int tid) {
+    const block_q3_K * x = (const block_q3_K *) vx;
+
+    const int64_t r = tid/4;
+    const int64_t t = r/2;
+    const int64_t is0 = r%2;
+    const int64_t l0 = 16*is0 + 4*(tid%4);
+    const int64_t n = t / 4;
+    const int64_t j = t - 4*n;
+
+    uint8_t m = 1 << (4*n + j);
+    int64_t is = 8*n + 2*j + is0;
+    int shift = 2*j;
+
+    int8_t us = is <  4 ? (x[ib].scales[is-0] & 0xF) | (((x[ib].scales[is+8] >> 0) & 3) << 4) :
+                is <  8 ? (x[ib].scales[is-0] & 0xF) | (((x[ib].scales[is+4] >> 2) & 3) << 4) :
+                is < 12 ? (x[ib].scales[is-8] >>  4) | (((x[ib].scales[is+0] >> 4) & 3) << 4) :
+                          (x[ib].scales[is-8] >>  4) | (((x[ib].scales[is-4] >> 6) & 3) << 4);
+    float d_all = x[ib].d;
+    float dl = d_all * (us - 32);
+
+    dst_t * y = yy + 128*n + 32*j;
+    const uint8_t * q = x[ib].qs + 32*n;
+    const uint8_t * hm = x[ib].hmask;
+
+    for (int l = l0; l < l0+4; ++l) {
+        y[l] = ggml_cuda_cast<dst_t>(dl * ((int8_t)((q[l] >> shift) & 3) - ((hm[l] & m) ? 0 : 4)));
+    }
+}
+
+static inline __device__ void get_scale_min_k4(int j, const uint8_t * q, uint8_t & d, uint8_t & m) {
+    if (j < 4) {
+        d = q[j] & 63; m = q[j + 4] & 63;
+    } else {
+        d = (q[j+4] & 0xF) | ((q[j-4] >> 6) << 4);
+        m = (q[j+4] >>  4) | ((q[j-0] >> 6) << 4);
+    }
+}
+
+template<typename dst_t>
+static __device__ __forceinline__ void dequantize_q4_K(const void * vx, const int64_t ib, dst_t * yy, const int tid) {
+    const block_q4_K * x = (const block_q4_K *) vx;
+
+    // assume 32 threads
+    const int64_t il  = tid/8;
+    const int64_t ir  = tid%8;
+    const int64_t is  = 2*il;
+    const int64_t n   = 4;
+
+    dst_t * y = yy + 64*il + n*ir;
+
+    const float dall = __low2half(x[ib].dm);
+    const float dmin = __high2half(x[ib].dm);
+
+    const uint8_t * q = x[ib].qs + 32*il + n*ir;
+
+    uint8_t sc, m;
+    get_scale_min_k4(is + 0, x[ib].scales, sc, m);
+    const float d1 = dall * sc; const float m1 = dmin * m;
+    get_scale_min_k4(is + 1, x[ib].scales, sc, m);
+    const float d2 = dall * sc; const float m2 = dmin * m;
+    for (int l = 0; l < n; ++l) {
+        y[l + 0] = ggml_cuda_cast<dst_t>(d1 * (q[l] & 0xF) - m1);
+        y[l +32] = ggml_cuda_cast<dst_t>(d2 * (q[l] >>  4) - m2);
+    }
+}
+
+template<typename dst_t>
+static __device__ __forceinline__ void dequantize_q5_K(const void * vx, const int64_t ib, dst_t * yy, const int tid) {
+    const block_q5_K * x = (const block_q5_K *) vx;
+
+    // assume 64 threads - this is very slightly better than the one below
+    const int64_t il  = tid/16;   // il is in 0...3
+    const int64_t ir  = tid%16;   // ir is in 0...15
+    const int64_t is  = 2*il;     // is is in 0...6
+
+    dst_t * y = yy + 64*il + 2*ir;
+
+    const float dall = __low2half(x[ib].dm);
+    const float dmin = __high2half(x[ib].dm);
+
+    const uint8_t * ql = x[ib].qs + 32*il + 2*ir;
+    const uint8_t * qh = x[ib].qh + 2*ir;
+
+    uint8_t sc, m;
+    get_scale_min_k4(is + 0, x[ib].scales, sc, m);
+    const float d1 = dall * sc; const float m1 = dmin * m;
+    get_scale_min_k4(is + 1, x[ib].scales, sc, m);
+    const float d2 = dall * sc; const float m2 = dmin * m;
+
+    uint8_t   hm  = 1 << (2*il);
+    y[ 0] = ggml_cuda_cast<dst_t>(d1 * ((ql[ 0] & 0xF) + (qh[ 0] & hm ? 16 : 0)) - m1);
+    y[ 1] = ggml_cuda_cast<dst_t>(d1 * ((ql[ 1] & 0xF) + (qh[ 1] & hm ? 16 : 0)) - m1);
+    hm <<= 1;
+    y[32] = ggml_cuda_cast<dst_t>(d2 * ((ql[ 0] >>  4) + (qh[ 0] & hm ? 16 : 0)) - m2);
+    y[33] = ggml_cuda_cast<dst_t>(d2 * ((ql[ 1] >>  4) + (qh[ 1] & hm ? 16 : 0)) - m2);
+}
+
+template<typename dst_t>
+static __device__ __forceinline__ void dequantize_q6_K(const void * vx, const int64_t ib, dst_t * yy, const int tid) {
+    const block_q6_K * x = (const block_q6_K *) vx;
+
+    // assume 64 threads - this is very slightly better than the one below
+    const int64_t ip  = tid/32;   // ip is 0 or 1
+    const int64_t il  = tid - 32*ip; // 0...32
+    const int64_t is  = 8*ip + il/16;
+
+    dst_t * y = yy + 128*ip + il;
+
+    const float d = x[ib].d;
+
+    const uint8_t * ql = x[ib].ql + 64*ip + il;
+    const uint8_t   qh = x[ib].qh[32*ip + il];
+    const int8_t  * sc = x[ib].scales + is;
+
+    y[ 0] = ggml_cuda_cast<dst_t>(d * sc[0] * ((int8_t)((ql[ 0] & 0xF) | (((qh >> 0) & 3) << 4)) - 32));
+    y[32] = ggml_cuda_cast<dst_t>(d * sc[2] * ((int8_t)((ql[32] & 0xF) | (((qh >> 2) & 3) << 4)) - 32));
+    y[64] = ggml_cuda_cast<dst_t>(d * sc[4] * ((int8_t)((ql[ 0]  >> 4) | (((qh >> 4) & 3) << 4)) - 32));
+    y[96] = ggml_cuda_cast<dst_t>(d * sc[6] * ((int8_t)((ql[32]  >> 4) | (((qh >> 6) & 3) << 4)) - 32));
+}
+
+//================================== i-quants
+
+// Each call dequantizes one super-block of QK_K values into y with 32
+// threads; iq4_nl packs QK_K/QK4_NL sub-blocks per super-block.
+
+template<typename dst_t>
+static __device__ __forceinline__ void dequantize_iq2_xxs(const void * vx, const int64_t ibs, dst_t * yy, const int tid) {
+
+    const block_iq2_xxs * x = (const block_iq2_xxs  *) vx;
+
+    const int64_t il = tid/8; // 0...3
+    const int64_t ib = tid%8; // 0...7
+    dst_t * y = yy + 32*ib + 8*il;
+    const uint16_t * q2 = x[ibs].qs + 4*ib;
+    const uint8_t  * aux8 = (const uint8_t *)q2;
+    const uint8_t  * grid = (const uint8_t *)(iq2xxs_grid + aux8[il]);
+    const uint32_t aux32 = q2[2] | (q2[3] << 16);
+    const float d = (float)x[ibs].d * (0.5f + (aux32 >> 28)) * 0.25f;
+    const uint8_t signs = ksigns_iq2xs[(aux32 >> 7*il) & 127];
+    for (int j = 0; j < 8; ++j) {
+        y[j] = ggml_cuda_cast<dst_t>(d * grid[j] * (signs & kmask_iq2xs[j] ? -1.f : 1.f));
+    }
+}
+
+template<typename dst_t>
+static __device__ __forceinline__ void dequantize_iq2_xs(const void * vx, const int64_t ibs, dst_t * yy, const int tid) {
+
+    const block_iq2_xs * x = (const block_iq2_xs *) vx;
+
+    const int64_t il = tid/8; // 0...3
+    const int64_t ib = tid%8; // 0...7
+    dst_t * y = yy + 32*ib + 8*il;
+    const uint16_t * q2 = x[ibs].qs + 4*ib;
+    const uint8_t  * grid = (const uint8_t *)(iq2xs_grid + (q2[il] & 511));
+    const float d = (float)x[ibs].d * (0.5f + ((x[ibs].scales[ib] >> 4*(il/2)) & 0xf)) * 0.25f;
+    const uint8_t signs = ksigns_iq2xs[q2[il] >> 9];
+    for (int j = 0; j < 8; ++j) {
+        y[j] = ggml_cuda_cast<dst_t>(d * grid[j] * (signs & kmask_iq2xs[j] ? -1.f : 1.f));
+    }
+}
+
+template<typename dst_t>
+static __device__ __forceinline__ void dequantize_iq2_s(const void * vx, const int64_t ibs, dst_t * yy, const int tid) {
+
+    const block_iq2_s * x = (const block_iq2_s *) vx;
+
+    const int64_t il = tid/8; // 0...3
+    const int64_t ib = tid%8; // 0...7
+    dst_t * y = yy + 32*ib + 8*il;
+    const uint8_t * grid = (const uint8_t *)(iq2s_grid + (x[ibs].qs[4*ib+il] | ((x[ibs].qh[ib] << (8-2*il)) & 0x300)));
+    const float d = (float)x[ibs].d * (0.5f + ((x[ibs].scales[ib] >> 4*(il/2)) & 0xf)) * 0.25f;
+    const uint8_t signs = x[ibs].qs[QK_K/8+4*ib+il];
+    for (int j = 0; j < 8; ++j) {
+        y[j] = ggml_cuda_cast<dst_t>(d * grid[j] * (signs & kmask_iq2xs[j] ? -1.f : 1.f));
+    }
+}
+
+template<typename dst_t>
+static __device__ __forceinline__ void dequantize_iq3_xxs(const void * vx, const int64_t ibs, dst_t * yy, const int tid) {
+
+    const block_iq3_xxs * x = (const block_iq3_xxs  *) vx;
+
+    const int64_t il = tid/8; // 0...3
+    const int64_t ib = tid%8; // 0...7
+    dst_t * y = yy + 32*ib + 8*il;
+    const uint8_t  * q3 = x[ibs].qs + 8*ib;
+    const uint16_t * gas = (const uint16_t *)(x[ibs].qs + QK_K/4) + 2*ib;
+    const uint8_t  * grid1 = (const uint8_t *)(iq3xxs_grid + q3[2*il+0]);
+    const uint8_t  * grid2 = (const uint8_t *)(iq3xxs_grid + q3[2*il+1]);
+    const uint32_t aux32 = gas[0] | (gas[1] << 16);
+    const float d = (float)x[ibs].d * (0.5f + (aux32 >> 28)) * 0.5f;
+    const uint8_t signs = ksigns_iq2xs[(aux32 >> 7*il) & 127];
+    for (int j = 0; j < 4; ++j) {
+        y[j+0] = ggml_cuda_cast<dst_t>(d * grid1[j] * (signs & kmask_iq2xs[j+0] ? -1.f : 1.f));
+        y[j+4] = ggml_cuda_cast<dst_t>(d * grid2[j] * (signs & kmask_iq2xs[j+4] ? -1.f : 1.f));
+    }
+}
+
+template<typename dst_t>
+static __device__ __forceinline__ void dequantize_iq3_s(const void * vx, const int64_t ibs, dst_t * yy, const int tid) {
+
+    const block_iq3_s * x = (const block_iq3_s *) vx;
+
+    const int64_t il = tid/8; // 0...3
+    const int64_t ib = tid%8; // 0...7
+    dst_t * y = yy + 32*ib + 8*il;
+    const uint8_t * qs = x[ibs].qs + 8*ib;
+    const uint8_t * grid1 = (const uint8_t *)(iq3s_grid + (qs[2*il+0] | ((x[ibs].qh[ib] << (8-2*il)) & 256)));
+    const uint8_t * grid2 = (const uint8_t *)(iq3s_grid + (qs[2*il+1] | ((x[ibs].qh[ib] << (7-2*il)) & 256)));
+    const float d = (float)x[ibs].d * (1 + 2*((x[ibs].scales[ib/2] >> 4*(ib%2)) & 0xf));
+    const uint8_t signs = x[ibs].signs[4*ib + il];
+    for (int j = 0; j < 4; ++j) {
+        y[j+0] = ggml_cuda_cast<dst_t>(d * grid1[j] * (signs & kmask_iq2xs[j+0] ? -1.f : 1.f));
+        y[j+4] = ggml_cuda_cast<dst_t>(d * grid2[j] * (signs & kmask_iq2xs[j+4] ? -1.f : 1.f));
+    }
+}
+
+template<typename dst_t>
+static __device__ __forceinline__ void dequantize_iq1_s(const void * vx, const int64_t ibs, dst_t * yy, const int tid) {
+
+    const block_iq1_s * x = (const block_iq1_s  *) vx;
+
+    const int64_t il = tid/8; // 0...3
+    const int64_t ib = tid%8; // 0...7
+    dst_t * y = yy + 32*ib + 8*il;
+    const float delta = x[ibs].qh[ib] & 0x8000 ? -1 - IQ1S_DELTA : -1 + IQ1S_DELTA;
+    const float d = (float)x[ibs].d * (2*((x[ibs].qh[ib] >> 12) & 7) + 1);
+    uint32_t grid32[2]; const int8_t * q = (const int8_t *)grid32;
+    grid32[0] = iq1s_grid_gpu[x[ibs].qs[4*ib+il] | (((x[ibs].qh[ib] >> 3*il) & 7) << 8)];
+    grid32[1] = (grid32[0] >> 4) & 0x0f0f0f0f;
+    grid32[0] &= 0x0f0f0f0f;
+    for (int j = 0; j < 8; ++j) {
+        y[j] = ggml_cuda_cast<dst_t>(d * (q[j] + delta));
+    }
+}
+
+template<typename dst_t>
+static __device__ __forceinline__ void dequantize_iq1_m(const void * vx, const int64_t ibs, dst_t * yy, const int tid) {
+
+    const block_iq1_m * x = (const block_iq1_m  *) vx;
+
+    const int64_t il = tid/8; // 0...3
+    const int64_t ib = tid%8; // 0...7
+    dst_t * y = yy + 32*ib + 8*il;
+    const uint16_t * sc = (const uint16_t *)x[ibs].scales;
+    iq1m_scale_t scale;
+    scale.u16 = (sc[0] >> 12) | ((sc[1] >> 8) & 0x00f0) | ((sc[2] >> 4) & 0x0f00) | (sc[3] & 0xf000);
+    const int64_t ib16 = 2*ib + il/2; // sc[ib16/4] >> 3*(ib16%4) -> sc[ib/2] >> 3*((2*ib+il/2)%4);
+    const float d = (float)scale.f16 * (2*((sc[ib16/4] >> 3*(ib16%4)) & 0x7) + 1);
+    const float delta = x[ibs].qh[2*ib+il/2] & (0x08 << 4*(il%2)) ? -1 - IQ1M_DELTA : -1 + IQ1M_DELTA;
+    uint32_t grid32[2]; const int8_t * q = (const int8_t *)grid32;
+    grid32[0] = iq1s_grid_gpu[x[ibs].qs[4*ib+il] | (((x[ibs].qh[2*ib+il/2] >> 4*(il%2)) & 7) << 8)];
+    grid32[1] = (grid32[0] >> 4) & 0x0f0f0f0f;
+    grid32[0] &= 0x0f0f0f0f;
+    for (int j = 0; j < 8; ++j) {
+        y[j] = ggml_cuda_cast<dst_t>(d * (q[j] + delta));
+    }
+}
+
+template<typename dst_t>
+static __device__ __forceinline__ void dequantize_iq4_nl(const void * vx, const int64_t ibs, dst_t * yy, const int tid) {
+
+    const block_iq4_nl * x = (const block_iq4_nl *) vx + ibs*(QK_K/QK4_NL);
+
+    const int64_t il = tid/8; // 0...3
+    const int64_t ib = tid%8; // 0...7
+    dst_t * y = yy + 32*ib + 4*il;
+    const uint8_t  * q4 = x[ib].qs + 4*il;
+    const float d = (float)x[ib].d;
+    for (int j = 0; j < 4; ++j) {
+        y[j+ 0] = ggml_cuda_cast<dst_t>(d * kvalues_iq4nl[q4[j] & 0xf]);
+        y[j+16] = ggml_cuda_cast<dst_t>(d * kvalues_iq4nl[q4[j] >>  4]);
+    }
+}
+
+template<typename dst_t>
+static __device__ __forceinline__ void dequantize_iq4_xs(const void * vx, const int64_t ibs, dst_t * yy, const int tid) {
+    const block_iq4_xs * x = (const block_iq4_xs *)vx;
+
+    const int64_t il = tid/8; // 0...3
+    const int64_t ib = tid%8; // 0...7
+    dst_t * y = yy + 32*ib + 4*il;
+    const uint8_t  * q4 = x[ibs].qs + 16*ib + 4*il;
+    const float d = (float)x[ibs].d * ((((x[ibs].scales_l[ib/2] >> 4*(ib%2)) & 0xf) | (((x[ibs].scales_h >> 2*ib) & 3) << 4)) - 32);
+    for (int j = 0; j < 4; ++j) {
+        y[j+ 0] = ggml_cuda_cast<dst_t>(d * kvalues_iq4nl[q4[j] & 0xf]);
+        y[j+16] = ggml_cuda_cast<dst_t>(d * kvalues_iq4nl[q4[j] >>  4]);
+    }
+}
+
+template<typename dst_t>
+static __device__ __forceinline__ void dequantize_mxfp4(const void * vx, const int64_t ibs, dst_t * yy, const int tid) {
+
+    const block_mxfp4 * x = (const block_mxfp4 *) vx + ibs*(QK_K/QK_MXFP4);
+
+    const int64_t il = tid/8; // 0...3
+    const int64_t ib = tid%8; // 0...7
+    dst_t * y = yy + 32*ib + 4*il;
+    const uint8_t  * q4 = x[ib].qs + 4*il;
+    const float d = ggml_cuda_e8m0_to_fp32(x[ib].e);
+    for (int j = 0; j < 4; ++j) {
+        y[j+ 0] = ggml_cuda_cast<dst_t>(d * kvalues_mxfp4[q4[j] & 0xf]*0.5f);
+        y[j+16] = ggml_cuda_cast<dst_t>(d * kvalues_mxfp4[q4[j] >>  4]*0.5f);
+    }
+}
@@ -40,6 +40,35 @@ static __global__ void k_get_rows(
    }
 }

+template<typename dst_t, dequantize_kq_t<dst_t> dequantize_kq>
+static __global__ void k_get_rows_kq(
+        const void * __restrict__ src0, const int32_t * __restrict__ src1, dst_t * __restrict__ dst,
+        const int64_t ne00, /*const int64_t ne01, const int64_t ne02, const int64_t ne03,*/
+        /*const int64_t ne10,*/ const int64_t ne11, const uint3 ne12_fdv, /*const int64_t ne13,*/
+        /*const size_t s0,*/ const size_t s1, const size_t s2, const size_t s3,
+        /*const size_t nb00,*/ const size_t nb01, const size_t nb02, const size_t nb03,
+        const size_t s10, const size_t s11, const size_t s12/*, const size_t s13*/) {
+
+    ggml_cuda_pdl_sync();
+    const int64_t nsb = ne00/QK_K; // super-blocks per row
+    for (int64_t z = blockIdx.z; z < ne11*(int64_t)ne12_fdv.z; z += gridDim.z) {
+        // The x and y dimensions of the grid are swapped because the maximum allowed grid size for x is higher.
+        const int i10 = blockIdx.x;
+        const uint2 dm  = fast_div_modulo((uint32_t)z, ne12_fdv);
+        const int i11 = dm.x;
+        const int i12 = dm.y;
+
+        const int i01 = src1[i10*s10 + i11*s11 + i12*s12];
+
+        dst_t * dst_row = dst + i10*s1 + i11*s2 + i12*s3;
+        const void * src0_row = (const char *) src0 + i01*nb01 + i11*nb02 + i12*nb03;
+
+        for (int64_t ib = blockIdx.y; ib < nsb; ib += gridDim.y) {
+            dequantize_kq(src0_row, ib, dst_row + ib*QK_K, threadIdx.x);
+        }
+    }
+}
+
 template<typename src0_t, typename dst_t>
 static __global__ void k_get_rows_float(
        const src0_t * src0_ptr, const int32_t * src1_ptr, dst_t * dst_ptr,
@@ -55,27 +84,51 @@ static __global__ void k_get_rows_float(
    dst_t         * GGML_CUDA_RESTRICT dst  = dst_ptr;
    ggml_cuda_pdl_sync();
    for (int64_t z = blockIdx.z; z < ne11*(int64_t)ne12_fdv.z; z += gridDim.z) {
+        // The x and y dimensions of the grid are swapped because the maximum allowed grid size for x is higher.
+        const int i10 = blockIdx.x;
+        const uint2 dm = fast_div_modulo((uint32_t)z, ne12_fdv);
+        const int i11 = dm.x;
+        const int i12 = dm.y;
+
+        const int i01 = src1[i10*s10 + i11*s11 + i12*s12];
+
+        dst_t * GGML_CUDA_RESTRICT dst_row = dst + i10*s1 + i11*s2 + i12*s3;
+        const src0_t * GGML_CUDA_RESTRICT src0_row = (const src0_t *)((const char *) src0 + i01*nb01 + i11*nb02 + i12*nb03);
+
        for (int64_t i00 = blockIdx.y*blockDim.x + threadIdx.x; i00 < ne00; i00 += gridDim.y*blockDim.x) {
-            // The x and y dimensions of the grid are swapped because the maximum allowed grid size for x is higher.
-            const int i10 = blockIdx.x;
-            const uint2 dm = fast_div_modulo((uint32_t)z, ne12_fdv);
-            const int i11 = dm.x;
-            const int i12 = dm.y;
-
-            if (i00 >= ne00) {
-                return;
-            }
-
-            const int i01 = src1[i10*s10 + i11*s11 + i12*s12];
-
-            dst_t * dst_row = dst + i10*s1 + i11*s2 + i12*s3;
-            const src0_t * src0_row = (const src0_t *)((const char *) src0 + i01*nb01 + i11*nb02 + i12*nb03);
-
            dst_row[i00] = ggml_cuda_cast<dst_t>(src0_row[i00]);
        }
    }
 }

+template<typename dst_t>
+static __global__ void k_get_rows_float_vec(
+        const dst_t * src0_ptr, const int32_t * src1_ptr, dst_t * dst_ptr,
+        const int64_t ne00v,
+        const int64_t ne11, const uint3 ne12_fdv,
+        const size_t s1, const size_t s2, const size_t s3,
+        const size_t nb01, const size_t nb02, const size_t nb03,
+        const size_t s10, const size_t s11, const size_t s12) {
+
+    ggml_cuda_pdl_lc();
+    ggml_cuda_pdl_sync();
+    for (int64_t z = blockIdx.z; z < ne11*(int64_t)ne12_fdv.z; z += gridDim.z) {
+        const int i10 = blockIdx.x;
+        const uint2 dm = fast_div_modulo((uint32_t)z, ne12_fdv);
+        const int i11 = dm.x;
+        const int i12 = dm.y;
+
+        const int i01 = src1_ptr[i10*s10 + i11*s11 + i12*s12];
+
+        int4       * GGML_CUDA_RESTRICT dst_row  = (int4 *)      (dst_ptr + i10*s1 + i11*s2 + i12*s3);
+        const int4 * GGML_CUDA_RESTRICT src0_row = (const int4 *)((const char *) src0_ptr + i01*nb01 + i11*nb02 + i12*nb03);
+
+        for (int64_t i = blockIdx.y*blockDim.x + threadIdx.x; i < ne00v; i += gridDim.y*blockDim.x) {
+            dst_row[i] = src0_row[i];
+        }
+    }
+}
+
 template<typename grad_t, typename dst_t>
 static __global__ void k_get_rows_back_float(
        const grad_t * __restrict__ grad, const int32_t * __restrict__ rows, dst_t * __restrict__ dst,
@@ -140,16 +193,18 @@ static void get_rows_cuda_q(
        s10, s11, s12/*, s13*/);
 }

-template<typename src0_t, typename dst_t>
-static void get_rows_cuda_float(
-        const src0_t * src0_d, const int32_t * src1_d, dst_t * dst_d,
+template<int block_dim, typename dst_t, dequantize_kq_t<dst_t> dequantize_kq>
+static void get_rows_cuda_kq(
+        const void * src0_d, const int32_t * src1_d, dst_t * dst_d,
        const int64_t ne00, const size_t nb01, const size_t nb02, const size_t nb03,
        const int64_t ne10, const int64_t ne11, const int64_t ne12, const size_t nb10, const size_t nb11, const size_t nb12,
        const size_t nb1, const size_t nb2, const size_t nb3,
        cudaStream_t stream) {
-    const dim3 block_dims(CUDA_GET_ROWS_BLOCK_SIZE, 1, 1);
-    const int block_num_y = (ne00 + CUDA_GET_ROWS_BLOCK_SIZE - 1) / CUDA_GET_ROWS_BLOCK_SIZE;
-    const dim3 block_nums(ne10, MIN(block_num_y, UINT16_MAX), MIN(ne11*ne12, UINT16_MAX));
+    GGML_ASSERT(ne00 % QK_K == 0);
+    const int64_t nsb = ne00/QK_K;
+
+    const dim3 block_dims(block_dim, 1, 1);
+    const dim3 block_nums(ne10, MIN(nsb, UINT16_MAX), MIN(ne11*ne12, UINT16_MAX));

    // strides in elements
    // const size_t s0 = nb0 / sizeof(dst_t);
@@ -166,6 +221,67 @@ static void get_rows_cuda_float(
    GGML_ASSERT(ne11 <= std::numeric_limits<uint32_t>::max() / ne12);
    const uint3 ne12_fdv = init_fastdiv_values(ne12);

+    k_get_rows_kq<dst_t, dequantize_kq><<<block_nums, block_dims, 0, stream>>>(
+        src0_d, src1_d, dst_d,
+        ne00, /*ne01, ne02, ne03,*/
+        /*ne10,*/ ne11, ne12_fdv, /*ne13,*/
+        /* s0,*/ s1, s2, s3,
+        /* nb00,*/ nb01, nb02, nb03,
+        s10, s11, s12/*, s13*/);
+}
+
+template<typename src0_t, typename dst_t>
+static void get_rows_cuda_float(
+        const src0_t * src0_d, const int32_t * src1_d, dst_t * dst_d,
+        const int64_t ne00, const size_t nb01, const size_t nb02, const size_t nb03,
+        const int64_t ne10, const int64_t ne11, const int64_t ne12, const size_t nb10, const size_t nb11, const size_t nb12,
+        const size_t nb1, const size_t nb2, const size_t nb3,
+        cudaStream_t stream) {
+    const dim3 block_dims(CUDA_GET_ROWS_BLOCK_SIZE, 1, 1);
+
+    // strides in elements
+    // const size_t s0 = nb0 / sizeof(dst_t);
+    const size_t s1 = nb1 / sizeof(dst_t);
+    const size_t s2 = nb2 / sizeof(dst_t);
+    const size_t s3 = nb3 / sizeof(dst_t);
+
+    const size_t s10 = nb10 / sizeof(int32_t);
+    const size_t s11 = nb11 / sizeof(int32_t);
+    const size_t s12 = nb12 / sizeof(int32_t);
+    // const size_t s13 = nb13 / sizeof(int32_t);
+
+    GGML_ASSERT(ne12 > 0);
+    GGML_ASSERT(ne11 <= std::numeric_limits<uint32_t>::max() / ne12);
+    const uint3 ne12_fdv = init_fastdiv_values(ne12);
+
+    if constexpr (std::is_same<src0_t, dst_t>::value) {
+        constexpr int VEC = 16 / sizeof(dst_t);
+        const int64_t ne00v = ne00 / VEC;
+        const int64_t vec_block_num_y = (ne00v + CUDA_GET_ROWS_BLOCK_SIZE - 1) / CUDA_GET_ROWS_BLOCK_SIZE;
+        const bool enough_blocks = vec_block_num_y * ne10 * ne11 * ne12 >= 128;
+        const bool can_vec = VEC > 1 && enough_blocks &&
+            (ne00 % VEC == 0) &&
+            (nb01 % 16 == 0) && (nb02 % 16 == 0) && (nb03 % 16 == 0) &&
+            (nb1  % 16 == 0) && (nb2  % 16 == 0) && (nb3  % 16 == 0) &&
+            (((uintptr_t) src0_d) % 16 == 0) && (((uintptr_t) dst_d) % 16 == 0);
+
+        if (can_vec) {
+            const int block_num_y = vec_block_num_y;
+            const dim3 block_nums(ne10, MIN(block_num_y, UINT16_MAX), MIN(ne11*ne12, UINT16_MAX));
+            const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params{block_nums, block_dims, 0, stream};
+            ggml_cuda_kernel_launch(k_get_rows_float_vec<dst_t>, launch_params,
+                (const dst_t *) src0_d, src1_d, dst_d,
+                ne00v, ne11, ne12_fdv,
+                s1, s2, s3,
+                nb01, nb02, nb03,
+                s10, s11, s12);
+            return;
+        }
+    }
+
+    const int block_num_y = (ne00 + CUDA_GET_ROWS_BLOCK_SIZE - 1) / CUDA_GET_ROWS_BLOCK_SIZE;
+    const dim3 block_nums(ne10, MIN(block_num_y, UINT16_MAX), MIN(ne11*ne12, UINT16_MAX));
+
    const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params{block_nums, block_dims, 0, stream};
    ggml_cuda_kernel_launch(k_get_rows_float<src0_t, dst_t>, launch_params,
        src0_d, src1_d, dst_d,
@@ -224,8 +340,67 @@ static void ggml_cuda_get_rows_switch_src0_type(
            get_rows_cuda_q<QK8_0, QR8_0, dequantize_q8_0>(src0_d, src1_d, dst_d,
                ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream);
            break;
+        case GGML_TYPE_Q2_K:
+            get_rows_cuda_kq<64, dst_t, dequantize_q2_K<dst_t>>(src0_d, src1_d, dst_d,
+                ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream);
+            break;
+        case GGML_TYPE_Q3_K:
+            get_rows_cuda_kq<64, dst_t, dequantize_q3_K<dst_t>>(src0_d, src1_d, dst_d,
+                ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream);
+            break;
+        case GGML_TYPE_Q4_K:
+            get_rows_cuda_kq<32, dst_t, dequantize_q4_K<dst_t>>(src0_d, src1_d, dst_d,
+                ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream);
+            break;
+        case GGML_TYPE_Q5_K:
+            get_rows_cuda_kq<64, dst_t, dequantize_q5_K<dst_t>>(src0_d, src1_d, dst_d,
+                ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream);
+            break;
+        case GGML_TYPE_Q6_K:
+            get_rows_cuda_kq<64, dst_t, dequantize_q6_K<dst_t>>(src0_d, src1_d, dst_d,
+                ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream);
+            break;
+        case GGML_TYPE_IQ2_XXS:
+            get_rows_cuda_kq<32, dst_t, dequantize_iq2_xxs<dst_t>>(src0_d, src1_d, dst_d,
+                ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream);
+            break;
+        case GGML_TYPE_IQ2_XS:
+            get_rows_cuda_kq<32, dst_t, dequantize_iq2_xs<dst_t>>(src0_d, src1_d, dst_d,
+                ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream);
+            break;
+        case GGML_TYPE_IQ2_S:
+            get_rows_cuda_kq<32, dst_t, dequantize_iq2_s<dst_t>>(src0_d, src1_d, dst_d,
+                ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream);
+            break;
+        case GGML_TYPE_IQ3_XXS:
+            get_rows_cuda_kq<32, dst_t, dequantize_iq3_xxs<dst_t>>(src0_d, src1_d, dst_d,
+                ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream);
+            break;
+        case GGML_TYPE_IQ3_S:
+            get_rows_cuda_kq<32, dst_t, dequantize_iq3_s<dst_t>>(src0_d, src1_d, dst_d,
+                ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream);
+            break;
+        case GGML_TYPE_IQ1_S:
+            get_rows_cuda_kq<32, dst_t, dequantize_iq1_s<dst_t>>(src0_d, src1_d, dst_d,
+                ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream);
+            break;
+        case GGML_TYPE_IQ1_M:
+            get_rows_cuda_kq<32, dst_t, dequantize_iq1_m<dst_t>>(src0_d, src1_d, dst_d,
+                ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream);
+            break;
+        case GGML_TYPE_IQ4_NL:
+            get_rows_cuda_kq<32, dst_t, dequantize_iq4_nl<dst_t>>(src0_d, src1_d, dst_d,
+                ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream);
+            break;
+        case GGML_TYPE_IQ4_XS:
+            get_rows_cuda_kq<32, dst_t, dequantize_iq4_xs<dst_t>>(src0_d, src1_d, dst_d,
+                ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream);
+            break;
+        case GGML_TYPE_MXFP4:
+            get_rows_cuda_kq<32, dst_t, dequantize_mxfp4<dst_t>>(src0_d, src1_d, dst_d,
+                ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream);
+            break;
        default:
-            // TODO: k-quants
            GGML_ABORT("%s: unsupported src0 type: %s\n", __func__, ggml_type_name(src0_type));
            break;
    }
@@ -1,6 +1,7 @@
 #include "ggml-cuda.h"
 #include "ggml-impl.h"
 #include "ggml-backend-impl.h"
+#include "ggml-profiler.h"

 #include "ggml-cuda/allreduce.cuh"
 #include "ggml-cuda/common.cuh"
@@ -91,6 +92,92 @@

 static_assert(sizeof(half) == sizeof(ggml_fp16_t), "wrong fp16 size");

+// CUDA profiler state
+struct ggml_cuda_profiler_state {
+    bool enabled = false;
+    int split_id = -1;
+    cudaStream_t stream = nullptr;
+
+    static constexpr int MAX_PENDING_EVENTS = 4096;
+    std::vector<cudaEvent_t> start_events;
+    std::vector<cudaEvent_t> end_events;
+    std::vector<uint64_t>    cpu_timestamps;  // CPU-side timestamps for global ordering
+    int event_count = 0;
+
+    std::vector<ggml_profile_record> records;
+    std::vector<int> record_event_indices;
+
+    void init(cudaStream_t stream) {
+        this->stream = stream;
+        start_events.reserve(MAX_PENDING_EVENTS);
+        end_events.reserve(MAX_PENDING_EVENTS);
+        cpu_timestamps.reserve(MAX_PENDING_EVENTS);
+    }
+
+    void reset() {
+        for (auto & ev : start_events) {
+            (void) cudaEventDestroy(ev);
+        }
+        for (auto & ev : end_events) {
+            (void) cudaEventDestroy(ev);
+        }
+        start_events.clear();
+        end_events.clear();
+        cpu_timestamps.clear();
+        event_count = 0;
+        records.clear();
+        record_event_indices.clear();
+    }
+
+    ~ggml_cuda_profiler_state() {
+        reset();
+    }
+
+    void record_start() {
+        cudaEvent_t ev;
+        (void) cudaEventCreate(&ev);
+        (void) cudaEventRecord(ev, stream);
+        start_events.push_back(ev);
+        cpu_timestamps.push_back(ggml_profiler_time_ns());
+        event_count++;
+    }
+
+    void record_end(const char * name, int backend_id, int split_id, uint64_t bytes, const char * extra,
+                    const ggml_tensor * node) {
+        cudaEvent_t ev;
+        (void) cudaEventCreate(&ev);
+        (void) cudaEventRecord(ev, stream);
+        end_events.push_back(ev);
+        record_event_indices.push_back(records.size());
+
+        ggml_profile_record rec;
+        rec.type = GGML_PROFILE_EVENT_OP;
+        rec.name = name;
+        rec.backend_id = backend_id;
+        rec.split_id = split_id;
+        rec.start_ns = 0;
+        rec.end_ns = 0;
+        rec.bytes = bytes;
+        rec.extra = extra;
+        ggml_profile_record_from_tensor(&rec, node);
+        records.push_back(rec);
+    }
+
+    void finalize() {
+        (void) cudaStreamSynchronize(stream);
+
+        for (int i = 0; i < (int)record_event_indices.size(); i++) {
+            float ms = 0.0f;
+            (void) cudaEventElapsedTime(&ms, start_events[i], end_events[i]);
+            uint64_t duration_ns = (uint64_t)(ms * 1e6f);
+            int rec_idx = record_event_indices[i];
+            // Use CPU-side timestamp for global ordering, GPU-measured duration for accuracy
+            records[rec_idx].start_ns = cpu_timestamps[i];
+            records[rec_idx].end_ns   = cpu_timestamps[i] + duration_ns;
+        }
+    }
+};
+
 #define GGML_LOG_WARN_ONCE(str) \
    { static std::once_flag warn_flag; std::call_once(warn_flag, []() { GGML_LOG_WARN(str); }); }

@@ -2703,6 +2790,7 @@ static int ggml_cuda_try_gdn_cache_fusion(

 static bool ggml_cuda_topk_moe_fusion(const struct ggml_cgraph * cgraph, int node_idx, ggml_cuda_topk_moe_args & args) {
    args.sigmoid         = false;
+    args.sqrt_softplus   = false;
    args.softmax         = false;
    args.delayed_softmax = false;
    args.prob_bias       = false;
@@ -2716,10 +2804,17 @@ static bool ggml_cuda_topk_moe_fusion(const struct ggml_cgraph * cgraph, int nod
    }

    if (nodes[node_idx]->op == GGML_OP_UNARY) {
-        if (ggml_get_unary_op(nodes[node_idx]) != GGML_UNARY_OP_SIGMOID) {
+        const ggml_unary_op unary_op = ggml_get_unary_op(nodes[node_idx]);
+        if (unary_op == GGML_UNARY_OP_SIGMOID) {
+            args.sigmoid = true;
+        } else if (unary_op == GGML_UNARY_OP_SOFTPLUS && node_idx + 1 < n_nodes &&
+                   nodes[node_idx + 1]->op == GGML_OP_SQRT && nodes[node_idx + 1]->src[0] == nodes[node_idx]) {
+            // sqrt(softplus(x)) scoring (DeepSeek-V4)
+            args.sqrt_softplus = true;
+            node_idx++;
+        } else {
            return false;
        }
-        args.sigmoid = true;
    }

    if (nodes[node_idx]->op == GGML_OP_ARGSORT) {
@@ -2728,7 +2823,7 @@ static bool ggml_cuda_topk_moe_fusion(const struct ggml_cgraph * cgraph, int nod

    node_idx++;

-    if (args.sigmoid || args.softmax) {
+    if (args.sigmoid || args.sqrt_softplus || args.softmax) {
        // SOFTMAX -> RESHAPE
        if (node_idx >= n_nodes || nodes[node_idx]->op != GGML_OP_RESHAPE ||
                nodes[node_idx]->src[0] != nodes[node_idx - 1]) {
@@ -3172,21 +3267,27 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph
            const ggml_tensor * scale   = nullptr;

            if (!args.delayed_softmax) {
-                ggml_op gating_op = args.sigmoid ? GGML_OP_UNARY : GGML_OP_SOFT_MAX;
-                int     out_nodes[2];  // nodes which can't be elided
+                int out_nodes[2];  // nodes which can't be elided
+
+                if (args.sigmoid) {
+                    ops.insert(ops.end(), { GGML_OP_UNARY });
+                } else if (args.sqrt_softplus) {
+                    ops.insert(ops.end(), { GGML_OP_UNARY, GGML_OP_SQRT });
+                } else {
+                    ops.insert(ops.end(), { GGML_OP_SOFT_MAX });
+                }
+                const int i_probs = i + (int) ops.size() - 1;  // last node of the gating activation

                if (args.prob_bias) {
-                    bias = cgraph->nodes[i + 2]->src[1];
-                    ops.insert(ops.end(), { gating_op, GGML_OP_RESHAPE, GGML_OP_ADD, GGML_OP_ARGSORT, GGML_OP_VIEW,
+                    bias = cgraph->nodes[i_probs + 2]->src[1];
+                    ops.insert(ops.end(), { GGML_OP_RESHAPE, GGML_OP_ADD, GGML_OP_ARGSORT, GGML_OP_VIEW,
                                            GGML_OP_GET_ROWS });
-                    out_nodes[0] = i + 4;
-                    ids          = cgraph->nodes[i + 4];
+                    out_nodes[0] = i_probs + 4;
                } else {
-                    ops.insert(ops.end(),
-                               { gating_op, GGML_OP_RESHAPE, GGML_OP_ARGSORT, GGML_OP_VIEW, GGML_OP_GET_ROWS });
-                    out_nodes[0] = i + 3;
-                    ids          = cgraph->nodes[i + 3];
+                    ops.insert(ops.end(), { GGML_OP_RESHAPE, GGML_OP_ARGSORT, GGML_OP_VIEW, GGML_OP_GET_ROWS });
+                    out_nodes[0] = i_probs + 3;
                }
+                ids = cgraph->nodes[out_nodes[0]];

                if (args.norm) {
                    ops.insert(ops.end(),
@@ -4016,8 +4117,23 @@ static void ggml_cuda_graph_evaluate_and_capture(ggml_backend_cuda_context * cud
 #else
                GGML_UNUSED(integrated);
 #endif  // NDEBUG
+                if (cuda_ctx->profiler_state != nullptr && cuda_ctx->profiler_state->enabled) {
+                    cuda_ctx->profiler_state->record_start();
+                }

                bool ok = ggml_cuda_compute_forward(*cuda_ctx, node);
+
+                if (cuda_ctx->profiler_state != nullptr && cuda_ctx->profiler_state->enabled) {
+                    cuda_ctx->profiler_state->record_end(
+                        ggml_op_name(node->op),
+                        -1,
+                        cuda_ctx->profiler_state->split_id,
+                        ggml_nbytes(node),
+                        nullptr,
+                        node
+                    );
+                }
+
                if (!ok) {
                    GGML_LOG_ERROR("%s: op not supported %s (%s)\n", __func__, node->name, ggml_op_name(node->op));
                }
@@ -4088,6 +4204,19 @@ static enum ggml_status ggml_backend_cuda_graph_compute(ggml_backend_t backend,

    ggml_cuda_set_device(cuda_ctx->device);

+    // Disable CUDA graphs when profiling (we need per-node timing)
+    bool was_graph_enabled = false;
+    if (cuda_ctx->profiler_state != nullptr && cuda_ctx->profiler_state->enabled) {
+#ifdef USE_CUDA_GRAPH
+        const void * graph_key = ggml_cuda_graph_get_key(cgraph);
+        ggml_cuda_graph * graph = cuda_ctx->cuda_graph(graph_key);
+        was_graph_enabled = graph->is_enabled();
+        if (was_graph_enabled) {
+            graph->disable_due_to_gpu_arch = true;
+        }
+#endif
+    }
+
    bool use_cuda_graph             = false;
    bool cuda_graph_update_required = false;
    const void * graph_key = nullptr;
@@ -4139,6 +4268,15 @@ static enum ggml_status ggml_backend_cuda_graph_compute(ggml_backend_t backend,

    ggml_cuda_graph_evaluate_and_capture(cuda_ctx, cgraph, use_cuda_graph, cuda_graph_update_required, graph_key);

+    // Restore CUDA graph enabled state after profiling
+    if (was_graph_enabled) {
+#ifdef USE_CUDA_GRAPH
+        const void * graph_key_prof = ggml_cuda_graph_get_key(cgraph);
+        ggml_cuda_graph * graph = cuda_ctx->cuda_graph(graph_key_prof);
+        graph->disable_due_to_gpu_arch = false;
+#endif
+    }
+
    return GGML_STATUS_SUCCESS;
 }

@@ -4831,7 +4969,25 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g
                    case GGML_TYPE_Q5_0:
                    case GGML_TYPE_Q5_1:
                    case GGML_TYPE_Q8_0:
+                    case GGML_TYPE_Q2_K:
+                    case GGML_TYPE_Q3_K:
+                    case GGML_TYPE_Q4_K:
+                    case GGML_TYPE_Q5_K:
+                    case GGML_TYPE_Q6_K:
+                    case GGML_TYPE_IQ2_XXS:
+                    case GGML_TYPE_IQ2_XS:
+                    case GGML_TYPE_IQ2_S:
+                    case GGML_TYPE_IQ3_XXS:
+                    case GGML_TYPE_IQ3_S:
+                    case GGML_TYPE_IQ1_S:
+                    case GGML_TYPE_IQ1_M:
+                    case GGML_TYPE_IQ4_XS:
                        return true;
+                    case GGML_TYPE_IQ4_NL:
+                    case GGML_TYPE_MXFP4:
+                        // 32-value sub-blocks, the row size does not guarantee
+                        // the QK_K super-blocks the get_rows kernel iterates on
+                        return op->src[0]->ne[0] % QK_K == 0;
                    default:
                        return false;
                }
@@ -5381,12 +5537,68 @@ ggml_backend_t ggml_backend_cuda_init(int device) {
    }

    ggml_backend_t cuda_backend = new ggml_backend {
-        /* .guid    = */ ggml_backend_cuda_guid(),
-        /* .iface   = */ ggml_backend_cuda_interface,
-        /* .device  = */ ggml_backend_reg_dev_get(ggml_backend_cuda_reg(), device),
-        /* .context = */ ctx,
+        /* .guid     = */ ggml_backend_cuda_guid(),
+        /* .iface    = */ ggml_backend_cuda_interface,
+        /* .device   = */ ggml_backend_reg_dev_get(ggml_backend_cuda_reg(), device),
+        /* .context  = */ ctx,
+        /* .profiler = */ nullptr,
    };

+    // Register profiler
+    auto * prof_state = new ggml_cuda_profiler_state();
+    prof_state->init(ctx->stream());
+    ctx->profiler_state = prof_state;
+
+    static auto cuda_prof_enable = [](void * ctx, bool enable) {
+        auto * cuda_ctx = (ggml_backend_cuda_context *)ctx;
+        if (cuda_ctx->profiler_state != nullptr) {
+            cuda_ctx->profiler_state->enabled = enable;
+            if (!enable) {
+                cuda_ctx->profiler_state->reset();
+            }
+        }
+    };
+    static auto cuda_prof_reset = [](void * ctx) {
+        auto * cuda_ctx = (ggml_backend_cuda_context *)ctx;
+        if (cuda_ctx->profiler_state != nullptr) {
+            cuda_ctx->profiler_state->reset();
+            cuda_ctx->profiler_state->split_id = -1;
+        }
+    };
+    static auto cuda_prof_set_split_id = [](void * ctx, int split_id) {
+        auto * cuda_ctx = (ggml_backend_cuda_context *)ctx;
+        if (cuda_ctx->profiler_state != nullptr) {
+            cuda_ctx->profiler_state->split_id = split_id;
+        }
+    };
+    static auto cuda_prof_get_records = [](void * ctx, const ggml_profile_record ** out) -> int {
+        auto * cuda_ctx = (ggml_backend_cuda_context *)ctx;
+        if (cuda_ctx->profiler_state != nullptr) {
+            cuda_ctx->profiler_state->finalize();
+            *out = cuda_ctx->profiler_state->records.data();
+            return (int)cuda_ctx->profiler_state->records.size();
+        }
+        *out = nullptr;
+        return 0;
+    };
+    static auto cuda_prof_free = [](void * ctx) {
+        auto * cuda_ctx = (ggml_backend_cuda_context *)ctx;
+        if (cuda_ctx->profiler_state != nullptr) {
+            delete cuda_ctx->profiler_state;
+            cuda_ctx->profiler_state = nullptr;
+        }
+    };
+
+    auto * profiler = new ggml_backend_profiler {
+        /* .context      = */ ctx,
+        /* .enable       = */ cuda_prof_enable,
+        /* .reset        = */ cuda_prof_reset,
+        /* .set_split_id = */ cuda_prof_set_split_id,
+        /* .get_records  = */ cuda_prof_get_records,
+        /* .free_context = */ cuda_prof_free,
+    };
+    ggml_backend_set_profiler(cuda_backend, profiler);
+
    return cuda_backend;
 }

@@ -25,12 +25,7 @@ static void ggml_cuda_mul_mat_q_switch_type(ggml_backend_cuda_context & ctx, con
        case GGML_TYPE_Q8_0:
            mul_mat_q_case<GGML_TYPE_Q8_0>(ctx, args, stream);
            break;
-        case GGML_TYPE_MXFP4:
-            mul_mat_q_case<GGML_TYPE_MXFP4>(ctx, args, stream);
-            break;
-        case GGML_TYPE_NVFP4:
-            mul_mat_q_case<GGML_TYPE_NVFP4>(ctx, args, stream);
-            break;
+// -----------------------------------------------------------------------
        case GGML_TYPE_Q2_K:
            mul_mat_q_case<GGML_TYPE_Q2_K>(ctx, args, stream);
            break;
@@ -46,6 +41,10 @@ static void ggml_cuda_mul_mat_q_switch_type(ggml_backend_cuda_context & ctx, con
        case GGML_TYPE_Q6_K:
            mul_mat_q_case<GGML_TYPE_Q6_K>(ctx, args, stream);
            break;
+// -----------------------------------------------------------------------
+        case GGML_TYPE_IQ1_S:
+            mul_mat_q_case<GGML_TYPE_IQ1_S>(ctx, args, stream);
+            break;
        case GGML_TYPE_IQ2_XXS:
            mul_mat_q_case<GGML_TYPE_IQ2_XXS>(ctx, args, stream);
            break;
@@ -61,15 +60,19 @@ static void ggml_cuda_mul_mat_q_switch_type(ggml_backend_cuda_context & ctx, con
        case GGML_TYPE_IQ3_S:
            mul_mat_q_case<GGML_TYPE_IQ3_S>(ctx, args, stream);
            break;
-        case GGML_TYPE_IQ1_S:
-            mul_mat_q_case<GGML_TYPE_IQ1_S>(ctx, args, stream);
-            break;
        case GGML_TYPE_IQ4_XS:
            mul_mat_q_case<GGML_TYPE_IQ4_XS>(ctx, args, stream);
            break;
        case GGML_TYPE_IQ4_NL:
            mul_mat_q_case<GGML_TYPE_IQ4_NL>(ctx, args, stream);
            break;
+// -----------------------------------------------------------------------
+        case GGML_TYPE_MXFP4:
+            mul_mat_q_case<GGML_TYPE_MXFP4>(ctx, args, stream);
+            break;
+        case GGML_TYPE_NVFP4:
+            mul_mat_q_case<GGML_TYPE_NVFP4>(ctx, args, stream);
+            break;
        default:
            GGML_ABORT("fatal error");
            break;
@@ -130,14 +133,20 @@ void ggml_cuda_mul_mat_q(
        const size_t nbytes_src1_q8_1 = ne13*ne12 * ne11*ne10_padded * y_block_size/y_values_per_block +
            ggml_cuda_mmq_get_J_max(src0->type, fallback, cc, ne11) * sizeof(block_q8_1_mmq);
        ggml_cuda_pool_alloc<char> src1_q8_1(ctx.pool(), nbytes_src1_q8_1);
+        ggml_cuda_pool_alloc<float> src1_scale(ctx.pool());
+        if (src0->type == GGML_TYPE_NVFP4 && use_native_fp4) {
+            src1_scale.alloc(ne13*ne12*ne11);
+        }

        {
            const int64_t s11 = src1->nb[1] / ts_src1;
            const int64_t s12 = src1->nb[2] / ts_src1;
            const int64_t s13 = src1->nb[3] / ts_src1;
            if (use_native_fp4) {
+                static constexpr size_t align_float8 = 32;
+                const bool use_aligned_float8 = ggml_cuda_is_aligned(src1, align_float8);
                static_assert(sizeof(block_fp4_mmq) == 4 * sizeof(block_q8_1));
-                quantize_mmq_fp4_cuda(src1_d, nullptr, src1_q8_1.get(), src0->type, ne10, s11, s12, s13, ne10_padded,
+                quantize_mmq_fp4_cuda(src1_d, nullptr, src1_q8_1.get(), src1_scale.ptr, src0->type, use_aligned_float8, ne10, s11, s12, s13, ne10_padded,
                                        ne11, ne12, ne13, stream);

            } else {
@@ -155,6 +164,7 @@ void ggml_cuda_mul_mat_q(

        const mmq_args args = {
            src0_d, src0->type, (const int *) src1_q8_1.ptr, nullptr, nullptr, dst_d,
+            src0->type == GGML_TYPE_NVFP4 && use_native_fp4 ? src1_scale.ptr : nullptr,
            ne00, ne01, ne1, s01, ne11, s1,
            ne02, ne12, s02, s12, s2,
            ne03, ne13, s03, s13, s3,
@@ -192,6 +202,10 @@ void ggml_cuda_mul_mat_q(
    const size_t nbytes_src1_q8_1 = ne12*n_expert_used*ne10_padded * y_block_size/y_values_per_block +
        ggml_cuda_mmq_get_J_max(src0->type, fallback, cc, ne11) * sizeof(block_q8_1_mmq);
    ggml_cuda_pool_alloc<char> src1_q8_1(ctx.pool(), nbytes_src1_q8_1);
+    ggml_cuda_pool_alloc<float> src1_scale(ctx.pool());
+    if (src0->type == GGML_TYPE_NVFP4 && use_native_fp4) {
+        src1_scale.alloc(ne12*n_expert_used);
+    }

    const int64_t ne11_flat = ne12*n_expert_used;
    const int64_t ne12_flat = 1;
@@ -202,18 +216,19 @@ void ggml_cuda_mul_mat_q(
        const int64_t s12 = src1->nb[2] / ts_src1;
        const int64_t s13 = src1->nb[3] / ts_src1;

-        if (dedup_bcast) {
-            // quantize each token once, scatter its block to all n_expert_used slots
-            if (use_native_fp4) {
-                quantize_scatter_mmq_fp4_cuda(src1_d, ids_src1.get(), src1_q8_1.get(), src0->type, ne10,
+        if (use_native_fp4) {
+            static constexpr size_t align_float8 = 32;
+            const bool use_aligned_float8 = ggml_cuda_is_aligned(src1, align_float8);
+            if (dedup_bcast) {
+                quantize_scatter_mmq_fp4_cuda(src1_d, ids_src1.get(), src1_q8_1.get(), src1_scale.ptr, src0->type, use_aligned_float8, ne10,
                                        /*stride_token=*/s12, ne10_padded, ne12, ne11_flat, n_expert_used, stream);
            } else {
-                quantize_scatter_mmq_q8_1_cuda(src1_d, ids_src1.get(), src1_q8_1.get(), src0->type, ne10,
-                                        /*stride_token=*/s12, ne10_padded, ne12, ne11_flat, n_expert_used, stream);
+                quantize_mmq_fp4_cuda(src1_d, ids_src1.get(), src1_q8_1.get(), src1_scale.ptr, src0->type, use_aligned_float8, ne10, s11, s12, s13,
+                                        ne10_padded, ne11_flat, ne12_flat, ne13_flat, stream);
            }
-        } else if (use_native_fp4) {
-            quantize_mmq_fp4_cuda(src1_d, ids_src1.get(), src1_q8_1.get(), src0->type, ne10, s11, s12, s13,
-                                    ne10_padded, ne11_flat, ne12_flat, ne13_flat, stream);
+        } else if (dedup_bcast) {
+            quantize_scatter_mmq_q8_1_cuda(src1_d, ids_src1.get(), src1_q8_1.get(), src0->type, ne10,
+                                    /*stride_token=*/s12, ne10_padded, ne12, ne11_flat, n_expert_used, stream);
        } else {
            quantize_mmq_q8_1_cuda(src1_d, ids_src1.get(), src1_q8_1.get(), src0->type, ne10, s11, s12, s13,
                                   ne10_padded, ne11_flat, ne12_flat, ne13_flat, stream);
@@ -229,6 +244,7 @@ void ggml_cuda_mul_mat_q(
    // Note that ne02 is used instead of ne12 because the number of y channels determines the z dimension of the CUDA grid.
    const mmq_args args = {
        src0_d, src0->type, (const int *) src1_q8_1.get(), ids_dst.get(), expert_bounds.get(), dst_d,
+        src1_scale.ptr,
        ne00, ne01, ne_get_rows, s01, ne_get_rows, s1,
        ne02, ne02, s02, s12, s2,
        ne03, ne13, s03, s13, s3,
@@ -251,21 +267,24 @@ bool ggml_cuda_should_use_mmq(enum ggml_type type, int cc, int64_t ne11, int64_t
        case GGML_TYPE_Q5_0:
        case GGML_TYPE_Q5_1:
        case GGML_TYPE_Q8_0:
-        case GGML_TYPE_MXFP4:
-        case GGML_TYPE_NVFP4:
+// -------------------------------------------------
        case GGML_TYPE_Q2_K:
        case GGML_TYPE_Q3_K:
        case GGML_TYPE_Q4_K:
        case GGML_TYPE_Q5_K:
        case GGML_TYPE_Q6_K:
+// -------------------------------------------------
+        case GGML_TYPE_IQ1_S:
        case GGML_TYPE_IQ2_XXS:
        case GGML_TYPE_IQ2_XS:
        case GGML_TYPE_IQ2_S:
        case GGML_TYPE_IQ3_XXS:
        case GGML_TYPE_IQ3_S:
-        case GGML_TYPE_IQ1_S:
        case GGML_TYPE_IQ4_XS:
        case GGML_TYPE_IQ4_NL:
+// -------------------------------------------------
+        case GGML_TYPE_MXFP4:
+        case GGML_TYPE_NVFP4:
            mmq_supported = true;
            break;
        default:
@@ -13,7 +13,7 @@
 typedef void (*ggml_cuda_mmq_load_tiles_t)(const char * __restrict__ x, int * x_tile, const int kbx0, const int i_max, const int stride);
 typedef void (*ggml_cuda_mmq_vec_dot_t)(const int * __restrict__ x, const int * __restrict__ y, float * __restrict__ sum, const int k00);
 typedef void (*ggml_cuda_mmq_write_back_t)(const float * __restrict__ sum, const int32_t * __restrict__ get_rows_to_sorted,
-    float * __restrict__ dst, const int stride, const int i_max, const int j_max);
+    float * __restrict__ dst, const float * __restrict__ y_scale, const int stride, const int i_max, const int j_max);

 enum mmq_q8_1_ds_layout {
    MMQ_Q8_1_DS_LAYOUT_D4,
@@ -413,11 +413,13 @@ static __host__ int ggml_cuda_mmq_get_nbytes_shared_x(const ggml_cuda_mmq_config

 template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_write_back_dp4a(
        const float * __restrict__ sum, const int32_t * __restrict__ ids_dst, float * __restrict__ dst,
-        const int stride, const int i_max, const int j_max) {
+        const float * __restrict__ y_scale, const int stride, const int i_max, const int j_max) {
    constexpr int warp_size = ggml_cuda_get_physical_warp_size();
    constexpr int nwarps    = ggml_cuda_mmq_get_nthreads(type, J, fallback) / warp_size;
    constexpr int I         = ggml_cuda_mmq_get_I(type, J, fallback);

+    const bool y_scale_used = y_scale != nullptr;
+
 #pragma unroll
    for (int j0 = 0; j0 < J; j0 += nwarps) {
        const int j = j0 + threadIdx.y;
@@ -434,7 +436,16 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
                continue;
            }

-            dst[ids_dst[j]*stride + i] = sum[(j0/nwarps) * (I/warp_size) + i0/warp_size];
+            if constexpr (type == GGML_TYPE_NVFP4) {
+                if (y_scale_used) {
+                    dst[ids_dst[j]*stride + i] = y_scale[j] * sum[(j0/nwarps) * (I/warp_size) + i0/warp_size];
+                } else {
+                    dst[ids_dst[j]*stride + i] = sum[(j0/nwarps) * (I/warp_size) + i0/warp_size];
+                }
+            } else {
+                dst[ids_dst[j]*stride + i] = sum[(j0/nwarps) * (I/warp_size) + i0/warp_size];
+                GGML_UNUSED(y_scale_used);
+            }
        }
    }
 }
@@ -442,7 +453,8 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
 template<ggml_type type, int J, bool fallback>
 static __device__ __forceinline__ void ggml_cuda_mmq_write_back_mma(
            const float * __restrict__ sum, const int * __restrict__ ids_dst, float * __restrict__ dst,
-            const int stride, const int i_max, const int j_max) {
+            const float * __restrict__ y_scale, const int stride, const int i_max, const int j_max) {
+
 #if defined(AMD_MFMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
    typedef tile<16, 16, int, DATA_LAYOUT_J_MAJOR> tile_C;
 #else
@@ -457,6 +469,8 @@ static __device__ __forceinline__ void ggml_cuda_mmq_write_back_mma(

    const int i0 = (threadIdx.y / ntx) * (ntx*tile_C::I);

+    const bool y_scale_used = y_scale != nullptr;
+
 #pragma unroll
    for (int j0 = 0; j0 < J; j0 += ntx*tile_C::J) {
 #pragma unroll
@@ -475,7 +489,16 @@ static __device__ __forceinline__ void ggml_cuda_mmq_write_back_mma(
                    continue;
                }

-                dst[ids_dst[j]*stride + i] = sum[(j0/tile_C::J + n)*tile_C::ne + l];
+                if constexpr (type == GGML_TYPE_NVFP4) {
+                    if (y_scale_used) {
+                        dst[ids_dst[j]*stride + i] = y_scale[j] * sum[(j0/tile_C::J + n)*tile_C::ne + l];
+                    } else {
+                        dst[ids_dst[j]*stride + i] = sum[(j0/tile_C::J + n)*tile_C::ne + l];
+                    }
+                } else {
+                    dst[ids_dst[j]*stride + i] = sum[(j0/tile_C::J + n)*tile_C::ne + l];
+                    GGML_UNUSED(y_scale_used);
+                }
            }
        }
    }
@@ -819,6 +842,7 @@ template <ggml_type type, int J, bool fallback, bool fixup>
 static __device__ __forceinline__ void mul_mat_q_process_tile(
        const char * __restrict__ x, const int offset_x, const int * __restrict__ y,
        const int * __restrict__ ids_dst, float * __restrict__ dst, float * __restrict__ tmp_fixup,
+        const float * __restrict__ y_scale,
        const int stride_row_x, const int ncols_y, const int stride_col_dst,
        const int tile_x_max_i, const int tile_y_max_j, const int kb0_start, const int kb0_stop) {

@@ -884,9 +908,9 @@ static __device__ __forceinline__ void mul_mat_q_process_tile(
    }

    if (fixup) {
-        write_back(sum, ids_dst, tmp_fixup + blockIdx.x*(J*I), I, I, J);
+        write_back(sum, ids_dst, tmp_fixup + blockIdx.x*(J*I), y_scale, I, I, J);
    } else {
-        write_back(sum, ids_dst, dst, stride_col_dst, tile_x_max_i, tile_y_max_j);
+        write_back(sum, ids_dst, dst, y_scale, stride_col_dst, tile_x_max_i, tile_y_max_j);
    }
 }

@@ -898,6 +922,7 @@ __launch_bounds__(ggml_cuda_mmq_get_nthreads(type, J, fallback), ggml_cuda_mmq_g
 static __global__ void mul_mat_q(
        const char * __restrict__ x, const int * __restrict__ y, const int32_t * __restrict__ ids_dst,
        const int32_t * __restrict__ expert_bounds, float * __restrict__ dst, float * __restrict__ tmp_fixup,
+        const float * __restrict__ y_scale,
        const uint3 blocks_per_ne00, const int nrows_x, const int ncols_dst, const int stride_row_x, const int ncols_y, const int stride_col_dst,
        const uint3 channel_ratio, const uint3 nchannels_y, const int stride_channel_x, const int stride_channel_y, const int stride_channel_dst,
        const uint3 sample_ratio, const uint3 nsamples_y, const int stride_sample_x, const int stride_sample_y, const int stride_sample_dst,
@@ -943,8 +968,14 @@ static __global__ void mul_mat_q(
        int col_low    = 0;
        int col_high   = ncols_dst;
        int col_diff   = ncols_dst;
-        int offset_y   = wt*stride_sample_y   + zt*stride_channel_y;
-        int offset_dst = wt*stride_sample_dst + zt*stride_channel_dst + jt*J*stride_col_dst;
+        int offset_y       = wt*stride_sample_y   + zt*stride_channel_y;
+        int offset_dst     = wt*stride_sample_dst + zt*stride_channel_dst + jt*J*stride_col_dst;
+        int offset_y_scale;
+        if constexpr (type == GGML_TYPE_NVFP4) {
+            offset_y_scale = wt*nchannels_y.z*ncols_y + zt*ncols_y;
+        } else {
+            GGML_UNUSED(offset_y_scale);
+        }

        if (ids_dst) {
            col_low  = expert_bounds[zt + 0];
@@ -953,6 +984,9 @@ static __global__ void mul_mat_q(

            offset_y   = 0;
            offset_dst = 0;
+            if constexpr (type == GGML_TYPE_NVFP4) {
+                offset_y_scale = 0;
+            }

            if (jt*J >= col_diff) {
                return;
@@ -974,6 +1008,11 @@ static __global__ void mul_mat_q(

        offset_y   += (col_low + jt*J)*(sizeof(block_q8_1_mmq)/sizeof(int));
        offset_dst += it*I;
+        const float * y_scale_tile = nullptr;
+        if constexpr (type == GGML_TYPE_NVFP4) {
+            offset_y_scale += col_low + jt*J;
+            y_scale_tile = y_scale ? y_scale + offset_y_scale : nullptr;
+        }

        const int tile_x_max_i = nrows_x  - it*I - 1;
        const int tile_y_max_j = col_diff - jt*J - 1;
@@ -982,7 +1021,8 @@ static __global__ void mul_mat_q(

        constexpr bool fixup = false;
        mul_mat_q_process_tile<type, J, fallback, fixup>
-            (x, offset_x, y + offset_y, ids_dst_shared, dst + offset_dst, tmp_fixup, stride_row_x, ncols_y, stride_col_dst,
+            (x, offset_x, y + offset_y, ids_dst_shared, dst + offset_dst, tmp_fixup, y_scale_tile,
+             stride_row_x, ncols_y, stride_col_dst,
             tile_x_max_i, tile_y_max_j, 0, blocks_per_ne00.z);
        return;
    }
@@ -1016,8 +1056,14 @@ static __global__ void mul_mat_q(
        int col_low    = 0;
        int col_high   = ncols_dst;
        int col_diff   = ncols_dst;
-        int offset_y   = wt*stride_sample_y   + zt*stride_channel_y;
-        int offset_dst = wt*stride_sample_dst + zt*stride_channel_dst + jt*J*stride_col_dst;
+        int offset_y       = wt*stride_sample_y   + zt*stride_channel_y;
+        int offset_dst     = wt*stride_sample_dst + zt*stride_channel_dst + jt*J*stride_col_dst;
+        int offset_y_scale;
+        if constexpr (type == GGML_TYPE_NVFP4) {
+            offset_y_scale = wt*nchannels_y.z*ncols_y + zt*ncols_y;
+        } else {
+            GGML_UNUSED(offset_y_scale);
+        }

        if (ids_dst) {
            col_low  = expert_bounds[zt + 0];
@@ -1026,6 +1072,9 @@ static __global__ void mul_mat_q(

            offset_y   = 0;
            offset_dst = 0;
+            if constexpr (type == GGML_TYPE_NVFP4) {
+                offset_y_scale = 0;
+            }

            if (jt*J >= col_diff) {
                kbc += blocks_per_ne00.z;
@@ -1053,6 +1102,11 @@ static __global__ void mul_mat_q(

        offset_y += (col_low + jt * J) * (sizeof(block_q8_1_mmq) / sizeof(int));
        offset_dst += it*I;
+        const float * y_scale_tile = nullptr;
+        if constexpr (type == GGML_TYPE_NVFP4) {
+            offset_y_scale += col_low + jt * J;
+            y_scale_tile = y_scale ? y_scale + offset_y_scale : nullptr;
+        }

        const int tile_x_max_i = nrows_x  - it*I - 1;
        const int tile_y_max_j = col_diff - jt*J - 1;
@@ -1061,7 +1115,8 @@ static __global__ void mul_mat_q(

        constexpr bool fixup = false; // All but (potentially) the last iterations write their data to dst rather than the fixup buffer.
        mul_mat_q_process_tile<type, J, fallback, fixup>
-            (x, offset_x, y + offset_y, ids_dst_shared, dst + offset_dst, tmp_fixup, stride_row_x, ncols_y, stride_col_dst,
+            (x, offset_x, y + offset_y, ids_dst_shared, dst + offset_dst, tmp_fixup, y_scale_tile,
+             stride_row_x, ncols_y, stride_col_dst,
             tile_x_max_i, tile_y_max_j, kb0_start, kb0_stop);

        kbc += blocks_per_ne00.z;
@@ -1090,8 +1145,14 @@ static __global__ void mul_mat_q(
    int col_low    = 0;
    int col_high   = ncols_dst;
    int col_diff   = ncols_dst;
-    int offset_y   = wt*stride_sample_y   + zt*stride_channel_y;
-    int offset_dst = wt*stride_sample_dst + zt*stride_channel_dst + jt*J*stride_col_dst;
+    int offset_y       = wt*stride_sample_y   + zt*stride_channel_y;
+    int offset_dst     = wt*stride_sample_dst + zt*stride_channel_dst + jt*J*stride_col_dst;
+    int offset_y_scale;
+    if constexpr (type == GGML_TYPE_NVFP4) {
+        offset_y_scale = wt*nchannels_y.z*ncols_y + zt*ncols_y;
+    } else {
+        GGML_UNUSED(offset_y_scale);
+    }

    if (ids_dst) {
        col_low  = expert_bounds[zt + 0];
@@ -1100,6 +1161,9 @@ static __global__ void mul_mat_q(

        offset_y   = 0;
        offset_dst = 0;
+        if constexpr (type == GGML_TYPE_NVFP4) {
+            offset_y_scale = 0;
+        }

        if (jt*J >= col_diff) {
            return;
@@ -1122,6 +1186,11 @@ static __global__ void mul_mat_q(

    offset_y += (col_low + jt * J) * (sizeof(block_q8_1_mmq) / sizeof(int));
    offset_dst += it*I;
+    const float * y_scale_tile = nullptr;
+    if constexpr (type == GGML_TYPE_NVFP4) {
+        offset_y_scale += col_low + jt * J;
+        y_scale_tile = y_scale ? y_scale + offset_y_scale : nullptr;
+    }

    const int tile_x_max_i = nrows_x  - it*I - 1;
    const int tile_y_max_j = col_diff - jt*J - 1;
@@ -1130,7 +1199,8 @@ static __global__ void mul_mat_q(

    constexpr bool fixup = true; // Last index writes its data to fixup buffer to avoid data races with other blocks.
    mul_mat_q_process_tile<type, J, fallback, fixup>
-        (x, offset_x, y + offset_y, ids_dst_shared, dst + offset_dst, tmp_fixup, stride_row_x, ncols_y, stride_col_dst,
+        (x, offset_x, y + offset_y, ids_dst_shared, dst + offset_dst, tmp_fixup, y_scale_tile,
+         stride_row_x, ncols_y, stride_col_dst,
         tile_x_max_i, tile_y_max_j, kb0_start, kb0_stop);
 }

@@ -1274,6 +1344,7 @@ static __global__ void mul_mat_q_stream_k_fixup(

 struct mmq_args {
    const char * x; ggml_type type_x; const int * y; const int32_t * ids_dst; const int32_t * expert_bounds; float * dst;
+    const float * y_scale;
    int64_t ncols_x; int64_t nrows_x; int64_t ncols_dst; int64_t stride_row_x; int64_t ncols_y; int64_t nrows_dst;
    int64_t nchannels_x; int64_t nchannels_y; int64_t stride_channel_x; int64_t stride_channel_y; int64_t stride_channel_dst;
    int64_t nsamples_x; int64_t nsamples_y; int64_t stride_sample_x; int64_t stride_sample_y; int64_t stride_sample_dst;
@@ -1323,7 +1394,7 @@ static void launch_mul_mat_q(ggml_backend_cuda_context & ctx, const mmq_args & a

    if (!ggml_cuda_mmq_get_stream_k(type, J, fallback, cc)) {
        mul_mat_q<type, J, fallback><<<block_nums_xy_tiling, block_dims, nbytes_shared, stream>>>
-            (args.x, args.y, args.ids_dst, args.expert_bounds, args.dst, nullptr,
+            (args.x, args.y, args.ids_dst, args.expert_bounds, args.dst, nullptr, args.y_scale,
             blocks_per_ne00_fd, args.nrows_x, args.ncols_dst, args.stride_row_x, args.ncols_y, args.nrows_dst,
             channel_ratio_fd, nchannels_y_fd, args.stride_channel_x, args.stride_channel_y, args.stride_channel_dst,
             sample_ratio_fd, nsamples_y_fd, args.stride_sample_x, args.stride_sample_y, args.stride_sample_dst,
@@ -1352,7 +1423,7 @@ static void launch_mul_mat_q(ggml_backend_cuda_context & ctx, const mmq_args & a
    const dim3 block_dims_fixup(block_dims.x, block_dims.y/2, block_dims.z);

    mul_mat_q<type, J, fallback><<<block_nums_stream_k, block_dims, nbytes_shared, stream>>>
-        (args.x, args.y, args.ids_dst, args.expert_bounds, args.dst, tmp_fixup.ptr,
+        (args.x, args.y, args.ids_dst, args.expert_bounds, args.dst, tmp_fixup.ptr, args.y_scale,
         blocks_per_ne00_fd, args.nrows_x, args.ncols_dst, args.stride_row_x, args.ncols_y, args.nrows_dst,
         channel_ratio_fd, nchannels_y_fd, args.stride_channel_x, args.stride_channel_y, args.stride_channel_dst,
         sample_ratio_fd, nsamples_y_fd, args.stride_sample_x, args.stride_sample_y, args.stride_sample_dst,
@@ -1466,26 +1537,30 @@ void mul_mat_q_case(ggml_backend_cuda_context & ctx, const mmq_args & args, cuda
 #define DECL_MMQ_CASE(type)                                                        \
    template void mul_mat_q_case<type>(ggml_backend_cuda_context & ctx, const mmq_args & args, cudaStream_t stream) \

+extern DECL_MMQ_CASE(GGML_TYPE_Q1_0);
 extern DECL_MMQ_CASE(GGML_TYPE_Q4_0);
 extern DECL_MMQ_CASE(GGML_TYPE_Q4_1);
 extern DECL_MMQ_CASE(GGML_TYPE_Q5_0);
 extern DECL_MMQ_CASE(GGML_TYPE_Q5_1);
 extern DECL_MMQ_CASE(GGML_TYPE_Q8_0);
-extern DECL_MMQ_CASE(GGML_TYPE_MXFP4);
-extern DECL_MMQ_CASE(GGML_TYPE_NVFP4);
+// -----------------------------------------
 extern DECL_MMQ_CASE(GGML_TYPE_Q2_K);
 extern DECL_MMQ_CASE(GGML_TYPE_Q3_K);
 extern DECL_MMQ_CASE(GGML_TYPE_Q4_K);
 extern DECL_MMQ_CASE(GGML_TYPE_Q5_K);
 extern DECL_MMQ_CASE(GGML_TYPE_Q6_K);
+// -----------------------------------------
+extern DECL_MMQ_CASE(GGML_TYPE_IQ1_S);
 extern DECL_MMQ_CASE(GGML_TYPE_IQ2_XXS);
 extern DECL_MMQ_CASE(GGML_TYPE_IQ2_XS);
 extern DECL_MMQ_CASE(GGML_TYPE_IQ2_S);
 extern DECL_MMQ_CASE(GGML_TYPE_IQ3_XXS);
 extern DECL_MMQ_CASE(GGML_TYPE_IQ3_S);
-extern DECL_MMQ_CASE(GGML_TYPE_IQ1_S);
 extern DECL_MMQ_CASE(GGML_TYPE_IQ4_NL);
 extern DECL_MMQ_CASE(GGML_TYPE_IQ4_XS);
+// -----------------------------------------
+extern DECL_MMQ_CASE(GGML_TYPE_MXFP4);
+extern DECL_MMQ_CASE(GGML_TYPE_NVFP4);

 // -------------------------------------------------------------------------------------------------------------------------

@@ -1,6 +1,55 @@
 #include "quantize.cuh"
 #include <cstdint>

+#if defined(BLACKWELL_MMA_AVAILABLE)
+// this maps to 256-bit loads in PTX on supported devices,
+// and otherwise falls back to 2 128-bit loads
+struct __builtin_align__(32) float8 {
+    float x; float y; float z; float w;
+    float p; float q; float r; float s;
+};
+#endif
+
+#if CUDART_VERSION >= 12080
+static __device__ __forceinline__ float nvfp4_native_scale_error(
+        const float vals[QK_NVFP4_SUB], const float inv_col_scale, const float inv_scale, const float scale) {
+    const float scale_dequant = 2.0f * scale;
+    float err = 0.0f;
+
+#pragma unroll
+    for (int k = 0; k < QK_NVFP4_SUB; k += 4) {
+        const float v0 = vals[k + 0] * inv_col_scale;
+        const float v1 = vals[k + 1] * inv_col_scale;
+        const float v2 = vals[k + 2] * inv_col_scale;
+        const float v3 = vals[k + 3] * inv_col_scale;
+
+        const __nv_fp4x4_e2m1 q(make_float4(v0 * inv_scale, v1 * inv_scale, v2 * inv_scale, v3 * inv_scale));
+        const __nv_fp4x4_storage_t q_storage = q.__x;
+        const __nv_fp4x2_storage_t q_lo = static_cast<__nv_fp4x2_storage_t>(q_storage);
+        const __nv_fp4x2_storage_t q_hi = static_cast<__nv_fp4x2_storage_t>(q_storage >> 8U);
+
+        const __half2_raw hraw2_lo = __nv_cvt_fp4x2_to_halfraw2(q_lo, __NV_E2M1);
+        const __half2_raw hraw2_hi = __nv_cvt_fp4x2_to_halfraw2(q_hi, __NV_E2M1);
+        const __half2 h2_lo = static_cast<__half2>(hraw2_lo);
+        const __half2 h2_hi = static_cast<__half2>(hraw2_hi);
+        const float2 dq_lo = __half22float2(h2_lo);
+        const float2 dq_hi = __half22float2(h2_hi);
+
+        const float err0 = fabsf(v0) - fabsf(dq_lo.x) * scale_dequant;
+        const float err1 = fabsf(v1) - fabsf(dq_lo.y) * scale_dequant;
+        const float err2 = fabsf(v2) - fabsf(dq_hi.x) * scale_dequant;
+        const float err3 = fabsf(v3) - fabsf(dq_hi.y) * scale_dequant;
+
+        err = fmaf(err0, err0, err);
+        err = fmaf(err1, err1, err);
+        err = fmaf(err2, err2, err);
+        err = fmaf(err3, err3, err);
+    }
+
+    return err;
+}
+#endif // CUDART_VERSION >= 12080
+
 __launch_bounds__(CUDA_QUANTIZE_BLOCK_SIZE, 1)
 static __global__ void quantize_q8_1(
        const float * x_ptr, void * vy_ptr,
@@ -74,115 +123,209 @@ __device__ __forceinline__ uint8_t compute_e8m0_scale(float amax) {
    return static_cast<uint8_t>(biased);
 }

-
 // scatter: grid over tokens, quantize once, write to all the token's compact rows
-template <bool scatter>
+template <bool scatter, bool use_aligned_float8>
 static __global__ void quantize_mmq_nvfp4(
-        const float * __restrict__ x, const int32_t * __restrict__ ids, void * __restrict__ vy,
+        const float * __restrict__ x, const int32_t * __restrict__ ids, void * __restrict__ vy, float * __restrict__ scale,
        const int64_t ne00, const int64_t s01, const int64_t s02, const int64_t s03,
        const int64_t ne0, const int64_t ne1, const int64_t ne2, const int n_expert_used) {
 #if defined(BLACKWELL_MMA_AVAILABLE)

-    const int64_t i0_base = ((int64_t) blockDim.x * blockIdx.y + threadIdx.x) * QK_NVFP4_SUB;
-    if (i0_base >= ne0) {
-        return;
-    }
-
-    const int64_t k_block = i0_base / QK_FP4_MMQ;
    const int64_t blocks_per_col = (ne0 + QK_FP4_MMQ - 1) / QK_FP4_MMQ;
-    if (k_block >= blocks_per_col) {
-        return;
-    }
-    const int sub = (i0_base % QK_FP4_MMQ) / QK_NVFP4_SUB;

    int64_t base_idx;
    if constexpr (scatter) {
        base_idx = (int64_t) blockIdx.x * s02; // one physical row per token
    } else {
-        const int64_t i2  = blockIdx.z % ne2;
-        const int64_t i3  = blockIdx.z / ne2;
+        const int64_t i2  = blockIdx.y % ne2;
+        const int64_t i3  = blockIdx.y / ne2;
        const int64_t i01 = ids ? ids[blockIdx.x] : blockIdx.x;
        base_idx = i3 * s03 + i2 * s02 + i01 * s01;
    }
+    const float * __restrict__  x_row = x + base_idx;

-    float vals_raw[QK_NVFP4_SUB];
-    float amax_raw = 0.0f;
+    float amax = 0.0f;
+    if constexpr (use_aligned_float8) {
+        for (int64_t i0 = 8 * threadIdx.x; i0 < ne00; i0 += 8 * blockDim.x) {
+            const float * x_base = x_row + i0;
+            const float8 v = reinterpret_cast<const float8 *>(x_base)[0];
+            amax = fmaxf(amax, fabsf(v.x));
+            amax = fmaxf(amax, fabsf(v.y));
+            amax = fmaxf(amax, fabsf(v.z));
+            amax = fmaxf(amax, fabsf(v.w));
+            amax = fmaxf(amax, fabsf(v.p));
+            amax = fmaxf(amax, fabsf(v.q));
+            amax = fmaxf(amax, fabsf(v.r));
+            amax = fmaxf(amax, fabsf(v.s));
+        }
+    } else {
+        for (int64_t i0 = threadIdx.x; i0 < ne00; i0 += blockDim.x) {
+            amax = fmaxf(amax, fabsf(x_row[i0]));
+        }
+    }
+
+    amax = warp_reduce_max<WARP_SIZE>(amax);
+
+    __shared__ float warp_amax[CUDA_QUANTIZE_BLOCK_SIZE_MMQ / WARP_SIZE];
+    const int lane = threadIdx.x % WARP_SIZE;
+    const int warp = threadIdx.x / WARP_SIZE;
+
+    if (lane == 0) {
+        warp_amax[warp] = amax;
+    }
+    __syncthreads();
+
+    if (warp == 0) {
+        amax = threadIdx.x < int(CUDA_QUANTIZE_BLOCK_SIZE_MMQ / WARP_SIZE) ? warp_amax[lane] : 0.0f;
+        amax = warp_reduce_max<WARP_SIZE>(amax);
+        if (lane == 0) {
+            warp_amax[0] = amax / (6.0f * 448.0f);
+            if constexpr (scatter) {
 #pragma unroll
-    for (int k = 0; k < QK_NVFP4_SUB; k++) {
-        const int64_t i00 = i0_base + k;
-        if (i00 < ne00) {
-            const float v = x[base_idx + i00];
-            vals_raw[k] = v;
-            amax_raw = fmaxf(amax_raw, fabsf(v));
-        } else {
-            vals_raw[k] = 0.0f;
+                for (int slot = 0; slot < n_expert_used; ++slot) {
+                    const int64_t i = ids[(int64_t) blockIdx.x * n_expert_used + slot];
+                    scale[i] = warp_amax[0];
+                }
+            } else {
+                scale[blockIdx.y * ne1 + blockIdx.x] = warp_amax[0];
+            }
        }
    }
-
-    static constexpr int test_offsets[5] = { 0, -1, 1, -2, 2};
-    const int first_fp8_code = (int) ggml_cuda_fp32_to_ue4m3(amax_raw / 6.0f);
-
-    float best_err = FLT_MAX;
-    uint8_t fp8_code = 0;
-    float subblock_scale = 0.0f;
-
-#pragma unroll // Check +/- 2 to find best code to reduce NVFP4 activation loss. Negligible overhead on Blackwell.
-    for (int i = 0; i < 5; i++) {
-        const int test_code = first_fp8_code + test_offsets[i];
-        if (test_code < 0 || test_code > 0x7e) {
-            continue;
-        }
-        const uint8_t code = (uint8_t) test_code;
-        const float test_scale = ggml_cuda_ue4m3_to_fp32(code);
-        const float test_inv_scale = test_scale > 0.0f ? 0.5f / test_scale : 0.0f;
-        float cur_err = 0.0f;
-#pragma unroll
-        for (int k = 0; k < QK_NVFP4_SUB; ++k) {
-            const float v = vals_raw[k];
-            const uint8_t q = ggml_cuda_float_to_fp4_e2m1(v, test_inv_scale);
-            const float err_diff = fabsf(v) - fabsf(kvalues_mxfp4[q & 0x7]) * test_scale;
-            cur_err = fmaf(err_diff, err_diff, cur_err);
-        }
-
-        if (cur_err < best_err) {
-            best_err = cur_err;
-            fp8_code = test_code;
-            subblock_scale = test_scale;
-        }
-    }
-
-    const float inv_scale = subblock_scale > 0.0f ? 0.5f / subblock_scale : 0.0f;
-    uint32_t q0 = 0;
-    uint32_t q1 = 0;
-#pragma unroll // this is faster than the previous __nv_fp4x4_e2m1
-    for (int k = 0; k < QK_NVFP4_SUB / 4; ++k) {
-        q0 |= (uint32_t) ggml_cuda_float_to_fp4_e2m1(vals_raw[k +  0], inv_scale) << (8 * k);
-        q0 |= (uint32_t) ggml_cuda_float_to_fp4_e2m1(vals_raw[k +  8], inv_scale) << (8 * k + 4);
-        q1 |= (uint32_t) ggml_cuda_float_to_fp4_e2m1(vals_raw[k +  4], inv_scale) << (8 * k);
-        q1 |= (uint32_t) ggml_cuda_float_to_fp4_e2m1(vals_raw[k + 12], inv_scale) << (8 * k + 4);
-    }
+    __syncthreads();

    block_fp4_mmq * y = (block_fp4_mmq *) vy;
-    if constexpr (scatter) {
+    const int64_t n_subblocks = (ne0 + QK_NVFP4_SUB - 1) / QK_NVFP4_SUB;
+
+    for (int64_t isb = threadIdx.x; isb < n_subblocks; isb += blockDim.x) {
+        const int64_t i0_base = isb * QK_NVFP4_SUB;
+        const int64_t k_block = i0_base / QK_FP4_MMQ;
+        const int sub = (i0_base % QK_FP4_MMQ) / QK_NVFP4_SUB;
+
+        const float row_scale = warp_amax[0];
+        const float inv_col_scale = row_scale > 0.0f ? 1.0f / row_scale : 0.0f;
+
+        float vals[QK_NVFP4_SUB];
+        if constexpr (use_aligned_float8) {
+            const float * x_base = x_row + i0_base;
+            const float8 v0 = i0_base +  7 < ne00 ? reinterpret_cast<const float8 *>(x_base)[0]     : float8{0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f};
+            const float8 v1 = i0_base + 15 < ne00 ? reinterpret_cast<const float8 *>(x_base + 8)[0] : float8{0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f};
+            vals[0] = v0.x; vals[1] = v0.y; vals[2] = v0.z; vals[3] = v0.w;
+            vals[4] = v0.p; vals[5] = v0.q; vals[6] = v0.r; vals[7] = v0.s;
+            vals[8] = v1.x; vals[9] = v1.y; vals[10] = v1.z; vals[11] = v1.w;
+            vals[12] = v1.p; vals[13] = v1.q; vals[14] = v1.r; vals[15] = v1.s;
+        } else {
 #pragma unroll
-        for (int slot = 0; slot < n_expert_used; ++slot) {
-            const int64_t i = ids[(int64_t) blockIdx.x * n_expert_used + slot];
-            block_fp4_mmq * yb = y + (k_block * ne1 + i);
+            for (int k = 0; k < QK_NVFP4_SUB; ++k) {
+                const int64_t i00 = i0_base + k;
+                vals[k] = i00 < ne00 ? x_row[i00] : 0.0f;
+            }
+        }
+
+        uint32_t q0 = 0;
+        uint32_t q1 = 0;
+
+        float amax_sub = 0.0f;
+#pragma unroll
+        for (int k = 0; k < QK_NVFP4_SUB; ++k) {
+            amax_sub = fmaxf(amax_sub, fabsf(vals[k] * inv_col_scale));
+        }
+
+        static constexpr int test_offsets[5] = { 0, -1, 1, -2, 2 };
+        const int first_fp8_code = (int) ggml_cuda_fp32_to_ue4m3(amax_sub / 6.0f);
+
+        uint8_t fp8_code = (uint8_t) first_fp8_code;
+        float subblock_scale = ggml_cuda_ue4m3_to_fp32(fp8_code);
+        float inv_scale_err = subblock_scale > 0.0f ? 0.5f / subblock_scale : 0.0f;
+#if CUDART_VERSION >= 12080
+        float best_err = nvfp4_native_scale_error(vals, inv_col_scale, inv_scale_err, subblock_scale);
+#else
+        float best_err = 0.0f;
+#pragma unroll
+        for (int k = 0; k < QK_NVFP4_SUB; ++k) {
+            const float v = vals[k] * inv_col_scale;
+            const uint8_t q = ggml_cuda_float_to_fp4_e2m1(v, inv_scale_err);
+            const float err_diff = fabsf(v) - fabsf(kvalues_fp4[q & 0x7]) * subblock_scale;
+            best_err = fmaf(err_diff, err_diff, best_err);
+        }
+#endif // CUDART_VERSION >= 12080
+
+#pragma unroll
+        for (int i = 1; i < 5; ++i) {
+            const int test_code = first_fp8_code + test_offsets[i];
+            if (test_code < 0 || test_code > 0x7e) {
+                continue;
+            }
+
+            const float test_scale = ggml_cuda_ue4m3_to_fp32((uint8_t) test_code);
+            const float test_inv_scale = test_scale > 0.0f ? 0.5f / test_scale : 0.0f;
+#if CUDART_VERSION >= 12080
+            const float cur_err = nvfp4_native_scale_error(vals, inv_col_scale, test_inv_scale, test_scale);
+#else
+            float cur_err = 0.0f;
+#pragma unroll
+            for (int k = 0; k < QK_NVFP4_SUB; ++k) {
+                const float v = vals[k] * inv_col_scale;
+                const uint8_t q = ggml_cuda_float_to_fp4_e2m1(v, test_inv_scale);
+                const float err_diff = fabsf(v) - fabsf(kvalues_fp4[q & 0x7]) * test_scale;
+                cur_err = fmaf(err_diff, err_diff, cur_err);
+            }
+#endif // CUDART_VERSION >= 12080
+
+            if (cur_err < best_err) {
+                best_err = cur_err;
+                fp8_code = (uint8_t) test_code;
+                subblock_scale = test_scale;
+            }
+        }
+#if CUDART_VERSION >= 12080
+        const float inv_scale = subblock_scale > 0.0f ? 0.5f / subblock_scale : 0.0f;
+        const float s = inv_col_scale * inv_scale;
+
+        __nv_fp4x4_e2m1 q0_lo(make_float4(vals[0] * s, vals[8]  * s, vals[1] * s, vals[9]  * s));
+        __nv_fp4x4_e2m1 q0_hi(make_float4(vals[2] * s, vals[10] * s, vals[3] * s, vals[11] * s));
+        __nv_fp4x4_e2m1 q1_lo(make_float4(vals[4] * s, vals[12] * s, vals[5] * s, vals[13] * s));
+        __nv_fp4x4_e2m1 q1_hi(make_float4(vals[6] * s, vals[14] * s, vals[7] * s, vals[15] * s));
+
+        const char2 q0_lo_c = *reinterpret_cast<char2 *>(&q0_lo);
+        const char2 q0_hi_c = *reinterpret_cast<char2 *>(&q0_hi);
+        const char2 q1_lo_c = *reinterpret_cast<char2 *>(&q1_lo);
+        const char2 q1_hi_c = *reinterpret_cast<char2 *>(&q1_hi);
+
+        q0 = uint32_t(uint8_t(q0_lo_c.x)) | (uint32_t(uint8_t(q0_lo_c.y)) <<  8) |
+            (uint32_t(uint8_t(q0_hi_c.x)) << 16) | (uint32_t(uint8_t(q0_hi_c.y)) << 24);
+        q1 = uint32_t(uint8_t(q1_lo_c.x)) | (uint32_t(uint8_t(q1_lo_c.y)) <<  8) |
+            (uint32_t(uint8_t(q1_hi_c.x)) << 16) | (uint32_t(uint8_t(q1_hi_c.y)) << 24);
+#else
+        const float inv_scale = subblock_scale > 0.0f ? 0.5f / subblock_scale : 0.0f;
+#pragma unroll
+        for (int k = 0; k < QK_NVFP4_SUB / 4; ++k) {
+            q0 |= uint32_t(ggml_cuda_float_to_fp4_e2m1(vals[k + 0] * inv_col_scale, inv_scale)) << (8 * k);
+            q0 |= uint32_t(ggml_cuda_float_to_fp4_e2m1(vals[k + 8] * inv_col_scale, inv_scale)) << (8 * k + 4);
+            q1 |= uint32_t(ggml_cuda_float_to_fp4_e2m1(vals[k + 4] * inv_col_scale, inv_scale)) << (8 * k);
+            q1 |= uint32_t(ggml_cuda_float_to_fp4_e2m1(vals[k + 12] * inv_col_scale, inv_scale)) << (8 * k + 4);
+        }
+#endif // CUDART_VERSION >= 12080
+
+        if constexpr (scatter) {
+#pragma unroll
+            for (int slot = 0; slot < n_expert_used; ++slot) {
+                const int64_t i = ids[(int64_t) blockIdx.x * n_expert_used + slot];
+                block_fp4_mmq * yb = y + (k_block * ne1 + i);
+                uint32_t * yqs = reinterpret_cast<uint32_t *>(yb->qs);
+                yqs[2 * sub + 0] = q0;
+                yqs[2 * sub + 1] = q1;
+                reinterpret_cast<uint8_t *>(yb->d4)[sub] = fp8_code;
+            }
+        } else {
+            block_fp4_mmq * yb = y + (blockIdx.y * ((int64_t) blocks_per_col * ne1) + k_block * ne1 + blockIdx.x);
            uint32_t * yqs = reinterpret_cast<uint32_t *>(yb->qs);
            yqs[2 * sub + 0] = q0;
            yqs[2 * sub + 1] = q1;
            reinterpret_cast<uint8_t *>(yb->d4)[sub] = fp8_code;
        }
-    } else {
-        block_fp4_mmq * yb = y + (blockIdx.z * ((int64_t) blocks_per_col * ne1) + k_block * ne1 + blockIdx.x);
-        uint32_t * yqs = reinterpret_cast<uint32_t *>(yb->qs);
-        yqs[2 * sub + 0] = q0;
-        yqs[2 * sub + 1] = q1;
-        reinterpret_cast<uint8_t *>(yb->d4)[sub] = fp8_code;
    }
-    GGML_UNUSED(n_expert_used);
 #else
-    GGML_UNUSED(n_expert_used);
+    GGML_UNUSED_VARS(x, ids, vy, scale, ne00, s01, s02, s03, ne0, ne1, ne2, n_expert_used);
    NO_DEVICE_CODE; // This is for Blackwell NVFP4 activations only.
 #endif // defined(BLACKWELL_MMA_AVAILABLE)

@@ -491,18 +634,22 @@ void quantize_scatter_mmq_q8_1_cuda(

 // scatter=true reuses the quant kernels: grid over tokens, ids = inverse map (token slot -> compact row)
 void quantize_scatter_mmq_fp4_cuda(
-        const float * x, const int32_t * ids_src1_inv, void * vy, const ggml_type type_src0,
+        const float * x, const int32_t * ids_src1_inv, void * vy, float * scale, const ggml_type type_src0, const bool use_aligned_float8,
        const int64_t ne00, const int64_t stride_token, const int64_t ne0,
        const int64_t n_tokens, const int64_t nrows_dst, const int n_expert_used, cudaStream_t stream) {
    GGML_ASSERT(ne0 > 0);
    if (type_src0 == GGML_TYPE_NVFP4) {
+        GGML_ASSERT(scale);
        GGML_ASSERT(ne00 % QK_NVFP4 == 0);
-        constexpr int nvfp4_block_size = 128;
-        const int64_t block_num_y = (ne0 + QK_NVFP4_SUB * nvfp4_block_size - 1) / (QK_NVFP4_SUB * nvfp4_block_size);
-        const dim3 block_size(nvfp4_block_size, 1, 1);
-        const dim3 num_blocks(n_tokens, block_num_y, 1);
-        quantize_mmq_nvfp4<true><<<num_blocks, block_size, 0, stream>>>(
-            x, ids_src1_inv, vy, ne00, /*s01=*/0, /*s02=*/stride_token, /*s03=*/0, ne0, /*ne1=*/nrows_dst, /*ne2=*/1, n_expert_used);
+        const dim3 block_size(CUDA_QUANTIZE_BLOCK_SIZE_MMQ, 1, 1);
+        const dim3 num_blocks(n_tokens, 1, 1);
+        if (use_aligned_float8) {
+            quantize_mmq_nvfp4<true, true><<<num_blocks, block_size, 0, stream>>>(
+                x, ids_src1_inv, vy, scale, ne00, /*s01=*/0, /*s02=*/stride_token, /*s03=*/0, ne0, /*ne1=*/nrows_dst, /*ne2=*/1, n_expert_used);
+        } else {
+            quantize_mmq_nvfp4<true, false><<<num_blocks, block_size, 0, stream>>>(
+                x, ids_src1_inv, vy, scale, ne00, /*s01=*/0, /*s02=*/stride_token, /*s03=*/0, ne0, /*ne1=*/nrows_dst, /*ne2=*/1, n_expert_used);
+        }
    } else {
        GGML_ASSERT(type_src0 == GGML_TYPE_MXFP4);
        constexpr int nwarps = 8;
@@ -516,20 +663,24 @@ void quantize_scatter_mmq_fp4_cuda(
 }

 void quantize_mmq_fp4_cuda(
-        const float * x, const int32_t * ids, void * vy, const ggml_type type_src0,
+        const float * x, const int32_t * ids, void * vy, float * scale, const ggml_type type_src0, const bool use_aligned_float8,
        const int64_t ne00, const int64_t s01, const int64_t s02, const int64_t s03,
        const int64_t ne0, const int64_t ne1, const int64_t ne2, const int64_t ne3, cudaStream_t stream) {
    GGML_ASSERT(type_src0 == GGML_TYPE_MXFP4 || type_src0 == GGML_TYPE_NVFP4);
    GGML_ASSERT(ne0 > 0);

    if (type_src0 == GGML_TYPE_NVFP4) {
+        GGML_ASSERT(scale);
        GGML_ASSERT(ne00 % QK_NVFP4 == 0);
-        constexpr int nvfp4_block_size = 128;
-        const int64_t block_num_y = (ne0 + QK_NVFP4_SUB * nvfp4_block_size - 1) / (QK_NVFP4_SUB * nvfp4_block_size);
-        const dim3 block_size(nvfp4_block_size, 1, 1);
-        const dim3 num_blocks(ne1, block_num_y, ne2 * ne3);
-        quantize_mmq_nvfp4<false><<<num_blocks, block_size, 0, stream>>>(
-            x, ids, vy, ne00, s01, s02, s03, ne0, ne1, ne2, /*n_expert_used=*/0);
+        const dim3 block_size(CUDA_QUANTIZE_BLOCK_SIZE_MMQ, 1, 1);
+        const dim3 num_blocks(ne1, ne2 * ne3, 1);
+        if (use_aligned_float8) {
+            quantize_mmq_nvfp4<false, true><<<num_blocks, block_size, 0, stream>>>(
+                x, ids, vy, scale, ne00, s01, s02, s03, ne0, ne1, ne2, /*n_expert_used=*/0);
+        } else {
+            quantize_mmq_nvfp4<false, false><<<num_blocks, block_size, 0, stream>>>(
+                x, ids, vy, scale, ne00, s01, s02, s03, ne0, ne1, ne2, /*n_expert_used=*/0);
+        }
    } else {
        GGML_ASSERT(ne0 % (2 * QK_MXFP4) == 0);

@@ -29,7 +29,9 @@ void quantize_mmq_q8_1_cuda(
 void quantize_mmq_fp4_cuda(const float *   x,
                             const int32_t * ids,
                             void *          vy,
+                             float *         scale,
                             ggml_type       type_src0,
+                             bool            use_aligned_float8,
                             int64_t         ne00,
                             int64_t         s01,
                             int64_t         s02,
@@ -44,7 +46,9 @@ void quantize_mmq_fp4_cuda(const float *   x,
 void quantize_scatter_mmq_fp4_cuda(const float *   x,
                                   const int32_t * ids_src1_inv,
                                   void *          vy,
+                                   float *         scale,
                                   ggml_type       type_src0,
+                                   bool            use_aligned_float8,
                                   int64_t         ne00,
                                   int64_t         stride_token,
                                   int64_t         ne0,
@@ -8,6 +8,7 @@
 // Kernel config struct - passed by value to CUDA kernel
 struct topk_moe_config {
    bool use_sigmoid;
+    bool use_sqrt_softplus;
    bool with_norm;
    bool delayed_softmax;
 };
@@ -67,6 +68,16 @@ __device__ void sigmoid_warp_inplace(float (&vals)[experts_per_thread], const in
    }
 }

+template <int experts_per_thread, bool use_limit>
+__device__ void sqrt_softplus_warp_inplace(float (&vals)[experts_per_thread], const int limit, const int lane) {
+#pragma unroll
+    for (int i = 0; i < experts_per_thread; i++) {
+        const int  idx    = lane + i * WARP_SIZE;
+        const bool active = !use_limit || (idx < limit);
+        vals[i]           = active ? sqrtf(vals[i] > 20.0f ? vals[i] : logf(1.0f + expf(vals[i]))) : -INFINITY;
+    }
+}
+
 /*
    This kernel does the following:
    1. optionally softmax over the logits per token [n_experts, n_tokens]
@@ -115,6 +126,8 @@ __launch_bounds__(4 * WARP_SIZE, 1) __global__ void topk_moe_cuda(const float *
    if (!config.delayed_softmax) {
        if (config.use_sigmoid) {
           sigmoid_warp_inplace<experts_per_thread, false>(wt, n_experts, threadIdx.x);
+        } else if (config.use_sqrt_softplus) {
+           sqrt_softplus_warp_inplace<experts_per_thread, false>(wt, n_experts, threadIdx.x);
        } else {
           softmax_warp_inplace<experts_per_thread, false>(wt, n_experts, threadIdx.x);
        }
@@ -364,9 +377,10 @@ void ggml_cuda_op_topk_moe(ggml_backend_cuda_context &     ctx,
    }

    topk_moe_config config;
-    config.use_sigmoid     = args.sigmoid;
-    config.with_norm       = with_norm;
-    config.delayed_softmax = args.delayed_softmax;
+    config.use_sigmoid       = args.sigmoid;
+    config.use_sqrt_softplus = args.sqrt_softplus;
+    config.with_norm         = with_norm;
+    config.delayed_softmax   = args.delayed_softmax;

    if (bias) {
        launch_topk_moe_cuda<true>(ctx, logits_d, weights_d, ids_d, bias_d, n_rows, n_experts, n_expert_used, clamp_val,
@@ -415,7 +429,7 @@ bool ggml_cuda_should_use_topk_moe(const ggml_tensor * gating_op,
    } else if (gating_op->op == GGML_OP_UNARY) {
        ggml_unary_op op = ggml_get_unary_op(gating_op);

-        if (op != GGML_UNARY_OP_SIGMOID) {
+        if (op != GGML_UNARY_OP_SIGMOID && op != GGML_UNARY_OP_SOFTPLUS) {
            return false;
        }
    }
@@ -5,6 +5,7 @@

 struct ggml_cuda_topk_moe_args {
    bool sigmoid{};
+    bool sqrt_softplus{};
    bool softmax{};
    bool delayed_softmax{};
    bool prob_bias{};
@@ -64,8 +64,10 @@
 #define cudaErrorMemoryAllocation hipErrorOutOfMemory
 #define cudaErrorPeerAccessAlreadyEnabled hipErrorPeerAccessAlreadyEnabled
 #define cudaErrorPeerAccessNotEnabled hipErrorPeerAccessNotEnabled
+#define cudaEventCreate hipEventCreate
 #define cudaEventCreateWithFlags hipEventCreateWithFlags
 #define cudaEventDisableTiming hipEventDisableTiming
+#define cudaEventElapsedTime hipEventElapsedTime
 #define cudaEventRecord hipEventRecord
 #define cudaEventSynchronize hipEventSynchronize
 #define cudaEvent_t hipEvent_t
@@ -48,8 +48,10 @@
 #define cudaErrorMemoryAllocation musaErrorMemoryAllocation
 #define cudaErrorPeerAccessAlreadyEnabled musaErrorPeerAccessAlreadyEnabled
 #define cudaErrorPeerAccessNotEnabled musaErrorPeerAccessNotEnabled
+#define cudaEventCreate musaEventCreate
 #define cudaEventCreateWithFlags musaEventCreateWithFlags
 #define cudaEventDisableTiming musaEventDisableTiming
+#define cudaEventElapsedTime musaEventElapsedTime
 #define cudaEventRecord musaEventRecord
 #define cudaEventSynchronize musaEventSynchronize
 #define cudaEvent_t musaEvent_t
@@ -143,12 +143,12 @@ static const char * htp_event_name(uint16_t id) {
        case HTP_TRACE_EVT_HMX_COMP:       return "HMX_COMP";
        case HTP_TRACE_EVT_L2FLUSH:        return "L2FLUSH";
        case HTP_TRACE_EVT_INIT:           return "INIT";
+        case HTP_TRACE_EVT_BUFF:           return "BUFF";
        default:                           return "UNKNOWN";
    }
 }

-static void ggml_hexagon_dump_op_prof(const std::string &sess_name, const htp_opnode & node,
-                                      const htp_prof_desc & pd) {
+static void ggml_hexagon_dump_op_prof(const std::string &sess_name, const htp_opnode & node, const htp_prof_desc & pd) {
    if (!opt_profile) return;

    uint32_t op_usec = pd.usecs;
@@ -168,6 +168,43 @@ static void ggml_hexagon_dump_op_prof(const std::string &sess_name, const htp_op
            node.op_name().c_str(), fmt.names, fmt.dims, fmt.types, fmt.strides, fmt.kparams, op_usec, op_cycles, pd.cycles_start, mhz, pmu_str);
 }

+static void ggml_hexagon_dump_batch_prof(const std::string & sess_name, const htp_opbatch_rsp & rsp) {
+    uint64_t batch_cycles = rsp.cycles_stop - rsp.cycles_start;
+    float batch_mhz = rsp.usecs > 0 ? (float) batch_cycles / rsp.usecs : 0.0f;
+
+    char evt_str[256] = "----";
+    if (opt_profile == 3) {
+        snprintf(evt_str, sizeof(evt_str), "evt-cnt %u,%u,%u,%u,%u,%u,%u,%u,%u,%u,%u",
+                rsp.n_traces[0], rsp.n_traces[1], rsp.n_traces[2], rsp.n_traces[3],
+                rsp.n_traces[4], rsp.n_traces[5], rsp.n_traces[6], rsp.n_traces[7],
+                rsp.n_traces[8], rsp.n_traces[9], rsp.n_traces[10]);
+    }
+
+    GGML_LOG_DEBUG("ggml-hex: %s profile-op OPBATCH|----|n-ops %u|%s|----|----|usec %u cycles %llu start %llu mhz %.1f\n",
+                   sess_name.c_str(), rsp.n_ops, evt_str, rsp.usecs, (unsigned long long) batch_cycles, (unsigned long long) rsp.cycles_start, batch_mhz);
+}
+
+static void ggml_hexagon_dump_trace_events(const std::string & sess_name, const htp_opbatch_rsp & rsp,
+                                           const htp_trace_desc * trace_events, uint32_t n_traces) {
+    if (opt_profile == 3 && trace_events) {
+        uint32_t valid_cnt[HTP_MAX_NTHREADS + 1] = {0};
+        for (uint32_t t = 0; t <= HTP_MAX_NTHREADS; t++) {
+            uint32_t count = rsp.n_traces[t];
+            valid_cnt[t] = count > n_traces ? n_traces : count;
+        }
+
+        for (uint32_t t = 0; t <= HTP_MAX_NTHREADS; t++) {
+            for (uint32_t idx = 0; idx < valid_cnt[t]; idx++) {
+                const auto & e = trace_events[t * n_traces + idx];
+                bool is_stop = (e.info & 0x8000) != 0;
+                uint16_t info = e.info & 0x7FFF;
+                GGML_LOG_DEBUG("ggml-hex: %s trace-evt %s: thread %u info %u %s %u\n",
+                               sess_name.c_str(), htp_event_name(e.id), t, info, is_stop ? "stop" : "start", e.cycles);
+            }
+        }
+    }
+}
+
 // **

 static inline bool ggml_hexagon_is_repack_type(enum ggml_type type) {
@@ -1128,13 +1165,7 @@ struct ggml_hexagon_opbatch {
    std::unordered_map<const ggml_tensor*, int> t_map; // tensor ptr  to index
    std::unordered_multimap<void*, int>         d_map; // tensor data to index

-    struct tensor_range {
-        uint64_t start;
-        uint64_t end;
-        int bi;
-        std::vector<int> tensors;
-    };
-    std::vector<tensor_range> ranges;
+

    unsigned int n_bufs;     // num buffers in the batch
    unsigned int n_tens;     // num tensors ...
@@ -1155,7 +1186,6 @@ struct ggml_hexagon_opbatch {
        b_map.clear();
        t_map.clear();
        d_map.clear();
-        ranges.clear();
    }

    ggml_hexagon_opbatch(ggml_hexagon_session *sess, size_t batch_size, size_t max_vmem) {
@@ -1209,70 +1239,7 @@ struct ggml_hexagon_opbatch {
        return bi;
    }

-    void add_range(const htp_tensor * h, int ti) {
-        uint64_t t_start = h->data;
-        uint64_t t_end   = t_start + h->size;
-        int      bi      = h->bi;

-        int first_match = -1;
-        int unused_idx  = -1;
-        for (size_t i = 0; i < ranges.size(); i++) {
-            if (ranges[i].bi == -1) {
-                unused_idx = i;
-                continue;
-            }
-            if (ranges[i].bi != bi) {
-                continue;
-            }
-            if (ranges[i].start >= t_end || ranges[i].end <= t_start) {
-                continue;
-            }
-
-            if (first_match == -1) {
-                first_match = i;
-                HEX_VERBOSE("ggml-hex: %s range-grow #%d : bi %d [%p, %p) + #%d [%p, %p) -> [%p, %p)\n",
-                    sess->c_name(), (int) i, ranges[i].bi,
-                    (void *) (h_bufs[ranges[i].bi].base + ranges[i].start),
-                    (void *) (h_bufs[ranges[i].bi].base + ranges[i].end),
-                    ti,
-                    (void *) (h_bufs[bi].base + t_start),
-                    (void *) (h_bufs[bi].base + t_end),
-                    (void *) (h_bufs[ranges[i].bi].base + std::min(ranges[i].start, t_start)),
-                    (void *) (h_bufs[ranges[i].bi].base + std::max(ranges[i].end, t_end)));
-
-                ranges[i].start = std::min(ranges[i].start, t_start);
-                ranges[i].end   = std::max(ranges[i].end, t_end);
-                ranges[i].tensors.push_back(ti);
-            } else {
-                HEX_VERBOSE("ggml-hex: %s range-merge #%d [%p, %p) + #%d [%p, %p) -> [%p, %p)\n",
-                    sess->c_name(), first_match,
-                    (void *) (h_bufs[bi].base + ranges[first_match].start),
-                    (void *) (h_bufs[bi].base + ranges[first_match].end),
-                    (int) i,
-                    (void *) (h_bufs[bi].base + ranges[i].start),
-                    (void *) (h_bufs[bi].base + ranges[i].end),
-                    (void *) (h_bufs[bi].base + std::min(ranges[first_match].start, ranges[i].start)),
-                    (void *) (h_bufs[bi].base + std::max(ranges[first_match].end, ranges[i].end)));
-
-                ranges[first_match].start = std::min(ranges[first_match].start, ranges[i].start);
-                ranges[first_match].end   = std::max(ranges[first_match].end, ranges[i].end);
-                ranges[first_match].tensors.insert(
-                    ranges[first_match].tensors.end(),
-                    ranges[i].tensors.begin(),
-                    ranges[i].tensors.end()
-                );
-                ranges[i].bi = -1;
-            }
-        }
-
-        if (first_match == -1) {
-            if (unused_idx != -1) {
-                ranges[unused_idx] = {t_start, t_end, bi, {ti}};
-            } else {
-                ranges.push_back({t_start, t_end, bi, {ti}});
-            }
-        }
-    }

    bool same_shape(const htp_tensor * h, const ggml_tensor * t) const {
        int64_t ne0 = t->ne[0];
@@ -1286,7 +1253,8 @@ struct ggml_hexagon_opbatch {
        int64_t nb2 = is_repack ? nb1 * ne1 : t->nb[2];
        int64_t nb3 = is_repack ? nb2 * t->ne[2] : t->nb[3];

-        return (h->ne[0] == ne0) && (h->ne[1] == ne1) && (h->ne[2] == t->ne[2]) && (h->ne[3] == t->ne[3]) &&
+        return (h->type == t->type) &&
+               (h->ne[0] == ne0) && (h->ne[1] == ne1) && (h->ne[2] == t->ne[2]) && (h->ne[3] == t->ne[3]) &&
               (h->nb[0] == t->nb[0]) && (h->nb[1] == nb1) && (h->nb[2] == nb2) && (h->nb[3] == nb3);
    }

@@ -1340,8 +1308,7 @@ struct ggml_hexagon_opbatch {
            h.nb[0] = t->nb[0]; h.nb[1] = t->nb[1]; h.nb[2] = t->nb[2]; h.nb[3] = t->nb[3];
        }

-        h.alias = ti;
-        add_range(&h, ti);
+

        h.flags = 0;
        if (ggml_backend_buffer_get_usage(t->buffer) != GGML_BACKEND_BUFFER_USAGE_WEIGHTS) {
@@ -1423,14 +1390,6 @@ struct ggml_hexagon_opbatch {
    }

    void finalize_ranges() {
-        for (const auto & r : ranges) {
-            if (r.bi == -1) {
-                continue;
-            }
-            for (size_t i = 0; i < r.tensors.size(); i++) {
-                h_tens[r.tensors[i]].alias = r.tensors[(i + 1) % r.tensors.size()];
-            }
-        }
    }
 };

@@ -1581,9 +1540,6 @@ struct ggml_hexagon_opqueue {
        if (opt_profile && rsp.n_ops > 0) {
            auto & ops = op_cache[rsp.id];

-            uint64_t batch_usec = ggml_time_us() - start_usec[rsp.id];
-            uint32_t htp_usec   = 0;
-
            GGML_ASSERT(rsp.n_ops <= ops.size());

            const htp_prof_desc * pd = (const htp_prof_desc *) p_ptr;
@@ -1594,55 +1550,13 @@ struct ggml_hexagon_opqueue {
                trace_events = (const htp_trace_desc *) (p_ptr + p_size);
            }

-            uint32_t trace_idx[HTP_MAX_NTHREADS + 1] = {0};
-            uint32_t valid_cnt[HTP_MAX_NTHREADS + 1] = {0};
-
-            if (opt_profile == 3) {
-                for (uint32_t t = 0; t <= HTP_MAX_NTHREADS; t++) {
-                    uint32_t count = rsp.n_traces[t];
-                    valid_cnt[t] = count > n_traces ? n_traces : count;
-                }
-            }
+            ggml_hexagon_dump_batch_prof(shm_buf->sess->name, rsp);

            for (uint32_t i = 0; i < rsp.n_ops; i++) {
-                htp_usec += pd[i].usecs;
-
                ggml_hexagon_dump_op_prof(shm_buf->sess->name, ops[i], pd[i]);
-
-                if (opt_profile == 3) {
-                    uint32_t op_duration = pd[i].cycles_stop - pd[i].cycles_start;
-
-                    for (uint32_t t = 0; t <= HTP_MAX_NTHREADS; t++) {
-                        while (trace_idx[t] < valid_cnt[t]) {
-                            const auto & e = trace_events[t * n_traces + trace_idx[t]];
-                            uint32_t offset = e.cycles - pd[i].cycles_start;
-                            if (offset >= 0x80000000) {
-                                trace_idx[t]++;
-                                continue;
-                            }
-                            if (offset > op_duration) {
-                                break;
-                            }
-                            bool is_stop = (e.info & 0x8000) != 0;
-                            uint16_t info = e.info & 0x7FFF;
-                            GGML_LOG_DEBUG("ggml-hex: %s trace-op %s: thread %u event %s info %u %s %u\n",
-                                           shm_buf->sess->c_name(), ops[i].op_name().c_str(), t, htp_event_name(e.id), info, is_stop ? "stop" : "start", e.cycles);
-                            trace_idx[t]++;
-                        }
-                    }
-                }
            }

-            char evt_str[256] = "";
-            if (opt_profile == 3) {
-                snprintf(evt_str, sizeof(evt_str), " evt [%u,%u,%u,%u,%u,%u,%u,%u,%u,%u,%u]",
-                        rsp.n_traces[0], rsp.n_traces[1], rsp.n_traces[2], rsp.n_traces[3],
-                        rsp.n_traces[4], rsp.n_traces[5], rsp.n_traces[6], rsp.n_traces[7],
-                        rsp.n_traces[8], rsp.n_traces[9], rsp.n_traces[10]);
-            }
-
-            GGML_LOG_DEBUG("ggml-hex: %s profile-batch n-ops %u batch-dur-usec %lld htp-ops-usec %u%s\n",
-                           shm_buf->sess->c_name(), rsp.n_ops, (long long) batch_usec, htp_usec, evt_str);
+            ggml_hexagon_dump_trace_events(shm_buf->sess->name, rsp, trace_events, n_traces);
        }
    }
 };
@@ -1661,7 +1575,7 @@ void ggml_hexagon_session::flush_pending(bool all) {
        const uint32_t timeo = opt_oppoll ? 0 : DSPQUEUE_TIMEOUT;

        int err = dspqueue_read(this->queue, &flags, 1, &n_dbufs, &dbuf, sizeof(rsp), &rsp_size, (uint8_t *) &rsp, timeo);
-        if (err == AEE_EEXPIRED) {
+        if (err == AEE_EEXPIRED || err == AEE_EWOULDBLOCK) {
            continue;
        }

@@ -2113,7 +2027,7 @@ static bool ggml_hexagon_precompute_flash_attn_params(
    const struct ggml_tensor * sinks = op->src[4];
    if (ggml_hexagon_flash_attn_is_hmx_eligible(sess, q, k, v, sinks)) {
        size_t Br = 0, Bc = 0;
-        int ret = hmx_fa_find_chunk_size(&Br, &Bc, G, DK, DV, neq1, nek1, sess->vtcm_size, sess->n_threads);
+        int ret = hmx_fa_find_chunk_size(&Br, &Bc, G, DK, DV, neq1, nek1, sess->vtcm_size, sess->n_threads, kparams->is_q_fp32 != 0);
        if (ret == 0) {
            kparams->kernel_type = HTP_FA_KERNEL_HMX;
            kparams->Br = Br;
@@ -2123,7 +2037,7 @@ static bool ggml_hexagon_precompute_flash_attn_params(

            kparams->u.hmx.g_br = hex_align_up(G * Br, 32);
            kparams->u.hmx.pipeline = (kparams->n_kv_blocks >= 3 && sess->n_threads >= 2) ? 1 : 0;
-            kparams->vtcm_size = hmx_fa_compute_vtcm_usage(G, DK, DV, Br, Bc, kparams->n_threads, kparams->u.hmx.pipeline != 0);
+            kparams->vtcm_size = hmx_fa_compute_vtcm_usage(G, DK, DV, Br, Bc, kparams->n_threads, kparams->u.hmx.pipeline != 0, kparams->is_q_fp32 != 0);

            const size_t row_vec_bytes = hex_align_up(Bc * sizeof(uint16_t), 256);
            kparams->u.hmx.row_buf_stride = row_vec_bytes / 128; // HVX vector is 128 bytes
@@ -2412,6 +2326,7 @@ static void ggml_hexagon_precompute_hvx_mm_params(
    int ne12,
    int ne13,
    bool is_matmul_id,
+    const size_t src2_row_size,
    size_t vtcm_budget,
    struct htp_mm_kernel_params * kparams
 ) {
@@ -2437,7 +2352,7 @@ static void ggml_hexagon_precompute_hvx_mm_params(
            for (uint32_t d = max_prefetch; d >= 2; d /= 2) {
                htp_mm_hvx_vtcm_layout_build(
                    &L, kparams->kernel_type, wtype, ne10, src1_nrows, sess->n_threads,
-                    0, src0->nb[1], 0, d, true, false, false
+                    0, src0->nb[1], 0, src2_row_size, d, true, false, false
                );
                if (L.total_bytes <= vtcm_budget) {
                    best_n_prefetch = d;
@@ -2447,7 +2362,7 @@ static void ggml_hexagon_precompute_hvx_mm_params(
            if (best_n_prefetch == 2 && L.total_bytes > vtcm_budget) {
                htp_mm_hvx_vtcm_layout_build(
                    &L, kparams->kernel_type, wtype, ne10, src1_nrows, sess->n_threads,
-                    0, src0->nb[1], 0, 2, true, false, false
+                    0, src0->nb[1], 0, src2_row_size, 2, true, false, false
                );
            }
            kparams->n_prefetch = best_n_prefetch;
@@ -2471,7 +2386,7 @@ static void ggml_hexagon_precompute_hvx_mm_params(
                for (uint32_t d = max_prefetch; d >= 2; d /= 2) {
                    htp_mm_hvx_vtcm_layout_build(
                        &L, kparams->kernel_type, wtype, ne10, src1_nrows, sess->n_threads,
-                        dst->nb[1], src0->nb[1], src1->nb[1], d, false, false, false
+                        dst->nb[1], src0->nb[1], src1->nb[1], src2_row_size, d, false, false, false
                    );
                    if (L.total_bytes <= vtcm_budget) {
                        best_n_prefetch = d;
@@ -2481,7 +2396,7 @@ static void ggml_hexagon_precompute_hvx_mm_params(
                if (best_n_prefetch == 2 && L.total_bytes > vtcm_budget) {
                    htp_mm_hvx_vtcm_layout_build(
                        &L, kparams->kernel_type, wtype, ne10, src1_nrows, sess->n_threads,
-                        dst->nb[1], src0->nb[1], src1->nb[1], 2, false, false, false
+                        dst->nb[1], src0->nb[1], src1->nb[1], src2_row_size, 2, false, false, false
                    );
                }

@@ -2505,7 +2420,7 @@ static void ggml_hexagon_precompute_hvx_mm_params(
                struct htp_mm_hvx_vtcm_layout L;
                htp_mm_hvx_vtcm_layout_build(
                    &L, kparams->kernel_type, wtype, ne10, src1_nrows, sess->n_threads,
-                    dst->nb[1], src0->nb[1], src1->nb[1], 16, false, false, false
+                    dst->nb[1], src0->nb[1], src1->nb[1], src2_row_size, 16, false, false, false
                );

                kparams->n_prefetch = 16;
@@ -2525,7 +2440,7 @@ static void ggml_hexagon_precompute_hvx_mm_params(
        struct htp_mm_hvx_vtcm_layout L;
        htp_mm_hvx_vtcm_layout_build(
            &L, HTP_MM_KERNEL_HVX_F16_F16_VTCM, wtype, ne10, src1_nrows, sess->n_threads,
-            dst->nb[1], src0->nb[1], src1->nb[1], 16, false, false, false
+            dst->nb[1], src0->nb[1], src1->nb[1], src2_row_size, 16, false, false, false
        );

        if (!is_batched && !is_permuted && L.total_bytes <= vtcm_budget) {
@@ -2545,7 +2460,7 @@ static void ggml_hexagon_precompute_hvx_mm_params(
            kparams->src1_row_size = src1->nb[1];
            htp_mm_hvx_vtcm_layout_build(
                &L, kparams->kernel_type, wtype, ne10, src1_nrows, sess->n_threads,
-                dst->nb[1], src0->nb[1], src1->nb[1], 16, false, false, false
+                dst->nb[1], src0->nb[1], src1->nb[1], src2_row_size, 16, false, false, false
            );
            kparams->vtcm_size = L.total_bytes;
            kparams->vtcm_src0_size = L.src0_bytes;
@@ -2561,7 +2476,7 @@ static void ggml_hexagon_precompute_hvx_mm_params(
        struct htp_mm_hvx_vtcm_layout L;
        htp_mm_hvx_vtcm_layout_build(
            &L, HTP_MM_KERNEL_HVX_F32_F32_VTCM, wtype, ne10, src1_nrows, sess->n_threads,
-            dst->nb[1], src0->nb[1], src1->nb[1], 16, false, false, false
+            dst->nb[1], src0->nb[1], src1->nb[1], src2_row_size, 16, false, false, false
        );

        if (!is_batched && !is_permuted && L.total_bytes <= vtcm_budget) {
@@ -2577,7 +2492,7 @@ static void ggml_hexagon_precompute_hvx_mm_params(
            kparams->src1_row_size = src1->nb[1];
            htp_mm_hvx_vtcm_layout_build(
                &L, kparams->kernel_type, wtype, ne10, src1_nrows, sess->n_threads,
-                dst->nb[1], src0->nb[1], src1->nb[1], 16, false, false, false
+                dst->nb[1], src0->nb[1], src1->nb[1], src2_row_size, 16, false, false, false
            );
            kparams->vtcm_size = L.total_bytes;
            kparams->vtcm_src0_size = L.src0_bytes;
@@ -2588,11 +2503,12 @@ static void ggml_hexagon_precompute_hvx_mm_params(
    }
 }

-static void ggml_hexagon_precompute_matmul_params(
+static void ggml_hexagon_precompute_matmul_params_impl(
    const struct ggml_hexagon_session * sess,
    const struct ggml_tensor * src0,
    const struct ggml_tensor * src1,
    const struct ggml_tensor * dst,
+    const size_t src2_row_size,
    struct htp_mm_kernel_params * kparams
 ) {
    memset(kparams, 0, sizeof(*kparams));
@@ -2627,7 +2543,7 @@ static void ggml_hexagon_precompute_matmul_params(
    }

    // Fallback to HVX parameter computation
-    ggml_hexagon_precompute_hvx_mm_params(sess, src0, src1, dst, wtype, ne02, ne03, ne10, ne11, ne12, ne13, is_matmul_id, vtcm_budget, kparams);
+    ggml_hexagon_precompute_hvx_mm_params(sess, src0, src1, dst, wtype, ne02, ne03, ne10, ne11, ne12, ne13, is_matmul_id, src2_row_size, vtcm_budget, kparams);

 finalize:
    kparams->div_ne12_ne1 = init_fastdiv_values(ne12 * ne11);
@@ -2637,6 +2553,27 @@ finalize:
    kparams->div_ne11     = init_fastdiv_values(ne11);
 }

+static void ggml_hexagon_precompute_matmul_params(
+    const struct ggml_hexagon_session * sess,
+    const struct ggml_tensor * src0,
+    const struct ggml_tensor * src1,
+    const struct ggml_tensor * dst,
+    struct htp_mm_kernel_params * kparams
+) {
+    ggml_hexagon_precompute_matmul_params_impl(sess, src0, src1, dst, 0, kparams);
+}
+
+static void ggml_hexagon_precompute_fused_matmul_add_params(
+    const struct ggml_hexagon_session * sess,
+    const struct ggml_tensor * src0,
+    const struct ggml_tensor * src1,
+    const struct ggml_tensor * src2,
+    const struct ggml_tensor * dst,
+    struct htp_mm_kernel_params * kparams
+) {
+    ggml_hexagon_precompute_matmul_params_impl(sess, src0, src1, dst, src2->nb[1], kparams);
+}
+
 static void ggml_hexagon_precompute_unary_params(
    const struct ggml_hexagon_session * sess,
    uint32_t op,
@@ -2730,7 +2667,7 @@ static void ggml_hexagon_precompute_fused_qkv_params(
            struct htp_mm_hvx_vtcm_layout L;
            htp_mm_hvx_vtcm_layout_build(
                &L, HTP_MM_KERNEL_HVX_QUANT_ROW, wtype, ne10, src1_nrows, sess->n_threads,
-                0, src0_row_size, src1_row_size, d, false, true, false
+                0, src0_row_size, src1_row_size, 0, d, false, true, false
            );
            if (L.total_bytes <= sess->vtcm_size) {
                best_n_prefetch = d;
@@ -2745,7 +2682,7 @@ static void ggml_hexagon_precompute_fused_qkv_params(
    // Test tiled first
    htp_mm_hvx_vtcm_layout_build(
        &L, HTP_MM_KERNEL_HVX_QUANT_ROW, wtype, ne10, src1_nrows, sess->n_threads,
-        0, src0_row_size, src1_row_size, best_n_prefetch, false, true, false
+        0, src0_row_size, src1_row_size, 0, best_n_prefetch, false, true, false
    );

    if (try_tiled && L.total_bytes <= sess->vtcm_size) {
@@ -2763,7 +2700,7 @@ static void ggml_hexagon_precompute_fused_qkv_params(

        htp_mm_hvx_vtcm_layout_build(
            &L, HTP_MM_KERNEL_HVX_QUANT_ROW_FLAT, wtype, ne10, src1_nrows, sess->n_threads,
-            0, src0_row_size, flat_src1_row_size, best_n_prefetch, false, true, false
+            0, src0_row_size, flat_src1_row_size, 0, best_n_prefetch, false, true, false
        );
        kparams->vtcm_src0_size = L.src0_bytes;
        kparams->vtcm_src1_size = L.src1_bytes;
@@ -2800,7 +2737,7 @@ static void ggml_hexagon_precompute_fused_ffn_params(
            struct htp_mm_hvx_vtcm_layout L;
            htp_mm_hvx_vtcm_layout_build(
                &L, HTP_MM_KERNEL_HVX_QUANT_ROW, wtype, ne10, src1_nrows, sess->n_threads,
-                0, src0_row_size, src1_row_size, d, false, false, true
+                0, src0_row_size, src1_row_size, 0, d, false, false, true
            );
            if (L.total_bytes <= sess->vtcm_size) {
                best_n_prefetch = d;
@@ -2815,7 +2752,7 @@ static void ggml_hexagon_precompute_fused_ffn_params(
    // Test tiled first
    htp_mm_hvx_vtcm_layout_build(
        &L, HTP_MM_KERNEL_HVX_QUANT_ROW, wtype, ne10, src1_nrows, sess->n_threads,
-        0, src0_row_size, src1_row_size, best_n_prefetch, false, false, true
+        0, src0_row_size, src1_row_size, 0, best_n_prefetch, false, false, true
    );

    if (try_tiled && L.total_bytes <= sess->vtcm_size) {
@@ -2832,7 +2769,7 @@ static void ggml_hexagon_precompute_fused_ffn_params(

        htp_mm_hvx_vtcm_layout_build(
            &L, HTP_MM_KERNEL_HVX_QUANT_ROW_FLAT, wtype, ne10, src1_nrows, sess->n_threads,
-            0, src0_row_size, flat_src1_row_size, best_n_prefetch, false, false, true
+            0, src0_row_size, flat_src1_row_size, 0, best_n_prefetch, false, false, true
        );
        kparams->vtcm_src0_size = L.src0_bytes;
        kparams->vtcm_src1_size = L.src1_bytes;
@@ -3083,7 +3020,10 @@ static bool ggml_hexagon_supported_activations(const struct ggml_hexagon_session
        return false;
    }

-    if (!ggml_is_contiguous(src0) || !ggml_is_contiguous(dst)) {
+    if (!ggml_is_contiguous_1(src0)) {
+        return false;
+    }
+    if (!ggml_is_contiguous(dst)) {
        return false;
    }

@@ -3094,7 +3034,7 @@ static bool ggml_hexagon_supported_activations(const struct ggml_hexagon_session
        if (!ggml_are_same_shape(src0, src1)) {
            return false;
        }
-        if (!ggml_is_contiguous(src1)) {
+        if (!ggml_is_contiguous_1(src1)) {
            return false;
        }
    }
@@ -3476,6 +3416,7 @@ static htp_op_code op_remap_to_htp(const ggml_tensor * t) {
        case GGML_OP_RMS_NORM:        return HTP_OP_RMS_NORM;
        case GGML_OP_CONCAT:          return HTP_OP_CONCAT;
        case GGML_OP_SCALE:           return HTP_OP_SCALE;
+        case GGML_OP_CLAMP:           return HTP_OP_CLAMP;
        case GGML_OP_SQR:             return HTP_OP_SQR;
        case GGML_OP_SQRT:            return HTP_OP_SQRT;
        case GGML_OP_SOFT_MAX:        return HTP_OP_SOFTMAX;
@@ -3651,16 +3592,19 @@ static bool try_fuse_node(const ggml_hexagon_session * sess, const ggml_cgraph *
    if (n->op == GGML_OP_MUL_MAT && next_node) {
        if (next_node->op == GGML_OP_ADD && op_is_compute(next_node) && ggml_can_fuse(graph, i, { GGML_OP_MUL_MAT, GGML_OP_ADD })) {
            if (next_node->src[0] == n || next_node->src[1] == n) {
+                const struct ggml_tensor * src2 = (next_node->src[0] == n) ? next_node->src[1] : next_node->src[0];
                struct htp_mm_kernel_params kparams;
-                ggml_hexagon_precompute_matmul_params(sess, n->src[0], n->src[1], next_node, &kparams);
-                if ((size_t)kparams.vtcm_size <= sess->vtcm_size) {
+                ggml_hexagon_precompute_fused_matmul_add_params(sess, n->src[0], n->src[1], src2, next_node, &kparams);
+                const int src1_nrows = n->src[1]->ne[1] * n->src[1]->ne[2] * n->src[1]->ne[3];
+                const bool can_fuse = (kparams.n_hmx > 0) || (src1_nrows == 1);
+                if (can_fuse && (size_t)kparams.vtcm_size <= sess->vtcm_size) {
                    htp_opnode node(n, {}, HTP_OP_MUL_MAT_ADD);
                    node.add_fused(next_node);
                    memcpy(node.kernel_params, &kparams, sizeof(kparams));
                    nodes.push_back(std::move(node));
                    i += 1;
                    return true;
-                } else {
+                } else if (can_fuse) {
                    HEX_VERBOSE("ggml-hex: skip MUL_MAT_ADD fusion because VTCM needed (%d) > budget (%zu)\n",
                                kparams.vtcm_size, sess->vtcm_size);
                }
@@ -3916,6 +3860,7 @@ static ggml_backend_t ggml_backend_hexagon_device_init(ggml_backend_dev_t dev, c
        /* .interface = */ hexagon_backend_i,
        /* .device    = */ dev,
        /* .context   = */ sess,
+        /* .profiler  = */ nullptr,
    };

    GGML_UNUSED(params);
@@ -4126,6 +4071,7 @@ static bool ggml_backend_hexagon_device_supports_op(ggml_backend_dev_t dev, cons
        case GGML_OP_L2_NORM:
        case GGML_OP_RMS_NORM:
        case GGML_OP_SCALE:
+        case GGML_OP_CLAMP:
            supp = ggml_hexagon_supported_unary(sess, op);
            break;

@@ -4149,12 +4095,10 @@ static bool ggml_backend_hexagon_device_supports_op(ggml_backend_dev_t dev, cons
                case GGML_UNARY_OP_SIGMOID:
                case GGML_UNARY_OP_SOFTPLUS:
                case GGML_UNARY_OP_TANH:
-                    supp = ggml_hexagon_supported_unary(sess, op);
-                    break;
                case GGML_UNARY_OP_SILU:
                case GGML_UNARY_OP_GELU:
                case GGML_UNARY_OP_GELU_QUICK:
-                    supp = ggml_hexagon_supported_activations(sess, op);
+                    supp = ggml_hexagon_supported_unary(sess, op);
                    break;
                default:
                    break;
@@ -4451,7 +4395,7 @@ static void ggml_hexagon_init(ggml_backend_reg * reg) {
    opt_opstage   = str_opstage  ? strtoul(str_opstage, NULL, 0)          : opt_opstage;
    opt_opbatch   = str_opbatch  ? strtoul(str_opbatch, NULL, 0)          : opt_opbatch;
    opt_opqueue   = str_opqueue  ? strtoul(str_opqueue, NULL, 0)          : opt_opqueue;
-    opt_optrace   = str_optrace  ? strtoul(str_optrace, NULL, 0)          : (opt_opbatch * 128);
+    opt_optrace   = str_optrace  ? strtoul(str_optrace, NULL, 0)          : (opt_opbatch * 256);
    opt_oppoll    = str_oppoll   ? strtoul(str_oppoll,  NULL, 0)          : opt_oppoll;
    opt_opfusion  = str_opfusion ? atoi(str_opfusion)                     : opt_opfusion;
    opt_profile   = str_profile  ? atoi(str_profile)                      : 0;
@@ -59,7 +59,11 @@ typedef AEEResult (*dspqueue_read_pfn_t)(dspqueue_t queue, uint32_t *flags,
                                         uint32_t max_message_length,
                                         uint32_t *message_length, uint8_t *message,
                                         uint32_t timeout_us);
-
+typedef AEEResult (*dspqueue_read_noblock_pfn_t)(dspqueue_t queue, uint32_t *flags,
+                                                 uint32_t max_buffers, uint32_t *num_buffers,
+                                                 struct dspqueue_buffer *buffers,
+                                                 uint32_t max_message_length,
+                                                 uint32_t *message_length, uint8_t *message);
 typedef int (*fastrpc_mmap_pfn_t)(int domain, int fd, void *addr, int offset, size_t length, enum fastrpc_map_flags flags);
 typedef int (*fastrpc_munmap_pfn_t)(int domain, int fd, void *addr, size_t length);

@@ -82,11 +86,12 @@ rpcmem_to_fd_pfn_t  rpcmem_to_fd_pfn  = nullptr;
 fastrpc_mmap_pfn_t   fastrpc_mmap_pfn   = nullptr;
 fastrpc_munmap_pfn_t fastrpc_munmap_pfn = nullptr;

-dspqueue_create_pfn_t dspqueue_create_pfn = nullptr;
-dspqueue_close_pfn_t  dspqueue_close_pfn  = nullptr;
-dspqueue_export_pfn_t dspqueue_export_pfn = nullptr;
-dspqueue_write_pfn_t  dspqueue_write_pfn  = nullptr;
-dspqueue_read_pfn_t   dspqueue_read_pfn   = nullptr;
+dspqueue_create_pfn_t       dspqueue_create_pfn       = nullptr;
+dspqueue_close_pfn_t        dspqueue_close_pfn        = nullptr;
+dspqueue_export_pfn_t       dspqueue_export_pfn       = nullptr;
+dspqueue_write_pfn_t        dspqueue_write_pfn        = nullptr;
+dspqueue_read_pfn_t         dspqueue_read_pfn         = nullptr;
+dspqueue_read_noblock_pfn_t dspqueue_read_noblock_pfn = nullptr;

 remote_handle64_open_pfn_t    remote_handle64_open_pfn    = nullptr;
 remote_handle64_invoke_pfn_t  remote_handle64_invoke_pfn  = nullptr;
@@ -167,6 +172,12 @@ AEEResult dspqueue_read(dspqueue_t               queue,
                        uint32_t *               message_length,
                        uint8_t *                message,
                        uint32_t                 timeout_us) {
+#ifdef _WIN32
+    if (timeout_us == 0) {
+        return dspqueue_read_noblock_pfn(queue, flags, max_buffers, num_buffers, buffers, max_message_length,
+                                     message_length, message);
+    }
+#endif
    return dspqueue_read_pfn(queue, flags, max_buffers, num_buffers, buffers, max_message_length, message_length,
                             message, timeout_us);
 }
@@ -349,6 +360,7 @@ int htpdrv_init() {
    dlsym(handle.get(), dspqueue_export_pfn_t, dspqueue_export_pfn, dspqueue_export, false);
    dlsym(handle.get(), dspqueue_write_pfn_t, dspqueue_write_pfn, dspqueue_write, false);
    dlsym(handle.get(), dspqueue_read_pfn_t, dspqueue_read_pfn, dspqueue_read, false);
+    dlsym(handle.get(), dspqueue_read_noblock_pfn_t, dspqueue_read_noblock_pfn, dspqueue_read_noblock, false);
    dlsym(handle.get(), remote_handle64_open_pfn_t, remote_handle64_open_pfn, remote_handle64_open, false);
    dlsym(handle.get(), remote_handle64_invoke_pfn_t, remote_handle64_invoke_pfn, remote_handle64_invoke, false);
    dlsym(handle.get(), remote_handle_control_pfn_t, remote_handle_control_pfn, remote_handle_control, false);
@@ -101,6 +101,4 @@ void dma_queue_alias_free(dma_queue_t q) {
    (void) q;
 }

-void dma_queue_flush(dma_queue_t q) {
-    while (dma_queue_pop(q).dst != NULL) ;
-}
+
@@ -106,7 +106,7 @@ struct dma_queue_s {
    bool                alias;     // When set, dma_queue_delete will not free the ring
 };

-void        dma_queue_flush(dma_queue_t q);
+

 size_t      dma_queue_sizeof(size_t capacity);
 size_t      dma_queue_alignof(void);
@@ -154,7 +154,6 @@ static inline bool dma_is_vtcm(const dma_queue * q, const void * ptr) {
 static inline bool dma_queue_push_single_1d(dma_queue * q, dma_ptr dptr, size_t size) {
    dma_ring * r = q->ring;
    if (((r->push_idx + 1) & r->idx_mask) == r->pop_idx) {
-        FARF(HIGH, "dma-push: queue full\n");
        return false;
    }

@@ -165,6 +164,8 @@ static inline bool dma_queue_push_single_1d(dma_queue * q, dma_ptr dptr, size_t

    r->dptr[r->push_idx] = dptr;

+    htp_trace_event_start(r->trace, HTP_TRACE_EVT_DMA, r->push_idx);
+
    if (size) {
        desc->next       = NULL;
        desc->desc_size  = 0; // 1D mode
@@ -173,7 +174,6 @@ static inline bool dma_queue_push_single_1d(dma_queue * q, dma_ptr dptr, size_t
        desc->order      = 0;
        desc->done       = 0;

-        htp_trace_event_start(r->trace, HTP_TRACE_EVT_DMA, r->push_idx);
        dmlink(r->tail, desc);
        r->tail = (dma_descriptor_2d *) desc;
    } else {
@@ -188,7 +188,6 @@ static inline bool dma_queue_push_single_1d(dma_queue * q, dma_ptr dptr, size_t
 static inline bool dma_queue_push_single_2d(dma_queue * q, dma_ptr dptr, size_t dst_stride, size_t src_stride, size_t row_size, size_t nrows) {
    dma_ring * r = q->ring;
    if (((r->push_idx + 1) & r->idx_mask) == r->pop_idx) {
-        FARF(HIGH, "dma-push: queue full\n");
        return false;
    }

@@ -224,8 +223,9 @@ static inline bool dma_queue_push_single_2d(dma_queue * q, dma_ptr dptr, size_t

    r->dptr[r->push_idx] = dptr;

+    htp_trace_event_start(r->trace, HTP_TRACE_EVT_DMA, r->push_idx);
+
    if (nrows) {
-        htp_trace_event_start(r->trace, HTP_TRACE_EVT_DMA, r->push_idx);
        dmlink(r->tail, desc);
        r->tail = desc;
    } else {
@@ -252,10 +252,11 @@ static inline dma_ptr dma_queue_pop(dma_queue * q) {
            dmpoll();
        }
    }
-    htp_trace_event_stop(r->trace, HTP_TRACE_EVT_DMA, r->pop_idx);

    dptr = r->dptr[r->pop_idx];

+    htp_trace_event_stop(r->trace, HTP_TRACE_EVT_DMA, r->pop_idx);
+
    r->pop_idx = (r->pop_idx + 1) & r->idx_mask;
    return dptr;
 }
@@ -270,6 +271,8 @@ static inline dma_ptr dma_queue_pop_nowait(dma_queue * q) {

    dptr = r->dptr[r->pop_idx];

+    htp_trace_event_stop(r->trace, HTP_TRACE_EVT_DMA, r->pop_idx);
+
    r->pop_idx = (r->pop_idx + 1) & r->idx_mask;
    return dptr;
 }
@@ -278,6 +281,10 @@ static inline bool dma_queue_empty(dma_queue * q) {
    return q->ring->push_idx == q->ring->pop_idx;
 }

+static inline void dma_queue_flush(dma_queue * q) {
+    while (dma_queue_pop(q).dst != NULL) ;
+}
+
 static inline uint32_t dma_queue_depth(dma_queue * q) {
    return (q->ring->push_idx - q->ring->pop_idx) & q->ring->idx_mask;
 }
@@ -314,14 +321,18 @@ static inline bool dma_queue_push(dma_queue *q, dma_ptr dptr, size_t dst_stride,
    {
        const uint8_t *src = (const uint8_t *) dptr.src;
        uint8_t       *dst = (uint8_t *)       dptr.dst;
-        for (size_t r = 0; r < nrows; ++r) {
+        size_t r = 0;
+        while (r + 1 < nrows) {
            dma_ptr p = dma_make_ptr(dst + r * dst_stride, src + r * src_stride);
-            if (!dma_queue_push_single_1d(q, p, row_size))
-                return false;
-            if (r + 1 < nrows)
-                dma_queue_pop(q);
+            if (!dma_queue_push_single_1d(q, p, row_size)) {
+                dma_queue_flush(q);
+            } else {
+                r++;
+            }
        }
-        return true;
+        dma_queue_flush(q);
+        dma_ptr p = dma_make_ptr(dst + r * dst_stride, src + r * src_stride);
+        return dma_queue_push_single_1d(q, p, row_size);
    }
 }

@@ -123,15 +123,17 @@ struct hmx_fa_context {
    uint32_t     g_br;  // hex_align_up(G * Br, 32) - actual tile row dim

    // VTCM buffers (allocated by vtcm_seq_alloc)
+    __fp16 *     vtcm_q_dma;           // Q DMA fetch buffer
    __fp16 *     vtcm_q_tiles;         // Q tile format [g_br, D]
    __fp16 *     vtcm_o_tiles[2];      // O ping-pong [g_br, D]
    __fp16 *     vtcm_k_fp16[2];       // K DMA double-buffer [Bc, D]
    __fp16 *     vtcm_v_fp16[2];       // V DMA double-buffer [Bc, D]
-    __fp16 *     vtcm_k_tiles;         // K tiles (transposed)
+    __fp16 *     vtcm_k_tiles[2];      // K tiles (transposed, double-buffered)
    __fp16 *     vtcm_v_tiles[2];      // V tiles (column-major, double-buffered)
-    __fp16 *     vtcm_s_tiles;         // S = QK^T [g_br, Bc]
-    __fp16 *     vtcm_p_tiles;         // P = softmax(S) [g_br, Bc]
+    __fp16 *     vtcm_s_tiles[2];      // S = QK^T [g_br, Bc] (double-buffered)
+    __fp16 *     vtcm_p_tiles[2];      // P = softmax(S) [g_br, Bc]
    __fp16 *     vtcm_d_tiles;         // Diagonal rescale [g_br, g_br]
+    __fp16 *     vtcm_d_inv_l;         // Diagonal rescale (1/l) [g_br, g_br]
    HVX_Vector * vtcm_m_vec;           // Row max [g_br]
    HVX_Vector * vtcm_l_vec;           // Row sum [g_br]
    HVX_Vector * vtcm_s_rowmax;        // Softmax intermediate [g_br]
@@ -236,10 +238,6 @@ static void flash_attn_ext_f16_thread(unsigned int nth, unsigned int ith, void *
        const uint32_t iv3 = fastdiv(iq3, &factx->broadcast_rv3);
        const uint32_t iv2 = fastdiv(iq2, &factx->broadcast_rv2);

-        // Fetch Q row
-        const uint8_t * q_row_ptr = (const uint8_t *) q->data + (iq1*nbq1 + iq2*nbq2 + iq3*nbq3);
-        dma_queue_push(dma, dma_make_ptr(spad_q, q_row_ptr), factx->size_q_row_padded, nbq1, size_q_row, 1);
-
        const __fp16 * mp_base = NULL;
        if (mask) {
            const uint32_t im2 = fastmodulo(iq2, mask->ne[2], &factx->src3_div2);
@@ -247,26 +245,91 @@ static void flash_attn_ext_f16_thread(unsigned int nth, unsigned int ith, void *
            mp_base = (const __fp16 *) ((const uint8_t *) mask->data + iq1*mask->nb[1] + im2*mask->nb[2] + im3*mask->nb[3]);
        }

-        // Prefetch first two blocks
-        for (uint32_t ib = 0; ib < MIN(factx->n_blocks, 2); ++ib) {
-            const uint32_t ic_start = ib * FLASH_ATTN_BLOCK_SIZE;
-            const uint32_t current_block_size = MIN(FLASH_ATTN_BLOCK_SIZE, nek1 - ic_start);
+        // Precalculate next row variables if there is a next row
+        bool has_next_ir = (ir + 1 < ir1);
+        uint32_t next_ik2 = 0, next_ik3 = 0, next_iv2 = 0, next_iv3 = 0;
+        const uint8_t * next_q_row_ptr = NULL;
+        const __fp16 * next_mp_base = NULL;

-            // K
-            const uint8_t * k_src = (const uint8_t *) k->data + (ic_start*nbk1 + ik2*nbk2 + ik3*nbk3);
-            uint8_t * k_dst = spad_k + (ib % 2) * factx->size_k_block;
-            dma_queue_push(dma, dma_make_ptr(k_dst, k_src), factx->size_k_row_padded, nbk1, size_k_row, current_block_size);
+        const uint8_t * next_k_src0 = NULL;
+        const uint8_t * next_v_src0 = NULL;
+        const uint8_t * next_m_src0 = NULL;
+        uint32_t next_block_size0 = 0;

-            // V
-            const uint8_t * v_src = (const uint8_t *) v->data + (ic_start*nbv1 + iv2*nbv2 + iv3*nbv3);
-            uint8_t * v_dst = spad_v + (ib % 2) * factx->size_v_block;
-            dma_queue_push(dma, dma_make_ptr(v_dst, v_src), factx->size_v_row_padded, nbv1, size_v_row, current_block_size);
+        const uint8_t * next_k_src1 = NULL;
+        const uint8_t * next_v_src1 = NULL;
+        const uint8_t * next_m_src1 = NULL;
+        uint32_t next_block_size1 = 0;
+
+        if (has_next_ir) {
+            const uint32_t next_ir = ir + 1;
+            const uint32_t next_iq3 = fastdiv(next_ir, &factx->src0_div21);
+            const uint32_t next_iq2 = fastdiv(next_ir - next_iq3*neq2*neq1, &factx->src0_div1);
+            const uint32_t next_iq1 = (next_ir - next_iq3*neq2*neq1 - next_iq2 * neq1);
+
+            next_ik3 = fastdiv(next_iq3, &factx->broadcast_rk3);
+            next_ik2 = fastdiv(next_iq2, &factx->broadcast_rk2);
+
+            next_iv3 = fastdiv(next_iq3, &factx->broadcast_rv3);
+            next_iv2 = fastdiv(next_iq2, &factx->broadcast_rv2);
+
+            next_q_row_ptr = (const uint8_t *) q->data + (next_iq1*nbq1 + next_iq2*nbq2 + next_iq3*nbq3);

-            // Mask
            if (mask) {
-                const uint8_t * m_src = (const uint8_t *) (mp_base + ic_start);
-                // Mask is 1D contiguous for this row
-                dma_cache_push(dma, &m_cache, m_src, current_block_size * 2, current_block_size * 2, current_block_size * 2, 1);
+                const uint32_t next_im2 = fastmodulo(next_iq2, mask->ne[2], &factx->src3_div2);
+                const uint32_t next_im3 = fastmodulo(next_iq3, mask->ne[3], &factx->src3_div3);
+                next_mp_base = (const __fp16 *) ((const uint8_t *) mask->data + next_iq1*mask->nb[1] + next_im2*mask->nb[2] + next_im3*mask->nb[3]);
+            }
+
+            // Precalculate next K/V block 0 source pointers
+            {
+                const uint32_t ic_start = 0;
+                next_block_size0 = MIN(FLASH_ATTN_BLOCK_SIZE, nek1 - ic_start);
+                next_k_src0 = (const uint8_t *) k->data + (ic_start*nbk1 + next_ik2*nbk2 + next_ik3*nbk3);
+                next_v_src0 = (const uint8_t *) v->data + (ic_start*nbv1 + next_iv2*nbv2 + next_iv3*nbv3);
+                if (mask) {
+                    next_m_src0 = (const uint8_t *) (next_mp_base + ic_start);
+                }
+            }
+
+            // Precalculate next K/V block 1 source pointers (if n_blocks > 1)
+            if (factx->n_blocks > 1) {
+                const uint32_t ic_start = 1 * FLASH_ATTN_BLOCK_SIZE;
+                next_block_size1 = MIN(FLASH_ATTN_BLOCK_SIZE, nek1 - ic_start);
+                next_k_src1 = (const uint8_t *) k->data + (ic_start*nbk1 + next_ik2*nbk2 + next_ik3*nbk3);
+                next_v_src1 = (const uint8_t *) v->data + (ic_start*nbv1 + next_iv2*nbv2 + next_iv3*nbv3);
+                if (mask) {
+                    next_m_src1 = (const uint8_t *) (next_mp_base + ic_start);
+                }
+            }
+        }
+
+        if (ir == ir0) {
+            // Fetch Q row
+            const uint8_t * q_row_ptr = (const uint8_t *) q->data + (iq1*nbq1 + iq2*nbq2 + iq3*nbq3);
+            dma_queue_push(dma, dma_make_ptr(spad_q, q_row_ptr), factx->size_q_row_padded, nbq1, size_q_row, 1);
+
+            // Prefetch first two blocks
+            for (uint32_t ib = 0; ib < MIN(factx->n_blocks, 2); ++ib) {
+                const uint32_t ic_start = ib * FLASH_ATTN_BLOCK_SIZE;
+                const uint32_t current_block_size = MIN(FLASH_ATTN_BLOCK_SIZE, nek1 - ic_start);
+
+                // K
+                const uint8_t * k_src = (const uint8_t *) k->data + (ic_start*nbk1 + ik2*nbk2 + ik3*nbk3);
+                uint8_t * k_dst = spad_k + (ib % 2) * factx->size_k_block;
+                dma_queue_push(dma, dma_make_ptr(k_dst, k_src), factx->size_k_row_padded, nbk1, size_k_row, current_block_size);
+
+                // V
+                const uint8_t * v_src = (const uint8_t *) v->data + (ic_start*nbv1 + iv2*nbv2 + iv3*nbv3);
+                uint8_t * v_dst = spad_v + (ib % 2) * factx->size_v_block;
+                dma_queue_push(dma, dma_make_ptr(v_dst, v_src), factx->size_v_row_padded, nbv1, size_v_row, current_block_size);
+
+                // Mask
+                if (mask) {
+                    const uint8_t * m_src = (const uint8_t *) (mp_base + ic_start);
+                    // Mask is 1D contiguous for this row
+                    dma_cache_push(dma, &m_cache, m_src, current_block_size * 2, current_block_size * 2, current_block_size * 2, 1);
+                }
            }
        }

@@ -287,6 +350,11 @@ static void flash_attn_ext_f16_thread(unsigned int nth, unsigned int ith, void *

        const HVX_Vector slope_vec = hvx_vec_splat_f16(slope);
        const HVX_Vector v_neg_inf = Q6_Vh_vsplat_R(0xfbff);
+        const HVX_Vector v_cap     = (factx->logit_softcap != 0.0f) ? hvx_vec_splat_f16(factx->logit_softcap) : Q6_V_vzero();
+        const HVX_Vector vinf      = Q6_Vh_vsplat_R(0xFC00);
+        const HVX_Vector vmin      = Q6_Vh_vsplat_R(0xFBFF);
+        const HVX_Vector v_log2e   = hvx_vec_splat_f16(EXP_LOG2E_F);
+        const uint32_t stride_v2   = factx->size_v_row_padded * 2;
        for (uint32_t ib = 0; ib < factx->n_blocks; ++ib) {
            const uint32_t ic_start = ib * FLASH_ATTN_BLOCK_SIZE;
            const uint32_t current_block_size = MIN(FLASH_ATTN_BLOCK_SIZE, nek1 - ic_start);
@@ -309,7 +377,6 @@ static void flash_attn_ext_f16_thread(unsigned int nth, unsigned int ith, void *

            // 2. Softcap (in FP16)
            if (factx->logit_softcap != 0.0f) {
-                const HVX_Vector v_cap = hvx_vec_splat_f16(factx->logit_softcap);
                scores_f16 = hvx_vec_tanh_f16(scores_f16);
                scores_f16 = hvx_vec_mul_f16_f16(scores_f16, v_cap);
            }
@@ -319,8 +386,6 @@ static void flash_attn_ext_f16_thread(unsigned int nth, unsigned int ith, void *
            // 3. Mask (in FP16)
            if (mask) {
                HVX_Vector m_vals_f16 = *(const HVX_UVector *) m_base;
-                HVX_Vector vinf = Q6_Vh_vsplat_R(0xFC00);
-                HVX_Vector vmin = Q6_Vh_vsplat_R(0xFBFF);
                HVX_VectorPred is_inf = Q6_Q_vcmp_eq_VhVh(m_vals_f16, vinf);
                m_vals_f16 = Q6_V_vmux_QVV(is_inf, vmin, m_vals_f16);

@@ -335,10 +400,30 @@ static void flash_attn_ext_f16_thread(unsigned int nth, unsigned int ith, void *
            HVX_Vector v_max     = Q6_V_lo_W(hvx_vec_f16_to_f32(v_max_f16)); // splat block max in FP32
            htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_FA_QK, ir);

+            if (ib + 1 == factx->n_blocks && has_next_ir) {
+                // Queue next row's Q row!
+                dma_queue_push(dma, dma_make_ptr(spad_q, next_q_row_ptr), factx->size_q_row_padded, nbq1, size_q_row, 1);
+
+                if (factx->n_blocks % 2 == 0) {
+                    // Queue next row's block 0 (into buffer slot 0)
+                    uint8_t * k_dst = spad_k + 0 * factx->size_k_block;
+                    uint8_t * v_dst = spad_v + 0 * factx->size_v_block;
+
+                    // K (block 0 of next row)
+                    dma_queue_push(dma, dma_make_ptr(k_dst, next_k_src0), factx->size_k_row_padded, nbk1, size_k_row, next_block_size0);
+
+                    // V (block 0 of next row)
+                    dma_queue_push(dma, dma_make_ptr(v_dst, next_v_src0), factx->size_v_row_padded, nbv1, size_v_row, next_block_size0);
+
+                    // Mask (block 0 of next row)
+                    if (mask) {
+                        dma_cache_push(dma, &m_cache, next_m_src0, next_block_size0 * 2, next_block_size0 * 2, next_block_size0 * 2, 1);
+                    }
+                }
+            }
+
            htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_FA_SFM, ir);
            {
-                const HVX_Vector v_log2e = hvx_vec_splat_f16(EXP_LOG2E_F);
-
                // 4. Online Softmax Update
                HVX_Vector M_new_vec = Q6_Vsf_vmax_VsfVsf(v_max, M_vec);
                HVX_Vector diff_vec  = HVX_OP_SUB_F32(M_vec, M_new_vec);
@@ -370,24 +455,20 @@ static void flash_attn_ext_f16_thread(unsigned int nth, unsigned int ith, void *
                S_vec = HVX_OP_ADD_F32(HVX_OP_MUL_F32(S_vec, ms_vec), p_sum_vec);

                // 5. Accumulate V (F16 * F16 -> F32 accumulator)
-                __fp16 __attribute__((aligned(128))) p_arr[VLEN_FP16];
-                hvx_vec_store_a(p_arr, 128, P);
+                const uint8_t * v_ptr = v_base;

                for (uint32_t j = 0; j < current_block_size; j += 2) {
                    if (j + 1 == current_block_size) {
-                        if (p_arr[j] != 0.0f) {
-                            const uint8_t * v_ptr = v_base + j * factx->size_v_row_padded;
-                            hvx_mad_f32_f16_aa(VKQ32, v_ptr, (p_arr + j), DV);
-                        }
+                        HVX_Vector S0 = hvx_vec_repl_f16(Q6_V_vror_VR(P, j * 2));
+                        hvx_mad_f32_f16_aa_vec(VKQ32, v_ptr, S0, DV);
                        break;
                    }

-                    if (p_arr[j] == 0.0f && p_arr[j + 1] == 0.0f) {
-                        continue;
-                    }
+                    HVX_Vector S0 = hvx_vec_repl_f16(Q6_V_vror_VR(P, j * 2));
+                    HVX_Vector S1 = hvx_vec_repl_f16(Q6_V_vror_VR(P, (j + 1) * 2));

-                    const uint8_t * v_ptr = v_base + j * factx->size_v_row_padded;
-                    hvx_mad_f32_f16_aa_rx2(VKQ32, v_ptr, v_ptr + factx->size_v_row_padded, (p_arr + j), (p_arr + j + 1), DV);
+                    hvx_mad_f32_f16_aa_rx2_vec(VKQ32, v_ptr, v_ptr + factx->size_v_row_padded, S0, S1, DV);
+                    v_ptr += stride_v2;
                }
            }
            htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_FA_SFM, ir);
@@ -414,6 +495,61 @@ static void flash_attn_ext_f16_thread(unsigned int nth, unsigned int ith, void *
            }
        }

+        if (has_next_ir) {
+            if (factx->n_blocks % 2 == 0) {
+                // Queue next row's block 1 (into buffer slot 1, if n_blocks > 1)
+                if (factx->n_blocks > 1) {
+                    uint8_t * k_dst = spad_k + 1 * factx->size_k_block;
+                    uint8_t * v_dst = spad_v + 1 * factx->size_v_block;
+
+                    // K (block 1 of next row)
+                    dma_queue_push(dma, dma_make_ptr(k_dst, next_k_src1), factx->size_k_row_padded, nbk1, size_k_row, next_block_size1);
+
+                    // V (block 1 of next row)
+                    dma_queue_push(dma, dma_make_ptr(v_dst, next_v_src1), factx->size_v_row_padded, nbv1, size_v_row, next_block_size1);
+
+                    // Mask (block 1 of next row)
+                    if (mask) {
+                        dma_cache_push(dma, &m_cache, next_m_src1, next_block_size1 * 2, next_block_size1 * 2, next_block_size1 * 2, 1);
+                    }
+                }
+            } else {
+                // Queue next row's block 0 (into buffer slot 0)
+                {
+                    uint8_t * k_dst = spad_k + 0 * factx->size_k_block;
+                    uint8_t * v_dst = spad_v + 0 * factx->size_v_block;
+
+                    // K (block 0 of next row)
+                    dma_queue_push(dma, dma_make_ptr(k_dst, next_k_src0), factx->size_k_row_padded, nbk1, size_k_row, next_block_size0);
+
+                    // V (block 0 of next row)
+                    dma_queue_push(dma, dma_make_ptr(v_dst, next_v_src0), factx->size_v_row_padded, nbv1, size_v_row, next_block_size0);
+
+                    // Mask (block 0 of next row)
+                    if (mask) {
+                        dma_cache_push(dma, &m_cache, next_m_src0, next_block_size0 * 2, next_block_size0 * 2, next_block_size0 * 2, 1);
+                    }
+                }
+
+                // Queue next row's block 1 (into buffer slot 1, if n_blocks > 1)
+                if (factx->n_blocks > 1) {
+                    uint8_t * k_dst = spad_k + 1 * factx->size_k_block;
+                    uint8_t * v_dst = spad_v + 1 * factx->size_v_block;
+
+                    // K (block 1 of next row)
+                    dma_queue_push(dma, dma_make_ptr(k_dst, next_k_src1), factx->size_k_row_padded, nbk1, size_k_row, next_block_size1);
+
+                    // V (block 1 of next row)
+                    dma_queue_push(dma, dma_make_ptr(v_dst, next_v_src1), factx->size_v_row_padded, nbv1, size_v_row, next_block_size1);
+
+                    // Mask (block 1 of next row)
+                    if (mask) {
+                        dma_cache_push(dma, &m_cache, next_m_src1, next_block_size1 * 2, next_block_size1 * 2, next_block_size1 * 2, 1);
+                    }
+                }
+            }
+        }
+
        htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_O_PROC, ir);
        // sinks
        float M = hvx_vec_get_f32(M_vec);
@@ -471,6 +607,7 @@ typedef struct {
    void *                  curr_k;
    uint32_t                kv_start;
    uint32_t                rows_per_t;
+    size_t                  buf_idx;
 } fa_k_int_args_t;

 static void fa_k_interleave_thread(unsigned int n, unsigned int i, void * data) {
@@ -488,19 +625,19 @@ static void fa_k_interleave_thread(unsigned int n, unsigned int i, void * data)

    struct htp_thread_trace * tr = &factx->octx->ctx->trace[i];
    htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_FA_K_PREP, (uint16_t) (args->kv_start + start));
-    hmx_interleave_rows_to_tiles(factx->vtcm_k_tiles, (const __fp16 *) args->curr_k, total_rows, factx->DK,
+    hmx_interleave_rows_to_tiles(factx->vtcm_k_tiles[args->buf_idx], (const __fp16 *) args->curr_k, total_rows, factx->DK,
                             args->src_stride, start, end);
    htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_FA_K_PREP, (uint16_t) (args->kv_start + start));
 }

-static void fa_phase_k_interleave(struct hmx_fa_context * factx, uint32_t kv_rows, size_t src_stride, void * curr_k, uint32_t kv_start) {
+static void fa_phase_k_interleave(struct hmx_fa_context * factx, uint32_t kv_rows, size_t src_stride, void * curr_k, uint32_t kv_start, size_t buf_idx) {
    work_queue_t wp = factx->octx->ctx->work_queue;
    uint32_t n = 1;
    if (factx->n_threads > 1 && kv_rows >= factx->n_threads * 2) {
        n = factx->n_threads;
    }
    uint32_t rows_per_t = hex_align_up(hmx_ceil_div(kv_rows, n), 2);
-    fa_k_int_args_t args = { factx, kv_rows, src_stride, curr_k, kv_start, rows_per_t };
+    fa_k_int_args_t args = { factx, kv_rows, src_stride, curr_k, kv_start, rows_per_t, buf_idx };
    if (n > 1) {
        work_queue_run(wp, fa_k_interleave_thread, &args, n);
    } else {
@@ -645,12 +782,13 @@ static void fa_q_load_thread(unsigned int n, unsigned int i, void * data) {
            }
        }

-        // Initialize vtcm_d_tiles to 0
+        // Initialize vtcm_d_tiles and vtcm_d_inv_l to 0
        const size_t d_bytes_per_t = hex_align_up(d_tile_bytes / n, 128);
        const size_t d_start       = i * d_bytes_per_t;
        const size_t d_end         = hex_smin(d_start + d_bytes_per_t, d_tile_bytes);
        if (d_start < d_tile_bytes) {
            hvx_splat_u8_a((char *) factx->vtcm_d_tiles + d_start, 0, d_end - d_start);
+            hvx_splat_u8_a((char *) factx->vtcm_d_inv_l + d_start, 0, d_end - d_start);
        }
    }

@@ -662,15 +800,14 @@ static void fa_q_load_thread(unsigned int n, unsigned int i, void * data) {

        assert(factx->DK == factx->DV);

-        const size_t o_tile_bytes = factx->o_tile_bytes;
-        const bool use_q_dma = (2 * o_tile_bytes >= factx->g_br * DK * (factx->is_q_fp32 ? 4 : 2));
+        const bool use_q_dma = (factx->vtcm_q_dma != NULL);

        __fp16 * q_tiles = factx->vtcm_q_tiles;
        if (use_q_dma) {
            const size_t g_rows_end = hex_smin(end, n_rows_g);
            const uint32_t d_limit = factx->is_q_fp32 ? DK / 32 : DK / 64;

-            uint8_t * q_flat  = (uint8_t *) factx->vtcm_o_tiles[0];
+            uint8_t * q_flat  = (uint8_t *) factx->vtcm_q_dma;
            if (factx->is_q_fp32) {
                switch (d_limit) {
                case 2:  hmx_fa_q_prep_fp32_d2(q_tiles, q_flat, start, end, g_rows_end, DK, G, args->n_rows_q, &factx->div_G, args->q_transposed); break;
@@ -781,10 +918,10 @@ static void fa_o_store_thread_f32(unsigned int n, unsigned int i, void * data) {
    const uint32_t            kv_head    = args->kv_head;
    const uint32_t            ib3        = args->ib3;

-    for (size_t r = start; r < end; ++r) {
-        const size_t q_idx = fastdiv(r, &factx->div_G);
-        const size_t h_idx = fastmodulo(r, G, &factx->div_G);
+    size_t q_idx = fastdiv(start, &factx->div_G);
+    size_t h_idx = fastmodulo(start, G, &factx->div_G);

+    for (size_t r = start; r < end; ++r) {
        float * out = (float *) ((uint8_t *) dst->data + (kv_head * G + h_idx) * dst->nb[1] +
                                 (q_start + q_idx) * dst->nb[2] + ib3 * dst->nb[3]);

@@ -801,6 +938,12 @@ static void fa_o_store_thread_f32(unsigned int n, unsigned int i, void * data) {
                *(HVX_UVector *) (out + d * 32) = Q6_V_hi_W(vp);
            }
        }
+
+        h_idx++;
+        if (h_idx == G) {
+            h_idx = 0;
+            q_idx++;
+        }
    }
    htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_O_PROC, (uint16_t) (args->q_start * G + start));
 }
@@ -829,10 +972,10 @@ static void fa_o_store_thread_f16(unsigned int n, unsigned int i, void * data) {
    const uint32_t            kv_head    = args->kv_head;
    const uint32_t            ib3        = args->ib3;

-    for (size_t r = start; r < end; ++r) {
-        const size_t q_idx = fastdiv(r, &factx->div_G);
-        const size_t h_idx = fastmodulo(r, G, &factx->div_G);
+    size_t q_idx = fastdiv(start, &factx->div_G);
+    size_t h_idx = fastmodulo(start, G, &factx->div_G);

+    for (size_t r = start; r < end; ++r) {
        __fp16 * out = (__fp16 *) ((uint8_t *) dst->data + (kv_head * G + h_idx) * dst->nb[1] +
                                   (q_start + q_idx) * dst->nb[2] + ib3 * dst->nb[3]);

@@ -851,6 +994,12 @@ static void fa_o_store_thread_f16(unsigned int n, unsigned int i, void * data) {
                *(HVX_UVector *) (out + d * 64) = Q6_V_hi_W(vp);
            }
        }
+
+        h_idx++;
+        if (h_idx == G) {
+            h_idx = 0;
+            q_idx++;
+        }
    }
    htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_O_PROC, (uint16_t) (args->q_start * G + start));
 }
@@ -879,6 +1028,7 @@ static void fa_phase_o_store(struct hmx_fa_context *   factx,

 typedef struct {
    struct hmx_fa_context *   factx;
+    size_t                    buf_idx;
    size_t                    kv_rows;
    size_t                    n_rows_g;
    size_t                    n_col_tiles;
@@ -960,8 +1110,8 @@ static inline void fa_softmax_impl(
            uint32_t r0 = r / HMX_FP16_TILE_N_ROWS;
            uint32_t r1 = r % HMX_FP16_TILE_N_ROWS;

-            const __fp16 * s_ld_base = factx->vtcm_s_tiles + r0 * HMX_FP16_TILE_N_ROWS * Bc;
-            __fp16 *       p_st_base = factx->vtcm_p_tiles + r0 * HMX_FP16_TILE_N_ROWS * Bc;
+            const __fp16 * s_ld_base = factx->vtcm_s_tiles[args->buf_idx] + r0 * HMX_FP16_TILE_N_ROWS * Bc;
+            __fp16 *       p_st_base = factx->vtcm_p_tiles[args->buf_idx] + r0 * HMX_FP16_TILE_N_ROWS * Bc;

            // Decode 2 rows from S tiles into per-thread row buffers
            if (has_softcap) {
@@ -983,7 +1133,26 @@ static inline void fa_softmax_impl(
                    my_row_buf1[ci] = hvx_vec_mul_f16_f16(t1, v_cap);
                }
            } else {
-                for (size_t c = 0; c < kv_rows; c += 64) {
+                size_t c = 0;
+                for (; c + 64 < kv_rows; c += 128) {
+                    size_t             ci0       = c / 64;
+                    size_t             ci1       = ci0 + 1;
+                    const __fp16 *     in_dtile0 = s_ld_base + ci0 * HMX_FP16_TILE_N_ELMS * 2;
+                    const __fp16 *     in_dtile1 = s_ld_base + ci1 * HMX_FP16_TILE_N_ELMS * 2;
+                    const HVX_Vector * pv_s_in0_0 = ((const HVX_Vector *) in_dtile0) + r1 / 2;
+                    const HVX_Vector * pv_s_in1_0 = pv_s_in0_0 + 16;
+                    const HVX_Vector * pv_s_in0_1 = ((const HVX_Vector *) in_dtile1) + r1 / 2;
+                    const HVX_Vector * pv_s_in1_1 = pv_s_in0_1 + 16;
+
+                    HVX_VectorPair vp_s_drow0 = Q6_W_vdeal_VVR(*pv_s_in1_0, *pv_s_in0_0, -2);
+                    my_row_buf0[ci0]          = Q6_V_lo_W(vp_s_drow0);
+                    my_row_buf1[ci0]          = Q6_V_hi_W(vp_s_drow0);
+
+                    HVX_VectorPair vp_s_drow1 = Q6_W_vdeal_VVR(*pv_s_in1_1, *pv_s_in0_1, -2);
+                    my_row_buf0[ci1]          = Q6_V_lo_W(vp_s_drow1);
+                    my_row_buf1[ci1]          = Q6_V_hi_W(vp_s_drow1);
+                }
+                for (; c < kv_rows; c += 64) {
                    size_t             ci       = c / 64;
                    const __fp16 *     in_dtile = s_ld_base + ci * HMX_FP16_TILE_N_ELMS * 2;
                    const HVX_Vector * pv_s_in0 = ((const HVX_Vector *) in_dtile) + r1 / 2;
@@ -1007,12 +1176,12 @@ static inline void fa_softmax_impl(

            HVX_Vector v_s_rowmax0 = v_neg_inf;
            HVX_Vector v_s_rowmax1 = v_neg_inf;
-            for (size_t c = 0; c < kv_rows; c += 64) {
-                size_t         ci          = c / 64;
-                const size_t   ne          = hex_smin(kv_rows - c, 64);
-                HVX_VectorPred q_tail_keep = Q6_Q_vsetq2_R(ne * sizeof(__fp16));
+            if (has_mask) {
+                for (size_t c = 0; c < kv_rows; c += 64) {
+                    size_t         ci          = c / 64;
+                    const size_t   ne          = hex_smin(kv_rows - c, 64);
+                    HVX_VectorPred q_tail_keep = Q6_Q_vsetq2_R(ne * sizeof(__fp16));

-                if (has_mask) {
                    HVX_Vector v_mask0, v_mask1;

                    if (mask_broadcast) {
@@ -1066,15 +1235,31 @@ static inline void fa_softmax_impl(
                        my_row_buf0[ci] = Q6_V_vmux_QVV(q_keep0, hvx_vec_add_f16_f16(my_row_buf0[ci], v_mask0_scaled), v_neg_inf);
                        my_row_buf1[ci] = Q6_V_vmux_QVV(q_keep1, hvx_vec_add_f16_f16(my_row_buf1[ci], v_mask1_scaled), v_neg_inf);
                    }
-                } else {
+
+                    v_s_rowmax0 = Q6_Vhf_vmax_VhfVhf(v_s_rowmax0, my_row_buf0[ci]);
+                    v_s_rowmax1 = Q6_Vhf_vmax_VhfVhf(v_s_rowmax1, my_row_buf1[ci]);
+                }
+            } else {
+                size_t c = 0;
+                for (; c + 64 < kv_rows; c += 128) {
+                    size_t ci0 = c / 64;
+                    size_t ci1 = ci0 + 1;
+                    v_s_rowmax0 = Q6_Vhf_vmax_VhfVhf(v_s_rowmax0, my_row_buf0[ci0]);
+                    v_s_rowmax1 = Q6_Vhf_vmax_VhfVhf(v_s_rowmax1, my_row_buf1[ci0]);
+                    v_s_rowmax0 = Q6_Vhf_vmax_VhfVhf(v_s_rowmax0, my_row_buf0[ci1]);
+                    v_s_rowmax1 = Q6_Vhf_vmax_VhfVhf(v_s_rowmax1, my_row_buf1[ci1]);
+                }
+                for (; c < kv_rows; c += 64) {
+                    size_t         ci          = c / 64;
+                    const size_t   ne          = hex_smin(kv_rows - c, 64);
+                    HVX_VectorPred q_tail_keep = Q6_Q_vsetq2_R(ne * sizeof(__fp16));
                    if (ne < 64) {
                        my_row_buf0[ci] = Q6_V_vmux_QVV(q_tail_keep, my_row_buf0[ci], v_neg_inf);
                        my_row_buf1[ci] = Q6_V_vmux_QVV(q_tail_keep, my_row_buf1[ci], v_neg_inf);
                    }
+                    v_s_rowmax0 = Q6_Vhf_vmax_VhfVhf(v_s_rowmax0, my_row_buf0[ci]);
+                    v_s_rowmax1 = Q6_Vhf_vmax_VhfVhf(v_s_rowmax1, my_row_buf1[ci]);
                }
-
-                v_s_rowmax0 = Q6_Vhf_vmax_VhfVhf(v_s_rowmax0, my_row_buf0[ci]);
-                v_s_rowmax1 = Q6_Vhf_vmax_VhfVhf(v_s_rowmax1, my_row_buf1[ci]);
            }

            v_s_rowmax0 = hvx_vec_reduce_max_f16(v_s_rowmax0);
@@ -1121,8 +1306,48 @@ static inline void fa_softmax_impl(
            HVX_Vector       v_p_rowsum0 = v_zero;
            HVX_Vector       v_p_rowsum1 = v_zero;

-            for (size_t c = 0; c < kv_rows; c += 64) {
-                size_t     ci           = c / 64;
+            size_t c = 0;
+            for (; c + 64 < kv_rows; c += 128) {
+                size_t     ci0          = c / 64;
+                size_t     ci1          = ci0 + 1;
+
+                HVX_Vector v_s_minus_m0_0 = Q6_Vqf16_vsub_VhfVhf(my_row_buf0[ci0], v_dup_m0);
+                HVX_Vector v_s_minus_m1_0 = Q6_Vqf16_vsub_VhfVhf(my_row_buf1[ci0], v_dup_m1);
+                HVX_Vector v_s_minus_m0_1 = Q6_Vqf16_vsub_VhfVhf(my_row_buf0[ci1], v_dup_m0);
+                HVX_Vector v_s_minus_m1_1 = Q6_Vqf16_vsub_VhfVhf(my_row_buf1[ci1], v_dup_m1);
+
+                HVX_Vector v_p_row0_hf_0  = hvx_vec_exp2_f16(Q6_Vhf_equals_Vqf16(v_s_minus_m0_0));
+                HVX_Vector v_p_row1_hf_0  = hvx_vec_exp2_f16(Q6_Vhf_equals_Vqf16(v_s_minus_m1_0));
+                HVX_Vector v_p_row0_hf_1  = hvx_vec_exp2_f16(Q6_Vhf_equals_Vqf16(v_s_minus_m0_1));
+                HVX_Vector v_p_row1_hf_1  = hvx_vec_exp2_f16(Q6_Vhf_equals_Vqf16(v_s_minus_m1_1));
+
+                __fp16 *     out_dtile0  = p_st_base + ci0 * HMX_FP16_TILE_N_ELMS * 2;
+                __fp16 *     out_dtile1  = p_st_base + ci1 * HMX_FP16_TILE_N_ELMS * 2;
+                HVX_Vector * pv_p_out0_0  = ((HVX_Vector *) out_dtile0) + r1 / 2;
+                HVX_Vector * pv_p_out1_0  = pv_p_out0_0 + 16;
+                HVX_Vector * pv_p_out0_1  = ((HVX_Vector *) out_dtile1) + r1 / 2;
+                HVX_Vector * pv_p_out1_1  = pv_p_out0_1 + 16;
+
+                HVX_VectorPair vp_p_dual0 = Q6_W_vshuff_VVR(v_p_row1_hf_0, v_p_row0_hf_0, -2);
+                *pv_p_out0_0               = Q6_V_lo_W(vp_p_dual0);
+                *pv_p_out1_0               = Q6_V_hi_W(vp_p_dual0);
+
+                HVX_VectorPair vp_p_dual1 = Q6_W_vshuff_VVR(v_p_row1_hf_1, v_p_row0_hf_1, -2);
+                *pv_p_out0_1               = Q6_V_lo_W(vp_p_dual1);
+                *pv_p_out1_1               = Q6_V_hi_W(vp_p_dual1);
+
+                HVX_VectorPair vp_p0_0 = hvx_vec_f16_to_f32_shuff(v_p_row0_hf_0);
+                HVX_VectorPair vp_p1_0 = hvx_vec_f16_to_f32_shuff(v_p_row1_hf_0);
+                HVX_VectorPair vp_p0_1 = hvx_vec_f16_to_f32_shuff(v_p_row0_hf_1);
+                HVX_VectorPair vp_p1_1 = hvx_vec_f16_to_f32_shuff(v_p_row1_hf_1);
+
+                v_p_rowsum0 = Q6_Vqf32_vadd_Vqf32Vqf32(v_p_rowsum0, Q6_Vqf32_vadd_VsfVsf(Q6_V_lo_W(vp_p0_0), Q6_V_hi_W(vp_p0_0)));
+                v_p_rowsum0 = Q6_Vqf32_vadd_Vqf32Vqf32(v_p_rowsum0, Q6_Vqf32_vadd_VsfVsf(Q6_V_lo_W(vp_p0_1), Q6_V_hi_W(vp_p0_1)));
+                v_p_rowsum1 = Q6_Vqf32_vadd_Vqf32Vqf32(v_p_rowsum1, Q6_Vqf32_vadd_VsfVsf(Q6_V_lo_W(vp_p1_0), Q6_V_hi_W(vp_p1_0)));
+                v_p_rowsum1 = Q6_Vqf32_vadd_Vqf32Vqf32(v_p_rowsum1, Q6_Vqf32_vadd_VsfVsf(Q6_V_lo_W(vp_p1_1), Q6_V_hi_W(vp_p1_1)));
+            }
+            for (size_t c_rem = c; c_rem < kv_rows; c_rem += 64) {
+                size_t     ci           = c_rem / 64;
                HVX_Vector v_s_minus_m0 = Q6_Vqf16_vsub_VhfVhf(my_row_buf0[ci], v_dup_m0);
                HVX_Vector v_s_minus_m1 = Q6_Vqf16_vsub_VhfVhf(my_row_buf1[ci], v_dup_m1);

@@ -1281,7 +1506,7 @@ static __attribute__((noinline)) void fa_build_d_diag_inv_l(struct hmx_fa_contex
            v_content = Q6_V_vror_VR(v_content, 64);
        }

-        __fp16 * out_base = factx->vtcm_d_tiles + i * (n_row_tiles_g_br + 1) * HMX_FP16_TILE_N_ELMS;
+        __fp16 * out_base = factx->vtcm_d_inv_l + i * (n_row_tiles_g_br + 1) * HMX_FP16_TILE_N_ELMS;
        Q6_vscatter_QRMVhV(q_32_mask, (size_t) out_base, HMX_FP16_TILE_SIZE - 1, v_offsets, v_content);
    }
 }
@@ -1514,6 +1739,27 @@ static void fa_pop_mask_dma_gqa(dma_queue * dma, uint32_t G) {
    }
 }

+static inline void fa_prefetch_block(dma_queue * dma, const struct htp_tensor * k, const struct htp_tensor * v, const struct htp_tensor * mask,
+                                     uint32_t b, size_t Bc, size_t size_k_row_padded, size_t size_k_row, size_t size_v_row_padded, size_t size_v_row,
+                                     uint32_t ik2, uint32_t ik3, uint32_t iv2, uint32_t iv3, uint32_t q_start, uint32_t im3, uint32_t kv_head, uint32_t G,
+                                     size_t m_line_bytes, size_t n_rows_q, size_t nek1, size_t prefetch_buf, struct hmx_fa_context * factx) {
+    const uint32_t  prefetch_start = b * Bc;
+    const uint32_t  prefetch_rows  = hex_smin(Bc, nek1 - prefetch_start);
+    const uint8_t * k_prefetch_src = (const uint8_t *) k->data + prefetch_start * k->nb[1] + ik2 * k->nb[2] + ik3 * k->nb[3];
+    dma_queue_push(dma, dma_make_ptr(factx->vtcm_k_fp16[prefetch_buf], k_prefetch_src), size_k_row_padded, k->nb[1], size_k_row, prefetch_rows);
+    const uint8_t * v_prefetch_src = (const uint8_t *) v->data + prefetch_start * v->nb[1] + iv2 * v->nb[2] + iv3 * v->nb[3];
+    dma_queue_push(dma, dma_make_ptr(factx->vtcm_v_fp16[prefetch_buf], v_prefetch_src), size_v_row_padded, v->nb[1], size_v_row, prefetch_rows);
+
+    if (mask) {
+        if (__builtin_expect(factx->mask_broadcast, true)) {
+            const uint8_t * ms_src = (const uint8_t *) mask->data + q_start * mask->nb[1] + im3 * mask->nb[3] + prefetch_start * sizeof(__fp16);
+            dma_cache_push(dma, &factx->m_cache, ms_src, m_line_bytes, mask->nb[1], prefetch_rows * sizeof(__fp16), n_rows_q);
+        } else {
+            fa_push_mask_dma_gqa(dma, mask, q_start, im3, prefetch_start, kv_head, G, m_line_bytes, prefetch_rows, n_rows_q, factx);
+        }
+    }
+}
+
 // ============================================================================
 // Core HMX flash attention algorithm (GQA-merged)
 // ============================================================================
@@ -1612,7 +1858,7 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
    // Build the VTCM layout once (shared with the host estimator) and place every
    // scratch buffer at its computed offset.
    struct hmx_fa_vtcm_layout L;
-    hmx_fa_vtcm_layout_build(&L, G, DK, DV, Br, Bc, n_threads, pipeline);
+    hmx_fa_vtcm_layout_build(&L, G, DK, DV, Br, Bc, n_threads, pipeline, factx.is_q_fp32);

    if (L.total_bytes > ctx->vtcm_size) {
        return HTP_STATUS_VTCM_TOO_SMALL;
@@ -1620,6 +1866,7 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {

    uint8_t * const base = ctx->vtcm_base;

+    factx.vtcm_q_dma          = VTCM_LAYOUT_PTR(__fp16, base, L.off_q_dma);
    factx.vtcm_q_tiles        = VTCM_LAYOUT_PTR(__fp16, base, L.off_q_tiles);
    factx.vtcm_o_tiles[0]     = VTCM_LAYOUT_PTR(__fp16, base, L.off_o_tiles[0]);
    factx.vtcm_o_tiles[1]     = VTCM_LAYOUT_PTR(__fp16, base, L.off_o_tiles[1]);
@@ -1627,12 +1874,16 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
    factx.vtcm_k_fp16[1]      = VTCM_LAYOUT_PTR(__fp16, base, L.off_k_fp16[1]);
    factx.vtcm_v_fp16[0]      = VTCM_LAYOUT_PTR(__fp16, base, L.off_v_fp16[0]);
    factx.vtcm_v_fp16[1]      = VTCM_LAYOUT_PTR(__fp16, base, L.off_v_fp16[1]);
-    factx.vtcm_k_tiles        = VTCM_LAYOUT_PTR(__fp16, base, L.off_k_tiles);
+    factx.vtcm_k_tiles[0]     = VTCM_LAYOUT_PTR(__fp16, base, L.off_k_tiles[0]);
+    factx.vtcm_k_tiles[1]     = VTCM_LAYOUT_PTR_OPTIONAL(__fp16, base, L.off_k_tiles[1], pipeline);
    factx.vtcm_v_tiles[0]     = VTCM_LAYOUT_PTR(__fp16, base, L.off_v_tiles[0]);
    factx.vtcm_v_tiles[1]     = VTCM_LAYOUT_PTR_OPTIONAL(__fp16, base, L.off_v_tiles[1], pipeline);
-    factx.vtcm_s_tiles        = VTCM_LAYOUT_PTR(__fp16, base, L.off_s_tiles);
-    factx.vtcm_p_tiles        = VTCM_LAYOUT_PTR(__fp16, base, L.off_p_tiles);
+    factx.vtcm_s_tiles[0]     = VTCM_LAYOUT_PTR(__fp16, base, L.off_s_tiles[0]);
+    factx.vtcm_s_tiles[1]     = VTCM_LAYOUT_PTR_OPTIONAL(__fp16, base, L.off_s_tiles[1], pipeline);
+    factx.vtcm_p_tiles[0]     = VTCM_LAYOUT_PTR(__fp16, base, L.off_p_tiles[0]);
+    factx.vtcm_p_tiles[1]     = VTCM_LAYOUT_PTR_OPTIONAL(__fp16, base, L.off_p_tiles[1], pipeline);
    factx.vtcm_d_tiles        = VTCM_LAYOUT_PTR(__fp16, base, L.off_d_tiles);
+    factx.vtcm_d_inv_l        = VTCM_LAYOUT_PTR(__fp16, base, L.off_d_inv_l);
    factx.vtcm_m_vec          = VTCM_LAYOUT_PTR(HVX_Vector, base, L.off_m_vec);
    factx.vtcm_l_vec          = VTCM_LAYOUT_PTR(HVX_Vector, base, L.off_l_vec);
    factx.vtcm_s_rowmax       = VTCM_LAYOUT_PTR(HVX_Vector, base, L.off_s_rowmax);
@@ -1670,6 +1921,12 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {

    const size_t qo_element_size = factx.is_q_fp32 ? sizeof(float) : sizeof(__fp16);

+    const bool q_transposed                 = q->nb[1] < q->nb[2];
+    const size_t q_src_stride               = q_transposed ? q->nb[2] : q->nb[1];
+    const size_t q_row_bytes_untransposed   = factx.G * factx.DK * qo_element_size;
+    const size_t q_row_bytes_trans_factor   = factx.DK * qo_element_size;
+    const uint32_t kv_rows0                 = hex_smin(Bc, nek1);
+
    // ======== Reusable job descriptors for pipeline ========
    hmx_fa_qk_job_t       qk_job;
    hmx_fa_o_update_job_t ou_job;
@@ -1690,34 +1947,34 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
                const uint32_t iv2 = kv_head;
                const uint32_t iv3 = fastdiv(ib3, &kparams->broadcast_rv3);

-                // 1. Push Q DMA (if Q DMA is used)
-                const size_t o_tile_bytes = factx.o_tile_bytes;
-                const bool use_q_dma = (2 * o_tile_bytes >= factx.g_br * factx.DK * (factx.is_q_fp32 ? 4 : 2));
-                if (use_q_dma) {
-                    const bool q_transposed = q->nb[1] < q->nb[2];
-                    const uint8_t * q_ptr = (const uint8_t *) q->data + q_start * q->nb[1] + (kv_head * factx.G) * q->nb[2] + ib3 * q->nb[3];
-                    const size_t el_size = factx.is_q_fp32 ? sizeof(float) : sizeof(__fp16);
-                    const size_t q_row_bytes = q_transposed ? n_rows_q * factx.DK * el_size : factx.G * factx.DK * el_size;
-                    const size_t src_stride  = q_transposed ? q->nb[2] : q->nb[1];
+                // 1. Push Q and KV DMAs for the very first iteration.
+                // Subsequent iterations are enqueued early at the end of the previous iteration.
+                if (ib3 == 0 && q_start == 0 && kv_head == 0) {
+                    const uint8_t * q_ptr = (const uint8_t *) q->data;
+                    const size_t q_row_bytes = q_transposed ? n_rows_q * q_row_bytes_trans_factor : q_row_bytes_untransposed;
                    const size_t n_rows      = q_transposed ? factx.G : n_rows_q;
-                    dma_queue_push(dma, dma_make_ptr(factx.vtcm_o_tiles[0], q_ptr), q_row_bytes, hex_smax(src_stride, q_row_bytes), q_row_bytes, n_rows);
+                    dma_queue_push(dma, dma_make_ptr(factx.vtcm_q_dma, q_ptr), q_row_bytes, hex_smax(q_src_stride, q_row_bytes), q_row_bytes, n_rows);
+
+                    if (factx.n_kv_blocks > 0) {
+                        const uint8_t * k_src = (const uint8_t *) k->data + ik2 * k->nb[2] + ik3 * k->nb[3];
+                        dma_queue_push(dma, dma_make_ptr(factx.vtcm_k_fp16[0], k_src), size_k_row_padded, k->nb[1], size_k_row, kv_rows0);
+
+                        const uint8_t * v_src = (const uint8_t *) v->data + iv2 * v->nb[2] + iv3 * v->nb[3];
+                        dma_queue_push(dma, dma_make_ptr(factx.vtcm_v_fp16[0], v_src), size_v_row_padded, v->nb[1], size_v_row, kv_rows0);
+
+                        if (factx.pipeline && mask) {
+                            if (__builtin_expect(factx.mask_broadcast, true)) {
+                                const uint8_t * ms_src = (const uint8_t *) mask->data + q_start * mask->nb[1] + im3 * mask->nb[3] + 0;
+                                dma_cache_push(dma, &factx.m_cache, ms_src, m_line_bytes, mask->nb[1], kv_rows0 * sizeof(__fp16), n_rows_q);
+                            } else {
+                                fa_push_mask_dma_gqa(dma, mask, q_start, im3, 0, kv_head, G, m_line_bytes, kv_rows0, n_rows_q, &factx);
+                            }
+                        }
+                    }
                }

-                // 2. Prefetch first KV block
-                if (factx.n_kv_blocks > 0) {
-                    const uint32_t kv_rows0 = hex_smin(Bc, nek1);
-
-                    const uint8_t * k_src = (const uint8_t *) k->data + ik2 * k->nb[2] + ik3 * k->nb[3];
-                    dma_queue_push(dma, dma_make_ptr(factx.vtcm_k_fp16[0], k_src), size_k_row_padded, k->nb[1], size_k_row, kv_rows0);
-
-                    const uint8_t * v_src = (const uint8_t *) v->data + iv2 * v->nb[2] + iv3 * v->nb[3];
-                    dma_queue_push(dma, dma_make_ptr(factx.vtcm_v_fp16[0], v_src), size_v_row_padded, v->nb[1], size_v_row, kv_rows0);
-                }
-
-                // 3. Pop Q DMA (blocks until Q is loaded)
-                if (use_q_dma) {
-                    dma_queue_pop(dma);
-                }
+                // 2. Pop Q DMA (blocks until Q is loaded)
+                dma_queue_pop(dma);

                // ---- Load Q block & Initialize per-block state ----
                fa_phase_q_load(&factx, q, q_start, kv_head, ib3, n_rows_g);
@@ -1738,76 +1995,40 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
                hmx_queue_t hmx_q = ctx->hmx_queue;

                if (factx.pipeline) {
-                    // Pipeline path
+                    // Double-buffered job structs because HMX queue runs asynchronously
+                    hmx_fa_qk_job_t qk_job[2];
+                    hmx_fa_o_update_job_t ou_job[2];
+
+                    // Prefetch block 1 early if there are multiple blocks
+                    if (factx.n_kv_blocks > 1) {
+                        fa_prefetch_block(dma, k, v, mask, 1, Bc, size_k_row_padded, size_k_row, size_v_row_padded, size_v_row,
+                                          ik2, ik3, iv2, iv3, q_start, im3, kv_head, G, m_line_bytes, n_rows_q, nek1, 1, &factx);
+                    }
+
+                    // Prep and start QK-dot(0)
+                    void * curr_k0 = dma_queue_pop(dma).dst;
+                    fa_phase_k_interleave(&factx, kv_rows0, k_src_stride, curr_k0, 0, 0);
+
+                    qk_job[0].q_tiles        = factx.vtcm_q_tiles;
+                    qk_job[0].k_tiles        = factx.vtcm_k_tiles[0];
+                    qk_job[0].s_tiles        = factx.vtcm_s_tiles[0];
+                    qk_job[0].n_row_tiles    = n_row_tiles;
+                    qk_job[0].n_col_tiles    = hmx_ceil_div(kv_rows0, HMX_FP16_TILE_N_COLS);
+                    qk_job[0].n_dot_tiles    = DK / 32;
+                    qk_job[0].n_tiles_per_bc = n_tiles_per_bc;
+                    qk_job[0].hmx_scales     = factx.vtcm_hmx_scales_qk;
+                    hmx_queue_push(hmx_q, hmx_queue_make_desc(hmx_fa_qk_dot_worker, &qk_job[0]));
+
                    for (uint32_t kv_blk = 0; kv_blk < factx.n_kv_blocks; ++kv_blk) {
                        const uint32_t kv_start    = kv_blk * Bc;
                        const uint32_t kv_rows     = hex_smin(Bc, nek1 - kv_start);
                        const size_t   n_col_tiles = hmx_ceil_div(kv_rows, HMX_FP16_TILE_N_COLS);

-                        // Push mask DMA
-                        if (mask) {
-                            if (__builtin_expect(factx.mask_broadcast, true)) {
-                                const uint8_t * ms_src = (const uint8_t *) mask->data + q_start * mask->nb[1] + im3 * mask->nb[3] + kv_start * sizeof(__fp16);
-                                dma_cache_push(dma, &factx.m_cache, ms_src, m_line_bytes, mask->nb[1], kv_rows * sizeof(__fp16), n_rows_q);
-                            } else {
-                                fa_push_mask_dma_gqa(dma, mask, q_start, im3, kv_start, kv_head, G, m_line_bytes, kv_rows, n_rows_q, &factx);
-                            }
-                        }
-
-                        // Prefetch next KV block early
-                        if (kv_blk + 1 < factx.n_kv_blocks) {
-                            const uint32_t  prefetch_start = (kv_blk + 1) * Bc;
-                            const uint32_t  prefetch_rows  = hex_smin(Bc, nek1 - prefetch_start);
-                            const size_t    prefetch_buf   = 1 - buf_idx;
-                            const uint8_t * k_prefetch_src = (const uint8_t *) k->data + prefetch_start * k->nb[1] + ik2 * k->nb[2] + ik3 * k->nb[3];
-                            dma_queue_push(dma, dma_make_ptr(factx.vtcm_k_fp16[prefetch_buf], k_prefetch_src), size_k_row_padded, k->nb[1], size_k_row, prefetch_rows);
-                            const uint8_t * v_prefetch_src = (const uint8_t *) v->data + prefetch_start * v->nb[1] + iv2 * v->nb[2] + iv3 * v->nb[3];
-                            dma_queue_push(dma, dma_make_ptr(factx.vtcm_v_fp16[prefetch_buf], v_prefetch_src), size_v_row_padded, v->nb[1], size_v_row, prefetch_rows);
-                        }
-
-                        // ---- Phase 1: K_int ----
-                        if (kv_blk > 0) {
-                            ou_job.o_curr           = o_tile_curr;
-                            ou_job.o_prev           = o_tile_prev;
-                            ou_job.p_tiles          = factx.vtcm_p_tiles;
-                            ou_job.v_tiles          = factx.vtcm_v_tiles[1 - buf_idx];
-                            ou_job.d_tiles          = factx.vtcm_d_tiles;
-                            ou_job.hmx_scales       = factx.vtcm_hmx_scales_id;
-                            ou_job.n_row_tiles      = n_row_tiles;
-                            ou_job.n_col_tiles      = hmx_ceil_div(hex_smin(Bc, nek1 - (kv_blk - 1) * Bc), HMX_FP16_TILE_N_COLS);
-                            ou_job.n_row_tiles_g_br = n_row_tiles_g_br;
-                            ou_job.n_tiles_per_bc   = n_tiles_per_bc;
-                            ou_job.DV               = DV;
-                            hmx_queue_push(hmx_q, hmx_queue_make_desc(hmx_fa_o_update_worker, &ou_job));
-                        }
-
-                        // Wait for current K DMA and interleave
-                        void * curr_k = dma_queue_pop(dma).dst;
-                        fa_phase_k_interleave(&factx, kv_rows, k_src_stride, curr_k, kv_start);
-
-                        // ---- Phase 2: qk_dot ----
-                        qk_job.q_tiles        = factx.vtcm_q_tiles;
-                        qk_job.k_tiles        = factx.vtcm_k_tiles;
-                        qk_job.s_tiles        = factx.vtcm_s_tiles;
-                        qk_job.n_row_tiles    = n_row_tiles;
-                        qk_job.n_col_tiles    = n_col_tiles;
-                        qk_job.n_dot_tiles    = DK / 32;
-                        qk_job.n_tiles_per_bc = n_tiles_per_bc;
-                        qk_job.hmx_scales     = factx.vtcm_hmx_scales_qk;
-                        hmx_queue_push(hmx_q, hmx_queue_make_desc(hmx_fa_qk_dot_worker, &qk_job));
-
-                        // Wait for current V DMA and interleave
+                        // ---- 1. Pop and run V-prep for current block ----
                        void * curr_v = dma_queue_pop(dma).dst;
                        fa_phase_v_interleave(&factx, kv_rows, v_src_stride, curr_v, factx.vtcm_v_tiles[buf_idx], n_tiles_per_bc, kv_start);

-                        if (kv_blk > 0) {
-                            hmx_queue_pop(hmx_q);
-                            hex_swap_ptr((void **) &o_tile_curr, (void **) &o_tile_prev);
-                        }
-
-                        hmx_queue_pop(hmx_q);
-
-                        // ---- Phase 3: softmax + build_D ----
+                        // ---- 2. Pop and run mask-prep for current block ----
                        __fp16 * current_mask_vtcm = NULL;
                        if (mask) {
                            if (__builtin_expect(factx.mask_broadcast, true)) {
@@ -1818,9 +2039,34 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
                            }
                        }

+                        // ---- 3. Pop and run K-prep for next block & push next QK-dot ----
+                        if (kv_blk + 1 < factx.n_kv_blocks) {
+                            const uint32_t next_start = (kv_blk + 1) * Bc;
+                            const uint32_t next_rows  = hex_smin(Bc, nek1 - next_start);
+                            const size_t   next_buf   = 1 - buf_idx;
+
+                            void * next_k = dma_queue_pop(dma).dst;
+                            fa_phase_k_interleave(&factx, next_rows, k_src_stride, next_k, next_start, next_buf);
+
+                            qk_job[next_buf].q_tiles        = factx.vtcm_q_tiles;
+                            qk_job[next_buf].k_tiles        = factx.vtcm_k_tiles[next_buf];
+                            qk_job[next_buf].s_tiles        = factx.vtcm_s_tiles[next_buf];
+                            qk_job[next_buf].n_row_tiles    = n_row_tiles;
+                            qk_job[next_buf].n_col_tiles    = hmx_ceil_div(next_rows, HMX_FP16_TILE_N_COLS);
+                            qk_job[next_buf].n_dot_tiles    = DK / 32;
+                            qk_job[next_buf].n_tiles_per_bc = n_tiles_per_bc;
+                            qk_job[next_buf].hmx_scales     = factx.vtcm_hmx_scales_qk;
+                            hmx_queue_push(hmx_q, hmx_queue_make_desc(hmx_fa_qk_dot_worker, &qk_job[next_buf]));
+                        }
+
+                        // ---- 4. Wait for current block's QK-dot to finish ----
+                        hmx_queue_pop(hmx_q);
+
+                        // ---- 5. Phase 2: softmax + build_D ----
                        fa_softmax_args_t sargs;
                        memset(&sargs, 0, sizeof(sargs));
                        sargs.factx                = &factx;
+                        sargs.buf_idx              = buf_idx;
                        sargs.kv_rows              = kv_rows;
                        sargs.n_rows_g             = n_rows_g;
                        sargs.n_col_tiles          = n_col_tiles;
@@ -1838,8 +2084,39 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
                        sargs.mask_vtcm            = current_mask_vtcm;
                        sargs.mask_vtcm_row_stride = factx.mask_buf_row_stride;
                        sargs.slopes               = factx.vtcm_slopes;
+
+                        // Start HMX O update for block kv_blk - 1 (reads P[1 - buf_idx], V[1 - buf_idx])
+                        if (kv_blk > 0) {
+                            const size_t prev_buf = 1 - buf_idx;
+                            ou_job[prev_buf].o_curr           = o_tile_curr;
+                            ou_job[prev_buf].o_prev           = o_tile_prev;
+                            ou_job[prev_buf].p_tiles          = factx.vtcm_p_tiles[prev_buf];
+                            ou_job[prev_buf].v_tiles          = factx.vtcm_v_tiles[prev_buf];
+                            ou_job[prev_buf].d_tiles          = factx.vtcm_d_tiles;
+                            ou_job[prev_buf].hmx_scales       = factx.vtcm_hmx_scales_id;
+                            ou_job[prev_buf].n_row_tiles      = n_row_tiles;
+                            ou_job[prev_buf].n_col_tiles      = hmx_ceil_div(hex_smin(Bc, nek1 - (kv_blk - 1) * Bc), HMX_FP16_TILE_N_COLS);
+                            ou_job[prev_buf].n_row_tiles_g_br = n_row_tiles_g_br;
+                            ou_job[prev_buf].n_tiles_per_bc   = n_tiles_per_bc;
+                            ou_job[prev_buf].DV               = DV;
+                            hmx_queue_push(hmx_q, hmx_queue_make_desc(hmx_fa_o_update_worker, &ou_job[prev_buf]));
+                        }
+
+                        // Run Softmax on HVX (blocking call)
                        fa_phase_softmax_and_build_d(&factx, &sargs, n_row_tiles, n_row_tiles_g_br);

+                        // Wait for HMX O update for block kv_blk - 1 to finish
+                        if (kv_blk > 0) {
+                            hmx_queue_pop(hmx_q);
+                            hex_swap_ptr((void **) &o_tile_curr, (void **) &o_tile_prev);
+                        }
+
+                        // Prefetch block kv_blk + 2
+                        if (kv_blk + 2 < factx.n_kv_blocks) {
+                            fa_prefetch_block(dma, k, v, mask, kv_blk + 2, Bc, size_k_row_padded, size_k_row, size_v_row_padded, size_v_row,
+                                              ik2, ik3, iv2, iv3, q_start, im3, kv_head, G, m_line_bytes, n_rows_q, nek1, buf_idx, &factx);
+                        }
+
                        buf_idx = 1 - buf_idx;
                    }

@@ -1847,18 +2124,23 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
                    if (factx.n_kv_blocks > 0) {
                        const uint32_t last_blk = factx.n_kv_blocks - 1;
                        const size_t last_cols  = hmx_ceil_div(hex_smin(Bc, nek1 - last_blk * Bc), HMX_FP16_TILE_N_COLS);
-                        ou_job.o_curr           = o_tile_curr;
-                        ou_job.o_prev           = o_tile_prev;
-                        ou_job.p_tiles          = factx.vtcm_p_tiles;
-                        ou_job.v_tiles          = factx.vtcm_v_tiles[1 - buf_idx];
-                        ou_job.d_tiles          = factx.vtcm_d_tiles;
-                        ou_job.hmx_scales       = factx.vtcm_hmx_scales_id;
-                        ou_job.n_row_tiles      = n_row_tiles;
-                        ou_job.n_col_tiles      = last_cols;
-                        ou_job.n_row_tiles_g_br = n_row_tiles_g_br;
-                        ou_job.n_tiles_per_bc   = n_tiles_per_bc;
-                        ou_job.DV               = DV;
-                        hmx_queue_push(hmx_q, hmx_queue_make_desc(hmx_fa_o_update_worker, &ou_job));
+                        ou_job[0].o_curr           = o_tile_curr;
+                        ou_job[0].o_prev           = o_tile_prev;
+                        ou_job[0].p_tiles          = factx.vtcm_p_tiles[1 - buf_idx];
+                        ou_job[0].v_tiles          = factx.vtcm_v_tiles[1 - buf_idx];
+                        ou_job[0].d_tiles          = factx.vtcm_d_tiles;
+                        ou_job[0].hmx_scales       = factx.vtcm_hmx_scales_id;
+                        ou_job[0].n_row_tiles      = n_row_tiles;
+                        ou_job[0].n_col_tiles      = last_cols;
+                        ou_job[0].n_row_tiles_g_br = n_row_tiles_g_br;
+                        ou_job[0].n_tiles_per_bc   = n_tiles_per_bc;
+                        ou_job[0].DV               = DV;
+                        hmx_queue_push(hmx_q, hmx_queue_make_desc(hmx_fa_o_update_worker, &ou_job[0]));
+
+                        // Overlapped: run HVX build diag inv L while HMX is busy executing the update
+                        htp_trace_event_start(tr_hvx, HTP_TRACE_EVT_HVX_O_PROC, (uint16_t) q_start);
+                        fa_build_d_diag_inv_l(&factx, n_row_tiles, n_row_tiles_g_br);
+                        htp_trace_event_stop(tr_hvx, HTP_TRACE_EVT_HVX_O_PROC, (uint16_t) q_start);
                        hmx_queue_pop(hmx_q);

                        hex_swap_ptr((void **) &o_tile_curr, (void **) &o_tile_prev);
@@ -1892,12 +2174,12 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {

                        // Wait for current K DMA and interleave
                        void * curr_k = dma_queue_pop(dma).dst;
-                        fa_phase_k_interleave(&factx, kv_rows, k_src_stride, curr_k, kv_start);
+                        fa_phase_k_interleave(&factx, kv_rows, k_src_stride, curr_k, kv_start, 0);

                        {
                            qk_job.q_tiles        = factx.vtcm_q_tiles;
-                            qk_job.k_tiles        = factx.vtcm_k_tiles;
-                            qk_job.s_tiles        = factx.vtcm_s_tiles;
+                            qk_job.k_tiles        = factx.vtcm_k_tiles[0];
+                            qk_job.s_tiles        = factx.vtcm_s_tiles[0];
                            qk_job.n_row_tiles    = n_row_tiles;
                            qk_job.n_col_tiles    = n_col_tiles;
                            qk_job.n_dot_tiles    = (size_t) (DK / 32);
@@ -1948,7 +2230,7 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
                        {
                            ou_job.o_curr           = o_tile_curr;
                            ou_job.o_prev           = o_tile_prev;
-                            ou_job.p_tiles          = factx.vtcm_p_tiles;
+                            ou_job.p_tiles          = factx.vtcm_p_tiles[0];
                            ou_job.v_tiles          = factx.vtcm_v_tiles[0];
                            ou_job.d_tiles          = factx.vtcm_d_tiles;
                            ou_job.hmx_scales       = factx.vtcm_hmx_scales_id;
@@ -1959,6 +2241,12 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
                            ou_job.DV               = DV;

                            hmx_queue_push(ctx->hmx_queue, hmx_queue_make_desc(hmx_fa_o_update_worker, &ou_job));
+                            if (kv_blk + 1 == factx.n_kv_blocks) {
+                                // Overlapped: run HVX build diag inv L while HMX is busy executing the update
+                                htp_trace_event_start(tr_hvx, HTP_TRACE_EVT_HVX_O_PROC, (uint16_t) q_start);
+                                fa_build_d_diag_inv_l(&factx, n_row_tiles, n_row_tiles_g_br);
+                                htp_trace_event_stop(tr_hvx, HTP_TRACE_EVT_HVX_O_PROC, (uint16_t) q_start);
+                            }
                            hmx_queue_pop(ctx->hmx_queue);

                            hex_swap_ptr((void **) &o_tile_curr, (void **) &o_tile_prev);
@@ -1968,15 +2256,63 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
                    }
                }

+                // Enqueue DMAs for the next iteration early so they overlap with O-PROC
+                uint32_t next_kv_head = kv_head + 1;
+                uint32_t next_q_start = q_start;
+                uint32_t next_ib3     = ib3;
+                if (next_kv_head >= n_kv_heads) {
+                    next_kv_head = 0;
+                    next_q_start = q_start + Br;
+                    if (next_q_start >= neq1) {
+                        next_q_start = 0;
+                        next_ib3     = ib3 + 1;
+                    }
+                }
+                bool has_next = (next_ib3 < neq3);
+
+                if (has_next) {
+                    const uint32_t next_n_rows_q = hex_smin(Br, neq1 - next_q_start);
+                    const uint8_t * next_q_ptr = (const uint8_t *) q->data + next_q_start * q->nb[1] + (next_kv_head * factx.G) * q->nb[2] + next_ib3 * q->nb[3];
+                    const size_t next_q_row_bytes = q_transposed ? next_n_rows_q * q_row_bytes_trans_factor : q_row_bytes_untransposed;
+                    const size_t next_n_rows      = q_transposed ? factx.G : next_n_rows_q;
+                    dma_queue_push(dma, dma_make_ptr(factx.vtcm_q_dma, next_q_ptr), next_q_row_bytes, hex_smax(q_src_stride, next_q_row_bytes), next_q_row_bytes, next_n_rows);
+
+                    if (factx.n_kv_blocks > 0) {
+                        const uint32_t next_ik2 = next_kv_head;
+                        const uint32_t next_iv2 = next_kv_head;
+                        uint32_t next_ik3 = ik3;
+                        uint32_t next_iv3 = iv3;
+                        if (next_ib3 != ib3) {
+                            next_ik3 = fastdiv(next_ib3, &kparams->broadcast_rk3);
+                            next_iv3 = fastdiv(next_ib3, &kparams->broadcast_rv3);
+                        }
+
+                        const uint8_t * next_k_src = (const uint8_t *) k->data + next_ik2 * k->nb[2] + next_ik3 * k->nb[3];
+                        dma_queue_push(dma, dma_make_ptr(factx.vtcm_k_fp16[0], next_k_src), size_k_row_padded, k->nb[1], size_k_row, kv_rows0);
+
+                        const uint8_t * next_v_src = (const uint8_t *) v->data + next_iv2 * v->nb[2] + next_iv3 * v->nb[3];
+                        dma_queue_push(dma, dma_make_ptr(factx.vtcm_v_fp16[0], next_v_src), size_v_row_padded, v->nb[1], size_v_row, kv_rows0);
+
+                        if (factx.pipeline && mask) {
+                            uint32_t next_im3 = im3;
+                            if (next_ib3 != ib3) {
+                                next_im3 = fastmodulo(next_ib3, mask->ne[3], &factx.src3_div3);
+                            }
+                            if (__builtin_expect(factx.mask_broadcast, true)) {
+                                const uint8_t * ms_src = (const uint8_t *) mask->data + next_q_start * mask->nb[1] + next_im3 * mask->nb[3] + 0;
+                                dma_cache_push(dma, &factx.m_cache, ms_src, m_line_bytes, mask->nb[1], kv_rows0 * sizeof(__fp16), next_n_rows_q);
+                            } else {
+                                fa_push_mask_dma_gqa(dma, mask, next_q_start, next_im3, 0, next_kv_head, G, m_line_bytes, kv_rows0, next_n_rows_q, &factx);
+                            }
+                        }
+                    }
+                }
+
                // ---- Final normalization ----
                {
-                    htp_trace_event_start(tr_hvx, HTP_TRACE_EVT_HVX_O_PROC, (uint16_t) q_start);
-                    fa_build_d_diag_inv_l(&factx, n_row_tiles, n_row_tiles_g_br);
-                    htp_trace_event_stop(tr_hvx, HTP_TRACE_EVT_HVX_O_PROC, (uint16_t) q_start);
-
                    on_job.o_curr           = o_tile_curr;
                    on_job.o_prev           = o_tile_prev;
-                    on_job.d_tiles          = factx.vtcm_d_tiles;
+                    on_job.d_tiles          = factx.vtcm_d_inv_l;
                    on_job.hmx_scales       = factx.vtcm_hmx_scales_id;
                    on_job.n_row_tiles      = n_row_tiles;
                    on_job.n_row_tiles_g_br = n_row_tiles_g_br;
@@ -101,14 +101,16 @@ static_assert(sizeof(struct htp_fa_kernel_params) <= 128, "htp_fa_kernel_params
 struct hmx_fa_vtcm_layout {
    // Byte offsets from vtcm_base for each region.
    size_t off_q_tiles;
+    size_t off_q_dma;
    size_t off_o_tiles[2];
    size_t off_k_fp16[2];
    size_t off_v_fp16[2];
-    size_t off_k_tiles;
-    size_t off_v_tiles[2];     // [1] allocated only when pipeline, else 0
-    size_t off_s_tiles;
-    size_t off_p_tiles;
+    size_t off_k_tiles[2];
+    size_t off_v_tiles[2];
+    size_t off_s_tiles[2];
+    size_t off_p_tiles[2];
    size_t off_d_tiles;
+    size_t off_d_inv_l;
    size_t off_m_vec;
    size_t off_l_vec;
    size_t off_s_rowmax;
@@ -140,7 +142,7 @@ struct hmx_fa_vtcm_layout {

 static inline void hmx_fa_vtcm_layout_build(struct hmx_fa_vtcm_layout * L,
                                       size_t gqa_factor, size_t DK, size_t DV,
-                                       size_t Br, size_t Bc, size_t n_threads, bool pipeline) {
+                                       size_t Br, size_t Bc, size_t n_threads, bool pipeline, bool is_q_fp32) {
    const size_t g_br         = hex_align_up(gqa_factor * Br, HMX_FP16_TILE_N_ROWS);
    const size_t q_tile_size  = hex_align_up(g_br * DK   * sizeof(__fp16), HTP_FA_HMX_TILE_SIZE);
    const size_t o_tile_size  = hex_align_up(g_br * DV   * sizeof(__fp16), HTP_FA_HMX_TILE_SIZE);
@@ -149,6 +151,7 @@ static inline void hmx_fa_vtcm_layout_build(struct hmx_fa_vtcm_layout * L,
    const size_t s_tile_size  = hex_align_up(g_br * Bc   * sizeof(__fp16), HTP_FA_HMX_TILE_SIZE);
    const size_t d_tile_size  = hex_align_up(g_br * g_br * sizeof(__fp16), HTP_FA_HMX_TILE_SIZE);

+    const size_t q_dma_size   = hex_align_up(g_br * DK * (is_q_fp32 ? sizeof(float) : sizeof(__fp16)), 128);
    const size_t k_dma_size   = hex_align_up(Bc * hex_round_up(DK * sizeof(__fp16), 128), 128);
    const size_t v_dma_size   = hex_align_up(Bc * hex_round_up(DV * sizeof(__fp16), 128), 128);
    const size_t col_vec_size = hex_align_up(g_br * sizeof(float),  256);
@@ -160,27 +163,47 @@ static inline void hmx_fa_vtcm_layout_build(struct hmx_fa_vtcm_layout * L,

    size_t off = 0;

-    // Section 1: HMX Tiled Buffers (FA_HMX_TILE_SIZE = 2KB Aligned)
+    // Group A (Part 1 - HMX Tiled buffers)
    VTCM_LAYOUT_ALLOC(off, off_q_tiles,       q_tile_size);
    VTCM_LAYOUT_ALLOC(off, off_o_tiles[0],    o_tile_size);
    VTCM_LAYOUT_ALLOC(off, off_o_tiles[1],    o_tile_size);
-    VTCM_LAYOUT_ALLOC(off, off_k_tiles,       k_tile_size);
-    VTCM_LAYOUT_ALLOC(off, off_v_tiles[0],    v_tile_size);
-    VTCM_LAYOUT_ALLOC_OPTIONAL(off, off_v_tiles[1], v_tile_size, pipeline);
-    VTCM_LAYOUT_ALLOC(off, off_s_tiles,       s_tile_size);
-    VTCM_LAYOUT_ALLOC(off, off_p_tiles,       s_tile_size);
    VTCM_LAYOUT_ALLOC(off, off_d_tiles,       d_tile_size);
+    VTCM_LAYOUT_ALLOC(off, off_d_inv_l,       d_tile_size);

-    // Section 2: HVX/DMA flat and vector buffers (128B / 256B Aligned)
+    // Group B & C share start offset (Group B tiles must be 2KB aligned)
+    size_t off_group_b_c = hex_align_up(off, HTP_FA_HMX_TILE_SIZE);
+
+    // Group B: Compute-only buffers
+    size_t off_group_b = off_group_b_c;
+    VTCM_LAYOUT_ALLOC(off_group_b, off_k_tiles[0],    k_tile_size);
+    VTCM_LAYOUT_ALLOC_OPTIONAL(off_group_b, off_k_tiles[1], k_tile_size, pipeline);
+    VTCM_LAYOUT_ALLOC(off_group_b, off_v_tiles[0],    v_tile_size);
+    VTCM_LAYOUT_ALLOC_OPTIONAL(off_group_b, off_v_tiles[1], v_tile_size, pipeline);
+    VTCM_LAYOUT_ALLOC(off_group_b, off_s_tiles[0],    s_tile_size);
+    VTCM_LAYOUT_ALLOC_OPTIONAL(off_group_b, off_s_tiles[1], s_tile_size, pipeline);
+    VTCM_LAYOUT_ALLOC(off_group_b, off_p_tiles[0],    s_tile_size);
+    VTCM_LAYOUT_ALLOC_OPTIONAL(off_group_b, off_p_tiles[1], s_tile_size, pipeline);
+    VTCM_LAYOUT_ALLOC(off_group_b, off_s_rowmax,      col_vec_size);
+    VTCM_LAYOUT_ALLOC(off_group_b, off_p_rowsum,      col_vec_size);
+    VTCM_LAYOUT_ALLOC(off_group_b, off_row_bufs,      row_vec_size * 2 * n_threads);
+
+    const size_t group_b_size = off_group_b - off_group_b_c;
+
+    // Group C: Q fetch DMA buffer
+    size_t off_group_c = off_group_b_c;
+    VTCM_LAYOUT_ALLOC(off_group_c, off_q_dma,         q_dma_size);
+
+    const size_t group_c_size = off_group_c - off_group_b_c;
+
+    off = off_group_b_c + hex_smax(group_b_size, group_c_size);
+
+    // Group A (Part 2 - remaining non-HMX buffers)
    VTCM_LAYOUT_ALLOC(off, off_k_fp16[0],     k_dma_size);
    VTCM_LAYOUT_ALLOC(off, off_k_fp16[1],     k_dma_size);
    VTCM_LAYOUT_ALLOC(off, off_v_fp16[0],     v_dma_size);
    VTCM_LAYOUT_ALLOC(off, off_v_fp16[1],     v_dma_size);
    VTCM_LAYOUT_ALLOC(off, off_m_vec,         col_vec_size);
    VTCM_LAYOUT_ALLOC(off, off_l_vec,         col_vec_size);
-    VTCM_LAYOUT_ALLOC(off, off_s_rowmax,      col_vec_size);
-    VTCM_LAYOUT_ALLOC(off, off_p_rowsum,      col_vec_size);
-    VTCM_LAYOUT_ALLOC(off, off_row_bufs,      row_vec_size * 2 * n_threads);
    VTCM_LAYOUT_ALLOC(off, off_hmx_scales_id, 256);
    VTCM_LAYOUT_ALLOC(off, off_hmx_scales_qk, 256);
    VTCM_LAYOUT_ALLOC(off, off_mask_buf,      m_buf_size);
@@ -200,9 +223,9 @@ static inline void hmx_fa_vtcm_layout_build(struct hmx_fa_vtcm_layout * L,
 }

 // Exact VTCM usage for a given (gqa_factor, DK, DV, Br, Bc) configuration.
-static inline size_t hmx_fa_compute_vtcm_usage(size_t gqa_factor, size_t DK, size_t DV, size_t Br, size_t Bc, size_t n_threads, bool pipeline) {
+static inline size_t hmx_fa_compute_vtcm_usage(size_t gqa_factor, size_t DK, size_t DV, size_t Br, size_t Bc, size_t n_threads, bool pipeline, bool is_q_fp32) {
    struct hmx_fa_vtcm_layout L;
-    hmx_fa_vtcm_layout_build(&L, gqa_factor, DK, DV, Br, Bc, n_threads, pipeline);
+    hmx_fa_vtcm_layout_build(&L, gqa_factor, DK, DV, Br, Bc, n_threads, pipeline, is_q_fp32);
    return L.total_bytes;
 }

@@ -239,7 +262,8 @@ static inline int hmx_fa_find_chunk_size(size_t * Br_out,
                                  size_t   qo_len,
                                  size_t   kv_len,
                                  size_t   vtcm_budget,
-                                  size_t   n_threads) {
+                                  size_t   n_threads,
+                                  bool     is_q_fp32) {
    const size_t T       = HMX_FP16_TILE_N_ROWS;  // 32
    const size_t br_unit = hmx_ceil_div(T, gqa_factor);
    const size_t bc_unit = HMX_FP16_TILE_N_COLS * 2;  // 64
@@ -253,8 +277,9 @@ static inline int hmx_fa_find_chunk_size(size_t * Br_out,
    const size_t Bc_limit     = can_pipeline ? hex_align_down(kv_len / FA_MIN_KV_BLOCKS, bc_unit) :
                                               (kv_len >= bc_unit ? hex_align_down(kv_len, bc_unit) : bc_unit);
    // Cost coefficients calibrated from profiling
-    const size_t c_q_fixed    = 1400;  // per-Q-block: q_load + epilogue o_update + o_norm + o_store
-    const size_t c_iter_fixed = 200;   // per-KV-iter: HMX queue push/pop + DMA pop + barriers
+    const size_t c_q_fixed    = 800;   // per-Q-block: q_load + epilogue o_update + o_norm + o_store
+    const size_t c_iter_base  = 200;   // per-KV-iter base (HMX dot/update + DMA)
+    const size_t c_softmax    = 600;   // per 64-row vector chunk on HVX

    size_t best_cost = SIZE_MAX, best_mn = 0;
    size_t best_Br = 0, best_Bc = 0;
@@ -262,13 +287,20 @@ static inline int hmx_fa_find_chunk_size(size_t * Br_out,
    for (size_t Br = Br_max; Br >= br_unit; Br -= br_unit) {
        // Try all Bc candidates from Bc_limit down to bc_unit
        for (size_t Bc = Bc_limit; Bc >= bc_unit; Bc -= bc_unit) {
-            size_t vtcm_needed = hmx_fa_compute_vtcm_usage(gqa_factor, DK, DV, Br, Bc, n_threads, can_pipeline);
+            size_t vtcm_needed = hmx_fa_compute_vtcm_usage(gqa_factor, DK, DV, Br, Bc, n_threads, can_pipeline, is_q_fp32);
            if (vtcm_needed <= vtcm_budget) {
                // This Bc fits for this Br!
-                const size_t q_blocks  = (qo_len + Br - 1) / Br;
-                const size_t kv_blocks = (kv_len + Bc - 1) / Bc;
-                const size_t cost      = q_blocks * (c_q_fixed + kv_blocks * c_iter_fixed);
-                const size_t mn        = Br * Bc;
+                const size_t q_blocks       = (qo_len + Br - 1) / Br;
+                const size_t kv_blocks      = (kv_len + Bc - 1) / Bc;
+                const size_t actual_threads = (kv_blocks >= 3 && n_threads >= 2) ? n_threads : 1;
+                const size_t n_rows_g       = Br * gqa_factor;
+                const size_t n_row_vec_cnt  = (n_rows_g + 63) / 64;
+                const size_t n_use          = n_row_vec_cnt < actual_threads ? n_row_vec_cnt : actual_threads;
+                const size_t vecs_per_t     = n_use > 0 ? (n_row_vec_cnt + n_use - 1) / n_use : 1;
+
+                const size_t c_iter_actual  = c_iter_base + c_softmax * vecs_per_t;
+                const size_t cost           = q_blocks * (c_q_fixed + kv_blocks * c_iter_actual);
+                const size_t mn             = Br * Bc;

                if (cost < best_cost || (cost == best_cost && mn > best_mn)) {
                    best_cost = cost;
@@ -767,23 +767,25 @@ static void core_mma_chunk_fp16(__fp16 *restrict c, const __fp16 *restrict a, co

 // output : fp16 -> f32p

-static void transfer_output_chunk_fp16_to_fp32(
+static void transfer_output_chunk_fp16_to_fp32_col_chunk(
    float *restrict dst,
    const float *restrict src2,
    const __fp16 *restrict vtcm_src,
    uint32_t start_row,
    uint32_t n_rows,
-    uint32_t n_cols,
+    uint32_t c_len,
+    uint32_t total_n_cols,
    uint32_t dst_stride,
    uint32_t src2_stride,
    uint32_t dst_cols
 ) {
-    assert(n_cols % HTP_MM_HMX_TILE_N_COLS == 0);
-    const size_t tile_row_stride = (n_cols / HTP_MM_HMX_TILE_N_COLS) * HTP_MM_HMX_TILE_N_ELMS;
+    assert(c_len % HTP_MM_HMX_TILE_N_COLS == 0);
+    assert(total_n_cols % HTP_MM_HMX_TILE_N_COLS == 0);
+    const size_t tile_row_stride = (total_n_cols / HTP_MM_HMX_TILE_N_COLS) * HTP_MM_HMX_TILE_N_ELMS;

    const HVX_Vector one = hvx_vec_splat_f16(1.0);

-    const size_t limit_c         = hex_smin(n_cols, dst_cols);
+    const size_t limit_c         = hex_smin(c_len, dst_cols);
    const size_t limit_c_aligned = (limit_c & ~31);

    for (size_t r = 0; r < n_rows; r += 2) {
@@ -848,6 +850,22 @@ static void transfer_output_chunk_fp16_to_fp32(
    }
 }

+static inline void transfer_output_chunk_fp16_to_fp32(
+    float *restrict dst,
+    const float *restrict src2,
+    const __fp16 *restrict vtcm_src,
+    uint32_t start_row,
+    uint32_t n_rows,
+    uint32_t n_cols,
+    uint32_t dst_stride,
+    uint32_t src2_stride,
+    uint32_t dst_cols
+) {
+    transfer_output_chunk_fp16_to_fp32_col_chunk(
+        dst, src2, vtcm_src, start_row, n_rows, n_cols, n_cols, dst_stride, src2_stride, dst_cols
+    );
+}
+
 typedef struct {
    const __fp16  *vtcm_src;
    float         *dst;
@@ -19,6 +19,8 @@
 #endif
 #define HTP_MAX_MMAPS    16

+#define HTP_MAX_DIRTY_RANGES 16
+
 // Memory mapping
 struct htp_mmap {
    uint64_t size;
@@ -95,7 +97,11 @@ struct htp_context {
    atomic_bool            vtcm_needs_release;

    uint64_t               max_vmem;
-    uint32_t               dirty_map[HTP_OP_MAX_TENSORS / 32];
+    struct htp_dirty_range {
+        uint32_t start;
+        uint32_t end;
+        uint32_t bi;
+    } dirty_ranges[HTP_MAX_DIRTY_RANGES];

    // Persistent DDR scratchpad for MUL_MAT_ID mappings
    void *                 ddr_spad_base;
@@ -97,6 +97,7 @@ enum htp_op_code {
    HTP_OP_PAD,
    HTP_OP_NORM,
    HTP_OP_CONCAT,
+    HTP_OP_CLAMP,

    HTP_OP_INVALID
 };
@@ -122,7 +123,7 @@ enum htp_tensor_flags {
 // Tensor descriptor
 struct htp_tensor {
    uint32_t data;                 // Buffer offset in the messages, and data pointer on the NPU
-    uint32_t alias;                // Index of the canonical tensor for this memory buffer
+    uint32_t reserved;             // Reserved for alignment padding (must be multiple of 8)
    uint32_t size;                 // Data size in bytes
    uint32_t flags;                // Buffer / tensor flags
    uint32_t type;                 // Data type
@@ -172,6 +173,7 @@ enum htp_trace_event_id {
    HTP_TRACE_EVT_DMA                 = 0,
    HTP_TRACE_EVT_L2FLUSH             = 1,
    HTP_TRACE_EVT_INIT                = 2,
+    HTP_TRACE_EVT_BUFF                = 3,

    HTP_TRACE_EVT_HVX_COMP            = 20,
    HTP_TRACE_EVT_HVX_A_QUANT         = 21,
@@ -224,7 +226,10 @@ struct htp_opbatch_rsp {
    uint32_t n_tensors;  // Number of tensors
    uint32_t n_ops;      // Number of op profile descriptors
    uint32_t n_traces[HTP_MAX_NTHREADS + 1];
-    uint8_t  pad[8];     // align to 8 bytes
+    uint32_t usecs;          // Number of usec
+    uint32_t pad;            // align to 8 bytes
+    uint64_t cycles_start;   // Start cycle counter
+    uint64_t cycles_stop;    // Stop cycle counter
    // struct htp_prof_desc profs[];  -- dspqueue buf 0
 };

@@ -2,6 +2,7 @@

 #include <qurt.h>
 #include <qurt_memory.h>
+#include <HAP_farf.h>

 #include "hex-common.h"
 #include "hex-utils.h"
@@ -10,84 +11,6 @@
 #include "htp-ctx.h"
 #include "work-queue.h"

-struct l2flush_task {
-    struct htp_thread_trace * trace;
-    uint32_t start;
-    uint32_t end;
-    uint32_t chunk_size;
-    uint32_t ti;
-};
-
-static void l2flush_thread_worker(unsigned int n, unsigned int i, void * data) {
-    struct l2flush_task * task = (struct l2flush_task *) data;
-    const uint32_t start = task->start;
-    const uint32_t end   = task->end;
-    const uint32_t ti    = task->ti;
-    const uint32_t chunk_size = task->chunk_size;
-
-    const uint32_t thread_s = start + i * chunk_size;
-    if (thread_s >= end) {
-        return;
-    }
-    uint32_t thread_e = thread_s + chunk_size;
-    if (thread_e > end) {
-        thread_e = end;
-    }
-
-    struct htp_thread_trace * tr = &task->trace[i];
-    htp_trace_event_start(tr, HTP_TRACE_EVT_L2FLUSH, ti);
-    hex_l2flush((void *) (uintptr_t) thread_s, thread_e - thread_s);
-    htp_trace_event_stop(tr, HTP_TRACE_EVT_L2FLUSH, ti);
-}
-
-static void flush_all_dcache(struct htp_context * ctx) {
-    struct htp_thread_trace * tr = &ctx->trace[0];
-    htp_trace_event_start(tr, HTP_TRACE_EVT_L2FLUSH, 0);
-    qurt_mem_cache_clean((qurt_addr_t) 0, 0, QURT_MEM_CACHE_FLUSH_INVALIDATE_ALL, QURT_MEM_DCACHE);
-    hex_l2fetch_block(ctx, ctx->footprint);
-    htp_trace_event_stop(tr, HTP_TRACE_EVT_L2FLUSH, 0);
-    bitmap_reset(ctx->dirty_map, HTP_OP_MAX_TENSORS);
-}
-
-static void flush_tensor_range(struct htp_context * ctx, const struct htp_tensor * t) {
-    struct htp_thread_trace * tr = &ctx->trace[0];
-
-    if (t->size > HEX_L2_FLUSH_WQ_THRESHOLD && ctx->n_threads > 1) {
-        struct l2flush_task task;
-        task.start = hex_align_down((size_t) t->data, HEX_L2_LINE_SIZE);
-        task.end   = hex_align_up((size_t) t->data + t->size, HEX_L2_LINE_SIZE);
-        task.ti    = t->ti;
-        task.trace = ctx->trace;
-
-        const uint32_t total_size = task.end - task.start;
-        const uint32_t n_blocks   = (total_size + HEX_L2_BLOCK_SIZE - 1) / HEX_L2_BLOCK_SIZE;
-        const uint32_t blocks_per_thread = fastdiv(n_blocks + ctx->n_threads - 1, &ctx->n_threads_div);
-        task.chunk_size = blocks_per_thread * HEX_L2_BLOCK_SIZE;
-
-        work_queue_run(ctx->work_queue, l2flush_thread_worker, &task, ctx->n_threads);
-    } else {
-        htp_trace_event_start(tr, HTP_TRACE_EVT_L2FLUSH, t->ti);
-        hex_l2flush((void *) t->data, t->size);
-        htp_trace_event_stop(tr, HTP_TRACE_EVT_L2FLUSH, t->ti);
-    }
-
-    htp_tensor_make_clean(t, ctx->dirty_map);
-}
-
-void htp_tensor_flush(struct htp_context * ctx, const struct htp_tensor * t) {
-    if (!bitmap_test(ctx->dirty_map, t->ti)) {
-        return;
-    }
-
-    if (t->size > HEX_L2_FLUSH_ALL_THRESHOLD) {
-        flush_all_dcache(ctx);
-        return;
-    }
-
-    flush_tensor_range(ctx, t);
-}
-
-// One dirty tensor's line-aligned range, placed in the flattened global block space.
 struct l2flush_range {
    uint32_t start;       // line-aligned start address
    uint32_t end;         // line-aligned end address
@@ -103,9 +26,18 @@ struct l2flush_multi_task {
    uint32_t                  blocks_per_thread;
 };

+static void flush_all_dcache(struct htp_context * ctx) {
+    struct htp_thread_trace * tr = &ctx->trace[0];
+    htp_trace_event_start(tr, HTP_TRACE_EVT_L2FLUSH, 0);
+    qurt_mem_cache_clean((qurt_addr_t) 0, 0, QURT_MEM_CACHE_FLUSH_INVALIDATE_ALL, QURT_MEM_DCACHE);
+    hex_l2fetch_block(ctx, ctx->footprint);
+    htp_trace_event_stop(tr, HTP_TRACE_EVT_L2FLUSH, 0);
+    memset(ctx->dirty_ranges, 0, sizeof(ctx->dirty_ranges));
+}
+
 static void l2flush_multi_worker(unsigned int n, unsigned int i, void * data) {
-    (void) n;
    struct l2flush_multi_task * task = (struct l2flush_multi_task *) data;
+    (void) n;

    const uint32_t gb_first = i * task->blocks_per_thread;
    uint32_t       gb_last  = gb_first + task->blocks_per_thread;
@@ -141,11 +73,177 @@ static void l2flush_multi_worker(unsigned int n, unsigned int i, void * data) {
    htp_trace_event_stop(tr, HTP_TRACE_EVT_L2FLUSH, gb_first);
 }

-void htp_tensor_flush_all(struct htp_context * ctx, const struct htp_tensor * const * tensors, uint32_t n) {
-    uint64_t total_dirty = 0;
+void htp_tensor_dirty_all(struct htp_context * ctx, const struct htp_tensor * const * tensors, uint32_t n) {
+    const struct htp_tensor * pending[HTP_OP_MAX_OUTPUTS];
+    uint32_t n_pending = 0;
+
    for (uint32_t i = 0; i < n; i++) {
        const struct htp_tensor * t = tensors[i];
-        if (t && bitmap_test(ctx->dirty_map, t->ti)) {
+        if (!t) continue;
+
+        uint32_t t_start = t->data;
+        uint32_t t_end   = t_start + t->size;
+
+        bool merged = false;
+        for (uint32_t j = 0; j < HTP_MAX_DIRTY_RANGES; j++) {
+            struct htp_dirty_range * r = &ctx->dirty_ranges[j];
+            if (!r->start) continue;
+
+            if (r->start <= t_end && t_start <= r->end) {
+                uint32_t new_start = (t_start < r->start) ? t_start : r->start;
+                uint32_t new_end   = (t_end > r->end)     ? t_end   : r->end;
+                r->start = new_start;
+                r->end   = new_end;
+                merged   = true;
+            }
+        }
+
+        if (!merged) {
+            pending[n_pending++] = t;
+        }
+    }
+
+    if (n_pending == 0) {
+        return;
+    }
+
+    uint32_t empty_indices[HTP_MAX_DIRTY_RANGES];
+    uint32_t active_indices[HTP_MAX_DIRTY_RANGES];
+    uint32_t n_active = 0;
+    uint32_t n_empty  = 0;
+    for (uint32_t j = 0; j < HTP_MAX_DIRTY_RANGES; j++) {
+        if (ctx->dirty_ranges[j].start) {
+            active_indices[n_active++] = j;
+        } else {
+            empty_indices[n_empty++]   = j;
+        }
+    }
+
+    if (n_pending <= n_empty) {
+        for (uint32_t i = 0; i < n_pending; i++) {
+            uint32_t idx = empty_indices[i];
+            struct htp_dirty_range * r = &ctx->dirty_ranges[idx];
+            r->start = pending[i]->data;
+            r->end   = pending[i]->data + pending[i]->size;
+            r->bi    = pending[i]->bi;
+        }
+        return;
+    }
+
+    uint32_t n_evict = n_pending - n_empty;
+    uint32_t total_evict_size = 0;
+    for (uint32_t i = 0; i < n_evict; i++) {
+        uint32_t idx = active_indices[i];
+        struct htp_dirty_range * r = &ctx->dirty_ranges[idx];
+        total_evict_size += r->end - r->start;
+    }
+
+    if (total_evict_size > HEX_L2_FLUSH_ALL_THRESHOLD) {
+        flush_all_dcache(ctx);
+        for (uint32_t i = 0; i < n_pending; i++) {
+            struct htp_dirty_range * r = &ctx->dirty_ranges[i];
+            r->start = pending[i]->data;
+            r->end   = pending[i]->data + pending[i]->size;
+            r->bi    = pending[i]->bi;
+        }
+        return;
+    }
+
+    if (total_evict_size > HEX_L2_FLUSH_WQ_THRESHOLD && ctx->n_threads > 1 && n_evict <= HTP_OP_MAX_INPUTS) {
+        struct l2flush_multi_task task;
+        task.trace    = ctx->trace;
+        task.n_ranges = n_evict;
+
+        uint32_t block_acc = 0;
+        for (uint32_t i = 0; i < n_evict; i++) {
+            uint32_t idx = active_indices[i];
+            struct htp_dirty_range * r = &ctx->dirty_ranges[idx];
+
+            struct l2flush_range * rg = &task.ranges[i];
+            rg->start = hex_align_down((size_t) r->start, HEX_L2_LINE_SIZE);
+            rg->end   = hex_align_up((size_t) r->end, HEX_L2_LINE_SIZE);
+            rg->block_first = block_acc;
+            rg->n_blocks = (rg->end - rg->start + HEX_L2_BLOCK_SIZE - 1) / HEX_L2_BLOCK_SIZE;
+            block_acc += rg->n_blocks;
+        }
+
+        task.total_blocks      = block_acc;
+        task.blocks_per_thread = fastdiv(block_acc + ctx->n_threads - 1, &ctx->n_threads_div);
+
+        work_queue_run(ctx->work_queue, l2flush_multi_worker, &task, ctx->n_threads);
+    } else {
+        struct htp_thread_trace * tr = &ctx->trace[0];
+        htp_trace_event_start(tr, HTP_TRACE_EVT_L2FLUSH, 0);
+        for (uint32_t i = 0; i < n_evict; i++) {
+            uint32_t idx = active_indices[i];
+            struct htp_dirty_range * r = &ctx->dirty_ranges[idx];
+            uint32_t size = r->end - r->start;
+            hex_l2flush((void *) (uintptr_t) r->start, size);
+        }
+        htp_trace_event_stop(tr, HTP_TRACE_EVT_L2FLUSH, 0);
+    }
+
+    for (uint32_t i = 0; i < n_evict; i++) {
+        uint32_t idx = active_indices[i];
+        struct htp_dirty_range * r = &ctx->dirty_ranges[idx];
+        r->start = pending[i]->data;
+        r->end   = pending[i]->data + pending[i]->size;
+        r->bi    = pending[i]->bi;
+    }
+
+    for (uint32_t i = 0; i < n_empty; i++) {
+        uint32_t idx = empty_indices[i];
+        struct htp_dirty_range * r = &ctx->dirty_ranges[idx];
+        r->start = pending[n_evict + i]->data;
+        r->end   = pending[n_evict + i]->data + pending[n_evict + i]->size;
+        r->bi    = pending[n_evict + i]->bi;
+    }
+}
+
+static void make_tensor_clean(struct htp_context * ctx, const struct htp_tensor * t) {
+    uint32_t t_start = t->data;
+    uint32_t t_end   = t_start + t->size;
+
+    for (uint32_t i = 0; i < HTP_MAX_DIRTY_RANGES; i++) {
+        struct htp_dirty_range * r = &ctx->dirty_ranges[i];
+        if (!r->start) continue;
+
+        if (r->start < t_end && t_start < r->end) {
+            if (t_start <= r->start && r->end <= t_end) {
+                r->start = 0;
+            } else if (t_start <= r->start) {
+                r->start = t_end;
+            } else if (r->end <= t_end) {
+                r->end = t_start;
+            }
+        }
+    }
+}
+
+static inline bool is_tensor_dirty(struct htp_context * ctx, const struct htp_tensor * t) {
+    uint32_t t_start = t->data;
+    uint32_t t_end   = t_start + t->size;
+
+    for (uint32_t i = 0; i < HTP_MAX_DIRTY_RANGES; i++) {
+        struct htp_dirty_range * r = &ctx->dirty_ranges[i];
+        if (!r->start) continue;
+
+        if (r->start < t_end && t_start < r->end) {
+            return true;
+        }
+    }
+    return false;
+}
+
+void htp_tensor_flush_all(struct htp_context * ctx, const struct htp_tensor * const * tensors, uint32_t n) {
+    const struct htp_tensor * dirty_tensors[HTP_OP_MAX_INPUTS];
+    uint32_t n_dirty = 0;
+    uint64_t total_dirty = 0;
+
+    for (uint32_t i = 0; i < n; i++) {
+        const struct htp_tensor * t = tensors[i];
+        if (t && (t->flags & HTP_TENSOR_COMPUTE) && is_tensor_dirty(ctx, t)) {
+            dirty_tensors[n_dirty++] = t;
            total_dirty += t->size;
        }
    }
@@ -159,21 +257,15 @@ void htp_tensor_flush_all(struct htp_context * ctx, const struct htp_tensor * co
        return;
    }

-    // Aggregate is small enough to walk. Thread it across all dirty ranges at once
-    // when it is worth the dispatch, otherwise flush sequentially.
-    if (total_dirty > HEX_L2_FLUSH_WQ_THRESHOLD && ctx->n_threads > 1) {
+    if (total_dirty >= HEX_L2_FLUSH_WQ_THRESHOLD && ctx->n_threads > 1) {
        struct l2flush_multi_task task;
        task.trace    = ctx->trace;
        task.n_ranges = 0;

        uint32_t block_acc = 0;
-        for (uint32_t i = 0; i < n; i++) {
-            const struct htp_tensor * t = tensors[i];
-            if (!t || !bitmap_test(ctx->dirty_map, t->ti)) {
-                continue;
-            }
-            // Clear as we go: dedups a tensor passed as multiple srcs (e.g. mul(x,x)).
-            htp_tensor_make_clean(t, ctx->dirty_map);
+        for (uint32_t i = 0; i < n_dirty; i++) {
+            const struct htp_tensor * t = dirty_tensors[i];
+            make_tensor_clean(ctx, t);

            struct l2flush_range * rg = &task.ranges[task.n_ranges++];
            rg->start = hex_align_down((size_t) t->data, HEX_L2_LINE_SIZE);
@@ -191,14 +283,11 @@ void htp_tensor_flush_all(struct htp_context * ctx, const struct htp_tensor * co
    }

    struct htp_thread_trace * tr = &ctx->trace[0];
-    for (uint32_t i = 0; i < n; i++) {
-        const struct htp_tensor * t = tensors[i];
-        if (!t || !bitmap_test(ctx->dirty_map, t->ti)) {
-            continue;
-        }
+    for (uint32_t i = 0; i < n_dirty; i++) {
+        const struct htp_tensor * t = dirty_tensors[i];
        htp_trace_event_start(tr, HTP_TRACE_EVT_L2FLUSH, t->ti);
-        hex_l2flush((void *) t->data, t->size);
+        hex_l2flush((void *) (uintptr_t) t->data, t->size);
        htp_trace_event_stop(tr, HTP_TRACE_EVT_L2FLUSH, t->ti);
-        htp_tensor_make_clean(t, ctx->dirty_map);
+        make_tensor_clean(ctx, t);
    }
 }
@@ -5,10 +5,6 @@
 #include "htp-ops.h"
 #include "hex-bitmap.h"

-static inline struct htp_tensor * htp_tensor_alias(const struct htp_tensor * t) {
-    return (struct htp_tensor *) (uintptr_t) t->alias;
-}
-
 static inline void * htp_tensor_data(const struct htp_tensor * t) {
    return (void *) (uintptr_t) t->data;
 }
@@ -17,20 +13,8 @@ static inline uint32_t * htp_tensor_flags(const struct htp_tensor * t) {
    return (uint32_t *) &t->flags;
 }

-static inline void htp_tensor_make_dirty(const struct htp_tensor * t, uint32_t * dirty_map) {
-    struct htp_tensor * curr = (struct htp_tensor *) t;
-    do {
-        bitmap_set(dirty_map, curr->ti);
-        curr = htp_tensor_alias(curr);
-    } while (curr != t);
-}
-
-static inline void htp_tensor_make_clean(const struct htp_tensor * t, uint32_t * dirty_map) {
-    bitmap_clear(dirty_map, t->ti);
-}
-
 struct htp_context;
-void htp_tensor_flush(struct htp_context * ctx, const struct htp_tensor * t);
 void htp_tensor_flush_all(struct htp_context * ctx, const struct htp_tensor * const * tensors, uint32_t n);
+void htp_tensor_dirty_all(struct htp_context * ctx, const struct htp_tensor * const * tensors, uint32_t n);

 #endif // HTP_TENSOR_H
@@ -208,6 +208,77 @@ static inline void hvx_mad_f32_f16_aa_rx2(float * restrict y, const void * restr
        }
    }
 }
+static inline void hvx_mad_f32_f16_aa_vec(float * restrict y, const void * restrict x, HVX_Vector S0, uint32_t n) {
+    const HVX_Vector * restrict vx0 = (const HVX_Vector *) x;
+
+    HVX_VectorPair * restrict vy_p = (HVX_VectorPair *) y;
+    HVX_Vector * restrict vy = (HVX_Vector *) y;
+
+    uint32_t nvec = n / VLEN_FP16; // num full fp16 hvx vectors
+    uint32_t nloe = n % VLEN_FP16; // leftover elements
+
+    uint32_t i = 0;
+
+    #pragma unroll(2)
+    for (i = 0; i < nvec; ++i) {
+        vy_p[i] = hvx_vec_mpyacc_f32_f16(vy_p[i], Q6_Vh_vshuff_Vh(vx0[i]), S0);
+    }
+
+    if (nloe) {
+        HVX_VectorPair xy_p = vy_p[i];
+        xy_p = hvx_vec_mpyacc_f32_f16(xy_p, Q6_Vh_vshuff_Vh(vx0[i]), S0);
+
+        HVX_Vector xy = Q6_V_lo_W(xy_p);
+        i = 2 * i;  // index for vy
+
+        if (nloe >= VLEN_FP32) {
+            vy[i] = xy;
+            nloe -= VLEN_FP32; ++i; xy = Q6_V_hi_W(xy_p);
+        }
+
+        if (nloe) {
+            hvx_vec_store_a(&vy[i], nloe * 4, xy);
+        }
+    }
+}
+
+static inline void hvx_mad_f32_f16_aa_rx2_vec(float * restrict y, const void * restrict x0, const void * restrict x1,
+                                          HVX_Vector S0, HVX_Vector S1, uint32_t n) {
+    const HVX_Vector * restrict vx0 = (const HVX_Vector *) x0;
+    const HVX_Vector * restrict vx1 = (const HVX_Vector *) x1;
+
+    HVX_VectorPair * restrict vy_p  = (HVX_VectorPair *) y;
+    HVX_Vector * restrict vy        = (HVX_Vector *) y;
+
+    uint32_t nvec = n / VLEN_FP16;  // num full fp16 hvx vectors
+    uint32_t nloe = n % VLEN_FP16;  // leftover elements
+
+    uint32_t i = 0;
+
+    #pragma unroll(2)
+    for (i = 0; i < nvec; ++i) {
+        vy_p[i] = hvx_vec_mpyacc_f32_f16(vy_p[i], Q6_Vh_vshuff_Vh(vx0[i]), S0);
+        vy_p[i] = hvx_vec_mpyacc_f32_f16(vy_p[i], Q6_Vh_vshuff_Vh(vx1[i]), S1);
+    }
+
+    if (nloe) {
+        HVX_VectorPair xy_p = vy_p[i];
+        xy_p = hvx_vec_mpyacc_f32_f16(xy_p, Q6_Vh_vshuff_Vh(vx0[i]), S0);
+        xy_p = hvx_vec_mpyacc_f32_f16(xy_p, Q6_Vh_vshuff_Vh(vx1[i]), S1);
+
+        HVX_Vector xy = Q6_V_lo_W(xy_p);
+        i = 2 * i;  // index for vy
+
+        if (nloe >= VLEN_FP32) {
+            vy[i] = xy;
+            nloe -= VLEN_FP32; ++i; xy = Q6_V_hi_W(xy_p);
+        }
+
+        if (nloe) {
+            hvx_vec_store_a(&vy[i], nloe * 4, xy);
+        }
+    }
+}

 static inline void hvx_scale_vec_f32_aa(uint8_t * restrict dst, const uint8_t * restrict src, const uint32_t n, HVX_Vector vs) {
    assert((size_t) dst % 128 == 0);
@@ -286,6 +286,46 @@ static inline float hvx_sum_of_squares_f32(const uint8_t * restrict src, const i
    }
 }

+// Signed 32-bit Integer Max variants
+
+static inline HVX_Vector hvx_vec_reduce_max_n_i32(HVX_Vector in, unsigned int n) {
+    unsigned int total = n * 4;  // total vec nbytes
+    unsigned int width = 4;      // int32 nbytes
+
+    HVX_Vector max_val = in, max_t;
+    while (width < total) {
+        max_t = Q6_V_vror_VR(max_val, width);          // rotate right
+        max_val = Q6_Vw_vmax_VwVw(max_t, max_val);      // elementwise signed max
+        width = width << 1;
+    }
+    return max_val;
+}
+
+static inline HVX_Vector hvx_vec_reduce_max_i32(HVX_Vector in) {
+    return hvx_vec_reduce_max_n_i32(in, 32);
+}
+
+static inline int32_t hvx_reduce_max_i32_a(const uint8_t * restrict src, const int num_elems) {
+    HVX_Vector init_vec = Q6_V_vsplat_R(((const int32_t *) src)[0]);
+    HVX_Vector pad_vec  = Q6_V_vsplat_R(0x80000000);
+    assert((uintptr_t) src % 128 == 0);
+    hvx_reduce_loop_body(HVX_Vector, init_vec, pad_vec, Q6_Vw_vmax_VwVw, hvx_vec_reduce_max_i32, hvx_vec_get_i32);
+}
+
+static inline int32_t hvx_reduce_max_i32_u(const uint8_t * restrict src, const int num_elems) {
+    HVX_Vector init_vec = Q6_V_vsplat_R(((const int32_t *) src)[0]);
+    HVX_Vector pad_vec  = Q6_V_vsplat_R(0x80000000);
+    hvx_reduce_loop_body(HVX_UVector, init_vec, pad_vec, Q6_Vw_vmax_VwVw, hvx_vec_reduce_max_i32, hvx_vec_get_i32);
+}
+
+static inline int32_t hvx_reduce_max_i32(const uint8_t * restrict src, const int num_elems) {
+    if (hex_is_aligned((void *) src, 128)) {
+        return hvx_reduce_max_i32_a(src, num_elems);
+    } else {
+        return hvx_reduce_max_i32_u(src, num_elems);
+    }
+}
+
 #undef hvx_reduce_loop_body
 #undef HVX_REDUCE_MAX_OP
 #undef HVX_REDUCE_SUM_OP
@@ -718,18 +718,19 @@ static int execute_op(struct htp_ops_context * octx) {
        case HTP_OP_RMS_NORM:
        case HTP_OP_RMS_NORM_MUL:
        case HTP_OP_SCALE:
+        case HTP_OP_CLAMP:
        case HTP_OP_SQR:
        case HTP_OP_SQRT:
        case HTP_OP_UNARY_SOFTPLUS:
        case HTP_OP_UNARY_SIGMOID:
+        case HTP_OP_UNARY_SILU:
+        case HTP_OP_UNARY_GELU:
        case HTP_OP_UNARY_NEG:
        case HTP_OP_UNARY_EXP:
        case HTP_OP_UNARY_TANH:
        case HTP_OP_L2_NORM:
            return op_unary(octx);

-        case HTP_OP_UNARY_SILU:
-        case HTP_OP_UNARY_GELU:
        case HTP_OP_GLU_SWIGLU:
        case HTP_OP_GLU_SWIGLU_OAI:
        case HTP_OP_GLU_GEGLU:
@@ -900,10 +901,8 @@ static void prep_tensor(struct htp_context *ctx, struct htp_buf_desc *bufs, stru
    uint32_t offset = t->data;
    uint32_t size   = t->size;
    uint32_t bi     = t->bi;
-    uint32_t alias  = t->alias;

    t->data  = (uint32_t) (bufs[bi].base + offset);  // update data to the actual pointer
-    t->alias = (uint32_t) (tens + alias);            // update alias to the actual pointer

    FARF(HIGH, "prep-tensor #%u: bi %u offset %u size %u data %p : %u:%u:%u:%u", idx, t->bi, offset, t->size, (void*) t->data,
        t->ne[0], t->ne[1], t->ne[3], t->ne[3]);
@@ -954,14 +953,14 @@ static int proc_op_req(struct htp_ops_context * octx, struct htp_tensor *tens, u
        octx->dsts[i]    = dst;
        octx->dst_dma[i] = octx->ctx->dma; // FIXME: ? octx->ctx->dma_cached : octx->ctx->dma;

-        htp_tensor_make_dirty(dst, octx->ctx->dirty_map);
-
        FARF(HIGH, "prep-dst[%u] #%u: data %p size %u : %u:%u:%u:%u", i, dst_idx, (void*) dst->data, dst->size,
            dst->ne[0], dst->ne[1], dst->ne[2], dst->ne[3]);
    }

    int status = execute_op(octx);

+    htp_tensor_dirty_all(octx->ctx, octx->dsts, HTP_OP_MAX_OUTPUTS);
+
    octx->src0_spad.src = NULL;
    octx->src1_spad.src = NULL;
    octx->src2_spad.src = NULL;
@@ -993,12 +992,6 @@ static void process_opbatch(struct htp_context * ctx, const struct htp_opbatch_r
    FARF(HIGH, "processing opbatch #%u: n-bufs %u n-tensors %u n-ops %u n-traces %u : m-size %u b-size %u t-size %u o-size %u", req->id,
            n_bufs, n_tens, n_ops, req->n_traces, dbuf->size, b_size, t_size, o_size);

-    // Clean cache at the start of the batch
-    // We cant trace this part because the trace buffer is setup later
-    qurt_mem_cache_clean((qurt_addr_t) 0, 0, QURT_MEM_CACHE_FLUSH_INVALIDATE_ALL, QURT_MEM_DCACHE);
-    hex_l2fetch_block(ctx, ctx->footprint);
-    bitmap_reset(ctx->dirty_map, HTP_OP_MAX_TENSORS);
-
    // Setup descriptor pointers
    uint8_t * m_ptr = dbuf->ptr;
    struct htp_buf_desc* bufs = (struct htp_buf_desc*)  m_ptr; m_ptr += b_size;
@@ -1006,13 +999,8 @@ static void process_opbatch(struct htp_context * ctx, const struct htp_opbatch_r
    struct htp_op_desc*   ops = (struct htp_op_desc*)   m_ptr; m_ptr += o_size;
    struct htp_prof_desc* pds = (struct htp_prof_desc*) m_ptr;

-    prep_op_bufs(ctx, bufs, n_bufs);
-    prep_tensors(ctx, bufs, tens, n_tens);
-
-    struct htp_ops_context *octx = &ctx->octx;
-    memset(octx, 0, sizeof(*octx));
-    octx->n_threads = ctx->n_threads;
-    octx->ctx       = ctx;
+    struct profile_data batch_prof;
+    profile_start(HTP_PROF_BASIC, &batch_prof);

    memset(ctx->trace, 0, sizeof(ctx->trace));
    if (ctx->profiler == HTP_PROF_TRACE) {
@@ -1023,6 +1011,24 @@ static void process_opbatch(struct htp_context * ctx, const struct htp_opbatch_r
        }
    }

+    // Clean cache at the start of the batch
+    htp_trace_event_start(&ctx->trace[0], HTP_TRACE_EVT_L2FLUSH, 0);
+    qurt_mem_cache_clean((qurt_addr_t) 0, 0, QURT_MEM_CACHE_FLUSH_INVALIDATE_ALL, QURT_MEM_DCACHE);
+    hex_l2fetch_block(ctx, ctx->footprint);
+    memset(ctx->dirty_ranges, 0, sizeof(ctx->dirty_ranges));
+    htp_trace_event_stop(&ctx->trace[0], HTP_TRACE_EVT_L2FLUSH, 0);
+
+    htp_trace_event_start(&ctx->trace[0], HTP_TRACE_EVT_BUFF, 0);
+    prep_op_bufs(ctx, bufs, n_bufs);
+    htp_trace_event_stop(&ctx->trace[0], HTP_TRACE_EVT_BUFF, 0);
+
+    prep_tensors(ctx, bufs, tens, n_tens);
+
+    struct htp_ops_context *octx = &ctx->octx;
+    memset(octx, 0, sizeof(*octx));
+    octx->n_threads = ctx->n_threads;
+    octx->ctx       = ctx;
+
    work_queue_wakeup(ctx->work_queue);
    if (ctx->hmx_queue) {
        hmx_queue_wakeup(ctx->hmx_queue);
@@ -1055,13 +1061,23 @@ static void process_opbatch(struct htp_context * ctx, const struct htp_opbatch_r
    }
    work_queue_suspend(ctx->work_queue);

+    // Flush remaining dirty tensors at the end of the batch
+    htp_trace_event_start(&ctx->trace[0], HTP_TRACE_EVT_L2FLUSH, 0);
+    qurt_mem_cache_clean((qurt_addr_t) 0, 0, QURT_MEM_CACHE_FLUSH_INVALIDATE_ALL, QURT_MEM_DCACHE);
+    htp_trace_event_stop(&ctx->trace[0], HTP_TRACE_EVT_L2FLUSH, 0);
+
+    profile_stop(HTP_PROF_BASIC, &batch_prof);
+
    struct htp_opbatch_rsp rsp;
    memset(&rsp, 0, sizeof(rsp));
-    rsp.id        = req->id;
-    rsp.status    = op_status;
-    rsp.n_bufs    = n_bufs;
-    rsp.n_tensors = n_tens;
-    rsp.n_ops     = n_ops;
+    rsp.id           = req->id;
+    rsp.status       = op_status;
+    rsp.n_bufs       = n_bufs;
+    rsp.n_tensors    = n_tens;
+    rsp.n_ops        = n_ops;
+    rsp.usecs        = batch_prof.usecs;
+    rsp.cycles_start = batch_prof.cycles_start;
+    rsp.cycles_stop  = batch_prof.cycles_stop;

    if (ctx->profiler == HTP_PROF_TRACE) {
        for (int t = 0; t <= HTP_MAX_NTHREADS; t++) {
@@ -1072,11 +1088,6 @@ static void process_opbatch(struct htp_context * ctx, const struct htp_opbatch_r
    struct dspqueue_buffer write_dbuf = *dbuf;
    write_dbuf.flags = DSPQUEUE_BUFFER_FLAG_FLUSH_SENDER | DSPQUEUE_BUFFER_FLAG_INVALIDATE_RECIPIENT;

-    // Flush remaining dirty tensors at the end of the batch
-    htp_trace_event_start(&ctx->trace[0], HTP_TRACE_EVT_L2FLUSH, 0);
-    qurt_mem_cache_clean((qurt_addr_t) 0, 0, QURT_MEM_CACHE_FLUSH_INVALIDATE_ALL, QURT_MEM_DCACHE);
-    htp_trace_event_stop(&ctx->trace[0], HTP_TRACE_EVT_L2FLUSH, 0);
-
    err = dspqueue_write(queue, 0, 1, &write_dbuf, sizeof(rsp), (const uint8_t *) &rsp, DSPQUEUE_TIMEOUT_NONE);
    if (err != 0) {
        FARF(ERROR, "dspqueue_write failed: 0x%08x", (unsigned) err);
@@ -14,6 +14,8 @@
 #include "hex-dma.h"
 #include "hvx-utils.h"
 #include "hvx-dump.h"
+#include "hvx-arith.h"
+#include "hvx-reduce.h"

 #define GGML_COMMON_DECL_C
 #include "ggml-common.h"
@@ -82,6 +84,8 @@ struct htp_mm_context {

    // Precomputed values
    uint32_t src0_nrows_per_thread;
+    uint32_t src0_row_size_padded;
+    uint32_t src1_nrows;

    struct fastdiv_values mm_div_ne12_ne1;
    struct fastdiv_values mm_div_ne1;
@@ -103,6 +107,7 @@ struct htp_mm_context {
    // Fields for scattered mapping & HMX support in MUL_MAT_ID
    const uint32_t * matrix_row_counts;
    const struct mmid_row_mapping * matrix_rows;
+    uint32_t mapping_stride;

    // Dynamic VTCM pointers allocated sequentially
    uint8_t * vtcm_src0;
@@ -154,8 +159,6 @@ static const uint8_t __attribute__((aligned(VLEN))) kvalues_mxfp4_lut[] = {
    0,    0, 0,    0, 0,    0, 0, 0, 0, 0, 0, 0, 0, 0, 0,  0, 0, 0, 0,    0, 0,    0, 0,    0,
 };

-
-
 #define htp_matmul_tensors_preamble                                 \
    const struct htp_tensor * restrict src0 = octx->src[0];         \
    const struct htp_tensor * restrict src1 = octx->src[1];         \
@@ -444,6 +447,16 @@ static void hvx_mv_2d_repacked_##SUFFIX(unsigned int nth, unsigned int ith, void
                                                                                                                                  \
    uint32_t push_ct = ct_start;                                                                                                  \
    if (src0_start_row < src0_end_row) {                                                                                          \
+        if (src2) {                                                                                                               \
+            float * vtcm_src2_ptr = (float *) mmctx->vtcm_src2 + src0_start_row;                                                  \
+            const float * src2_ptr = (const float *) src2->data + src0_start_row;                                                 \
+            int slice_size = (int)MIN(src0_end_row, ne0) - (int)src0_start_row;                                                   \
+            if (slice_size > 0) {                                                                                                 \
+                dma_queue_push(dma_queue, dma_make_ptr(vtcm_src2_ptr, src2_ptr),                                                  \
+                               slice_size * sizeof(float), slice_size * sizeof(float), slice_size * sizeof(float), 1);            \
+                dma_queue_pop_nowait(dma_queue);                                                                                  \
+            }                                                                                                                     \
+        }                                                                                                                         \
        for (uint32_t d = 0; d < n_prefetch && push_ct < ct_end; d++, push_ct++) {                                                \
            dma_queue_push(dma_queue, dma_make_ptr(vtcm_src0_ptr + d * tile_row_transfer_size_aligned,                            \
                           src0_row + push_ct * tile_row_stride), aligned_tile_size, tile_size, tile_size, n_k_tiles_a);          \
@@ -465,7 +478,7 @@ static void hvx_mv_2d_repacked_##SUFFIX(unsigned int nth, unsigned int ith, void
                                                                                                                                  \
        htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, ct);                                                                    \
        DOT_2X1(ne10, dst_ptr, w_tile, src1_col, valid_rows, NULL);                                                               \
-        htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, ct);                                                                     \
+        htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, ct);                                                                    \
                                                                                                                                  \
        if (push_ct < ct_end) {                                                                                                   \
            dma_queue_push(dma_queue, dma_make_ptr((uint8_t *)w_tile, src0_row + push_ct * tile_row_stride),                      \
@@ -476,24 +489,16 @@ static void hvx_mv_2d_repacked_##SUFFIX(unsigned int nth, unsigned int ith, void
                                                                                                                                  \
    int copy_cnt = (int)MIN(src0_end_row, ne0) - (int)src0_start_row;                                                             \
    if (copy_cnt > 0) {                                                                                                           \
+        htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, ct_end);                                                                \
        if (src2) {                                                                                                               \
-            float * dst_ptr = &dst_col[src0_start_row];                                                                           \
-            const float * src2_ptr = (const float *) src2->data + src0_start_row;                                                 \
-            float * tmp_ptr = tmp;                                                                                                \
-            int remaining = copy_cnt;                                                                                             \
-            while (remaining > 0) {                                                                                               \
-                int n = MIN(remaining, 32);                                                                                       \
-                HVX_Vector v_out = hvx_vmemu(tmp_ptr);                                                                            \
-                HVX_Vector v_z   = hvx_vmemu(src2_ptr);                                                                           \
-                hvx_vec_store_u(dst_ptr, n * sizeof(float), hvx_vec_add_f32_f32(v_out, v_z));                                     \
-                dst_ptr += n;                                                                                                     \
-                src2_ptr += n;                                                                                                    \
-                tmp_ptr += n;                                                                                                     \
-                remaining -= n;                                                                                                   \
-            }                                                                                                                     \
+            hvx_add_f32_uaa((uint8_t *) &dst_col[src0_start_row],                                                                 \
+                            (const uint8_t *) tmp,                                                                                \
+                            (const uint8_t *) ((const float *) mmctx->vtcm_src2 + src0_start_row),                                \
+                            copy_cnt);                                                                                            \
        } else {                                                                                                                  \
            hvx_copy_f32_ua((uint8_t *) &dst_col[src0_start_row], (uint8_t *) tmp, copy_cnt);                                     \
        }                                                                                                                         \
+        htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, ct_end);                                                                 \
    }                                                                                                                             \
 }

@@ -1069,6 +1074,16 @@ static void hvx_mv_2d(unsigned int nth, unsigned int ith, void * data) {

    // Prefill vtcm with 2x src0 rows
    if (src0_start_row < src0_end_row) {
+        if (src2) {
+            float * vtcm_src2_ptr = (float *) mmctx->vtcm_src2 + src0_start_row;
+            const float * src2_ptr = (const float *) src2->data + src0_start_row;
+            int slice_size = (int)src0_end_row - (int)src0_start_row;
+            if (slice_size > 0) {
+                dma_queue_push(dma_queue, dma_make_ptr(vtcm_src2_ptr, src2_ptr),
+                               slice_size * sizeof(float), slice_size * sizeof(float), slice_size * sizeof(float), 1);
+                dma_queue_pop_nowait(dma_queue);
+            }
+        }
        for (uint32_t ir0 = src0_start_row; ir0 < src0_end_row_x2; ir0 += 2) {
            const uint32_t is0 = (ir0 - src0_start_row);
            if (is0 >= n_prefetch) {
@@ -1114,27 +1129,21 @@ static void hvx_mv_2d(unsigned int nth, unsigned int ith, void * data) {
    }

    int copy_cnt = src0_end_row - src0_start_row;
-    if (src2) {
-        float * dst_ptr = &dst_col[src0_start_row];
-        const float * src2_ptr = (const float *) src2->data + src0_start_row;
-        float * tmp_ptr = tmp;
-        int remaining = copy_cnt;
-        while (remaining > 0) {
-            int n = MIN(remaining, 32);
-            HVX_Vector v_out = hvx_vmemu(tmp_ptr);
-            HVX_Vector v_z   = hvx_vmemu(src2_ptr);
-            hvx_vec_store_u(dst_ptr, n * sizeof(float), hvx_vec_add_f32_f32(v_out, v_z));
-            dst_ptr += n;
-            src2_ptr += n;
-            tmp_ptr += n;
-            remaining -= n;
+    if (copy_cnt > 0) {
+        htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, src0_end_row);
+        if (src2) {
+            hvx_add_f32_uaa((uint8_t *) &dst_col[src0_start_row],
+                            (const uint8_t *) tmp,
+                            (const uint8_t *) ((const float *) mmctx->vtcm_src2 + src0_start_row),
+                            copy_cnt);
+        } else {
+            hvx_copy_f32_ua((uint8_t *) &dst_col[src0_start_row], (uint8_t *) tmp, copy_cnt);
        }
-    } else {
-        hvx_copy_f32_ua((uint8_t *) &dst_col[src0_start_row], (uint8_t *) tmp, copy_cnt);
+        htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, src0_end_row);
    }
 }

-#define MMID_MATRIX_ROW(row_id, i1) matrix_rows[(row_id) * ids->ne[0] * ids->ne[1] + (i1)]
+#define MMID_MATRIX_ROW(row_id, i1) matrix_rows[(row_id) * mmctx->mapping_stride + (i1)]

 static void hvx_mm_id(unsigned int nth, unsigned int ith, void * data) {
    htp_matmul_preamble;
@@ -1519,7 +1528,7 @@ static int hvx_mm_matmul(struct htp_ops_context * octx) {

    struct htp_mm_hvx_vtcm_layout L;
    htp_mm_hvx_vtcm_layout_build(&L, kparams->kernel_type, src0->type, ne10, src1_nrows, octx->n_threads,
-                                 dst_row_size, src0_row_size, src1_row_size, kparams->n_prefetch, false, false, false);
+                                 dst_row_size, src0_row_size, src1_row_size, src2 ? src2->nb[1] : 0, kparams->n_prefetch, false, false, false);

    if (kparams->kernel_type == HTP_MM_KERNEL_HVX_F16_F16_VTCM ||
        kparams->kernel_type == HTP_MM_KERNEL_HVX_F32_F32_VTCM ||
@@ -1551,6 +1560,7 @@ static int hvx_mm_matmul(struct htp_ops_context * octx) {
    uint8_t * const base = (uint8_t *) octx->ctx->vtcm_base;
    mmctx->vtcm_src1 = VTCM_LAYOUT_PTR(uint8_t, base, L.off_src1);
    mmctx->vtcm_src0 = VTCM_LAYOUT_PTR(uint8_t, base, L.off_src0);
+    mmctx->vtcm_src2 = VTCM_LAYOUT_PTR(uint8_t, base, L.off_src2);
    mmctx->vtcm_dst  = VTCM_LAYOUT_PTR(uint8_t, base, L.off_dst);

    octx->src1_spad.src  = NULL;
@@ -2346,12 +2356,77 @@ static void dequantize_tiled_weight_chunk_to_fp16_tiles(
    }
 }

+typedef struct {
+    float *dst;
+    const float *src2;
+    const __fp16 *vtcm_src;
+    uint32_t n_rows;
+    uint32_t n_cols;
+    uint32_t dst_stride;
+    uint32_t src2_stride;
+    uint32_t dst_cols;
+    struct fastdiv_values n_threads_div;
+    struct htp_thread_trace *traces;
+    struct htp_context *ctx;
+} output_transfer_col_chunk_state_t;
+
+static void transfer_output_chunk_col_chunk_worker_fn(unsigned int n, unsigned int i, void *data) {
+    (void) n;
+    output_transfer_col_chunk_state_t *st = (output_transfer_col_chunk_state_t *) data;
+    struct htp_thread_trace * tr = &st->traces[i];
+
+    uint32_t n_blocks = st->n_cols / 32;
+    uint32_t b_first = fastdiv(n_blocks * i, &st->n_threads_div);
+    uint32_t b_last  = fastdiv(n_blocks * (i + 1), &st->n_threads_div);
+    uint32_t c_first = b_first * 32;
+    uint32_t c_last  = b_last * 32;
+    uint32_t c_len   = c_last - c_first;
+
+    if (c_len == 0) return;
+
+    htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_O_PROC, c_first);
+
+    float *dst = st->dst + c_first;
+    const float *src2 = st->src2 ? (st->src2 + c_first) : NULL;
+    const __fp16 *vtcm_src = st->vtcm_src + b_first * HTP_MM_HMX_TILE_N_ELMS;
+
+    int chunk_dst_cols = (int)st->dst_cols - (int)c_first;
+    if (chunk_dst_cols > 0) {
+        transfer_output_chunk_fp16_to_fp32_col_chunk(
+            dst, src2, vtcm_src, 0, st->n_rows, c_len, st->n_cols,
+            st->dst_stride, st->src2_stride, (uint32_t)chunk_dst_cols
+        );
+    }
+
+    htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_O_PROC, c_first);
+}
+
 static void transfer_output_chunk_threaded(struct htp_context *ctx, float *dst, const float *src2, const __fp16 *vtcm_src,
                                              int n_rows, int n_cols, int dst_stride, uint32_t src2_stride, int dst_cols, int n_threads) {
    assert(n_cols % HTP_MM_HMX_TILE_N_COLS == 0);

    if (n_rows <= 0) return;

+    uint32_t n_blocks = (uint32_t)n_cols / 32;
+    if (n_threads > 1 && n_blocks >= (uint32_t)n_threads) {
+        struct fastdiv_values n_threads_div = init_fastdiv_values(n_threads);
+        output_transfer_col_chunk_state_t col_state;
+        col_state.dst = dst;
+        col_state.src2 = src2;
+        col_state.vtcm_src = vtcm_src;
+        col_state.n_rows = (uint32_t)n_rows;
+        col_state.n_cols = (uint32_t)n_cols;
+        col_state.dst_stride = (uint32_t)dst_stride;
+        col_state.src2_stride = src2_stride;
+        col_state.dst_cols = (uint32_t)dst_cols;
+        col_state.n_threads_div = n_threads_div;
+        col_state.traces = ctx->trace;
+        col_state.ctx = ctx;
+
+        worker_pool_run_func(ctx->worker_pool, transfer_output_chunk_col_chunk_worker_fn, &col_state, n_threads);
+        return;
+    }
+
    size_t n_tot_chunks      = n_rows;
    size_t n_chunks_per_task = (n_threads == 1) ? n_tot_chunks : hmx_ceil_div(n_rows, n_threads);
    n_chunks_per_task        = hex_align_up(n_chunks_per_task, 2);
@@ -3338,12 +3413,10 @@ int op_matmul(struct htp_ops_context * octx) {

 static int hmx_mm_op_matmul_id(
    struct htp_ops_context * octx,
-    struct htp_mm_context * mmctx,
-    const uint32_t * matrix_row_counts,
-    const struct mmid_row_mapping * matrix_rows,
-    void * mapping_buf,
-    bool must_free_mapping
+    struct htp_mm_context * mmctx
 ) {
+    const uint32_t *                matrix_row_counts = mmctx->matrix_row_counts;
+    const struct mmid_row_mapping * matrix_rows       = mmctx->matrix_rows;
    htp_matmul_tensors_preamble;
    const struct htp_mm_kernel_params * kparams = (const struct htp_mm_kernel_params *) octx->kernel_params;
    const int n_ids = octx->src[2]->ne[0];
@@ -3361,28 +3434,24 @@ static int hmx_mm_op_matmul_id(
                                   nb11, nb12,
                                   nb1, nb2,
                                   (int) src0->nb[1], (int) src0->type,
-                                   matrix_rows, cur_a, n_ids * octx->src[2]->ne[1]);
+                                   matrix_rows, cur_a, mmctx->mapping_stride);
        if (ret != 0) {
            FARF(ERROR, "HMX matmul failed for expert %u, error %d\n", cur_a, ret);
-            if (must_free_mapping) free(mapping_buf);
            return HTP_STATUS_NO_SUPPORT;
        }
    }

-    if (must_free_mapping) free(mapping_buf);
    return HTP_STATUS_OK;
 }

 static int hvx_mm_matmul_id(
    struct htp_ops_context * octx,
    struct htp_mm_context * mmctx,
-    size_t src0_row_size_padded,
-    uint32_t src1_nrows,
-    worker_callback_t matmul_id_job_func,
-    void * mapping_buf,
-    bool must_free_mapping
+    work_queue_func_t hvx_mmid_task_func
 ) {
    htp_matmul_tensors_preamble;
+    const uint32_t src0_row_size_padded = mmctx->src0_row_size_padded;
+    const uint32_t src1_nrows           = mmctx->src1_nrows;

    struct htp_thread_trace * tr = &octx->ctx->trace[0];
    htp_trace_event_start(tr, HTP_TRACE_EVT_INIT, 0);
@@ -3395,7 +3464,7 @@ static int hvx_mm_matmul_id(
    const uint32_t nb = (ne10 + qk - 1) / qk;
    const uint32_t total_nb = src1_nrows * nb;

-    worker_callback_t quant_task_func;
+    work_queue_func_t quant_task_func;
    uint32_t n_quant_tasks = 1;
    if (src1_nrows < octx->n_threads) {
        n_quant_tasks = MIN(total_nb, octx->n_threads);
@@ -3416,7 +3485,7 @@ static int hvx_mm_matmul_id(

    struct htp_mm_hvx_vtcm_layout L;
    htp_mm_hvx_vtcm_layout_build(&L, kparams->kernel_type, src0->type, ne10, src1_nrows, octx->n_threads,
-                                 0, src0_row_size, src1_row_size, kparams->n_prefetch, true, false, false);
+                                 0, src0_row_size, src1_row_size, 0, kparams->n_prefetch, true, false, false);

    size_t vtcm_size = kparams->vtcm_size > 0 ? (size_t)kparams->vtcm_size : L.total_bytes;

@@ -3431,7 +3500,6 @@ static int hvx_mm_matmul_id(
    // Make sure the reserved vtcm size is sufficient
    if (octx->ctx->vtcm_size < vtcm_size) {
        FARF(ERROR, "matmul-id-%s : current VTCM reservation %zu is too small, needed %zu\n", mmctx->type, octx->ctx->vtcm_size, vtcm_size);
-        if (must_free_mapping) free(mapping_buf);
        return HTP_STATUS_VTCM_TOO_SMALL;
    }

@@ -3461,12 +3529,78 @@ static int hvx_mm_matmul_id(

    htp_trace_event_stop(tr, HTP_TRACE_EVT_INIT, 0);

-    worker_pool_run_func(octx->ctx->worker_pool, matmul_id_job_func, mmctx, octx->n_threads);
+    worker_pool_run_func(octx->ctx->worker_pool, hvx_mmid_task_func, mmctx, octx->n_threads);

-    if (must_free_mapping) free(mapping_buf);
    return HTP_STATUS_OK;
 }

+static inline void scan_expert_ids_n(
+    const struct htp_tensor * ids,
+    const uint32_t n_ids,
+    uint32_t n_as,
+    uint32_t * counts,
+    struct mmid_row_mapping * matrix_rows,
+    uint32_t mapping_stride
+) {
+    const size_t ids_nb1 = ids->nb[1];
+    const uint8_t * ids_data = (const uint8_t *) ids->data;
+
+    for (uint32_t iid1 = 0; iid1 < ids->ne[1]; ++iid1) {
+        const int32_t * row_ptr = (const int32_t *) (ids_data + iid1 * ids_nb1);
+        for (uint32_t id = 0; id < n_ids; ++id) {
+            const int32_t i02 = row_ptr[id];
+            if (i02 < 0) {
+                continue;
+            }
+            assert(i02 < n_as);
+
+            if (matrix_rows) {
+                matrix_rows[i02 * mapping_stride + counts[i02]] = (struct mmid_row_mapping) { id, iid1 };
+            }
+            counts[i02] += 1;
+        }
+    }
+}
+
+static inline void scan_expert_ids(
+    const struct htp_tensor * ids,
+    uint32_t n_ids,
+    uint32_t n_as,
+    uint32_t * counts,
+    struct mmid_row_mapping * matrix_rows,
+    uint32_t mapping_stride
+) {
+    const size_t ids_nb0 = ids->nb[0];
+
+    if (ids_nb0 == 4) {
+        switch (n_ids) {
+            case 8:  scan_expert_ids_n(ids, 8,     n_as, counts, matrix_rows, mapping_stride); break;
+            case 4:  scan_expert_ids_n(ids, 4,     n_as, counts, matrix_rows, mapping_stride); break;
+            case 2:  scan_expert_ids_n(ids, 2,     n_as, counts, matrix_rows, mapping_stride); break;
+            default: scan_expert_ids_n(ids, n_ids, n_as, counts, matrix_rows, mapping_stride); break;
+        }
+    } else {
+        // Strided fallback
+        const size_t ids_nb1 = ids->nb[1];
+        const uint8_t * ids_data = (const uint8_t *) ids->data;
+        for (uint32_t iid1 = 0; iid1 < ids->ne[1]; ++iid1) {
+            const int32_t * row_ptr = (const int32_t *) (ids_data + iid1 * ids_nb1);
+            for (uint32_t id = 0; id < n_ids; ++id) {
+                const int32_t i02 = *(const int32_t *) ((const uint8_t *) row_ptr + id * ids_nb0);
+                if (i02 < 0) {
+                    continue;
+                }
+                assert(i02 < n_as);
+
+                if (matrix_rows) {
+                    matrix_rows[i02 * mapping_stride + counts[i02]] = (struct mmid_row_mapping) { id, iid1 };
+                }
+                counts[i02] += 1;
+            }
+        }
+    }
+}
+
 int op_matmul_id(struct htp_ops_context * octx) {
    htp_matmul_tensors_preamble;

@@ -3489,74 +3623,72 @@ int op_matmul_id(struct htp_ops_context * octx) {
    const uint32_t src0_nrows = ne01;  // per expert
    const uint32_t src1_nrows = ne11 * ne12 * ne13;

-    worker_callback_t quant_task_func;
-    worker_callback_t matmul_id_job_func = src1_nrows > 1 ? hvx_mm_id : hvx_mv_id;
-
-    // Compute src0_nrows_per_thread
-    mmctx->src0_nrows_per_thread  = (src0_nrows + octx->n_threads - 1) / octx->n_threads;
-    mmctx->src0_nrows_per_thread  = hex_round_up(mmctx->src0_nrows_per_thread, 32);
+    mmctx->src0_nrows_per_thread = (src0_nrows + octx->n_threads - 1) / octx->n_threads;
+    mmctx->src0_nrows_per_thread = hex_round_up(mmctx->src0_nrows_per_thread, 32);

    // row groups
    const int n_ids = ids->ne[0];  // n_expert_used
    const int n_as  = ne02;        // n_expert

-    size_t matrix_row_counts_size = n_as * sizeof(uint32_t);
-    size_t matrix_row_map_size    = n_as * ids->ne[0] * ids->ne[1] * sizeof(struct mmid_row_mapping);
-    const size_t total_map_size   = matrix_row_counts_size + matrix_row_map_size;
-
-    void * mapping_buf = NULL;
-    bool must_free_mapping = false;
-
-    if (octx->ctx->ddr_spad_base && total_map_size <= octx->ctx->ddr_spad_size) {
-        mapping_buf = octx->ctx->ddr_spad_base;
-    } else {
-        mapping_buf = memalign(128, total_map_size);
-        if (mapping_buf) {
-            must_free_mapping = true;
-        } else {
-            return HTP_STATUS_INTERNAL_ERR;
-        }
-    }
-
-    uint32_t *                matrix_row_counts = (uint32_t *) mapping_buf;
-    struct mmid_row_mapping * matrix_rows       = (struct mmid_row_mapping *) ((uint8_t *) mapping_buf + matrix_row_counts_size);
-
-    mmctx->matrix_row_counts = matrix_row_counts;
-    mmctx->matrix_rows       = matrix_rows;
-    mmctx->mm_div_ne11       = kparams->div_ne11;
-
-    if (hvx_mm_init_vec_dot(mmctx, src0->type) != 0) {
-        if (must_free_mapping) free(mapping_buf);
-        return HTP_STATUS_NO_SUPPORT;
-    }
+    uint8_t  * mapping_buf       = octx->ctx->ddr_spad_base;
+    uint32_t   mapping_stride    = 1;
+    uint32_t * matrix_row_counts = (uint32_t *) mapping_buf;
+    struct mmid_row_mapping * matrix_rows = NULL;

    if (src1_nrows > 1) {
-        // initialize matrix_row_counts and map
-        memset(matrix_row_counts, 0, n_as * sizeof(uint32_t));
+        const size_t matrix_row_counts_size = n_as * sizeof(uint32_t);
+        assert(octx->ctx->ddr_spad_size >= matrix_row_counts_size);

-        // group rows by src0 matrix
-        for (uint32_t iid1 = 0; iid1 < ids->ne[1]; ++iid1) {  // token idx
-            for (uint32_t id = 0; id < n_ids; ++id) {         // expert idx
-                const int32_t i02 = *(const int32_t *) ((const uint8_t *) ids->data + iid1 * ids->nb[1] + id * ids->nb[0]);
+        hex_l2fetch_block((const void *) ids->data, ids->ne[1] * ids->nb[1]);

-                if (i02 < 0) {
-                    continue;
-                }
-                assert(i02 < n_as);
+        memset(matrix_row_counts, 0, matrix_row_counts_size);
+        scan_expert_ids(ids, n_ids, n_as, matrix_row_counts, NULL, 0);

-                matrix_rows[i02 * n_ids * ids->ne[1] + matrix_row_counts[i02]] = (struct mmid_row_mapping) { id, iid1 };
-                matrix_row_counts[i02] += 1;
+        uint32_t max_count = hvx_reduce_max_i32((const uint8_t *) matrix_row_counts, n_as);
+        mapping_stride = max_count > 0 ? max_count : 1;
+
+        size_t matrix_row_map_size  = n_as * mapping_stride * sizeof(struct mmid_row_mapping);
+        const size_t total_map_size = matrix_row_counts_size + matrix_row_map_size;
+
+        if (total_map_size > octx->ctx->ddr_spad_size) {
+            mapping_buf = memalign(128, total_map_size);
+            if (!mapping_buf) {
+                return HTP_STATUS_INTERNAL_ERR;
            }
        }
+
+        matrix_row_counts = (uint32_t *) mapping_buf;
+        matrix_rows       = (struct mmid_row_mapping *) (mapping_buf + matrix_row_counts_size);
+
+        memset(matrix_row_counts, 0, n_as * sizeof(uint32_t));
+        scan_expert_ids(ids, n_ids, n_as, matrix_row_counts, matrix_rows, mapping_stride);
    }

+    mmctx->matrix_row_counts    = matrix_row_counts;
+    mmctx->matrix_rows          = matrix_rows;
+    mmctx->mapping_stride       = mapping_stride;
+    mmctx->mm_div_ne11          = kparams->div_ne11;
+    mmctx->src0_row_size_padded = src0_row_size_padded;
+    mmctx->src1_nrows           = src1_nrows;
+
    htp_trace_event_stop(tr, HTP_TRACE_EVT_INIT, 0);

+    int s;
    if (kparams->n_hmx) {
-        return hmx_mm_op_matmul_id(octx, mmctx, matrix_row_counts, matrix_rows, mapping_buf, must_free_mapping);
+        s = hmx_mm_op_matmul_id(octx, mmctx);
+    } else {
+        if (hvx_mm_init_vec_dot(mmctx, src0->type) == 0) {
+            s = hvx_mm_matmul_id(octx, mmctx, src1_nrows > 1 ? hvx_mm_id : hvx_mv_id);
+        } else {
+            s = HTP_STATUS_NO_SUPPORT;
+        }
    }

-    return hvx_mm_matmul_id(octx, mmctx, src0_row_size_padded, src1_nrows, matmul_id_job_func, mapping_buf, must_free_mapping);
+    if (mapping_buf != octx->ctx->ddr_spad_base) {
+        free(mapping_buf);
+    }
+
+    return s;
 }

 int op_matmul_qkv(struct htp_ops_context * octx) {
@@ -3633,7 +3765,7 @@ int op_matmul_qkv(struct htp_ops_context * octx) {

    struct htp_mm_hvx_vtcm_layout L;
    htp_mm_hvx_vtcm_layout_build(&L, kparams->kernel_type, src0->type, src1->ne[0], src1_nrows, octx->n_threads,
-                                 0, src0_row_size, src1_row_size, kparams->n_prefetch, false, true, false);
+                                 0, src0_row_size, src1_row_size, 0, kparams->n_prefetch, false, true, false);

    size_t vtcm_size = kparams->vtcm_size > 0 ? (size_t)kparams->vtcm_size : L.total_bytes;

@@ -3778,7 +3910,7 @@ int op_matmul_ffn(struct htp_ops_context * octx) {

    struct htp_mm_hvx_vtcm_layout L;
    htp_mm_hvx_vtcm_layout_build(&L, kparams->kernel_type, src0->type, src1->ne[0], src1_nrows, octx->n_threads,
-                                 0, src0_row_size, src1_row_size, kparams->n_prefetch, false, false, true);
+                                 0, src0_row_size, src1_row_size, 0, kparams->n_prefetch, false, false, true);

    size_t vtcm_size = kparams->vtcm_size > 0 ? (size_t)kparams->vtcm_size : L.total_bytes;

@@ -460,6 +460,7 @@ static inline void htp_mm_hvx_vtcm_layout_build(
    size_t dst_row_size,
    size_t src0_row_size,
    size_t src1_row_size,
+    size_t src2_row_size,
    uint32_t n_prefetch,
    bool is_matmul_id,
    bool is_fused_qkv,
@@ -467,7 +468,7 @@ static inline void htp_mm_hvx_vtcm_layout_build(
 ) {
    size_t src0_sz = 0;
    size_t src1_sz = 0;
-    size_t src2_sz = 0;
+    size_t src2_sz = src2_row_size > 0 ? htp_mm_round_up(src2_row_size, 128) : 0;
    size_t src3_sz = 0;
    size_t dst_sz  = 0;

@@ -138,6 +138,24 @@ static void scale_f32(const float * restrict src,
    }
 }

+static void clamp_f32(const float * restrict src,
+                      float * restrict dst,
+                      const uint32_t num_rows,
+                      const struct htp_unary_context * uctx) {
+    htp_unary_op_preamble;
+    float min = 0.f;
+    float max = 0.f;
+    memcpy(&min, &op_params[0], sizeof(float));
+    memcpy(&max, &op_params[1], sizeof(float));
+
+    for (uint32_t ir = 0; ir < num_rows; ir++) {
+        const uint8_t * restrict src_local = (const uint8_t *)src + (ir * src0_row_size_aligned);
+        uint8_t * restrict dst_local       = (uint8_t *)dst + (ir * dst_row_size_aligned);
+
+        hvx_clamp_scalar_f32(dst_local, src_local, min, max, ne0);
+    }
+}
+
 static void rms_norm_f32(const float * restrict src,
                         float * restrict dst,
                         const uint32_t num_rows,
@@ -258,6 +276,39 @@ static void sigmoid_f32(const float * restrict src,
    }
 }

+// silu(x) = x * sigmoid(x)
+static void silu_f32(const float * restrict src,
+                     float * restrict dst,
+                     const uint32_t num_rows,
+                     const struct htp_unary_context * uctx) {
+    htp_unary_op_preamble;
+
+    for (uint32_t ir = 0; ir < num_rows; ir++) {
+        const uint8_t * restrict src_local = (const uint8_t *)src + (ir * src0_row_size_aligned);
+        uint8_t * restrict dst_local       = (uint8_t *)dst + (ir * dst_row_size_aligned);
+
+        hvx_sigmoid_f32_aa(dst_local, src_local, ne0);
+        hvx_mul_f32_aaa(dst_local, src_local, dst_local, ne0);
+    }
+}
+
+// gelu(x) = x * sigmoid(1.702 * x)  (quick/sigmoid approximation, matches CPU GELU_QUICK reference)
+static void gelu_f32(const float * restrict src,
+                     float * restrict dst,
+                     const uint32_t num_rows,
+                     const struct htp_unary_context * uctx) {
+    htp_unary_op_preamble;
+
+    for (uint32_t ir = 0; ir < num_rows; ir++) {
+        const uint8_t * restrict src_local = (const uint8_t *)src + (ir * src0_row_size_aligned);
+        uint8_t * restrict dst_local       = (uint8_t *)dst + (ir * dst_row_size_aligned);
+
+        hvx_mul_scalar_f32(dst_local, src_local, 1.702f, ne0);
+        hvx_sigmoid_f32_aa(dst_local, dst_local, ne0);
+        hvx_mul_f32_aaa(dst_local, src_local, dst_local, ne0);
+    }
+}
+
 static void tri_f32(const float * restrict src,
                    float * restrict dst,
                    const uint32_t num_rows,
@@ -542,11 +593,14 @@ DEFINE_UNARY_TASK(norm,           false, false, norm_f32(src0_vtcm, dst_vtcm, bl
 DEFINE_UNARY_TASK(rms_norm,       false, false, rms_norm_f32(src0_vtcm, dst_vtcm, block_size, uctx))
 DEFINE_UNARY_TASK(rms_norm_mul,   true,  false, rms_norm_mul_f32(src0_vtcm, uctx->broadcast_weight ? (const float *) src1_vtcm_data : src1_vtcm, dst_vtcm, block_size, uctx))
 DEFINE_UNARY_TASK(scale,          false, false, scale_f32(src0_vtcm, dst_vtcm, block_size, uctx))
+DEFINE_UNARY_TASK(clamp,          false, false, clamp_f32(src0_vtcm, dst_vtcm, block_size, uctx))
 DEFINE_UNARY_TASK(sqr,            false, false, sqr_f32(src0_vtcm, dst_vtcm, block_size, uctx))
 DEFINE_UNARY_TASK(sqrt,           false, false, sqrt_f32(src0_vtcm, dst_vtcm, block_size, uctx))
 DEFINE_UNARY_TASK(unary_neg,      false, false, neg_f32(src0_vtcm, dst_vtcm, block_size, uctx))
 DEFINE_UNARY_TASK(unary_exp,      false, false, exp_f32(src0_vtcm, dst_vtcm, block_size, uctx))
 DEFINE_UNARY_TASK(unary_sigmoid,  false, false, sigmoid_f32(src0_vtcm, dst_vtcm, block_size, uctx))
+DEFINE_UNARY_TASK(unary_silu,     false, false, silu_f32(src0_vtcm, dst_vtcm, block_size, uctx))
+DEFINE_UNARY_TASK(unary_gelu,     false, false, gelu_f32(src0_vtcm, dst_vtcm, block_size, uctx))
 DEFINE_UNARY_TASK(unary_softplus, false, false, softplus_f32(src0_vtcm, dst_vtcm, block_size, uctx))
 DEFINE_UNARY_TASK(unary_tanh,     false, false, tanh_f32(src0_vtcm, dst_vtcm, block_size, uctx))
 DEFINE_UNARY_TASK(l2_norm,        false, false, l2_norm_f32(src0_vtcm, dst_vtcm, block_size, uctx))
@@ -681,6 +735,14 @@ static inline void tile_scale_f32(uint8_t * dst_vtcm, const uint8_t * src_vtcm,
    hvx_scale_offset_f32_aa(dst_vtcm, src_vtcm, tw, scale, bias);
 }

+static inline void tile_clamp_f32(uint8_t * dst_vtcm, const uint8_t * src_vtcm, uint32_t tw, const int32_t * op_params) {
+    float min = 0.f;
+    float max = 0.f;
+    memcpy(&min, &op_params[0], sizeof(float));
+    memcpy(&max, &op_params[1], sizeof(float));
+    hvx_clamp_scalar_f32(dst_vtcm, src_vtcm, min, max, tw);
+}
+
 static inline void tile_unary_softplus_f32(uint8_t * dst_vtcm, const uint8_t * src_vtcm, uint32_t tw) {
    const float * restrict sf = (const float *) src_vtcm;
    float * restrict df       = (float *) dst_vtcm;
@@ -690,6 +752,19 @@ static inline void tile_unary_softplus_f32(uint8_t * dst_vtcm, const uint8_t * s
    }
 }

+// silu(x) = x * sigmoid(x)
+static inline void tile_silu_f32(uint8_t * dst_vtcm, const uint8_t * src_vtcm, uint32_t tw) {
+    hvx_sigmoid_f32_aa(dst_vtcm, src_vtcm, tw);
+    hvx_mul_f32_aaa(dst_vtcm, src_vtcm, dst_vtcm, tw);
+}
+
+// gelu(x) = x * sigmoid(1.702 * x)  (quick/sigmoid approximation, matches CPU GELU_QUICK reference)
+static inline void tile_gelu_f32(uint8_t * dst_vtcm, const uint8_t * src_vtcm, uint32_t tw) {
+    hvx_mul_scalar_f32(dst_vtcm, src_vtcm, 1.702f, tw);
+    hvx_sigmoid_f32_aa(dst_vtcm, dst_vtcm, tw);
+    hvx_mul_f32_aaa(dst_vtcm, src_vtcm, dst_vtcm, tw);
+}
+
 // Triangular mask applied to one column tile. Boundary is an absolute column index, so
 // each vector compares against its absolute column position (col_start + i*VLEN_FP32).
 static inline void tri_apply_tile_f32(const uint8_t * restrict src, uint8_t * restrict dst,
@@ -765,11 +840,14 @@ static inline void tri_apply_tile_f32(const uint8_t * restrict src, uint8_t * re
 }

 DEFINE_UNARY_TILED_TASK(scale,          false, tile_scale_f32(dst_vtcm, src_vtcm, tw, op_params))
+DEFINE_UNARY_TILED_TASK(clamp,          false, tile_clamp_f32(dst_vtcm, src_vtcm, tw, op_params))
 DEFINE_UNARY_TILED_TASK(sqr,            false, hvx_sqr_f32_aa(dst_vtcm, src_vtcm, tw))
 DEFINE_UNARY_TILED_TASK(sqrt,           false, hvx_sqrt_f32_aa(dst_vtcm, src_vtcm, tw))
 DEFINE_UNARY_TILED_TASK(unary_neg,      false, hvx_scale_f32_aa(dst_vtcm, src_vtcm, tw, -1.0f))
 DEFINE_UNARY_TILED_TASK(unary_exp,      false, hvx_exp_f32(dst_vtcm, src_vtcm, tw, false))
 DEFINE_UNARY_TILED_TASK(unary_sigmoid,  false, hvx_sigmoid_f32_aa(dst_vtcm, src_vtcm, tw))
+DEFINE_UNARY_TILED_TASK(unary_silu,     false, tile_silu_f32(dst_vtcm, src_vtcm, tw))
+DEFINE_UNARY_TILED_TASK(unary_gelu,     false, tile_gelu_f32(dst_vtcm, src_vtcm, tw))
 DEFINE_UNARY_TILED_TASK(unary_softplus, false, tile_unary_softplus_f32(dst_vtcm, src_vtcm, tw))
 DEFINE_UNARY_TILED_TASK(unary_tanh,     false, hvx_tanh_f32_aa(dst_vtcm, src_vtcm, tw))
 DEFINE_UNARY_TILED_TASK(tri,            true,  tri_apply_tile_f32(src_vtcm, dst_vtcm, tw, col, i01, ne0, tri_ttype))
@@ -787,11 +865,14 @@ static int execute_op_unary_f32(struct htp_ops_context * octx) {
        case HTP_OP_RMS_NORM:        op_type = "rmsnorm-f32";      break;
        case HTP_OP_RMS_NORM_MUL:    op_type = "rmsnorm-mul-f32";  break;
        case HTP_OP_SCALE:           op_type = "scale-f32";        break;
+        case HTP_OP_CLAMP:           op_type = "clamp-f32";        break;
        case HTP_OP_SQR:             op_type = "sqr-f32";          break;
        case HTP_OP_SQRT:            op_type = "sqrt-f32";         break;
        case HTP_OP_UNARY_NEG:       op_type = "neg-f32";          break;
        case HTP_OP_UNARY_EXP:       op_type = "exp-f32";          break;
        case HTP_OP_UNARY_SIGMOID:   op_type = "sigmoid-f32";      break;
+        case HTP_OP_UNARY_SILU:      op_type = "silu-f32";         break;
+        case HTP_OP_UNARY_GELU:      op_type = "gelu-f32";         break;
        case HTP_OP_UNARY_SOFTPLUS:  op_type = "softplus-f32";     break;
        case HTP_OP_UNARY_TANH:      op_type = "tanh-f32";         break;
        case HTP_OP_L2_NORM:         op_type = "l2norm-f32";       break;
@@ -882,11 +963,14 @@ static int execute_op_unary_f32(struct htp_ops_context * octx) {
        if (col_tile) {
            switch (octx->op) {
                case HTP_OP_SCALE:           task_func = unary_task_f32_tiled_scale;          break;
+                case HTP_OP_CLAMP:           task_func = unary_task_f32_tiled_clamp;          break;
                case HTP_OP_SQR:             task_func = unary_task_f32_tiled_sqr;            break;
                case HTP_OP_SQRT:            task_func = unary_task_f32_tiled_sqrt;           break;
                case HTP_OP_UNARY_NEG:       task_func = unary_task_f32_tiled_unary_neg;      break;
                case HTP_OP_UNARY_EXP:       task_func = unary_task_f32_tiled_unary_exp;      break;
                case HTP_OP_UNARY_SIGMOID:   task_func = unary_task_f32_tiled_unary_sigmoid;  break;
+                case HTP_OP_UNARY_SILU:      task_func = unary_task_f32_tiled_unary_silu;     break;
+                case HTP_OP_UNARY_GELU:      task_func = unary_task_f32_tiled_unary_gelu;     break;
                case HTP_OP_UNARY_SOFTPLUS:  task_func = unary_task_f32_tiled_unary_softplus; break;
                case HTP_OP_UNARY_TANH:      task_func = unary_task_f32_tiled_unary_tanh;     break;
                case HTP_OP_TRI:             task_func = unary_task_f32_tiled_tri;            break;
@@ -898,11 +982,14 @@ static int execute_op_unary_f32(struct htp_ops_context * octx) {
                case HTP_OP_RMS_NORM:        task_func = unary_task_f32_rms_norm;             break;
                case HTP_OP_RMS_NORM_MUL:    task_func = unary_task_f32_rms_norm_mul;         break;
                case HTP_OP_SCALE:           task_func = unary_task_f32_scale;                break;
+                case HTP_OP_CLAMP:           task_func = unary_task_f32_clamp;                break;
                case HTP_OP_SQR:             task_func = unary_task_f32_sqr;                  break;
                case HTP_OP_SQRT:            task_func = unary_task_f32_sqrt;                 break;
                case HTP_OP_UNARY_NEG:       task_func = unary_task_f32_unary_neg;            break;
                case HTP_OP_UNARY_EXP:       task_func = unary_task_f32_unary_exp;            break;
                case HTP_OP_UNARY_SIGMOID:   task_func = unary_task_f32_unary_sigmoid;        break;
+                case HTP_OP_UNARY_SILU:      task_func = unary_task_f32_unary_silu;           break;
+                case HTP_OP_UNARY_GELU:      task_func = unary_task_f32_unary_gelu;           break;
                case HTP_OP_UNARY_SOFTPLUS:  task_func = unary_task_f32_unary_softplus;       break;
                case HTP_OP_UNARY_TANH:      task_func = unary_task_f32_unary_tanh;           break;
                case HTP_OP_L2_NORM:         task_func = unary_task_f32_l2_norm;              break;
@@ -41,6 +41,7 @@ _Static_assert(sizeof(struct htp_unary_kernel_params) <= 128, "htp_unary_kernel_

 static inline bool htp_op_is_unary(uint32_t opcode) {
    switch (opcode) {
+        case HTP_OP_CLAMP:
        case HTP_OP_NORM:
        case HTP_OP_RMS_NORM:
        case HTP_OP_RMS_NORM_MUL:
@@ -50,6 +51,8 @@ static inline bool htp_op_is_unary(uint32_t opcode) {
        case HTP_OP_UNARY_NEG:
        case HTP_OP_UNARY_EXP:
        case HTP_OP_UNARY_SIGMOID:
+        case HTP_OP_UNARY_SILU:
+        case HTP_OP_UNARY_GELU:
        case HTP_OP_UNARY_SOFTPLUS:
        case HTP_OP_UNARY_TANH:
        case HTP_OP_L2_NORM:
@@ -36,6 +36,32 @@ void ggml_metal_set_abort_callback  (ggml_metal_t ctx, ggml_abort_callback abort
 bool ggml_metal_supports_family     (ggml_metal_t ctx, int family);
 void ggml_metal_capture_next_compute(ggml_metal_t ctx);

+//
+// Profiling
+//
+
+// Opaque profiler state, owned by the C++ backend layer (ggml-metal.cpp). Holds the std::vector of
+// records returned via the ggml_backend_profiler interface. ggml_metal_t keeps a borrowed pointer
+// so that graph_compute can push records when sampling is active.
+struct ggml_metal_profiler_state;
+
+// Inject (or clear, with NULL) the profiler state pointer. Called once at backend init.
+void ggml_metal_set_profiler_state(ggml_metal_t ctx, struct ggml_metal_profiler_state * state);
+
+// Bridge function implemented in ggml-metal.cpp. Used by graph_compute (in .m) to push records.
+void ggml_metal_profiler_push_record(
+        struct ggml_metal_profiler_state * state,
+        const struct ggml_tensor * node,
+        uint64_t start_ns,
+        uint64_t end_ns);
+
+// Query whether the injected profiler state is currently enabled.
+// (Avoids exposing the C++ struct layout to the .m file.)
+bool ggml_metal_profiler_is_enabled(struct ggml_metal_profiler_state * state);
+
+// Query the split-id currently set on the profiler state.
+int  ggml_metal_profiler_get_split_id(struct ggml_metal_profiler_state * state);
+
 #ifdef __cplusplus
 }
 #endif
@@ -6,6 +6,7 @@
 #import "ggml-metal-impl.h"
 #import "ggml-metal-common.h"
 #import "ggml-metal-ops.h"
+#import "ggml-profiler.h"

 #import <Foundation/Foundation.h>

@@ -79,6 +80,19 @@ struct ggml_metal {
    // error state - set when a command buffer fails during synchronize
    // once set, graph_compute will return GGML_STATUS_FAILED until the backend is recreated
    bool has_error;
+
+    // Profiling
+    // Borrowed; owned by the C++ backend layer (ggml-metal.cpp). NULL when profiling is unavailable.
+    struct ggml_metal_profiler_state * profiler_state;
+
+    // Per-graph-compute scratch state populated when profiling is active for the current invocation.
+    // Lifetime: from start of ggml_metal_graph_compute until its end.
+    bool                              profiler_active;
+    ggml_metal_sample_buf_t           profiler_sample_buf;
+    size_t                            profiler_total_slots;
+    uint64_t                          profiler_cpu_anchor_ns;
+    uint64_t                          profiler_gpu_anchor_ns;
+    struct ggml_metal_op_sample_slot * profiler_slot_map;  // size = gf->n_nodes
 };

 ggml_metal_t ggml_metal_init(ggml_metal_device_t dev) {
@@ -450,6 +464,34 @@ enum ggml_status ggml_metal_graph_compute(ggml_metal_t ctx, struct ggml_cgraph *
    // keep the memory wired
    ggml_metal_device_rsets_keep_alive(ctx->dev);

+    // Decide whether profiling is active for this invocation. Activation is sticky for the whole
+    // call: allocate sample buffer + slot map up front so the encode_async block can use them.
+    ctx->profiler_active = false;
+    if (ctx->profiler_state != NULL && ggml_metal_profiler_is_enabled(ctx->profiler_state)) {
+        if (ggml_metal_device_supports_profiling(ctx->dev) && gf->n_nodes > 0) {
+            const size_t total_slots = 2 * (size_t) gf->n_nodes;
+            ctx->profiler_sample_buf = ggml_metal_device_create_sample_buf(ctx->dev, total_slots);
+            if (ctx->profiler_sample_buf != NULL) {
+                ctx->profiler_total_slots = total_slots;
+                ctx->profiler_slot_map = (struct ggml_metal_op_sample_slot *) calloc(
+                    (size_t) gf->n_nodes, sizeof(struct ggml_metal_op_sample_slot));
+                if (ctx->profiler_slot_map != NULL) {
+                    // Mark all entries unused (node_idx < 0).
+                    for (int i = 0; i < gf->n_nodes; ++i) {
+                        ctx->profiler_slot_map[i].node_idx = -1;
+                    }
+                    ggml_metal_device_sample_timestamps(ctx->dev,
+                                                       &ctx->profiler_cpu_anchor_ns,
+                                                       &ctx->profiler_gpu_anchor_ns);
+                    ctx->profiler_active = true;
+                } else {
+                    ggml_metal_sample_buf_free(ctx->profiler_sample_buf);
+                    ctx->profiler_sample_buf = NULL;
+                }
+            }
+        }
+    }
+
    // submit the ggml compute graph to the GPU by creating command buffers and encoding the ops in them
    // the first n_nodes_0 are encoded and submitted for processing directly by the calling thread
    // while these nodes are processing, we start n_cb threads to enqueue the rest of the nodes
@@ -609,6 +651,66 @@ enum ggml_status ggml_metal_graph_compute(ggml_metal_t ctx, struct ggml_cgraph *

            ctx->capture_started = false;
        }
+
+        // Profiling drain: wait for all command buffers to complete, resolve timestamps, push records.
+        // This forces a synchronous wait — Vulkan does the same when its profiler is active.
+        if (ctx->profiler_active) {
+            {
+                id<MTLCommandBuffer> cmd_buf = ctx->cmd_bufs[n_cb].obj;
+                if (cmd_buf) {
+                    [cmd_buf waitUntilCompleted];
+                }
+            }
+            for (int i = 0; i < n_cb; ++i) {
+                id<MTLCommandBuffer> cmd_buf = ctx->cmd_bufs[i].obj;
+                if (cmd_buf) {
+                    // Ensure cmd_bufs that were not auto-enqueued get committed.
+                    if ([cmd_buf status] == MTLCommandBufferStatusNotEnqueued) {
+                        [cmd_buf commit];
+                    }
+                    [cmd_buf waitUntilCompleted];
+                }
+            }
+
+            uint64_t * ns = (uint64_t *) calloc(ctx->profiler_total_slots, sizeof(uint64_t));
+            if (ns != NULL) {
+                ggml_metal_sample_buf_resolve(ctx->profiler_sample_buf,
+                                              /*base=*/0,
+                                              ctx->profiler_total_slots,
+                                              ctx->profiler_cpu_anchor_ns,
+                                              ctx->profiler_gpu_anchor_ns,
+                                              ns);
+
+                for (int i = 0; i < gf->n_nodes; ++i) {
+                    const struct ggml_metal_op_sample_slot * slot = &ctx->profiler_slot_map[i];
+                    if (slot->node_idx < 0) {
+                        continue;
+                    }
+                    if (slot->slot_start >= ctx->profiler_total_slots ||
+                        slot->slot_end   >= ctx->profiler_total_slots) {
+                        continue;
+                    }
+                    const uint64_t t0 = ns[slot->slot_start];
+                    const uint64_t t1 = ns[slot->slot_end];
+                    if (t0 == 0 || t1 == 0 || t1 < t0) {
+                        continue;
+                    }
+                    ggml_metal_profiler_push_record(ctx->profiler_state,
+                                                    ggml_graph_node(gf, slot->node_idx),
+                                                    t0,
+                                                    t1);
+                }
+
+                free(ns);
+            }
+
+            free(ctx->profiler_slot_map);
+            ctx->profiler_slot_map = NULL;
+            ggml_metal_sample_buf_free(ctx->profiler_sample_buf);
+            ctx->profiler_sample_buf = NULL;
+            ctx->profiler_total_slots = 0;
+            ctx->profiler_active = false;
+        }
    }

    return GGML_STATUS_SUCCESS;
@@ -660,6 +762,10 @@ ggml_metal_event_t ggml_metal_get_ev_cpy(ggml_metal_t ctx) {
    return ctx->ev_cpy;
 }

+void ggml_metal_set_profiler_state(ggml_metal_t ctx, struct ggml_metal_profiler_state * state) {
+    ctx->profiler_state = state;
+}
+
 void ggml_metal_set_n_cb(ggml_metal_t ctx, int n_cb) {
    if (ctx->n_cb != n_cb) {
        ctx->n_cb = MIN(n_cb, GGML_METAL_MAX_COMMAND_BUFFERS);
@@ -704,6 +810,17 @@ void ggml_metal_set_n_cb(ggml_metal_t ctx, int n_cb) {
            ctx->debug_graph,
            ctx->debug_fusion);

+        if (ctx->profiler_active) {
+            // Base slot for this command buffer is 2 * (first graph-node index it processes).
+            // The encoder uses one (start, end) pair per encoded op group; we over-allocate so
+            // that empty/filtered nodes simply leave gaps in the sample buffer.
+            const size_t base_slot = 2 * (size_t) idx_start;
+            ggml_metal_op_enable_profiling(ctx_op,
+                                           ctx->profiler_sample_buf,
+                                           base_slot,
+                                           ctx->profiler_slot_map);
+        }
+
        for (int idx = 0; idx < ggml_metal_op_n_nodes(ctx_op); ++idx) {
            const int res = ggml_metal_op_encode(ctx_op, idx);
            if (res == 0) {
@@ -91,6 +91,17 @@ void ggml_metal_encoder_memory_barrier(ggml_metal_encoder_t encoder);

 void ggml_metal_encoder_end_encoding(ggml_metal_encoder_t encoder);

+//
+// MTLCounterSampleBuffer wrapper (used by the profiler)
+//
+
+typedef struct ggml_metal_sample_buf * ggml_metal_sample_buf_t;
+
+// Insert a GPU timestamp sample on the encoder at the given slot.
+// Caller must ensure (a) the encoder belongs to a command buffer using a counter sample buffer that
+// supports MTLCounterSamplingPointAtDispatchBoundary, and (b) `index` is unique within the buffer.
+void ggml_metal_encoder_sample_timestamp(ggml_metal_encoder_t encoder, ggml_metal_sample_buf_t buf, size_t index);
+
 //
 // MTLLibrary wrapper
 //
@@ -296,6 +307,35 @@ bool ggml_metal_device_supports_op(ggml_metal_device_t dev, const struct ggml_te

 const struct ggml_metal_device_props * ggml_metal_device_get_props(ggml_metal_device_t dev);

+//
+// Profiling helpers
+//
+
+// Returns true if the device supports MTLCounterSamplingPointAtDispatchBoundary on compute encoders
+// AND exposes the MTLCommonCounterSetTimestamp counter set.
+bool ggml_metal_device_supports_profiling(ggml_metal_device_t dev);
+
+// Allocate a counter sample buffer with `sample_count` slots backed by the timestamp counter set.
+// Returns NULL on failure (e.g. unsupported device).
+ggml_metal_sample_buf_t ggml_metal_device_create_sample_buf(ggml_metal_device_t dev, size_t sample_count);
+
+void ggml_metal_sample_buf_free(ggml_metal_sample_buf_t buf);
+
+// Capture correlated CPU/GPU timestamps (both in Mach-absolute units, i.e. nanoseconds on
+// current Apple Silicon; the conversion via mach_timebase_info is applied internally).
+// Used to anchor the GPU timestamp domain to the CPU clock returned by ggml_profiler_time_ns().
+void ggml_metal_device_sample_timestamps(ggml_metal_device_t dev, uint64_t * cpu_ns, uint64_t * gpu_ns);
+
+// Resolve `count` consecutive samples starting at `base` into nanosecond timestamps anchored against
+// `cpu_anchor_ns` / `gpu_anchor_ns` (obtained from ggml_metal_device_sample_timestamps).
+// `out_ns` must hold at least `count` uint64_t entries.
+void ggml_metal_sample_buf_resolve(ggml_metal_sample_buf_t buf,
+                                   size_t   base,
+                                   size_t   count,
+                                   uint64_t cpu_anchor_ns,
+                                   uint64_t gpu_anchor_ns,
+                                   uint64_t * out_ns);
+
 //
 // device buffers
 //
@@ -8,6 +8,7 @@
 #include <Metal/Metal.h>

 #include <stdatomic.h>
+#include <mach/mach_time.h>

 #ifndef TARGET_OS_VISION
 #define TARGET_OS_VISION 0
@@ -517,6 +518,21 @@ void ggml_metal_encoder_end_encoding(ggml_metal_encoder_t encoder) {
    [encoder->obj endEncoding];
 }

+//
+// MTLCounterSampleBuffer wrapper
+//
+
+struct ggml_metal_sample_buf {
+    id<MTLCounterSampleBuffer> obj;
+    size_t sample_count;
+};
+
+void ggml_metal_encoder_sample_timestamp(ggml_metal_encoder_t encoder, ggml_metal_sample_buf_t buf, size_t index) {
+    if (@available(macOS 11.0, iOS 14.0, *)) {
+        [encoder->obj sampleCountersInBuffer:buf->obj atSampleIndex:index withBarrier:YES];
+    }
+}
+
 struct ggml_metal_device {
    id<MTLDevice> mtl_device;

@@ -1218,8 +1234,9 @@ bool ggml_metal_device_supports_op(ggml_metal_device_t dev, const struct ggml_te
                   (ggml_get_op_params_i32(op, 4) == 0) && (ggml_get_op_params_i32(op, 6) == 0);
        case GGML_OP_PAD_REFLECT_1D:
        case GGML_OP_TIMESTEP_EMBEDDING:
-        case GGML_OP_LEAKY_RELU:
            return op->src[0]->type == GGML_TYPE_F32;
+        case GGML_OP_LEAKY_RELU:
+            return op->src[0]->type == GGML_TYPE_F32 || op->src[0]->type == GGML_TYPE_F16;
        case GGML_OP_ARGSORT:
        case GGML_OP_TOP_K:
        case GGML_OP_ARANGE:
@@ -1377,6 +1394,139 @@ const struct ggml_metal_device_props * ggml_metal_device_get_props(ggml_metal_de
    return &dev->props;
 }

+//
+// Profiling helpers
+//
+
+// Look up the MTLCommonCounterSetTimestamp counter set, or nil if the device doesn't expose it.
+static id<MTLCounterSet> ggml_metal_device_get_timestamp_counter_set(ggml_metal_device_t dev) {
+    if (@available(macOS 10.15, iOS 14.0, *)) {
+        NSArray<id<MTLCounterSet>> * sets = [dev->mtl_device counterSets];
+        for (id<MTLCounterSet> cs in sets) {
+            if ([[cs name] isEqualToString:MTLCommonCounterSetTimestamp]) {
+                return cs;
+            }
+        }
+    }
+    return nil;
+}
+
+bool ggml_metal_device_supports_profiling(ggml_metal_device_t dev) {
+    if (@available(macOS 11.0, iOS 14.0, *)) {
+        if (![dev->mtl_device supportsCounterSampling:MTLCounterSamplingPointAtDispatchBoundary]) {
+            return false;
+        }
+        return ggml_metal_device_get_timestamp_counter_set(dev) != nil;
+    }
+    return false;
+}
+
+ggml_metal_sample_buf_t ggml_metal_device_create_sample_buf(ggml_metal_device_t dev, size_t sample_count) {
+    if (sample_count == 0) {
+        return NULL;
+    }
+    if (@available(macOS 10.15, iOS 14.0, *)) {
+        id<MTLCounterSet> cs = ggml_metal_device_get_timestamp_counter_set(dev);
+        if (cs == nil) {
+            return NULL;
+        }
+
+        MTLCounterSampleBufferDescriptor * desc = [[MTLCounterSampleBufferDescriptor alloc] init];
+        desc.counterSet = cs;
+        desc.label = @"ggml-metal-profiler";
+        desc.storageMode = MTLStorageModeShared;
+        desc.sampleCount = sample_count;
+
+        NSError * error = nil;
+        id<MTLCounterSampleBuffer> sb = [dev->mtl_device newCounterSampleBufferWithDescriptor:desc error:&error];
+        [desc release];
+
+        if (sb == nil) {
+            GGML_LOG_WARN("%s: failed to create counter sample buffer: %s\n", __func__,
+                          error ? [[error localizedDescription] UTF8String] : "unknown");
+            return NULL;
+        }
+
+        ggml_metal_sample_buf_t res = calloc(1, sizeof(struct ggml_metal_sample_buf));
+        res->obj = sb;
+        res->sample_count = sample_count;
+        return res;
+    }
+    return NULL;
+}
+
+void ggml_metal_sample_buf_free(ggml_metal_sample_buf_t buf) {
+    if (buf == NULL) {
+        return;
+    }
+    [buf->obj release];
+    free(buf);
+}
+
+void ggml_metal_device_sample_timestamps(ggml_metal_device_t dev, uint64_t * cpu_ns, uint64_t * gpu_ns) {
+    if (@available(macOS 10.15, iOS 14.0, *)) {
+        MTLTimestamp cpu_t = 0;
+        MTLTimestamp gpu_t = 0;
+        [dev->mtl_device sampleTimestamps:&cpu_t gpuTimestamp:&gpu_t];
+
+        // Apple docs: both timestamps are in Mach-absolute time units. Convert to nanoseconds.
+        static mach_timebase_info_data_t tb = {0, 0};
+        if (tb.denom == 0) {
+            mach_timebase_info(&tb);
+        }
+        if (cpu_ns) *cpu_ns = (uint64_t)((__uint128_t)cpu_t * tb.numer / tb.denom);
+        if (gpu_ns) *gpu_ns = (uint64_t)((__uint128_t)gpu_t * tb.numer / tb.denom);
+    } else {
+        if (cpu_ns) *cpu_ns = 0;
+        if (gpu_ns) *gpu_ns = 0;
+    }
+}
+
+void ggml_metal_sample_buf_resolve(ggml_metal_sample_buf_t buf,
+                                   size_t   base,
+                                   size_t   count,
+                                   uint64_t cpu_anchor_ns,
+                                   uint64_t gpu_anchor_ns,
+                                   uint64_t * out_ns) {
+    if (buf == NULL || count == 0 || out_ns == NULL) {
+        return;
+    }
+    if (@available(macOS 10.15, iOS 14.0, *)) {
+        NSRange range = NSMakeRange(base, count);
+        NSData * data = [buf->obj resolveCounterRange:range];
+        if (data == nil || [data length] < count * sizeof(MTLCounterResultTimestamp)) {
+            // Failed resolve: fill with zeros so the caller can detect dropped samples.
+            for (size_t i = 0; i < count; ++i) {
+                out_ns[i] = 0;
+            }
+            return;
+        }
+
+        static mach_timebase_info_data_t tb = {0, 0};
+        if (tb.denom == 0) {
+            mach_timebase_info(&tb);
+        }
+
+        const MTLCounterResultTimestamp * src = (const MTLCounterResultTimestamp *) [data bytes];
+        for (size_t i = 0; i < count; ++i) {
+            // MTLCounterErrorValue (~0ULL) marks an unrecorded sample slot; map to 0.
+            if (src[i].timestamp == MTLCounterErrorValue) {
+                out_ns[i] = 0;
+                continue;
+            }
+            const uint64_t gpu_ns = (uint64_t)((__uint128_t) src[i].timestamp * tb.numer / tb.denom);
+            // Anchor to CPU clock returned by ggml_profiler_time_ns().
+            out_ns[i] = (gpu_ns >= gpu_anchor_ns)
+                      ? cpu_anchor_ns + (gpu_ns - gpu_anchor_ns)
+                      : cpu_anchor_ns - (gpu_anchor_ns - gpu_ns);
+        }
+    } else {
+        for (size_t i = 0; i < count; ++i) {
+            out_ns[i] = 0;
+        }
+    }
+}
+
 //
 // device buffers
 //
@@ -77,6 +77,11 @@ struct ggml_metal_op {
        return ggml_graph_node(gf, idxs[i]);
    }

+    int node_global_idx(int i) const {
+        assert(i >= 0 && i < (int) idxs.size());
+        return idxs[i];
+    }
+
    bool can_fuse(int i0, const ggml_op * ops, int n_ops) const {
        assert(use_fusion);
        assert(i0 >= 0 && i0 < n_nodes());
@@ -100,6 +105,12 @@ struct ggml_metal_op {
    int debug_graph;
    int debug_fusion;

+    // Profiling: when sample_buf is non-null, ggml_metal_op_encode brackets each impl call with
+    // two timestamp samples. The (node_idx, slot_start, slot_end) tuple is recorded in slot_map.
+    ggml_metal_sample_buf_t sample_buf      = nullptr;
+    size_t                  next_slot       = 0;
+    ggml_metal_op_sample_slot * slot_map    = nullptr;
+
 private:
    ggml_cgraph * gf;

@@ -144,6 +155,16 @@ int ggml_metal_op_n_nodes(ggml_metal_op_t ctx) {
    return ctx->n_nodes();
 }

+void ggml_metal_op_enable_profiling(
+        ggml_metal_op_t ctx,
+        ggml_metal_sample_buf_t sample_buf,
+        size_t slot_base,
+        ggml_metal_op_sample_slot * slot_map) {
+    ctx->sample_buf = sample_buf;
+    ctx->next_slot  = slot_base;
+    ctx->slot_map   = slot_map;
+}
+
 static bool ggml_metal_op_concurrency_reset(ggml_metal_op_t ctx) {
    if (!ctx->mem_ranges) {
        return true;
@@ -509,12 +530,31 @@ int ggml_metal_op_encode(ggml_metal_op_t ctx, int idx) {
        ggml_metal_encoder_debug_group_push(ctx->enc, ggml_op_desc(ctx->node(idx)));
    }

+    const size_t slot_start = ctx->next_slot;
+    const bool   profiling  = (ctx->sample_buf != nullptr && ctx->slot_map != nullptr);
+    if (profiling) {
+        ggml_metal_encoder_sample_timestamp(ctx->enc, ctx->sample_buf, slot_start);
+        ctx->next_slot++;
+    }
+
    int res = ggml_metal_op_encode_impl(ctx, idx);
    if (idx + res > ctx->n_nodes()) {
        GGML_ABORT("fusion error: nodes spanning multiple encoders have been fused. this indicates a bug in the fusion logic %s",
                "https://github.com/ggml-org/llama.cpp/pull/14849");
    }

+    if (profiling) {
+        const size_t slot_end = ctx->next_slot;
+        ggml_metal_encoder_sample_timestamp(ctx->enc, ctx->sample_buf, slot_end);
+        ctx->next_slot++;
+
+        const int gidx = ctx->node_global_idx(idx);
+        ctx->slot_map[gidx].node_idx   = gidx;
+        ctx->slot_map[gidx].n_fused    = res;
+        ctx->slot_map[gidx].slot_start = slot_start;
+        ctx->slot_map[gidx].slot_end   = slot_end;
+    }
+
    if (ctx->use_capture) {
        ggml_metal_encoder_debug_group_pop(ctx->enc);
    }
@@ -26,6 +26,29 @@ int ggml_metal_op_n_nodes(ggml_metal_op_t ctx);

 int ggml_metal_op_encode(ggml_metal_op_t ctx, int idx);

+//
+// Profiling: per-op timestamp sampling
+//
+
+// (node_idx, slot_start, slot_end) triple recorded for each encoded op group.
+// node_idx < 0 marks an unused slot.
+struct ggml_metal_op_sample_slot {
+    int    node_idx;
+    int    n_fused;
+    size_t slot_start;
+    size_t slot_end;
+};
+
+// Enable profiling on this op encoder. `slot_base` is the first index in `sample_buf` available
+// to this encoder; `slot_map` is an array (size = number of graph nodes) into which
+// the encoder will record one entry per non-empty encoded op group.
+// Must be called before any ggml_metal_op_encode().
+void ggml_metal_op_enable_profiling(
+        ggml_metal_op_t ctx,
+        ggml_metal_sample_buf_t sample_buf,
+        size_t slot_base,
+        struct ggml_metal_op_sample_slot * slot_map);
+
 //
 // available ops:
 //
@@ -2,6 +2,7 @@

 #include "ggml-impl.h"
 #include "ggml-backend-impl.h"
+#include "ggml-profiler.h"

 #include "ggml-metal-device.h"
 #include "ggml-metal-context.h"
@@ -9,6 +10,7 @@

 #include <mutex>
 #include <string>
+#include <vector>

 #define GGML_METAL_NAME "MTL"
 #define GGML_METAL_MAX_DEVICES 16
@@ -20,6 +22,59 @@ static int g_devices = 1;
 // forward declaration
 static bool ggml_backend_buffer_is_metal(ggml_backend_buffer_t buffer);

+//
+// Profiler state (mirrors the per-backend profiler state from the Vulkan integration)
+//
+// Owned by ggml_backend_metal_init / device_init_backend; lifetime is the backend's lifetime.
+// The Metal context holds a borrowed pointer and calls into the bridge functions below to push
+// records during graph_compute.
+//
+
+struct ggml_metal_profiler_state {
+    bool enabled  = false;
+    int  split_id = -1;
+    std::vector<ggml_profile_record> records;
+
+    void reset() {
+        records.clear();
+        split_id = -1;
+    }
+};
+
+extern "C" {
+
+void ggml_metal_profiler_push_record(
+        ggml_metal_profiler_state * state,
+        const ggml_tensor * node,
+        uint64_t start_ns,
+        uint64_t end_ns) {
+    if (state == nullptr || node == nullptr) {
+        return;
+    }
+
+    ggml_profile_record rec;
+    rec.type       = GGML_PROFILE_EVENT_OP;
+    rec.name       = ggml_op_name(node->op);
+    rec.backend_id = -1;
+    rec.split_id   = state->split_id;
+    rec.start_ns   = start_ns;
+    rec.end_ns     = end_ns;
+    rec.bytes      = ggml_nbytes(node);
+    rec.extra      = nullptr;
+    ggml_profile_record_from_tensor(&rec, node);
+    state->records.push_back(rec);
+}
+
+bool ggml_metal_profiler_is_enabled(ggml_metal_profiler_state * state) {
+    return state != nullptr && state->enabled;
+}
+
+int ggml_metal_profiler_get_split_id(ggml_metal_profiler_state * state) {
+    return state != nullptr ? state->split_id : -1;
+}
+
+} // extern "C"
+
 ////////////////////////////////////////////////////////////////////////////////
 // backend interface
 ////////////////////////////////////////////////////////////////////////////////
@@ -589,6 +644,48 @@ static ggml_guid_t ggml_backend_metal_guid(void) {
    return &guid;
 }

+// Register the per-backend ggml_backend_profiler interface on `backend` and inject the borrowed
+// state pointer into the Metal context so graph_compute can push records.
+static void ggml_backend_metal_register_profiler(ggml_backend_t backend, ggml_metal_t ctx) {
+    auto * state = new ggml_metal_profiler_state();
+    ggml_metal_set_profiler_state(ctx, state);
+
+    static auto metal_prof_enable = [](void * context, bool enable) {
+        auto * state = (ggml_metal_profiler_state *) context;
+        state->enabled = enable;
+        if (!enable) {
+            state->reset();
+        }
+    };
+    static auto metal_prof_reset = [](void * context) {
+        auto * state = (ggml_metal_profiler_state *) context;
+        state->reset();
+    };
+    static auto metal_prof_set_split_id = [](void * context, int split_id) {
+        auto * state = (ggml_metal_profiler_state *) context;
+        state->split_id = split_id;
+    };
+    static auto metal_prof_get_records = [](void * context, const ggml_profile_record ** out) -> int {
+        auto * state = (ggml_metal_profiler_state *) context;
+        *out = state->records.data();
+        return (int) state->records.size();
+    };
+    static auto metal_prof_free = [](void * context) {
+        auto * state = (ggml_metal_profiler_state *) context;
+        delete state;
+    };
+
+    auto * profiler = new ggml_backend_profiler {
+        /* .context      = */ state,
+        /* .enable       = */ metal_prof_enable,
+        /* .reset        = */ metal_prof_reset,
+        /* .set_split_id = */ metal_prof_set_split_id,
+        /* .get_records  = */ metal_prof_get_records,
+        /* .free_context = */ metal_prof_free,
+    };
+    ggml_backend_set_profiler(backend, profiler);
+}
+
 ggml_backend_t ggml_backend_metal_init(void) {
    ggml_backend_dev_t dev = ggml_backend_reg_dev_get(ggml_backend_metal_reg(), 0);
    ggml_metal_device_t ctx_dev = (ggml_metal_device_t)dev->context;
@@ -606,10 +703,13 @@ ggml_backend_t ggml_backend_metal_init(void) {
        /* .interface = */ ggml_backend_metal_i,
        /* .device    = */ dev,
        /* .context   = */ ctx,
+        /* .profiler  = */ NULL,
    };

    ggml_backend_metal_set_n_cb(backend, 1);

+    ggml_backend_metal_register_profiler(backend, ctx);
+
    return backend;
 }

@@ -700,10 +800,13 @@ static ggml_backend_t ggml_backend_metal_device_init_backend(ggml_backend_dev_t
        /* .interface = */ ggml_backend_metal_i,
        /* .device    = */ dev,
        /* .context   = */ ctx,
+        /* .profiler  = */ NULL,
    };

    ggml_backend_metal_set_n_cb(backend, 1);

+    ggml_backend_metal_register_profiler(backend, ctx);
+
    return backend;

    GGML_UNUSED(params);
@@ -7491,7 +7491,8 @@ ggml_backend_t ggml_backend_opencl_init(void) {
        /* .guid    = */ ggml_backend_opencl_guid(),
        /* .iface   = */ ggml_backend_opencl_i,
        /* .device  = */ dev,
-        /* .context = */ backend_ctx
+        /* .context  = */ backend_ctx,
+        /* .profiler = */ nullptr,
    };

    return backend;
@@ -10767,6 +10768,7 @@ static ggml_backend_t ggml_backend_opencl_device_init(ggml_backend_dev_t dev, co
        /* .interface = */ ggml_backend_opencl_i,
        /* .device    = */ dev,
        /* .context   = */ backend_ctx,
+        /* .profiler  = */ nullptr,
    };

    ggml_backend_opencl_device_context * dev_ctx = (ggml_backend_opencl_device_context *) dev->context;
@@ -703,6 +703,7 @@ GGML_BACKEND_API ggml_backend_t ggml_backend_openvino_init(int device) {
        /* .interface = */ ggml_backend_openvino_interface,
        /* .device    = */ ggml_backend_reg_dev_get(ggml_backend_openvino_reg(), device),
        /* .context   = */ ctx,
+        /* .profiler  = */ nullptr,
    };

    return openvino_backend;
@@ -1378,3 +1379,5 @@ GGML_BACKEND_API ggml_backend_reg_t ggml_backend_openvino_reg(void) {

    return &reg;
 }
+
+GGML_BACKEND_DL_IMPL(ggml_backend_openvino_reg)
@@ -0,0 +1,121 @@
+#include "ggml-profiler.h"
+
+#include "ggml-backend-impl.h"
+#include "ggml-impl.h"
+#include "ggml.h"
+
+#include <stdio.h>
+#include <string.h>
+
+#ifdef _WIN32
+#    define WIN32_LEAN_AND_MEAN
+#    ifndef NOMINMAX
+#        define NOMINMAX
+#    endif
+#    include <windows.h>
+#else
+#    include <time.h>
+#    include <unistd.h>
+#endif
+
+//
+// Time utilities
+//
+
+uint64_t ggml_profiler_time_ns(void) {
+#ifdef _WIN32
+    LARGE_INTEGER freq, count;
+    QueryPerformanceFrequency(&freq);
+    QueryPerformanceCounter(&count);
+    return (uint64_t) (count.QuadPart * 1000000000ULL / freq.QuadPart);
+#elif defined(CLOCK_MONOTONIC_RAW)
+    struct timespec ts;
+    clock_gettime(CLOCK_MONOTONIC_RAW, &ts);
+    return (uint64_t) ts.tv_sec * 1000000000ULL + (uint64_t) ts.tv_nsec;
+#else
+    struct timespec ts;
+    clock_gettime(CLOCK_MONOTONIC, &ts);
+    return (uint64_t) ts.tv_sec * 1000000000ULL + (uint64_t) ts.tv_nsec;
+#endif
+}
+
+//
+// Record helpers
+//
+
+void ggml_profile_record_from_tensor(ggml_profile_record * rec, const struct ggml_tensor * node) {
+    if (rec == NULL) {
+        return;
+    }
+
+    // Output tensor info
+    if (node != NULL) {
+        memcpy(rec->ne, node->ne, sizeof(rec->ne));
+        rec->out_type = (int) node->type;
+        memcpy(rec->op_params, node->op_params, sizeof(rec->op_params));
+        // Copy the tensor name (rather than aliasing node->name) so the record stays
+        // valid after the graph's meta context is reused on the next build.
+        snprintf(rec->tensor_name, sizeof(rec->tensor_name), "%s", node->name);
+    } else {
+        memset(rec->ne, 0, sizeof(rec->ne));
+        rec->out_type = -1;
+        memset(rec->op_params, 0, sizeof(rec->op_params));
+        rec->tensor_name[0] = '\0';
+    }
+
+    // Sub-op (UNARY/GLU)
+    rec->sub_op = -1;
+    if (node != NULL) {
+        if (node->op == GGML_OP_UNARY) {
+            rec->sub_op = (int) ggml_get_unary_op(node);
+        } else if (node->op == GGML_OP_GLU) {
+            rec->sub_op = (int) ggml_get_glu_op(node);
+        }
+    }
+
+    // Source tensors
+    rec->n_src = 0;
+    for (int i = 0; i < GGML_MAX_SRC; i++) {
+        const struct ggml_tensor * src = (node != NULL) ? node->src[i] : NULL;
+        if (src == NULL) {
+            memset(rec->ne_src[i], 0, sizeof(rec->ne_src[i]));
+            memset(rec->nb_src[i], 0, sizeof(rec->nb_src[i]));
+            rec->type_src[i] = -1;
+        } else {
+            memcpy(rec->ne_src[i], src->ne, sizeof(rec->ne_src[i]));
+            for (int d = 0; d < 4; d++) {
+                rec->nb_src[i][d] = (int64_t) src->nb[d];
+            }
+            rec->type_src[i] = (int) src->type;
+            rec->n_src = i + 1;
+        }
+    }
+}
+
+//
+// Backend profiler registration
+//
+
+void ggml_backend_set_profiler(ggml_backend_t backend, ggml_backend_profiler_t profiler) {
+    if (backend == NULL) {
+        return;
+    }
+
+    // Free any existing profiler
+    if (backend->profiler != NULL) {
+        if (backend->profiler->free_context != NULL) {
+            backend->profiler->free_context(backend->profiler->context);
+        }
+        delete backend->profiler;
+        backend->profiler = NULL;
+    }
+
+    backend->profiler = profiler;
+}
+
+ggml_backend_profiler_t ggml_backend_get_profiler(ggml_backend_t backend) {
+    if (backend == NULL) {
+        return NULL;
+    }
+    return backend->profiler;
+}
@@ -785,7 +785,8 @@ ggml_backend_t ggml_backend_rpc_init(const char * endpoint, uint32_t device) {
        /* .guid    = */ ggml_backend_rpc_guid(),
        /* .iface   = */ ggml_backend_rpc_interface,
        /* .device  = */ ggml_backend_reg_dev_get(reg, device),
-        /* .context = */ ctx
+        /* .context  = */ ctx,
+        /* .profiler = */ nullptr,
    };
    return backend;
 }
@@ -6486,7 +6486,8 @@ ggml_backend_t ggml_backend_sycl_init(int device) {
        /* .guid    = */ ggml_backend_sycl_guid(),
        /* .iface   = */ ggml_backend_sycl_interface,
        /* .device  = */ ggml_backend_reg_dev_get(ggml_backend_sycl_reg(), device),
-        /* .context = */ ctx
+        /* .context  = */ ctx,
+        /* .profiler = */ nullptr,
    };

    return sycl_backend;
@@ -65,6 +65,7 @@ ggml_backend_t ggml_backend_remoting_device_init(ggml_backend_dev_t dev, const c
        /* .interface = */ ggml_backend_remoting_interface,
        /* .device    = */ ggml_backend_reg_dev_get(ggml_backend_virtgpu_reg(), ctx->device),
        /* .context   = */ ctx,
+        /* .profiler  = */ nullptr,
    };

    return remoting_backend;
@@ -1,4 +1,5 @@
 #include "ggml-vulkan.h"
+#include "ggml-profiler.h"
 #include <vulkan/vulkan_core.h>
 #if defined(GGML_VULKAN_RUN_TESTS) || defined(GGML_VULKAN_CHECK_RESULTS)
 #include <chrono>
@@ -156,6 +157,24 @@ typedef struct VkPhysicalDeviceShaderFloat8FeaturesEXT {
 } VkPhysicalDeviceShaderFloat8FeaturesEXT;
 #endif

+#ifndef VK_KHR_INTERNALLY_SYNCHRONIZED_QUEUES_EXTENSION_NAME
+#define VK_KHR_INTERNALLY_SYNCHRONIZED_QUEUES_EXTENSION_NAME "VK_KHR_internally_synchronized_queues"
+#define VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_INTERNALLY_SYNCHRONIZED_QUEUES_FEATURES_KHR ((VkStructureType)1000504000)
+#define VK_DEVICE_QUEUE_CREATE_INTERNALLY_SYNCHRONIZED_BIT_KHR ((VkDeviceQueueCreateFlagBits)0x00000004)
+
+// Compile-time constant guaranteed; no runtime initialization overhead
+static constexpr vk::DeviceQueueCreateFlagBits eInternallySynchronizedKHR =
+    static_cast<vk::DeviceQueueCreateFlagBits>(0x00000004);
+
+typedef struct VkPhysicalDeviceInternallySynchronizedQueuesFeaturesKHR {
+    VkStructureType    sType;
+    void* pNext;
+    VkBool32           internallySynchronizedQueues;
+} VkPhysicalDeviceInternallySynchronizedQueuesFeaturesKHR;
+#else
+static constexpr vk::DeviceQueueCreateFlagBits eInternallySynchronizedKHR = vk::DeviceQueueCreateFlagBits::eInternallySynchronizedKHR;
+#endif
+
 #define ROUNDUP_POW2(M, N) (((M) + (N) - 1) & ~((N) - 1))
 #define CEIL_DIV(M, N) (((M) / (N)) + (((M) % (N)) != 0))
 static bool is_pow2(uint32_t x) { return x > 1 && (x & (x-1)) == 0; }
@@ -285,27 +304,41 @@ struct vk_command_pool {
 };

 // Prevent simultaneous submissions to the same queue.
-// This could be per vk_queue if we stopped having two vk_queue structures
-// sharing the same vk::Queue.
-static std::mutex queue_mutex;
+struct vk_queue_handle {
+    vk::Queue queue;
+    virtual void submit(vk::ArrayProxy<const vk::SubmitInfo> submits, vk::Fence fence) = 0;
+    virtual void lock()   {}   // no-op by default (internally synchronized case)
+    virtual void unlock() {}
+    virtual ~vk_queue_handle() = default;
+};
+
+struct vk_queue_handle_synchronized : vk_queue_handle {
+    std::mutex mutex;
+    void submit(vk::ArrayProxy<const vk::SubmitInfo> submits, vk::Fence fence) override {
+        std::lock_guard<std::mutex> guard(mutex);
+        queue.submit(submits, fence);
+    }
+    void lock()   override { mutex.lock(); }
+    void unlock() override { mutex.unlock(); }
+};
+
+struct vk_queue_handle_unsynchronized : vk_queue_handle {
+    void submit(vk::ArrayProxy<const vk::SubmitInfo> submits, vk::Fence fence) override {
+        // Driver guarantees internal synchronization via VK_KHR_internally_synchronized_queues
+        queue.submit(submits, fence);
+    }
+    // lock()/unlock() inherited no-ops
+};

 struct vk_queue {
    uint32_t queue_family_index;
-    vk::Queue queue;
+    std::shared_ptr<vk_queue_handle> handle;

    vk_command_pool cmd_pool;

    vk::PipelineStageFlags stage_flags;

    bool transfer_only;
-
-    // copy everything except the cmd_pool
-    void copyFrom(vk_queue &other) {
-        queue_family_index = other.queue_family_index;
-        queue = other.queue;
-        stage_flags = other.stage_flags;
-        transfer_only = other.transfer_only;
-    }
 };

 static const char * ggml_backend_vk_buffer_type_name(ggml_backend_buffer_type_t buft);
@@ -712,11 +745,12 @@ struct vk_device_struct {
    uint32_t vendor_id;
    vk::DriverId driver_id;
    vk_device_architecture architecture;
-    vk_queue compute_queue;
-    vk_queue transfer_queue;
+    std::unique_ptr<vk_queue> compute_queue;
+    std::unique_ptr<vk_queue> transfer_queue;
    bool single_queue;
    bool support_async;
    bool async_use_transfer_queue;
+    bool has_internally_synchronized_queues = false;
    uint32_t subgroup_size;
    uint32_t subgroup_size_log2;
    uint32_t shader_core_count;
@@ -1019,8 +1053,13 @@ struct vk_device_struct {

        ggml_vk_destroy_buffer(sync_staging);

-        compute_queue.cmd_pool.destroy(device);
-        transfer_queue.cmd_pool.destroy(device);
+        if (compute_queue) compute_queue->cmd_pool.destroy(device);
+        if (transfer_queue) transfer_queue->cmd_pool.destroy(device);
+
+        // Explicitly clear to ensure queues drop their shared_ptrs to handles
+        // before the Vulkan logical device instance is destroyed
+        compute_queue.reset();
+        transfer_queue.reset();

        for (auto& pipeline : all_pipelines) {
            if (pipeline.expired()) {
@@ -1932,8 +1971,8 @@ private:

 std::mutex vk_memory_logger::log_mutex;

-static bool vk_perf_logger_enabled = false;
-static bool vk_perf_logger_concurrent = false;
+static bool vk_perf_logger_enabled = false;       // deprecated: use --profile instead
+static bool vk_perf_logger_concurrent = false;     // GGML_VK_PERF_LOGGER_CONCURRENT: use concurrent timestamp mode
 static bool vk_enable_sync_logger = false;
 // number of calls between perf logger prints
 static uint32_t vk_perf_logger_frequency = 1;
@@ -2123,6 +2162,21 @@ class vk_perf_logger {
    uint32_t print_count {};
 };

+// Profiler state for the new ggml_backend_profiler interface
+struct ggml_vk_profiler_state {
+    bool enabled = false;
+    int  split_id = -1;
+
+    std::vector<ggml_profile_record> records;
+    std::vector<uint64_t>            cpu_timestamps;  // CPU-side timestamps for global ordering
+
+    void reset() {
+        records.clear();
+        cpu_timestamps.clear();
+        split_id = -1;
+    }
+};
+
 struct ggml_backend_vk_context {
    std::string name;

@@ -2183,8 +2237,9 @@ struct ggml_backend_vk_context {
    topk_moe_mode fused_topk_moe_mode {};
    bool fused_topk_moe_scale {};

-    // for GGML_VK_PERF_LOGGER
-    std::unique_ptr<vk_perf_logger> perf_logger;
+    // Profiling
+    std::unique_ptr<vk_perf_logger> perf_logger;   // legacy env-var profiler
+    ggml_vk_profiler_state * profiler_state = nullptr;
    vk::QueryPool query_pool;
    std::vector<const char *> query_fusion_names;
    std::vector<int> query_fusion_node_count;
@@ -2909,8 +2964,7 @@ static vk_command_buffer* ggml_vk_create_cmd_buffer(vk_device& device, vk_comman
 static void ggml_vk_submit(vk_context& ctx, vk::Fence fence) {
    if (ctx->seqs.empty()) {
        if (fence) {
-            std::lock_guard<std::mutex> guard(queue_mutex);
-            ctx->p->q->queue.submit({}, fence);
+            ctx->p->q->handle->submit({}, fence);
        }
        return;
    }
@@ -2979,8 +3033,7 @@ static void ggml_vk_submit(vk_context& ctx, vk::Fence fence) {
        }
    }

-    std::lock_guard<std::mutex> guard(queue_mutex);
-    ctx->p->q->queue.submit(submit_infos, fence);
+    ctx->p->q->handle->submit(submit_infos, fence);

    ctx->seqs.clear();
 }
@@ -3031,18 +3084,44 @@ static uint32_t ggml_vk_find_queue_family_index(std::vector<vk::QueueFamilyPrope
    abort();
 }

-static void ggml_vk_create_queue(vk_device& device, vk_queue& q, uint32_t queue_family_index, uint32_t queue_index, vk::PipelineStageFlags&& stage_flags, bool transfer_only) {
+static std::unique_ptr<vk_queue> ggml_vk_create_queue(vk_device& device, uint32_t queue_family_index, uint32_t queue_index, vk::PipelineStageFlags&& stage_flags, bool transfer_only) {
    VK_LOG_DEBUG("ggml_vk_create_queue()");
    std::lock_guard<std::recursive_mutex> guard(device->mutex);

-    q.queue_family_index = queue_family_index;
-    q.transfer_only = transfer_only;
+    auto q = std::make_unique<vk_queue>();
+    q->queue_family_index = queue_family_index;
+    q->transfer_only = transfer_only;

-    q.cmd_pool.init(device, &q);
+    std::shared_ptr<vk_queue_handle> h;
+    vk::DeviceQueueInfo2 queue_info2{};
+    queue_info2.queueFamilyIndex = queue_family_index;
+    queue_info2.queueIndex = queue_index;

-    q.queue = device->device.getQueue(queue_family_index, queue_index);
+    if (device->has_internally_synchronized_queues) {
+        h = std::make_shared<vk_queue_handle_unsynchronized>();
+        queue_info2.flags = eInternallySynchronizedKHR;
+    } else {
+        h = std::make_shared<vk_queue_handle_synchronized>();
+    }

-    q.stage_flags = stage_flags;
+    h->queue = device->device.getQueue2(queue_info2);
+    q->handle = h;
+
+    q->cmd_pool.init(device, q.get());
+
+    q->stage_flags = stage_flags;
+    return q;
+}
+
+static std::unique_ptr<vk_queue> ggml_vk_create_aliased_queue(vk_device& device, const std::unique_ptr<vk_queue>& source) {
+    std::lock_guard<std::recursive_mutex> guard(device->mutex);
+    auto q = std::make_unique<vk_queue>();
+    q->handle = source->handle;
+    q->queue_family_index = source->queue_family_index;
+    q->stage_flags = source->stage_flags;
+    q->transfer_only = source->transfer_only;
+    q->cmd_pool.init(device, q.get());
+    return q;
 }

 static vk_context ggml_vk_create_context(ggml_backend_vk_context * ctx, vk_command_pool& p) {
@@ -3107,11 +3186,11 @@ static void ggml_vk_queue_command_pools_cleanup(vk_device& device) {
    // Arbitrary frequency to cleanup/reuse command buffers
    static constexpr uint32_t cleanup_frequency = 10;

-    if (device->compute_queue.cmd_pool.buffers_in_use() >= cleanup_frequency) {
-        ggml_vk_command_pool_cleanup(device, device->compute_queue.cmd_pool);
+    if (device->compute_queue->cmd_pool.buffers_in_use() >= cleanup_frequency) {
+        ggml_vk_command_pool_cleanup(device, device->compute_queue->cmd_pool);
    }
-    if (device->transfer_queue.cmd_pool.buffers_in_use() >= cleanup_frequency) {
-        ggml_vk_command_pool_cleanup(device, device->transfer_queue.cmd_pool);
+    if (device->transfer_queue->cmd_pool.buffers_in_use() >= cleanup_frequency) {
+        ggml_vk_command_pool_cleanup(device, device->transfer_queue->cmd_pool);
    }
 }

@@ -5886,6 +5965,7 @@ static vk_device ggml_vk_get_device(size_t idx) {
        bool coopmat2_support = false;
        bool coopmat2_decode_vector_support = false;
        bool pipeline_executable_properties_support = false;
+        bool internally_sync_support = false;
        device->coopmat_support = false;
        device->integer_dot_product = false;
        device->shader_64b_indexing = false;
@@ -5957,6 +6037,8 @@ static vk_device ggml_vk_get_device(size_t idx) {
            } else if (strcmp("VK_EXT_shader_64bit_indexing", properties.extensionName) == 0) {
                device->shader_64b_indexing = true;
 #endif
+            } else if (strcmp(VK_KHR_INTERNALLY_SYNCHRONIZED_QUEUES_EXTENSION_NAME, properties.extensionName) == 0) {
+                internally_sync_support = true;
            }
        }

@@ -6143,14 +6225,6 @@ static vk_device ggml_vk_get_device(size_t idx) {
        device->single_queue = compute_queue_family_index == transfer_queue_family_index && queue_family_props[compute_queue_family_index].queueCount == 1;

        std::vector<vk::DeviceQueueCreateInfo> device_queue_create_infos;
-        if (compute_queue_family_index != transfer_queue_family_index) {
-            device_queue_create_infos.push_back({vk::DeviceQueueCreateFlags(), compute_queue_family_index, 1, priorities});
-            device_queue_create_infos.push_back({vk::DeviceQueueCreateFlags(), transfer_queue_family_index, 1, priorities + 1});
-        } else if(!device->single_queue) {
-            device_queue_create_infos.push_back({vk::DeviceQueueCreateFlags(), compute_queue_family_index, 2, priorities});
-        } else {
-            device_queue_create_infos.push_back({vk::DeviceQueueCreateFlags(), compute_queue_family_index, 1, priorities});
-        }
        vk::DeviceCreateInfo device_create_info{};
        std::vector<const char *> device_extensions;
        vk::PhysicalDeviceFeatures device_features = device->physical_device.getFeatures();
@@ -6172,6 +6246,17 @@ static vk_device ggml_vk_get_device(size_t idx) {

        last_struct = (VkBaseOutStructure *)&vk12_features;

+        VkPhysicalDeviceInternallySynchronizedQueuesFeaturesKHR internally_synchronized_queues_features{};
+        internally_synchronized_queues_features.sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_INTERNALLY_SYNCHRONIZED_QUEUES_FEATURES_KHR;
+        internally_synchronized_queues_features.pNext = nullptr;
+        internally_synchronized_queues_features.internallySynchronizedQueues = VK_FALSE;
+
+        if (internally_sync_support) {
+            last_struct->pNext = (VkBaseOutStructure *)&internally_synchronized_queues_features;
+            last_struct = (VkBaseOutStructure *)&internally_synchronized_queues_features;
+            device_extensions.push_back(VK_KHR_INTERNALLY_SYNCHRONIZED_QUEUES_EXTENSION_NAME);
+        }
+
        VkPhysicalDevicePipelineRobustnessFeaturesEXT pl_robustness_features;
        pl_robustness_features.pNext = nullptr;
        pl_robustness_features.sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_PIPELINE_ROBUSTNESS_FEATURES_EXT;
@@ -6310,6 +6395,23 @@ static vk_device ggml_vk_get_device(size_t idx) {

        vkGetPhysicalDeviceFeatures2(device->physical_device, &device_features2);

+        device->has_internally_synchronized_queues = internally_synchronized_queues_features.internallySynchronizedQueues;
+
+        // Build queue create infos only after querying whether internally synchronized queues are enabled.
+        // getQueue2() later uses the same flag, so creation/retrieval must stay consistent.
+        vk::DeviceQueueCreateFlags queue_flags = device->has_internally_synchronized_queues ?
+                                                eInternallySynchronizedKHR :
+                                                vk::DeviceQueueCreateFlags();
+
+        if (compute_queue_family_index != transfer_queue_family_index) {
+            device_queue_create_infos.push_back({queue_flags, compute_queue_family_index, 1, priorities});
+            device_queue_create_infos.push_back({queue_flags, transfer_queue_family_index, 1, priorities + 1});
+        } else if(!device->single_queue) {
+            device_queue_create_infos.push_back({queue_flags, compute_queue_family_index, 2, priorities});
+        } else {
+            device_queue_create_infos.push_back({queue_flags, compute_queue_family_index, 1, priorities});
+        }
+
        device->pipeline_executable_properties_support = pipeline_executable_properties_support;

        device->fp16 = device->fp16 && vk12_features.shaderFloat16;
@@ -6592,7 +6694,7 @@ static vk_device ggml_vk_get_device(size_t idx) {
        device->device = device->physical_device.createDevice(device_create_info);

        // Queues
-        ggml_vk_create_queue(device, device->compute_queue, compute_queue_family_index, 0, { vk::PipelineStageFlagBits::eComputeShader | vk::PipelineStageFlagBits::eTransfer }, false);
+        device->compute_queue = ggml_vk_create_queue(device, compute_queue_family_index, 0, { vk::PipelineStageFlagBits::eComputeShader | vk::PipelineStageFlagBits::eTransfer }, false);

        // Shaders
        // Disable matmul tile sizes early if performance low or not supported
@@ -6694,13 +6796,11 @@ static vk_device ggml_vk_get_device(size_t idx) {

        if (!device->single_queue) {
            const uint32_t transfer_queue_index = compute_queue_family_index == transfer_queue_family_index ? 1 : 0;
-            ggml_vk_create_queue(device, device->transfer_queue, transfer_queue_family_index, transfer_queue_index, { vk::PipelineStageFlagBits::eTransfer }, true);
+            device->transfer_queue = ggml_vk_create_queue(device, transfer_queue_family_index, transfer_queue_index, { vk::PipelineStageFlagBits::eTransfer }, true);

            device->async_use_transfer_queue = prefers_transfer_queue || (getenv("GGML_VK_ASYNC_USE_TRANSFER_QUEUE") != nullptr);
        } else {
-            // TODO: Use pointer or reference to avoid copy
-            device->transfer_queue.copyFrom(device->compute_queue);
-            device->transfer_queue.cmd_pool.init(device, &device->transfer_queue);
+            device->transfer_queue = ggml_vk_create_aliased_queue(device, device->compute_queue);

            device->async_use_transfer_queue = false;
        }
@@ -7263,7 +7363,7 @@ static void ggml_vk_init(ggml_backend_vk_context * ctx, size_t idx) {
    ctx->fence = ctx->device->device.createFence({});
    ctx->almost_ready_fence = ctx->device->device.createFence({});

-    ctx->compute_cmd_pool.init(ctx->device, &ctx->device->compute_queue);
+    ctx->compute_cmd_pool.init(ctx->device, ctx->device->compute_queue.get());
    if (ctx->device->async_use_transfer_queue) {
        vk::SemaphoreTypeCreateInfo tci{ vk::SemaphoreType::eTimeline, 0 };
        vk::SemaphoreCreateInfo ci{};
@@ -7271,7 +7371,7 @@ static void ggml_vk_init(ggml_backend_vk_context * ctx, size_t idx) {
        ctx->transfer_semaphore.s = ctx->device->device.createSemaphore(ci);
        ctx->transfer_semaphore.value = 0;

-        ctx->transfer_cmd_pool.init(ctx->device, &ctx->device->transfer_queue);
+        ctx->transfer_cmd_pool.init(ctx->device, ctx->device->transfer_queue.get());
    }

    if (vk_perf_logger_enabled) {
@@ -8126,7 +8226,7 @@ static void ggml_vk_buffer_write_2d(vk_buffer& dst, size_t offset, const void *
    } else {
        std::lock_guard<std::recursive_mutex> guard(dst->device->mutex);

-        vk_context subctx = ggml_vk_create_temporary_context(dst->device->transfer_queue.cmd_pool);
+        vk_context subctx = ggml_vk_create_temporary_context(dst->device->transfer_queue->cmd_pool);
        ggml_vk_ctx_begin(dst->device, subctx);
        bool ret = ggml_vk_buffer_write_2d_async(subctx, dst, offset, src, spitch, dpitch, width, height, true);
        GGML_ASSERT(ret);
@@ -8241,7 +8341,7 @@ static void ggml_vk_buffer_read_2d(vk_buffer& src, size_t offset, void * dst, si
        GGML_ASSERT(src->memory_property_flags & vk::MemoryPropertyFlagBits::eHostCoherent);

        std::lock_guard<std::recursive_mutex> guard(src->device->mutex);
-        vk_context subctx = ggml_vk_create_temporary_context(src->device->compute_queue.cmd_pool);
+        vk_context subctx = ggml_vk_create_temporary_context(src->device->compute_queue->cmd_pool);
        ggml_vk_ctx_begin(src->device, subctx);
        subctx->s->buffer->buf.pipelineBarrier(
            vk::PipelineStageFlagBits::eComputeShader | vk::PipelineStageFlagBits::eTransfer,
@@ -8267,7 +8367,7 @@ static void ggml_vk_buffer_read_2d(vk_buffer& src, size_t offset, void * dst, si
    } else {
        std::lock_guard<std::recursive_mutex> guard(src->device->mutex);

-        vk_context subctx = ggml_vk_create_temporary_context(src->device->transfer_queue.cmd_pool);
+        vk_context subctx = ggml_vk_create_temporary_context(src->device->transfer_queue->cmd_pool);
        ggml_vk_ctx_begin(src->device, subctx);
        bool ret = ggml_vk_buffer_read_2d_async(subctx, src, offset, dst, spitch, dpitch, width, height, true);
        GGML_ASSERT(ret);
@@ -8304,7 +8404,7 @@ static void ggml_vk_buffer_copy(vk_buffer& dst, size_t dst_offset, vk_buffer& sr
        std::lock_guard<std::recursive_mutex> guard(src->device->mutex);
        VK_LOG_DEBUG("ggml_vk_buffer_copy(SINGLE_DEVICE, " << size << ")");
        // Copy within the device
-        vk_context subctx = ggml_vk_create_temporary_context(src->device->transfer_queue.cmd_pool);
+        vk_context subctx = ggml_vk_create_temporary_context(src->device->transfer_queue->cmd_pool);
        ggml_vk_ctx_begin(src->device, subctx);
        ggml_vk_buffer_copy_async(subctx, dst, dst_offset, src, src_offset, size);
        ggml_vk_ctx_end(subctx);
@@ -8347,7 +8447,7 @@ static void ggml_vk_buffer_memset(vk_buffer& dst, size_t offset, uint32_t c, siz
    }

    std::lock_guard<std::recursive_mutex> guard(dst->device->mutex);
-    vk_context subctx = ggml_vk_create_temporary_context(dst->device->transfer_queue.cmd_pool);
+    vk_context subctx = ggml_vk_create_temporary_context(dst->device->transfer_queue->cmd_pool);
    ggml_vk_ctx_begin(dst->device, subctx);
    subctx->s->buffer->buf.fillBuffer(dst->buffer, offset, size, c);
    ggml_vk_ctx_end(subctx);
@@ -14812,10 +14912,14 @@ static bool ggml_vk_build_graph(ggml_backend_vk_context * ctx, ggml_cgraph * cgr
            ctx->unsynced_nodes_read.clear();
            ggml_vk_sync_buffers(ctx, compute_ctx);

-            if (vk_perf_logger_enabled && vk_perf_logger_concurrent) {
+            if ((vk_perf_logger_enabled || (ctx->profiler_state != nullptr && ctx->profiler_state->enabled))
+                    && vk_perf_logger_concurrent) {
                ctx->query_node_idx[ctx->query_idx] = node_idx;
                compute_ctx->s->buffer->buf.writeTimestamp(vk::PipelineStageFlagBits::eAllCommands, ctx->query_pool, ctx->query_idx++);
                ggml_vk_sync_buffers(ctx, compute_ctx);
+                if (ctx->profiler_state != nullptr && ctx->profiler_state->enabled) {
+                    ctx->profiler_state->cpu_timestamps.push_back(ggml_profiler_time_ns());
+                }
            }
        }
        // Add all fused nodes to the unsynchronized lists.
@@ -15362,7 +15466,7 @@ static void ggml_vk_cleanup(ggml_backend_vk_context * ctx) {

        ctx->transfer_cmd_pool.destroy(ctx->device->device);
    }
-    if (vk_perf_logger_enabled) {
+    if (ctx->perf_logger) {
        ctx->perf_logger->print_timings(true);
    }
 }
@@ -15875,19 +15979,17 @@ static void ggml_vk_synchronize(ggml_backend_vk_context * ctx) {
                1, &ctx->transfer_semaphore.value,
                0, nullptr,
            };
-            vk::PipelineStageFlags stage = ctx->device->transfer_queue.stage_flags;
+            vk::PipelineStageFlags stage = ctx->device->transfer_queue->stage_flags;
            vk::SubmitInfo si{
                1, &ctx->transfer_semaphore.s, &stage,
                0, nullptr,
                0, nullptr,
            };
            si.setPNext(&tl_info);
-            std::lock_guard<std::mutex> guard(queue_mutex);
-            ctx->device->compute_queue.queue.submit({ si }, ctx->fence);
+            ctx->device->compute_queue->handle->submit({ si }, ctx->fence);
            ctx->transfer_semaphore_last_submitted = ctx->transfer_semaphore.value;
        } else {
-            std::lock_guard<std::mutex> guard(queue_mutex);
-            ctx->device->compute_queue.queue.submit({}, ctx->fence);
+            ctx->device->compute_queue->handle->submit({}, ctx->fence);
        }
        ggml_vk_wait_for_fence(ctx);
        ctx->submit_pending = false;
@@ -16449,7 +16551,9 @@ static ggml_status ggml_backend_vk_graph_compute(ggml_backend_t backend, ggml_cg
        vk::DebugUtilsLabelEXT dul = {};
        dul.pLabelName = "ggml_backend_vk_graph_compute";
        dul.color = std::array<float,4>{1.0f, 1.0f, 1.0f, 1.0f};
-        vk_instance.pfn_vkQueueBeginDebugUtilsLabelEXT(ctx->device->compute_queue.queue, reinterpret_cast<VkDebugUtilsLabelEXT*>(&dul));
+
+        std::lock_guard<vk_queue_handle> guard(*ctx->device->compute_queue->handle);
+        vk_instance.pfn_vkQueueBeginDebugUtilsLabelEXT(ctx->device->compute_queue->handle->queue, reinterpret_cast<VkDebugUtilsLabelEXT*>(&dul));
    }

    ctx->prealloc_size_add_rms_partials_offset = 0;
@@ -16472,7 +16576,9 @@ static ggml_status ggml_backend_vk_graph_compute(ggml_backend_t backend, ggml_cg
    ggml_vk_submit_transfer_ctx(ctx);

    vk_context compute_ctx;
-    if (vk_perf_logger_enabled) {
+    bool profiling = vk_perf_logger_enabled ||
+                     (ctx->profiler_state != nullptr && ctx->profiler_state->enabled);
+    if (profiling) {
        // allocate/resize the query pool
        if (ctx->num_queries < cgraph->n_nodes + 1) {
            if (ctx->query_pool) {
@@ -16500,6 +16606,10 @@ static ggml_status ggml_backend_vk_graph_compute(ggml_backend_t backend, ggml_cg
        ctx->query_idx = 0;
        compute_ctx->s->buffer->buf.writeTimestamp(vk::PipelineStageFlagBits::eAllCommands, ctx->query_pool, ctx->query_idx++);
        ggml_vk_sync_buffers(ctx, compute_ctx);
+        if (ctx->profiler_state != nullptr && ctx->profiler_state->enabled) {
+            ctx->profiler_state->cpu_timestamps.clear();
+            ctx->profiler_state->cpu_timestamps.push_back(ggml_profiler_time_ns());
+        }
    }

    ctx->prealloc_y_last_pipeline_used = nullptr;
@@ -16761,7 +16871,7 @@ static ggml_status ggml_backend_vk_graph_compute(ggml_backend_t backend, ggml_cg

        bool enqueued = ggml_vk_build_graph(ctx, cgraph, i, cgraph->nodes[submit_node_idx], submit_node_idx, i + ctx->num_additional_fused_ops >= last_node, almost_ready, submit);

-        if (vk_perf_logger_enabled && enqueued) {
+        if (profiling && enqueued) {
            compute_ctx = ggml_vk_get_compute_ctx(ctx);
            if (!vk_perf_logger_concurrent) {
                // track a single node/fusion for the current query
@@ -16769,6 +16879,9 @@ static ggml_status ggml_backend_vk_graph_compute(ggml_backend_t backend, ggml_cg
                ctx->query_fusion_names[ctx->query_idx] = fusion_string;
                compute_ctx->s->buffer->buf.writeTimestamp(vk::PipelineStageFlagBits::eAllCommands, ctx->query_pool, ctx->query_idx++);
                ggml_vk_sync_buffers(ctx, compute_ctx);
+                if (ctx->profiler_state != nullptr && ctx->profiler_state->enabled) {
+                    ctx->profiler_state->cpu_timestamps.push_back(ggml_profiler_time_ns());
+                }
            } else {
                // track a fusion string and number of fused ops for the current node_idx
                ctx->query_fusion_names[i] = fusion_string;
@@ -16802,7 +16915,7 @@ static ggml_status ggml_backend_vk_graph_compute(ggml_backend_t backend, ggml_cg

    ctx->last_total_flops = total_flops;

-    if (vk_perf_logger_enabled) {
+    if (profiling) {
        // End the command buffer and submit/wait
        GGML_ASSERT(!ctx->compute_ctx.expired());
        compute_ctx = ctx->compute_ctx.lock();
@@ -16816,15 +16929,40 @@ static ggml_status ggml_backend_vk_graph_compute(ggml_backend_t backend, ggml_cg
        // Get the results and pass them to the logger
        std::vector<uint64_t> timestamps(cgraph->n_nodes + 1);
        VK_CHECK(ctx->device->device.getQueryPoolResults(ctx->query_pool, 0, ctx->query_idx, (cgraph->n_nodes + 1)*sizeof(uint64_t), timestamps.data(), sizeof(uint64_t), vk::QueryResultFlagBits::e64 | vk::QueryResultFlagBits::eWait), "get timestamp results");
+
+        const double ts_period = ctx->device->properties.limits.timestampPeriod;
+        const bool has_profiler = ctx->profiler_state != nullptr && ctx->profiler_state->enabled;
+
        if (!vk_perf_logger_concurrent) {
            // Log each op separately
            for (int i = 1; i < ctx->query_idx; i++) {
                auto node = ctx->query_nodes[i];
                auto name = ctx->query_fusion_names[i];
-                ctx->perf_logger->log_timing(node, name, uint64_t((timestamps[i] - timestamps[i-1]) * ctx->device->properties.limits.timestampPeriod));
+                uint64_t duration_ns = uint64_t((timestamps[i] - timestamps[i-1]) * ts_period);
+
+                if (ctx->perf_logger) {
+                    ctx->perf_logger->log_timing(node, name, duration_ns);
+                }
+
+                if (has_profiler && node != nullptr) {
+                    uint64_t cpu_ts = (i < (int)ctx->profiler_state->cpu_timestamps.size())
+                                    ? ctx->profiler_state->cpu_timestamps[i] : 0;
+
+                    ggml_profile_record rec;
+                    rec.type       = GGML_PROFILE_EVENT_OP;
+                    rec.name       = ggml_op_name(node->op);
+                    rec.backend_id = -1;
+                    rec.split_id   = ctx->profiler_state->split_id;
+                    rec.start_ns   = cpu_ts;
+                    rec.end_ns     = cpu_ts + duration_ns;
+                    rec.bytes      = ggml_nbytes(node);
+                    rec.extra      = name;  // fusion name or NULL
+                    ggml_profile_record_from_tensor(&rec, node);
+                    ctx->profiler_state->records.push_back(rec);
+                }
            }
        } else {
-            // Log each group of nodes
+            // Log each group of nodes (concurrent mode)
            int prev_node_idx = 0;
            for (int i = 1; i < ctx->query_idx; i++) {
                auto cur_node_idx = ctx->query_node_idx[i];
@@ -16839,10 +16977,40 @@ static ggml_status ggml_backend_vk_graph_compute(ggml_backend_t backend, ggml_cg
                    node_idx += ctx->query_fusion_node_count[node_idx];
                }
                prev_node_idx = cur_node_idx;
-                ctx->perf_logger->log_timing(nodes, names, uint64_t((timestamps[i] - timestamps[i-1]) * ctx->device->properties.limits.timestampPeriod));
+                uint64_t duration_ns = uint64_t((timestamps[i] - timestamps[i-1]) * ts_period);
+
+                if (ctx->perf_logger) {
+                    ctx->perf_logger->log_timing(nodes, names, duration_ns);
+                }
+
+                if (has_profiler && !nodes.empty()) {
+                    uint64_t cpu_ts = (i < (int)ctx->profiler_state->cpu_timestamps.size())
+                                    ? ctx->profiler_state->cpu_timestamps[i] : 0;
+                    // In concurrent mode, report the group as a single combined operation
+                    auto * node = nodes[0];
+
+                    uint64_t total_bytes = 0;
+                    for (size_t j = 0; j < nodes.size(); j++) {
+                        total_bytes += ggml_nbytes(nodes[j]);
+                    }
+
+                    ggml_profile_record rec;
+                    rec.type       = GGML_PROFILE_EVENT_OP;
+                    rec.name       = ggml_op_name(node->op);
+                    rec.backend_id = -1;
+                    rec.split_id   = ctx->profiler_state->split_id;
+                    rec.start_ns   = cpu_ts;
+                    rec.end_ns     = cpu_ts + duration_ns;
+                    rec.bytes      = total_bytes;
+                    rec.extra      = names[0];  // fusion name of first op, or NULL
+                    ggml_profile_record_from_tensor(&rec, node);
+                    ctx->profiler_state->records.push_back(rec);
+                }
            }
        }
-        ctx->perf_logger->print_timings();
+        if (ctx->perf_logger) {
+            ctx->perf_logger->print_timings();
+        }
    }

    if (!ctx->device->support_async) {
@@ -17203,13 +17371,66 @@ ggml_backend_t ggml_backend_vk_init(size_t dev_num) {
        /* .guid    = */ ggml_backend_vk_guid(),
        /* .iface   = */ ggml_backend_vk_interface,
        /* .device  = */ ggml_backend_reg_dev_get(ggml_backend_vk_reg(), dev_num),
-        /* .context = */ ctx,
+        /* .context  = */ ctx,
+        /* .profiler = */ nullptr,
    };

    if (!ctx->device->support_async) {
        vk_backend->iface.get_tensor_async = nullptr;
    }

+    // Register profiler
+    auto * prof_state = new ggml_vk_profiler_state();
+    ctx->profiler_state = prof_state;
+
+    static auto vk_prof_enable = [](void * context, bool enable) {
+        auto * vk_ctx = (ggml_backend_vk_context *)context;
+        if (vk_ctx->profiler_state != nullptr) {
+            vk_ctx->profiler_state->enabled = enable;
+            if (!enable) {
+                vk_ctx->profiler_state->reset();
+            }
+        }
+    };
+    static auto vk_prof_reset = [](void * context) {
+        auto * vk_ctx = (ggml_backend_vk_context *)context;
+        if (vk_ctx->profiler_state != nullptr) {
+            vk_ctx->profiler_state->reset();
+        }
+    };
+    static auto vk_prof_set_split_id = [](void * context, int split_id) {
+        auto * vk_ctx = (ggml_backend_vk_context *)context;
+        if (vk_ctx->profiler_state != nullptr) {
+            vk_ctx->profiler_state->split_id = split_id;
+        }
+    };
+    static auto vk_prof_get_records = [](void * context, const ggml_profile_record ** out) -> int {
+        auto * vk_ctx = (ggml_backend_vk_context *)context;
+        if (vk_ctx->profiler_state != nullptr) {
+            *out = vk_ctx->profiler_state->records.data();
+            return (int)vk_ctx->profiler_state->records.size();
+        }
+        *out = nullptr;
+        return 0;
+    };
+    static auto vk_prof_free = [](void * context) {
+        auto * vk_ctx = (ggml_backend_vk_context *)context;
+        if (vk_ctx->profiler_state != nullptr) {
+            delete vk_ctx->profiler_state;
+            vk_ctx->profiler_state = nullptr;
+        }
+    };
+
+    auto * profiler = new ggml_backend_profiler {
+        /* .context      = */ ctx,
+        /* .enable       = */ vk_prof_enable,
+        /* .reset        = */ vk_prof_reset,
+        /* .set_split_id = */ vk_prof_set_split_id,
+        /* .get_records  = */ vk_prof_get_records,
+        /* .free_context = */ vk_prof_free,
+    };
+    ggml_backend_set_profiler(vk_backend, profiler);
+
    return vk_backend;
 }

@@ -6,6 +6,7 @@
 #extension GL_EXT_shader_explicit_arithmetic_types_int16 : require
 #extension GL_EXT_shader_explicit_arithmetic_types_int8 : require
 #extension GL_EXT_shader_16bit_storage : require
+#extension GL_EXT_control_flow_attributes : require

 #ifdef USE_OCP_FP4
 #extension GL_EXT_float_e2m1 : require
@@ -1792,7 +1793,9 @@ shared FLOAT_TYPE kvalues_iq4nl[16];
 void init_iq_shmem(uvec3 wgsize)
 {
    // copy the table into shared memory and sync
-    for (uint i = gl_LocalInvocationIndex.x; i < kvalues_iq4nl.length(); i += wgsize.x) {
+    // Use gl_LocalInvocationIndex (uint scalar) directly - each thread loads its portion
+    // with stride equal to workgroup size for coalesced shared memory writes
+    [[unroll]] for (uint i = gl_LocalInvocationIndex; i < kvalues_iq4nl.length(); i += wgsize.x) {
        kvalues_iq4nl[i] = FLOAT_TYPE(kvalues_iq4nl_const[i]);
    }
    barrier();
@@ -1832,11 +1835,12 @@ float ue4m3_to_fp32_build(uint u) {
 void init_iq_shmem(uvec3 wgsize)
 {
    // copy the table into shared memory and sync
-    for (uint i = gl_LocalInvocationIndex.x; i < kvalues_mxfp4.length(); i += wgsize.x) {
+    // Use gl_LocalInvocationIndex (uint scalar) directly for thread-strided initialization
+    [[unroll]] for (uint i = gl_LocalInvocationIndex; i < kvalues_mxfp4.length(); i += wgsize.x) {
        kvalues_mxfp4[i] = kvalues_mxfp4_const[i];
    }
 #if defined(DATA_A_NVFP4)
-    for (uint i = gl_LocalInvocationIndex.x; i < 128u; i += wgsize.x) {
+    [[unroll]] for (uint i = gl_LocalInvocationIndex; i < 128u; i += wgsize.x) {
        ue4m3_fp32_lut[i] = ue4m3_to_fp32_build(i);
    }
 #endif
@@ -355,6 +355,30 @@ struct ggml_webgpu_conv2d_pipeline_key_hash {
    }
 };

+// Same type fields as conv2d plus the input layout (WHCN vs CWHN).
+struct ggml_webgpu_conv2d_dw_pipeline_key {
+    ggml_type weight_type;
+    ggml_type input_type;
+    ggml_type output_type;
+    bool      whcn;
+
+    bool operator==(const ggml_webgpu_conv2d_dw_pipeline_key & other) const {
+        return weight_type == other.weight_type && input_type == other.input_type && output_type == other.output_type &&
+               whcn == other.whcn;
+    }
+};
+
+struct ggml_webgpu_conv2d_dw_pipeline_key_hash {
+    size_t operator()(const ggml_webgpu_conv2d_dw_pipeline_key & key) const {
+        size_t seed = 0;
+        ggml_webgpu_hash_combine(seed, key.weight_type);
+        ggml_webgpu_hash_combine(seed, key.input_type);
+        ggml_webgpu_hash_combine(seed, key.output_type);
+        ggml_webgpu_hash_combine(seed, key.whcn);
+        return seed;
+    }
+};
+
 /** Im2Col **/
 struct ggml_webgpu_im2col_pipeline_key {
    ggml_type input_type;
@@ -1210,6 +1234,8 @@ class ggml_webgpu_shader_lib {
        soft_max_pipelines;
    std::unordered_map<ggml_webgpu_conv2d_pipeline_key, webgpu_pipeline, ggml_webgpu_conv2d_pipeline_key_hash>
        conv2d_pipelines;
+    std::unordered_map<ggml_webgpu_conv2d_dw_pipeline_key, webgpu_pipeline, ggml_webgpu_conv2d_dw_pipeline_key_hash>
+        conv2d_dw_pipelines;
    std::unordered_map<ggml_webgpu_im2col_pipeline_key, webgpu_pipeline, ggml_webgpu_im2col_pipeline_key_hash>
        im2col_pipelines;

@@ -3172,6 +3198,50 @@ class ggml_webgpu_shader_lib {
        return conv2d_pipelines[key];
    }

+    // whcn selects the input layout: contiguous WHCN vs contiguous-channels CWHN
+    webgpu_pipeline get_conv2d_dw_pipeline(const ggml_webgpu_shader_lib_context & context, bool whcn) {
+        ggml_webgpu_conv2d_dw_pipeline_key key = {};
+        key.weight_type                        = context.src0->type;
+        key.input_type                         = context.src1->type;
+        key.output_type                        = context.dst->type;
+        key.whcn                               = whcn;
+
+        auto it = conv2d_dw_pipelines.find(key);
+        if (it != conv2d_dw_pipelines.end()) {
+            return it->second;
+        }
+
+        std::vector<std::string> defines;
+        std::string              variant = whcn ? "conv_2d_dw_whcn" : "conv_2d_dw_cwhn";
+
+        auto push_type_defines = [&](const char * prefix, ggml_type type) {
+            std::string s_prefix = prefix;
+            if (type == GGML_TYPE_F32) {
+                defines.push_back(s_prefix + "_F32");
+            } else if (type == GGML_TYPE_F16) {
+                defines.push_back(s_prefix + "_F16");
+            } else {
+                GGML_ABORT("Unsupported type for CONV_2D_DW shader");
+            }
+        };
+
+        push_type_defines("WEIGHT", key.weight_type);
+        push_type_defines("INPUT", key.input_type);
+        push_type_defines("OUTPUT", key.output_type);
+        if (whcn) {
+            defines.push_back("WHCN");
+        }
+        defines.push_back(std::string("WG_SIZE=") + std::to_string(context.max_wg_size));
+
+        auto processed           = preprocessor.preprocess(wgsl_conv2d_dw, defines);
+        auto decisions           = std::make_shared<ggml_webgpu_generic_shader_decisions>();
+        decisions->wg_size       = context.max_wg_size;
+        webgpu_pipeline pipeline = ggml_webgpu_create_pipeline(device, processed, variant);
+        pipeline.context         = decisions;
+        conv2d_dw_pipelines[key] = pipeline;
+        return conv2d_dw_pipelines[key];
+    }
+
    webgpu_pipeline get_im2col_pipeline(const ggml_webgpu_shader_lib_context & context) {
        ggml_webgpu_im2col_pipeline_key key = {};
        key.input_type                      = context.src1->type;
@@ -978,6 +978,67 @@ static webgpu_encoded_op ggml_webgpu_conv_2d(webgpu_context & ctx,
    return ggml_backend_webgpu_build(ctx, pipeline, params, entries, wg_x, wg_y);
 }

+// Same param/binding layout as conv_2d; the shader differs
+static webgpu_encoded_op ggml_webgpu_conv_2d_dw(webgpu_context & ctx,
+                                                ggml_tensor *    src0,
+                                                ggml_tensor *    src1,
+                                                ggml_tensor *    dst) {
+    const int32_t s0 = ggml_get_op_params_i32(dst, 0);
+    const int32_t s1 = ggml_get_op_params_i32(dst, 1);
+    const int32_t p0 = ggml_get_op_params_i32(dst, 2);
+    const int32_t p1 = ggml_get_op_params_i32(dst, 3);
+    const int32_t d0 = ggml_get_op_params_i32(dst, 4);
+    const int32_t d1 = ggml_get_op_params_i32(dst, 5);
+
+    // Scalar params matching conv2d_dw.wgsl (weight src0 [KW,KH,1,C], input src1, output dst).
+    std::vector<uint32_t> params = {
+        (uint32_t) (ggml_webgpu_tensor_misalignment(ctx, src0) / ggml_type_size(src0->type)),
+        (uint32_t) (ggml_webgpu_tensor_misalignment(ctx, src1) / ggml_type_size(src1->type)),
+        (uint32_t) (ggml_webgpu_tensor_misalignment(ctx, dst) / ggml_type_size(dst->type)),
+
+        (uint32_t) ggml_nelements(dst),
+        (uint32_t) dst->ne[2],
+        (uint32_t) dst->ne[3],
+        (uint32_t) dst->ne[0],
+        (uint32_t) dst->ne[1],
+        (uint32_t) src1->ne[0],
+        (uint32_t) src1->ne[1],
+        (uint32_t) src0->ne[0],
+        (uint32_t) src0->ne[1],
+
+        (uint32_t) s0,
+        (uint32_t) s1,
+        (uint32_t) p0,
+        (uint32_t) p1,
+        (uint32_t) d0,
+        (uint32_t) d1,
+    };
+
+    std::vector<wgpu::BindGroupEntry> entries = {
+        ggml_webgpu_make_tensor_bind_group_entry(ctx, 0, src0),
+        ggml_webgpu_make_tensor_bind_group_entry(ctx, 1, src1),
+        ggml_webgpu_make_tensor_bind_group_entry(ctx, 2, dst),
+    };
+
+    ggml_webgpu_shader_lib_context shader_lib_ctx = {};
+    shader_lib_ctx.src0                           = src0;
+    shader_lib_ctx.src1                           = src1;
+    shader_lib_ctx.dst                            = dst;
+    shader_lib_ctx.max_wg_size = ctx->global_ctx->capabilities.limits.maxComputeInvocationsPerWorkgroup;
+
+    // Input layout: contiguous -> WHCN, contiguous-channels -> CWHN
+    const bool      whcn      = ggml_is_contiguous(src1);
+    webgpu_pipeline pipeline  = ctx->shader_lib->get_conv2d_dw_pipeline(shader_lib_ctx, whcn);
+    auto *          decisions = static_cast<ggml_webgpu_generic_shader_decisions *>(pipeline.context.get());
+
+    uint32_t wg_x;
+    uint32_t wg_y;
+    uint32_t total_wg = CEIL_DIV((uint32_t) ggml_nelements(dst), decisions->wg_size);
+    compute_2d_workgroups(total_wg, ctx->global_ctx->capabilities.limits.maxComputeWorkgroupsPerDimension, wg_x, wg_y);
+
+    return ggml_backend_webgpu_build(ctx, pipeline, params, entries, wg_x, wg_y);
+}
+
 static webgpu_encoded_op ggml_webgpu_im2col(webgpu_context & ctx,
                                            ggml_tensor *    src0,
                                            ggml_tensor *    src1,
@@ -3164,6 +3225,8 @@ static std::optional<webgpu_encoded_op> ggml_webgpu_encode(webgpu_context ctx,
            return ggml_webgpu_sum_rows(ctx, src0, node);
        case GGML_OP_CONV_2D:
            return ggml_webgpu_conv_2d(ctx, src0, src1, node);
+        case GGML_OP_CONV_2D_DW:
+            return ggml_webgpu_conv_2d_dw(ctx, src0, src1, node);
        case GGML_OP_IM2COL:
            return ggml_webgpu_im2col(ctx, src0, src1, node);
        case GGML_OP_UPSCALE:
@@ -4004,6 +4067,7 @@ static ggml_backend_t ggml_backend_webgpu_backend_init(ggml_backend_dev_t dev, c
        /* .interface = */ ggml_backend_webgpu_i,
        /* .device    = */ dev,
        /* .context   = */ backend_ctx,
+        /* .profiler  = */ nullptr,
    };
    return backend;
 }
@@ -4349,6 +4413,12 @@ static bool ggml_backend_webgpu_device_supports_op(ggml_backend_dev_t dev, const
                          (src0->type == GGML_TYPE_F32 || src0->type == GGML_TYPE_F16) &&
                          (src1->type == GGML_TYPE_F32 || src1->type == GGML_TYPE_F16);
            break;
+        case GGML_OP_CONV_2D_DW:
+            supports_op = (op->type == GGML_TYPE_F32 || op->type == GGML_TYPE_F16) &&
+                          (src0->type == GGML_TYPE_F32 || src0->type == GGML_TYPE_F16) &&
+                          (src1->type == GGML_TYPE_F32 || src1->type == GGML_TYPE_F16) &&
+                          (ggml_is_contiguous(src1) || ggml_is_contiguous_channels(src1));
+            break;
        case GGML_OP_IM2COL:
            supports_op = (op->type == GGML_TYPE_F32 || op->type == GGML_TYPE_F16) &&
                          (src0->type == GGML_TYPE_F32 || src0->type == GGML_TYPE_F16);
@@ -0,0 +1,137 @@
+#include "common_decls.tmpl"
+enable f16;
+
+// Ported from the Vulkan backend's conv2d_dw.comp. Two variants (based on WHCN)
+// selected by the input (src1) layout: contiguous -> WHCN, else CWHN.
+// weight (src0) is [KW,KH,1,C]; output matches the input layout.
+
+@group(0) @binding(0)
+#if defined(WEIGHT_F32)
+var<storage, read_write> weights: array<f32>;
+#elif defined(WEIGHT_F16)
+var<storage, read_write> weights: array<f16>;
+#endif
+
+@group(0) @binding(1)
+#if defined(INPUT_F32)
+var<storage, read_write> input: array<f32>;
+#elif defined(INPUT_F16)
+var<storage, read_write> input: array<f16>;
+#endif
+
+@group(0) @binding(2)
+#if defined(OUTPUT_F32)
+var<storage, read_write> output: array<f32>;
+#elif defined(OUTPUT_F16)
+var<storage, read_write> output: array<f16>;
+#endif
+
+struct Params {
+    offset_w: u32,
+    offset_i: u32,
+    offset_o: u32,
+
+    ne: u32,
+    channels: u32,
+    batches: u32,
+    dst_w: u32, dst_h: u32,
+    src_w: u32, src_h: u32,
+    knl_w: u32, knl_h: u32,
+
+    stride_x: i32, stride_y: i32,
+    pad_x: i32, pad_y: i32,
+    dilation_x: i32, dilation_y: i32,
+};
+
+@group(0) @binding(3)
+var<uniform> params: Params;
+
+fn load_weight(idx: u32) -> f32 {
+    #if defined(WEIGHT_F32)
+        return weights[idx];
+    #elif defined(WEIGHT_F16)
+        return f32(weights[idx]);
+    #endif
+}
+fn load_input(idx: u32) -> f32 {
+    #if defined(INPUT_F32)
+        return input[idx];
+    #elif defined(INPUT_F16)
+        return f32(input[idx]);
+    #endif
+}
+fn store_output(idx: u32, val: f32) {
+    #if defined(OUTPUT_F32)
+        output[idx] = val;
+    #elif defined(OUTPUT_F16)
+        output[idx] = f16(val);
+    #endif
+}
+
+#if defined(WHCN)
+// Input/output/kernel contiguous in [W, H, C, N] order (kernel [KW,KH,C]).
+fn conv_2d_dw(idx: u32) -> f32 {
+    let i0    = idx / params.dst_w;
+    let dst_x = idx - i0 * params.dst_w;
+    let i1    = i0 / params.dst_h;
+    let dst_y = i0 - i1 * params.dst_h;
+    let n     = i1 / params.channels;
+    let c     = i1 - n * params.channels;
+
+    let src_i = params.offset_i + n * params.channels * params.src_h * params.src_w
+                                + c * params.src_h * params.src_w;
+    let knl_i = params.offset_w + c * params.knl_h * params.knl_w;
+
+    var sum: f32 = 0.0;
+    for (var ky: u32 = 0u; ky < params.knl_h; ky += 1u) {
+        let src_y = i32(dst_y) * params.stride_y + i32(ky) * params.dilation_y - params.pad_y;
+        if (src_y < 0 || src_y >= i32(params.src_h)) { continue; }
+        for (var kx: u32 = 0u; kx < params.knl_w; kx += 1u) {
+            let src_x = i32(dst_x) * params.stride_x + i32(kx) * params.dilation_x - params.pad_x;
+            if (src_x < 0 || src_x >= i32(params.src_w)) { continue; }
+            let v = load_input(src_i + u32(src_y) * params.src_w + u32(src_x));
+            let k = load_weight(knl_i + ky * params.knl_w + kx);
+            sum += v * k;
+        }
+    }
+    return sum;
+}
+#else
+// Channels contiguous (CWHN): channel is the innermost axis.
+fn conv_2d_dw(idx: u32) -> f32 {
+    let i0    = idx / params.channels;
+    let c     = idx - i0 * params.channels;
+    let i1    = i0 / params.dst_w;
+    let dst_x = i0 - i1 * params.dst_w;
+    let n     = i1 / params.dst_h;
+    let dst_y = i1 - n * params.dst_h;
+
+    let src_i   = params.offset_i + n * params.channels * params.src_h * params.src_w;
+    let src_row = params.src_w * params.channels;
+    let knl_row = params.knl_w * params.channels;
+
+    var sum: f32 = 0.0;
+    for (var ky: u32 = 0u; ky < params.knl_h; ky += 1u) {
+        let src_y = i32(dst_y) * params.stride_y + i32(ky) * params.dilation_y - params.pad_y;
+        if (src_y < 0 || src_y >= i32(params.src_h)) { continue; }
+        for (var kx: u32 = 0u; kx < params.knl_w; kx += 1u) {
+            let src_x = i32(dst_x) * params.stride_x + i32(kx) * params.dilation_x - params.pad_x;
+            if (src_x < 0 || src_x >= i32(params.src_w)) { continue; }
+            let v = load_input(src_i + u32(src_y) * src_row + u32(src_x) * params.channels + c);
+            let k = load_weight(params.offset_w + ky * knl_row + kx * params.channels + c);
+            sum += v * k;
+        }
+    }
+    return sum;
+}
+#endif
+
+@compute @workgroup_size(WG_SIZE)
+fn main(
+    @builtin(global_invocation_id) gid: vec3<u32>,
+    @builtin(num_workgroups) num_wg: vec3<u32>
+) {
+    let idx = gid.x + (num_wg.x * u32(WG_SIZE)) * gid.y;
+    if (idx >= params.ne) { return; }
+    store_output(params.offset_o + idx, conv_2d_dw(idx));
+}
@@ -503,7 +503,8 @@ static ggml_backend_t ggml_backend_zdnn_device_init(ggml_backend_dev_t dev, cons
        /* .guid       = */ ggml_backend_zdnn_guid(),
        /* .iface      = */ ggml_backend_zdnn_i,
        /* .device     = */ dev,
-        /* .context    = */ ctx
+        /* .context    = */ ctx,
+        /* .profiler   = */ NULL,
    };

    return backend;
@@ -467,7 +467,8 @@ ggml_backend_t ggml_backend_zendnn_init(void) {
        /* .guid    = */ ggml_backend_zendnn_guid(),
        /* .iface   = */ ggml_backend_zendnn_i,
        /* .device  = */ ggml_backend_reg_dev_get(ggml_backend_zendnn_reg(), 0),
-        /* .context = */ ctx,
+        /* .context  = */ ctx,
+        /* .profiler = */ nullptr,
    };

    return backend;
@@ -507,6 +507,7 @@ class MODEL_ARCH(IntEnum):
    DOTS1            = auto()
    ARCEE            = auto()
    AFMOE            = auto()
+    LAGUNA           = auto()
    ERNIE4_5         = auto()
    ERNIE4_5_MOE     = auto()
    HUNYUAN_MOE      = auto()
@@ -1088,6 +1089,7 @@ MODEL_ARCH_NAMES: dict[MODEL_ARCH, str] = {
    MODEL_ARCH.DOTS1:            "dots1",
    MODEL_ARCH.ARCEE:            "arcee",
    MODEL_ARCH.AFMOE:            "afmoe",
+    MODEL_ARCH.LAGUNA:           "laguna",
    MODEL_ARCH.ERNIE4_5:         "ernie4_5",
    MODEL_ARCH.ERNIE4_5_MOE:     "ernie4_5-moe",
    MODEL_ARCH.FALCON_H1:        "falcon-h1",
@@ -3823,6 +3825,31 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = {
        MODEL_TENSOR.FFN_POST_NORM,
        MODEL_TENSOR.FFN_EXP_PROBS_B,
    ],
+    MODEL_ARCH.LAGUNA: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_Q,
+        MODEL_TENSOR.ATTN_K,
+        MODEL_TENSOR.ATTN_V,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.ATTN_Q_NORM,
+        MODEL_TENSOR.ATTN_K_NORM,
+        MODEL_TENSOR.ATTN_GATE,
+        MODEL_TENSOR.FFN_NORM,
+        MODEL_TENSOR.FFN_GATE,
+        MODEL_TENSOR.FFN_DOWN,
+        MODEL_TENSOR.FFN_UP,
+        MODEL_TENSOR.FFN_GATE_INP,
+        MODEL_TENSOR.FFN_EXP_PROBS_B,
+        MODEL_TENSOR.FFN_GATE_EXP,
+        MODEL_TENSOR.FFN_DOWN_EXP,
+        MODEL_TENSOR.FFN_UP_EXP,
+        MODEL_TENSOR.FFN_GATE_SHEXP,
+        MODEL_TENSOR.FFN_UP_SHEXP,
+        MODEL_TENSOR.FFN_DOWN_SHEXP,
+    ],
    MODEL_ARCH.ERNIE4_5: [
        MODEL_TENSOR.TOKEN_EMBD,
        MODEL_TENSOR.OUTPUT_NORM,
@@ -479,6 +479,7 @@ class TensorNameMap:
            "model.layers.{bid}.mlp.e_score_correction",                    # exaone-moe
            "model.layers.{bid}.block_sparse_moe.gate.e_score_correction",  # kimi
            "model.layers.{bid}.moe.router_bias",                           # step3.5 expert selection bias
+            "model.layers.{bid}.mlp.experts.e_score_correction",       # laguna
        ),

        # Feed-forward up
@@ -202,6 +202,16 @@ extern "C" {
        LLAMA_SPLIT_MODE_TENSOR = 3,
    };

+    enum llama_load_mode {
+        LLAMA_LOAD_MODE_NONE      = 0, // no special loading mode
+        LLAMA_LOAD_MODE_MMAP      = 1, // memory map the model
+        LLAMA_LOAD_MODE_MLOCK     = 2, // mmap + force system to keep model in RAM rather than swapping or compressing
+        LLAMA_LOAD_MODE_DIRECT_IO = 3, // use direct I/O if available
+    };
+
+    LLAMA_API const char * llama_load_mode_name(enum llama_load_mode load_mode);
+    LLAMA_API enum llama_load_mode llama_load_mode_from_str(const char * str);
+
    enum llama_context_type {
        LLAMA_CONTEXT_TYPE_DEFAULT = 0,
        LLAMA_CONTEXT_TYPE_MTP     = 1,
@@ -301,6 +311,7 @@ extern "C" {

        int32_t n_gpu_layers; // number of layers to store in VRAM, a negative value means all layers
        enum llama_split_mode split_mode; // how to split the model across multiple GPUs
+        enum llama_load_mode  load_mode;  // how to load the model

        // the GPU that is used for the entire model when split_mode is LLAMA_SPLIT_MODE_NONE
        int32_t main_gpu;
@@ -321,9 +332,6 @@ extern "C" {

        // Keep the booleans together to avoid misalignment during copy-by-value.
        bool vocab_only;      // only load the vocabulary, no weights
-        bool use_mmap;        // use mmap if possible
-        bool use_direct_io;   // use direct io, takes precedence over use_mmap when supported
-        bool use_mlock;       // force system to keep model in RAM
        bool check_tensors;   // validate model tensor data
        bool use_extra_bufts; // use extra buffer types (used for weight repacking)
        bool no_host;         // bypass host buffer allowing extra buffers to be used
@@ -559,6 +567,8 @@ extern "C" {
    LLAMA_API           llama_memory_t   llama_get_memory  (const struct llama_context * ctx);
    LLAMA_API  enum llama_pooling_type   llama_pooling_type(const struct llama_context * ctx); // TODO: rename to llama_get_pooling_type

+    LLAMA_API struct ggml_backend_sched * llama_context_get_sched(const struct llama_context * ctx);
+
    LLAMA_API const struct llama_vocab * llama_model_get_vocab(const struct llama_model * model);
    LLAMA_API enum llama_rope_type       llama_model_rope_type(const struct llama_model * model);

@@ -8,12 +8,15 @@
    {%- set thinking = false -%}
  {%- endif -%}
 {%- endif -%}
+{%- if not drop_thinking is defined -%}
+  {%- set drop_thinking = false -%}
+{%- endif -%}
 {%- set dsml_token = '｜DSML｜' -%}
 {%- set thinking_start_token = '<think>' -%}
 {%- set thinking_end_token = '</think>' -%}
 {%- set tools_header = '## Tools\n\nYou have access to a set of tools to help answer the user\'s question. You can invoke tools by writing a "<' + dsml_token + 'tool_calls>" block like the following:\n\n<' + dsml_token + 'tool_calls>\n<' + dsml_token + 'invoke name="$TOOL_NAME">\n<' + dsml_token + 'parameter name="$PARAMETER_NAME" string="true|false">$PARAMETER_VALUE</' + dsml_token + 'parameter>\n...\n</' + dsml_token + 'invoke>\n<' + dsml_token + 'invoke name="$TOOL_NAME2">\n...\n</' + dsml_token + 'invoke>\n</' + dsml_token + 'tool_calls>\n\nString parameters should be specified as is and set `string="true"`. For all other types (numbers, booleans, arrays, objects), pass the value in JSON format and set `string="false"`.\n\nIf thinking_mode is enabled (triggered by ' + thinking_start_token + '), you MUST output your complete reasoning inside ' + thinking_start_token + '...' + thinking_end_token + ' BEFORE any tool calls or final response.\n\nOtherwise, output directly after ' + thinking_end_token + ' with tool calls or final response.\n\n### Available Tool Schemas\n\n' -%}
 {%- set tools_footer = '\nYou MUST strictly follow the above defined tool name and parameter schemas to invoke tool calls.\n' -%}
-{%- set ns = namespace(system_prompt='', is_first_sp=true) -%}
+{%- set ns = namespace(system_prompt='', is_first_sp=true, has_tool_calls=false) -%}
 {%- for message in messages -%}
  {%- if message['role'] == 'system' -%}
    {%- if ns.is_first_sp -%}
@@ -46,6 +49,11 @@
  {%- endif -%}
 {%- endfor -%}
 {%- set state = namespace(in_user=false) -%}
+{%- for message in messages -%}
+  {%- if message['role'] == 'tool' -%}
+    {%- set ns.has_tool_calls = true -%}
+  {%- endif -%}
+{%- endfor -%}
 {%- for message in messages -%}
  {%- if message['role'] == 'user' or message['role'] == 'developer' -%}
    {%- if state.in_user -%}
@@ -67,7 +75,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 retain_reasoning = (not drop_thinking) or (is_after_last_user or ns.has_tool_calls) -%}
+    {%- if retain_reasoning and thinking -%}
      {{- thinking_start_token -}}
      {%- if message['reasoning_content'] is defined and message['reasoning_content'] -%}
        {{- message['reasoning_content'] -}}
@@ -0,0 +1,93 @@
+{#- Iteration on laguna_glm_thinking_v8/chat_template.jinja -#}
+{#- No formatting instructions -#}
+{{- "〈|EOS|〉" -}}
+{%- set enable_thinking = enable_thinking | default(false) -%}
+{%- set add_generation_prompt = add_generation_prompt | default(false) -%}
+
+{#- ───── header (system message) ───── -#}
+{#- A caller-supplied system message with empty content opts out of the default below, producing no <system> block — used to train without a system message. -#}
+{%- set system_message = "You are a helpful, conversationally-fluent assistant made by Poolside. You are here to be helpful to users through natural language conversations." -%}
+{%- if messages and messages[0].role == "system" -%}
+  {%- set system_message = messages[0].content -%}
+  {%- set messages = messages[1:] -%}
+{%- endif -%}
+
+{%- set has_sys = system_message and system_message.strip() -%}
+{%- if has_sys or tools or enable_thinking -%}
+  {{- "<system>" -}}
+
+  {%- if has_sys -%}
+    {{- system_message.rstrip() -}}
+    {%- if tools -%}{{- "\n\n" -}}{%- endif -%}
+  {%- endif -%}
+
+  {%- if tools -%}
+    {{- "### Tools\n\n" -}}
+    {{- "You may call functions to assist with the user query.\n" -}}
+    {{- "All available function signatures are listed below:\n" -}}
+    {{- "<available_tools>\n" -}}
+    {%- for tool in tools -%}
+      {{- (tool | tojson) ~ "\n" -}}
+    {%- endfor -%}
+    {{- "</available_tools>" -}}
+  {%- endif -%}
+
+  {{- "</system>\n" -}}
+{%- endif -%}
+
+{#- ───── main loop ───── -#}
+{%- for message in messages -%}
+  {%- set content = message.content if message.content is string else "" -%}
+  {%- if message.role == "user" -%}
+    {{- "<user>" + content + "</user>\n" -}}
+  {%- elif message.role == "assistant" -%}
+    {%- generation -%}
+      {{- "<assistant>" -}}
+      {#- Extract reasoning content from message.reasoning (vLLM field name) or message.reasoning_content -#}
+      {%- set reasoning_content = '' -%}
+      {%- if message.reasoning is string -%}
+        {%- set reasoning_content = message.reasoning -%}
+      {%- elif message.reasoning_content is string -%}
+        {%- set reasoning_content = message.reasoning_content -%}
+      {%- endif -%}
+      {#- Display reasoning content for all messages if enable_thinking -#}
+      {%- if enable_thinking -%}
+        {{- '<think>' + reasoning_content + '</think>' -}}
+      {%- else -%}
+        {{- '</think>' -}}
+      {%- endif -%}
+      {#- Display main content (trailing newline only when no tool_calls follow) -#}
+      {%- if content -%}
+        {{- content -}}
+      {%- endif -%}
+      {%- if message.tool_calls -%}
+        {%- for tool_call in message.tool_calls -%}
+          {%- set function_data = tool_call.function -%}
+          {{- '<tool_call>' + function_data.name -}}
+          {%- set _args = function_data.arguments -%}
+          {%- for k, v in _args.items() -%}
+            {{- "<arg_key>" ~ k ~ "</arg_key>" -}}
+            {{- "<arg_value>" -}}{{- v | tojson(ensure_ascii=False) if v is not string else v -}}{{- "</arg_value>" -}}
+          {%- endfor -%}
+          {{- "</tool_call>" -}}
+        {%- endfor -%}
+      {%- endif -%}
+      {{- "</assistant>\n" -}}
+    {%- endgeneration -%}
+  {%- elif message.role == "tool" -%}
+    {{- "<tool_response>" + content + "</tool_response>\n" -}}
+  {%- elif message.role == "system" -%}
+    {#- Render additional system messages (the first one, if any, is handled separately in the header and was sliced off above) -#}
+    {{- "<system>" + content + "</system>\n" -}}
+  {%- endif -%}
+{%- endfor -%}
+{#- ───── generation prompt ───── -#}
+{%- if add_generation_prompt -%}
+  {{- "<assistant>" -}}
+  {#- ───── Include reasoning mode directive ───── -#}
+  {%- if enable_thinking -%}
+    {{- '<think>' -}}
+  {%- else -%}
+    {{- '</think>' -}}
+  {%- endif -%}
+{%- endif -%}
@@ -0,0 +1,132 @@
+{#- Copied from laguna_glm_thinking_v4/chat_template.jinja -#}
+{#- Removes prefix that references <think> token, and replaces message.reasoning_content reference with message.reasoning -#}
+{{- "〈|EOS|〉" -}}
+{%- set enable_thinking = enable_thinking | default(false) -%}
+{%- set render_assistant_messages_raw = render_assistant_messages_raw | default(false) -%}
+{%- set add_generation_prompt = add_generation_prompt | default(false) -%}
+
+{#- ───── header (system message) ───── -#}
+{%- set system_message = "" -%}
+{%- if messages and messages[0].role == "system" -%}
+  {%- set system_message = messages[0].content -%}
+{%- endif -%}
+
+{%- if (system_message and system_message.strip()) or tools -%}
+  {{- "<system>\n" -}}
+
+  {%- if system_message and system_message.strip() -%}
+    {{- "\n" -}}
+    {{- system_message.rstrip() -}}
+  {%- endif -%}
+
+  {%- if tools -%}
+    {{- "\n\n### Tools\n\n" -}}
+    {%- set ns = namespace(tool_string="You may call functions to assist with the user query.\n"
+          ~ "All available function signatures are listed below:\n"
+          ~ "<available_tools>\n") -%}
+    {%- for tool in tools -%}
+      {%- set ns.tool_string = ns.tool_string ~ (tool | tojson) ~ "\n" -%}
+    {%- endfor -%}
+    {%- if enable_thinking -%}
+      {%- set tool_string = ns.tool_string + "</available_tools>\n\n" ~
+            "Wrap your thinking in '<think>', '</think>' tags, followed by a function call. For each function call, return an unescaped XML-like object with function name and arguments within '<tool_call>' and '</tool_call>' tags, like here:\n" ~
+            "<think> your thoughts here </think>\n" ~
+            "<tool_call>function-name\n<arg_key>argument-key</arg_key>\n<arg_value>value-of-argument-key</arg_value>\n" ~
+            "</tool_call>" -%}
+    {%- else -%}
+      {%- set tool_string = ns.tool_string + "</available_tools>\n\n" ~
+            "For each function call, return an unescaped XML-like object " ~
+            "with function name and arguments within '<tool_call>' and '</tool_call>' tags, like here:\n" ~
+            "<tool_call>function-name\n<arg_key>argument-key</arg_key>\n<arg_value>value-of-argument-key</arg_value>\n" ~
+            "</tool_call>" -%}
+    {%- endif -%}
+    {{- tool_string -}}
+  {%- endif -%}
+
+  {{- "\n</system>\n" -}}
+{%- endif -%}
+
+{#- ───── main loop ───── -#}
+{%- for message in messages -%}
+  {%- set content = message.content if message.content is string else "" -%}
+  {%- if message.role == "user" -%}
+    {{- "<user>\n" + content + "\n</user>\n" -}}
+  {%- elif message.role == "assistant" -%}
+    {%- generation -%}
+      {{- "<assistant>\n" -}}
+      {%- if render_assistant_messages_raw -%}
+        {#- Raw mode: prepend the generation prompt token, then dump content verbatim. -#}
+        {#- The generation prompt is <think> when enable_thinking, </think> otherwise. -#}
+        {#- Only prepend if content doesn't already start with it. -#}
+        {%- if enable_thinking -%}
+          {%- if not content.startswith('<think>') -%}
+            {{- '<think>' -}}
+          {%- endif -%}
+        {%- else -%}
+          {%- if not content.startswith('</think>') -%}
+            {{- '</think>' -}}
+          {%- endif -%}
+        {%- endif -%}
+        {{- content -}}
+        {#- Append closing tag if content doesn't already end with it. -#}
+        {%- if not content.endswith('</assistant>\n') and not content.endswith('</assistant>') -%}
+          {{- '\n</assistant>' -}}
+        {%- endif -%}
+        {{- "\n" -}}
+      {%- else -%}
+        {#- Extract reasoning content from message.reasoning (vLLM field name) or message.reasoning_content, or from <think> tags -#}
+        {%- set reasoning_content = '' %}
+        {%- if message.reasoning is string %}
+          {%- set reasoning_content = message.reasoning %}
+        {%- elif message.reasoning_content is string %}
+          {%- set reasoning_content = message.reasoning_content %}
+        {%- endif %}
+        {#- Always strip <think> tags from content if present to avoid duplication -#}
+        {%- if '</think>' in content %}
+          {%- if not reasoning_content %}
+            {%- set reasoning_content = content.split('</think>')[0].rstrip('\n').split('<think>')[-1].lstrip('\n') %}
+          {%- endif %}
+          {%- set content = content.split('</think>')[-1].lstrip('\n') %}
+        {%- endif %}
+        {#- Display reasoning content for all messages -#}
+        {%- if reasoning_content -%}
+          {{- '<think>\n' + reasoning_content.strip() + '\n</think>\n' -}}
+        {%- else -%}
+          {{- '</think>\n' -}}
+        {%- endif -%}
+        {#- Display main content -#}
+        {%- if content.strip() -%}
+          {{- content.strip() ~ "\n" -}}
+        {%- endif -%}
+        {%- if message.tool_calls -%}
+          {%- for tool_call in message.tool_calls -%}
+            {%- set function_data = tool_call.function -%}
+            {{- '<tool_call>' + function_data.name }}
+            {% set _args = function_data.arguments %}
+            {%- for k, v in _args.items() -%}
+              {{- "<arg_key>" ~ k ~ "</arg_key>\n" -}}
+              {{- "<arg_value>"}}{{ v | tojson(ensure_ascii=False) if v is not string else v }}{{ "</arg_value>\n" -}}
+            {%- endfor -%}
+            {{- "</tool_call>\n" -}}
+          {%- endfor -%}
+        {%- endif -%}
+        {{- "</assistant>\n" -}}
+      {%- endif -%}
+    {%- endgeneration -%}
+  {%- elif message.role == "tool" -%}
+    {{- "<tool_response>\n" + content + "\n</tool_response>\n" -}}
+  {%- elif message.role == "system" and loop.index0 != 0 -%}
+    {#- Render additional system messages (skip the first one which is handled separately in the header) -#}
+    {{- "<system>\n" + content + "\n</system>\n" -}}
+  {%- endif -%}
+{%- endfor -%}
+{#- ───── generation prompt ───── -#}
+{%- if add_generation_prompt -%}
+  {{- "<assistant>\n" -}}
+  {#- ───── Include reasoning mode directive ───── -#}
+  {%- if not enable_thinking %}
+    {{- '</think>' -}}
+  {%- else %}
+    {{- '<think>' -}}
+  {%- endif %}
+{%- endif -%}
@@ -0,0 +1,132 @@
+{#- Iteration on laguna_glm_thinking_v5/chat_template.jinja -#}
+{#- Adds a default system message (used when no system message is provided in `messages`). -#}
+{{- "〈|EOS|〉" -}}
+{%- set enable_thinking = enable_thinking | default(false) -%}
+{%- set render_assistant_messages_raw = render_assistant_messages_raw | default(false) -%}
+{%- set add_generation_prompt = add_generation_prompt | default(false) -%}
+
+{#- ───── header (system message) ───── -#}
+{%- set system_message = "You are a helpful, conversationally-fluent assistant made by Poolside. You are here to be helpful to users through natural language conversations." -%}
+{%- if messages and messages[0].role == "system" -%}
+  {%- set system_message = messages[0].content -%}
+{%- endif -%}
+
+{%- if (system_message and system_message.strip()) or tools -%}
+  {{- "<system>\n" -}}
+
+  {%- if system_message and system_message.strip() -%}
+    {{- "\n" -}}
+    {{- system_message.rstrip() -}}
+  {%- endif -%}
+
+  {%- if tools -%}
+    {{- "\n\n### Tools\n\n" -}}
+    {%- set ns = namespace(tool_string="You may call functions to assist with the user query.\n"
+          ~ "All available function signatures are listed below:\n"
+          ~ "<available_tools>\n") -%}
+    {%- for tool in tools -%}
+      {%- set ns.tool_string = ns.tool_string ~ (tool | tojson) ~ "\n" -%}
+    {%- endfor -%}
+    {%- if enable_thinking -%}
+      {%- set tool_string = ns.tool_string + "</available_tools>\n\n" ~
+            "Wrap your thinking in '<think>', '</think>' tags, followed by a function call. For each function call, return an unescaped XML-like object with function name and arguments within '<tool_call>' and '</tool_call>' tags, like here:\n" ~
+            "<think> your thoughts here </think>\n" ~
+            "<tool_call>function-name\n<arg_key>argument-key</arg_key>\n<arg_value>value-of-argument-key</arg_value>\n" ~
+            "</tool_call>" -%}
+    {%- else -%}
+      {%- set tool_string = ns.tool_string + "</available_tools>\n\n" ~
+            "For each function call, return an unescaped XML-like object " ~
+            "with function name and arguments within '<tool_call>' and '</tool_call>' tags, like here:\n" ~
+            "<tool_call>function-name\n<arg_key>argument-key</arg_key>\n<arg_value>value-of-argument-key</arg_value>\n" ~
+            "</tool_call>" -%}
+    {%- endif -%}
+    {{- tool_string -}}
+  {%- endif -%}
+
+  {{- "\n</system>\n" -}}
+{%- endif -%}
+
+{#- ───── main loop ───── -#}
+{%- for message in messages -%}
+  {%- set content = message.content if message.content is string else "" -%}
+  {%- if message.role == "user" -%}
+    {{- "<user>\n" + content + "\n</user>\n" -}}
+  {%- elif message.role == "assistant" -%}
+    {%- generation -%}
+      {{- "<assistant>\n" -}}
+      {%- if render_assistant_messages_raw -%}
+        {#- Raw mode: prepend the generation prompt token, then dump content verbatim. -#}
+        {#- The generation prompt is <think> when enable_thinking, </think> otherwise. -#}
+        {#- Only prepend if content doesn't already start with it. -#}
+        {%- if enable_thinking -%}
+          {%- if not content.startswith('<think>') -%}
+            {{- '<think>' -}}
+          {%- endif -%}
+        {%- else -%}
+          {%- if not content.startswith('</think>') -%}
+            {{- '</think>' -}}
+          {%- endif -%}
+        {%- endif -%}
+        {{- content -}}
+        {#- Append closing tag if content doesn't already end with it. -#}
+        {%- if not content.endswith('</assistant>\n') and not content.endswith('</assistant>') -%}
+          {{- '\n</assistant>' -}}
+        {%- endif -%}
+        {{- "\n" -}}
+      {%- else -%}
+        {#- Extract reasoning content from message.reasoning (vLLM field name) or message.reasoning_content, or from <think> tags -#}
+        {%- set reasoning_content = '' %}
+        {%- if message.reasoning is string %}
+          {%- set reasoning_content = message.reasoning %}
+        {%- elif message.reasoning_content is string %}
+          {%- set reasoning_content = message.reasoning_content %}
+        {%- endif %}
+        {#- Always strip <think> tags from content if present to avoid duplication -#}
+        {%- if '</think>' in content %}
+          {%- if not reasoning_content %}
+            {%- set reasoning_content = content.split('</think>')[0].rstrip('\n').split('<think>')[-1].lstrip('\n') %}
+          {%- endif %}
+          {%- set content = content.split('</think>')[-1].lstrip('\n') %}
+        {%- endif %}
+        {#- Display reasoning content for all messages -#}
+        {%- if reasoning_content -%}
+          {{- '<think>\n' + reasoning_content.strip() + '\n</think>\n' -}}
+        {%- else -%}
+          {{- '</think>\n' -}}
+        {%- endif -%}
+        {#- Display main content -#}
+        {%- if content.strip() -%}
+          {{- content.strip() ~ "\n" -}}
+        {%- endif -%}
+        {%- if message.tool_calls -%}
+          {%- for tool_call in message.tool_calls -%}
+            {%- set function_data = tool_call.function -%}
+            {{- '<tool_call>' + function_data.name }}
+            {% set _args = function_data.arguments %}
+            {%- for k, v in _args.items() -%}
+              {{- "<arg_key>" ~ k ~ "</arg_key>\n" -}}
+              {{- "<arg_value>"}}{{ v | tojson(ensure_ascii=False) if v is not string else v }}{{ "</arg_value>\n" -}}
+            {%- endfor -%}
+            {{- "</tool_call>\n" -}}
+          {%- endfor -%}
+        {%- endif -%}
+        {{- "</assistant>\n" -}}
+      {%- endif -%}
+    {%- endgeneration -%}
+  {%- elif message.role == "tool" -%}
+    {{- "<tool_response>\n" + content + "\n</tool_response>\n" -}}
+  {%- elif message.role == "system" and loop.index0 != 0 -%}
+    {#- Render additional system messages (skip the first one which is handled separately in the header) -#}
+    {{- "<system>\n" + content + "\n</system>\n" -}}
+  {%- endif -%}
+{%- endfor -%}
+{#- ───── generation prompt ───── -#}
+{%- if add_generation_prompt -%}
+  {{- "<assistant>\n" -}}
+  {#- ───── Include reasoning mode directive ───── -#}
+  {%- if not enable_thinking %}
+    {{- '</think>' -}}
+  {%- else %}
+    {{- '<think>' -}}
+  {%- endif %}
+{%- endif -%}
@@ -28,7 +28,7 @@ LLAMA_BENCH_DB_FIELDS = [
    "model_type",   "model_size",   "model_n_params", "n_batch",    "n_ubatch",     "n_threads",
    "cpu_mask",     "cpu_strict",   "poll",           "type_k",     "type_v",       "n_gpu_layers",
    "split_mode",   "main_gpu",     "no_kv_offload",  "flash_attn", "tensor_split", "tensor_buft_overrides",
-    "use_mmap",     "embeddings",   "no_op_offload",  "n_prompt",   "n_gen",        "n_depth",
+    "load_mode",    "embeddings",   "no_op_offload",  "n_prompt",   "n_gen",        "n_depth",
    "test_time",    "avg_ns",       "stddev_ns",      "avg_ts",     "stddev_ts",    "n_cpu_moe",
    "fit_target",   "fit_min_ctx"
 ]
@@ -38,7 +38,7 @@ LLAMA_BENCH_DB_TYPES = [
    "TEXT",    "INTEGER", "INTEGER", "INTEGER", "INTEGER", "INTEGER",
    "TEXT",    "INTEGER", "INTEGER", "TEXT",    "TEXT",    "INTEGER",
    "TEXT",    "INTEGER", "INTEGER", "INTEGER", "TEXT",    "TEXT",
-    "INTEGER", "INTEGER", "INTEGER", "INTEGER", "INTEGER", "INTEGER",
+    "TEXT", "INTEGER", "INTEGER", "INTEGER", "INTEGER", "INTEGER",
    "TEXT",    "INTEGER", "INTEGER", "REAL",    "REAL",    "INTEGER",
    "INTEGER", "INTEGER"
 ]
@@ -63,7 +63,7 @@ assert len(TEST_BACKEND_OPS_DB_FIELDS) == len(TEST_BACKEND_OPS_DB_TYPES)
 LLAMA_BENCH_KEY_PROPERTIES = [
    "cpu_info", "gpu_info", "backends", "n_gpu_layers", "n_cpu_moe", "tensor_buft_overrides", "model_filename", "model_type",
    "n_batch", "n_ubatch", "embeddings", "cpu_mask", "cpu_strict", "poll", "n_threads", "type_k", "type_v",
-    "use_mmap", "no_kv_offload", "split_mode", "main_gpu", "tensor_split", "flash_attn", "n_prompt", "n_gen", "n_depth",
+    "load_mode", "no_kv_offload", "split_mode", "main_gpu", "tensor_split", "flash_attn", "n_prompt", "n_gen", "n_depth",
    "fit_target", "fit_min_ctx"
 ]

@@ -73,7 +73,7 @@ TEST_BACKEND_OPS_KEY_PROPERTIES = [
 ]

 # Properties that are boolean and are converted to Yes/No for the table:
-LLAMA_BENCH_BOOL_PROPERTIES = ["embeddings", "cpu_strict", "use_mmap", "no_kv_offload", "flash_attn"]
+LLAMA_BENCH_BOOL_PROPERTIES = ["embeddings", "cpu_strict", "no_kv_offload", "flash_attn"]
 TEST_BACKEND_OPS_BOOL_PROPERTIES = ["supported", "passed"]

 # Header names for the table (llama-bench):
@@ -82,7 +82,7 @@ LLAMA_BENCH_PRETTY_NAMES = {
    "tensor_buft_overrides": "Tensor overrides", "model_filename": "File", "model_type": "Model", "model_size": "Model size [GiB]",
    "model_n_params": "Num. of par.", "n_batch": "Batch size", "n_ubatch": "Microbatch size", "embeddings": "Embeddings",
    "cpu_mask": "CPU mask", "cpu_strict": "CPU strict", "poll": "Poll", "n_threads": "Threads", "type_k": "K type", "type_v": "V type",
-    "use_mmap": "Use mmap", "no_kv_offload": "NKVO", "split_mode": "Split mode", "main_gpu": "Main GPU", "tensor_split": "Tensor split",
+    "load_mode": "Load mode", "no_kv_offload": "NKVO", "split_mode": "Split mode", "main_gpu": "Main GPU", "tensor_split": "Tensor split",
    "flash_attn": "FlashAttention",
 }

@@ -6,6 +6,7 @@ import re
 import argparse
 import statistics
 import logging
+import bisect
 from typing import Any, Dict, List, Optional

 from collections import defaultdict
@@ -30,7 +31,7 @@ op_pattern = re.compile(
 )

 trace_pattern = re.compile(
-    r"trace-op\s+(?P<op_name>[A-Z_0-9+]+):\s+thread\s+(?P<thread>\d+)\s+event\s+(?P<event>[A-Z_0-9\-]+)\s+info\s+(?P<info>\d+)\s+(?P<state>start|stop)\s+(?P<cycles>\d+)"
+    r"trace-evt\s+(?P<event>[A-Z_0-9\-]+):\s+thread\s+(?P<thread>\d+)\s+info\s+(?P<info>\d+)\s+(?P<state>start|stop)\s+(?P<cycles>\d+)"
 )

 logger = logging.getLogger("ggml-hexagon-profile")
@@ -50,9 +51,13 @@ def normalize_event_name(evt_type):


 class CycleUnwrapper:
-    def __init__(self):
-        self.last_raw = None
-        self.high_part = 0
+    def __init__(self, initial_val=None):
+        if initial_val is not None:
+            self.last_raw = initial_val & 0xFFFFFFFF
+            self.high_part = initial_val & 0xFFFFFFFF00000000
+        else:
+            self.last_raw = None
+            self.high_part = 0

    def unwrap(self, raw):
        if self.last_raw is None:
@@ -78,10 +83,12 @@ def parse_log(file_path, pmu_index=None):
        sys.exit(1)

    all_ops: List[Dict[str, Any]] = []
+    all_traces: List[Dict[str, Any]] = []
    current_op: Optional[Dict[str, Any]] = None

    timestamp_pattern = re.compile(r"^(?P<min>\d+)\.(?P<sec>\d+)\.(?P<ms>\d+)\.(?P<us>\d+)\s+[A-Z]\s+")
-    unwrapper = CycleUnwrapper()
+    unwrapper = None
+    trace_unwrapper = None

    for line in f:
        ts_match = timestamp_pattern.match(line)
@@ -100,6 +107,7 @@ def parse_log(file_path, pmu_index=None):
            if not prefix_match:
                continue

+            names = parts[1]
            if len(parts) == 7:
                dims, types, timings = parts[2], parts[3], parts[6]
            elif len(parts) == 6:
@@ -120,6 +128,7 @@ def parse_log(file_path, pmu_index=None):
            op_match = op_pattern.search(line)
            if op_match:
                op_name = op_match.group('op_name')
+                names = ""
                dims = op_match.group('dims').strip()
                types = op_match.group('types').strip()
            else:
@@ -136,24 +145,31 @@ def parse_log(file_path, pmu_index=None):
                except (ValueError, IndexError):
                    pmu_val = None

-            evt_raw = op_match.group('evt') if 'evt' in op_match.groupdict() else None
            evt_val = None
-            if evt_raw:
+            evt_val = None
+            if types.startswith("evt-cnt "):
                try:
-                    evt_val = [int(x.strip()) for x in evt_raw.split(',')]
+                    evt_val = [int(x.strip()) for x in types[8:].split(',')]
                except ValueError:
                    evt_val = None

            cycles_start_raw = op_match.group('start')
            unwrapped_cycles_start = None
-            if cycles_start_raw:
-                unwrapped_cycles_start = unwrapper.unwrap(int(cycles_start_raw))
+            if op_name == "OPBATCH":
+                if cycles_start_raw:
+                    unwrapped_cycles_start = int(cycles_start_raw)
+                    unwrapper = CycleUnwrapper(unwrapped_cycles_start)
+                    trace_unwrapper = CycleUnwrapper(unwrapped_cycles_start)
+            else:
+                if cycles_start_raw and unwrapper is not None:
+                    unwrapped_cycles_start = unwrapper.unwrap(int(cycles_start_raw))

            idx = line.find("profile-op ")
            op_text = line[idx + 11:].strip() if idx != -1 else line.strip()

            current_op = {
                'name':         op_name,
+                'names':        names,
                'dims':         dims,
                'types':        types,
                'op_text':      op_text,
@@ -170,110 +186,239 @@ def parse_log(file_path, pmu_index=None):
            continue

        trace_match = trace_pattern.search(line)
-        if trace_match and current_op:
-            if trace_match.group('op_name') == current_op['name']:
-                raw_cyc = int(trace_match.group('cycles'))
-                current_op['trace_events'].append({
-                    'thread': int(trace_match.group('thread')),
-                    'event':  trace_match.group('event'),
-                    'info':   int(trace_match.group('info')),
-                    'cycles': raw_cyc,
-                    'unwrapped_cycles': unwrapper.unwrap(raw_cyc),
-                    'state':  trace_match.group('state')
-                })
+        if trace_match:
+            raw_cyc = int(trace_match.group('cycles'))
+            unwrapped_cyc = None
+            if trace_unwrapper is not None:
+                unwrapped_cyc = trace_unwrapper.unwrap(raw_cyc)
+            all_traces.append({
+                'thread': int(trace_match.group('thread')),
+                'event':  trace_match.group('event'),
+                'info':   int(trace_match.group('info')),
+                'cycles': raw_cyc,
+                'unwrapped_cycles': unwrapped_cyc,
+                'state':  trace_match.group('state')
+            })

    f.close()
+
+    # Assign start/end cycles to all ops
+    for op in all_ops:
+        op['start_cycles'] = op['unwrapped_cycles_start']
+        op['end_cycles'] = op['start_cycles'] + op['cycles'] if op['start_cycles'] is not None else None
+
+    # Filter ops with valid start_cycles
+    valid_ops = [op for op in all_ops if op['start_cycles'] is not None and op['end_cycles'] is not None]
+
+    # Separate OPBATCH ops from other ops
+    opbatch_ops = [op for op in valid_ops if op['name'] == "OPBATCH"]
+    other_ops = [op for op in valid_ops if op['name'] != "OPBATCH"]
+
+    # Sort them by start_cycles to enable binary search
+    opbatch_ops.sort(key=lambda op: op['start_cycles'])
+    other_ops.sort(key=lambda op: op['start_cycles'])
+
+    opbatch_starts = [op['start_cycles'] for op in opbatch_ops]
+    other_starts = [op['start_cycles'] for op in other_ops]
+
+    # Map trace events to any operator whose cycles contain them
+    for e in all_traces:
+        cyc = e['unwrapped_cycles']
+        if cyc is None:
+            continue
+
+        # Map to OPBATCH
+        idx = bisect.bisect_right(opbatch_starts, cyc) - 1
+        if idx >= 0:
+            op = opbatch_ops[idx]
+            if op['start_cycles'] <= cyc <= op['end_cycles']:
+                op['trace_events'].append(e)
+
+        # Map to other ops
+        idx = bisect.bisect_right(other_starts, cyc) - 1
+        if idx >= 0:
+            op = other_ops[idx]
+            if op['start_cycles'] <= cyc <= op['end_cycles']:
+                op['trace_events'].append(e)
+
    return all_ops


-def print_ascii_timeline(op_name, dims, types, usec, cycles, events, evt_val=None):
-    evt_str = ""
-    if evt_val:
-        evt_str = " - evt [" + ",".join(str(x) for x in evt_val) + "]"
+def print_bubbles_timeline(op):
+    op_name = op['name']
+    dims = op['dims']
+    types = op['types']
+    usec = op['usec']
+    cycles = op['cycles']
+    events = op['trace_events']
    logger.info("=" * 100)
-    logger.info(f"{op_name} ({dims} : {types}) - {usec} usec {cycles} cycles{evt_str}")
+    logger.info(f"{op_name} ({dims} : {types}) - {usec} usec {cycles} cycles")
    logger.info("=" * 100)

-    events = sorted(events, key=lambda e: e['cycles'])
    if not events:
        logger.info("  No trace events recorded.")
        return

-    min_cycles = events[0]['cycles']
+    # Identify start and end cycles for this operator
+    op_start = op['start_cycles']
+    op_end = op['end_cycles']
+    if op_start is None or op_end is None:
+        logger.info("  Cannot analyze bubbles: missing start/end cycle counts.")
+        return

-    logger.info("Cycles      %-30s" % "EventDetails" + " ".join(f"T{i:<2}" for i in range(10)) + " HMX")
-    logger.info("-" * 100)
-
-    thread_stacks = [[] for _ in range(11)]
+    batch_duration = op_end - op_start
+    if batch_duration <= 0:
+        logger.info("  Cannot analyze bubbles: batch duration is 0.")
+        return

+    # Group events by (thread, track_type)
+    tracks = defaultdict(list)
    for e in events:
        t = e['thread']
-        if t < 0 or t > 10:
-            continue
+        is_dma = (normalize_event_name(e['event']) == 'DMA')
+        track_type = 'dma' if is_dma else 'compute'
+        tracks[(t, track_type)].append(e)

-        if e['cycles'] >= min_cycles:
-            rel_cycles = e['cycles'] - min_cycles
-        else:
-            rel_cycles = (e['cycles'] + 0x100000000) - min_cycles
+    active_threads = sorted(list(set(t for (t, track_type) in tracks.keys())))
+    if not active_threads:
+        logger.info("  No active threads in trace.")
+        return

-        state = e['state']
-        evt_type = e['event']
+    bubble_threshold = 10000  # 10k cycles

-        # Determine char representing the event
-        norm_evt = normalize_event_name(evt_type)
-        char = '?'
-        if norm_evt == 'V-COMP':
-            char = 'V'
-        elif norm_evt == 'M-COMP':
-            char = 'H'
-        elif norm_evt == 'A-QUANT':
-            char = 'Q'
-        elif norm_evt == 'A-PREP':
-            char = 'A'
-        elif norm_evt == 'Q-PREP':
-            char = 'q'
-        elif norm_evt == 'K-PREP':
-            char = 'k'
-        elif norm_evt == 'V-PREP':
-            char = 'v'
-        elif norm_evt == 'W-DEQUANT':
-            char = 'D'
-        elif norm_evt == 'O-PROC':
-            char = 'O'
-        elif norm_evt == 'W-PREP':
-            char = 'P'
-        elif norm_evt == 'DMA':
-            char = 'M'
+    thread_stats = {}
+    for t in active_threads:
+        thread_stats[t] = {
+            'compute_idle_cycles': batch_duration,
+            'compute_idle_pct': 100.0,
+            'compute_bubbles': [],

-        if state == 'start':
-            thread_stacks[t].append(char)
-        elif state == 'stop':
-            if thread_stacks[t]:
-                if thread_stacks[t][-1] == char:
-                    thread_stacks[t].pop()
-                elif char in thread_stacks[t]:
-                    thread_stacks[t].remove(char)
-                else:
-                    thread_stacks[t].pop()
+            'dma_idle_cycles': batch_duration,
+            'dma_idle_pct': 100.0,
+            'dma_bubbles': []
+        }

-        cols = []
-        for i in range(11):
-            if thread_stacks[i]:
-                cols.append(f"[{thread_stacks[i][-1]}]")
+    total_compute_idle_pct = 0.0
+    total_dma_idle_pct = 0.0
+
+    for t in active_threads:
+        for track_type in ['compute', 'dma']:
+            key = (t, track_type)
+            track_events = tracks.get(key, [])
+
+            if not track_events:
+                gaps = [(op_start, op_end)]
+                idle_cycles = batch_duration
            else:
-                cols.append(" | ")
+                track_events = sorted(track_events, key=lambda e: e.get('unwrapped_cycles') or e['cycles'])

-        evt_desc = f"T{t}: {evt_type} {state} ({e['info']})"
-        logger.info(f"{rel_cycles:10d}  %-30s" % evt_desc + " ".join(cols[:10]) + "  " + cols[10])
+                active_intervals = []
+                active_count = 0
+                curr_start = None
+
+                for e in track_events:
+                    cyc = e.get('unwrapped_cycles') or e['cycles']
+                    cyc = max(op_start, min(op_end, cyc))
+                    state = e['state']
+
+                    if state == 'start':
+                        if active_count == 0:
+                            curr_start = cyc
+                        active_count += 1
+                    elif state == 'stop':
+                        if active_count > 0:
+                            active_count -= 1
+                            if active_count == 0:
+                                active_intervals.append((curr_start, cyc))
+                        else:
+                            active_intervals.append((op_start, cyc))
+
+                if active_count > 0 and curr_start is not None:
+                    active_intervals.append((curr_start, op_end))
+
+                # Merge intervals
+                active_intervals.sort(key=lambda x: x[0])
+                merged_intervals = []
+                for start, end in active_intervals:
+                    if not merged_intervals:
+                        merged_intervals.append([start, end])
+                    else:
+                        last_start, last_end = merged_intervals[-1]
+                        if start <= last_end:
+                            merged_intervals[-1][1] = max(last_end, end)
+                        else:
+                            merged_intervals.append([start, end])
+
+                # Calculate gaps
+                gaps = []
+                curr_time = op_start
+                for start, end in merged_intervals:
+                    if start > curr_time:
+                        gaps.append((curr_time, start))
+                    curr_time = max(curr_time, end)
+                if curr_time < op_end:
+                    gaps.append((curr_time, op_end))
+
+                idle_cycles = sum(end - start for start, end in gaps)
+
+            idle_pct = (idle_cycles / batch_duration) * 100.0
+
+            bubbles = []
+            for start, end in gaps:
+                dur = end - start
+                if dur >= bubble_threshold:
+                    bubbles.append((start, end, dur))
+
+            if track_type == 'compute':
+                thread_stats[t]['compute_idle_cycles'] = idle_cycles
+                thread_stats[t]['compute_idle_pct'] = idle_pct
+                thread_stats[t]['compute_bubbles'] = bubbles
+                total_compute_idle_pct += idle_pct
+            else:
+                thread_stats[t]['dma_idle_cycles'] = idle_cycles
+                thread_stats[t]['dma_idle_pct'] = idle_pct
+                thread_stats[t]['dma_bubbles'] = bubbles
+                total_dma_idle_pct += idle_pct
+
+    avg_compute_idle = total_compute_idle_pct / len(active_threads)
+    avg_dma_idle = total_dma_idle_pct / len(active_threads)
+
+    logger.info("  Combined Idle Statistics:")
+    logger.info(f"    Active Threads   : {', '.join(str(t) for t in active_threads)}")
+    logger.info(f"    Avg Thread Compute IDLE : {avg_compute_idle:.1f}%")
+    logger.info(f"    Avg Thread DMA IDLE     : {avg_dma_idle:.1f}%")
    logger.info("-" * 100)

+    logger.info("  Per-Thread Idle Analysis:")
+    for t in active_threads:
+        stats = thread_stats[t]
+        thread_name = f"Thread {t:<2} (HVX)" if t != 10 else "Thread 10 (HMX)"
+        logger.info(f"    {thread_name} -> Compute Idle: {stats['compute_idle_pct']:.1f}% | DMA Idle: {stats['dma_idle_pct']:.1f}%")

-def print_ascii_summary(op_name, dims, types, usec, cycles, events, evt_val=None):
-    evt_str = ""
-    if evt_val:
-        evt_str = " - evt [" + ",".join(str(x) for x in evt_val) + "]"
+    all_bubbles = []
+    for t in active_threads:
+        stats = thread_stats[t]
+        for start, end, dur in stats['compute_bubbles']:
+            pct = (dur / batch_duration) * 100.0
+            all_bubbles.append((dur, f"Thread {t} Compute: bubble of {dur} cycles ({pct:.1f}%) at {start - op_start} to {end - op_start}"))
+        for start, end, dur in stats['dma_bubbles']:
+            pct = (dur / batch_duration) * 100.0
+            all_bubbles.append((dur, f"Thread {t} DMA    : bubble of {dur} cycles ({pct:.1f}%) at {start - op_start} to {end - op_start}"))
+
+    if all_bubbles:
+        logger.info("-" * 100)
+        logger.info(f"  Significant Bubbles (>= {bubble_threshold} cycles):")
+        all_bubbles.sort(key=lambda x: x[0], reverse=True)
+        for dur, desc in all_bubbles[:15]:
+            logger.info(f"    {desc}")
+    else:
+        logger.info("-" * 100)
+        logger.info(f"  No significant bubbles detected (all idle gaps < {bubble_threshold} cycles).")
+
+
+def print_ascii_summary(op_name, dims, types, usec, cycles, events):
    logger.info("=" * 100)
-    logger.info(f"{op_name} ({dims} : {types}) - {usec} usec {cycles} cycles{evt_str}")
+    logger.info(f"{op_name} ({dims} : {types}) - {usec} usec {cycles} cycles")
    logger.info("=" * 100)

    events = sorted(events, key=lambda e: e['cycles'])
@@ -415,8 +560,8 @@ def main():
    parser.add_argument("--pmu-index", type=int)
    parser.add_argument("--pmu-name", type=str)
    parser.add_argument("--width", action='append', default=['dims:40'], help="Override column width, e.g. --width dims:50")
-    parser.add_argument("--timeline", type=str, nargs='?', const='summary', choices=["summary", "diagram"],
-                        help="Output ASCII art event summary or timing diagram (default: summary)")
+    parser.add_argument("--timeline", type=str, nargs='?', const='summary', choices=["summary", "bubbles"],
+                        help="Output ASCII art event summary or thread idle bubble analysis (default: summary)")
    parser.add_argument("--filter", type=str, help="Regex filter matching against the original profile-op line")

    group = parser.add_mutually_exclusive_group()
@@ -457,12 +602,11 @@ def main():
        ops = ops[-args.tail:]

    if args.timeline:
-        logger.info(f"\n# ASCII Timing {args.timeline.capitalize()}\n")
        for op in ops:
            if args.timeline == "summary":
-                print_ascii_summary(op['name'], op['dims'], op['types'], op['usec'], op['cycles'], op['trace_events'], op.get('evt_val'))
-            elif args.timeline == "diagram":
-                print_ascii_timeline(op['name'], op['dims'], op['types'], op['usec'], op['cycles'], op['trace_events'], op.get('evt_val'))
+                print_ascii_summary(op['name'], op['dims'], op['types'], op['usec'], op['cycles'], op['trace_events'])
+            elif args.timeline == "bubbles":
+                print_bubbles_timeline(op)
    else:
        generate_report(ops, args.top, overrides, args.sort, pmu_name=final_pmu_name)

@@ -6,6 +6,7 @@ import re
 import argparse
 import statistics
 import logging
+import bisect
 from typing import Any, Dict, List, Optional
 from collections import defaultdict

@@ -16,11 +17,11 @@ op_pattern = re.compile(
 )

 trace_pattern = re.compile(
-    r"trace-op\s+(?P<op_name>[A-Z_0-9+]+):\s+thread\s+(?P<thread>\d+)\s+event\s+(?P<event>[A-Z_0-9\-]+)\s+info\s+(?P<info>\d+)\s+(?P<state>start|stop)\s+(?P<cycles>\d+)"
+    r"trace-evt\s+(?P<event>[A-Z_0-9\-]+):\s+thread\s+(?P<thread>\d+)\s+info\s+(?P<info>\d+)\s+(?P<state>start|stop)\s+(?P<cycles>\d+)"
 )


-def normalize_event_name(evt_type):
+def normalize_event_name(evt_type, info=0):
    if evt_type == "HVX_COMP":
        return "V-COMP"
    if evt_type == "HMX_COMP":
@@ -32,9 +33,13 @@ def normalize_event_name(evt_type):


 class CycleUnwrapper:
-    def __init__(self):
-        self.last_raw = None
-        self.high_part = 0
+    def __init__(self, initial_val=None):
+        if initial_val is not None:
+            self.last_raw = initial_val & 0xFFFFFFFF
+            self.high_part = initial_val & 0xFFFFFFFF00000000
+        else:
+            self.last_raw = None
+            self.high_part = 0

    def unwrap(self, raw):
        if self.last_raw is None:
@@ -60,8 +65,10 @@ def parse_log(file_path):
        sys.exit(1)

    all_ops: List[Dict[str, Any]] = []
+    all_traces: List[Dict[str, Any]] = []
    current_op: Optional[Dict[str, Any]] = None
-    unwrapper = CycleUnwrapper()
+    unwrapper = None
+    trace_unwrapper = None
    line_idx = 0

    for line in f:
@@ -73,6 +80,7 @@ def parse_log(file_path):
            if not prefix_match:
                continue

+            names = parts[1]
            if len(parts) == 7:
                dims, types, strides, params, timings = parts[2], parts[3], parts[4], parts[5], parts[6]
            elif len(parts) == 6:
@@ -93,6 +101,7 @@ def parse_log(file_path):
            op_match = op_pattern.search(line)
            if op_match:
                op_name = op_match.group('op_name')
+                names = ""
                dims = op_match.group('dims').strip() if op_match.group('dims') else ''
                types = op_match.group('types').strip() if op_match.group('types') else ''
                strides = op_match.group('strides').strip() if op_match.group('strides') else ''
@@ -103,18 +112,30 @@ def parse_log(file_path):
        if op_match:
            cycles_start_raw = op_match.group('start')
            unwrapped_cycles_start = None
-            if cycles_start_raw:
-                unwrapped_cycles_start = unwrapper.unwrap(int(cycles_start_raw))
+            if op_name == "OPBATCH":
+                if cycles_start_raw:
+                    unwrapped_cycles_start = int(cycles_start_raw)
+                    unwrapper = CycleUnwrapper(unwrapped_cycles_start)
+                    trace_unwrapper = CycleUnwrapper(unwrapped_cycles_start)
+            else:
+                if cycles_start_raw and unwrapper is not None:
+                    unwrapped_cycles_start = unwrapper.unwrap(int(cycles_start_raw))

            idx = line.find("profile-op ")
            op_text = line[idx + 11:].strip() if idx != -1 else line.strip()

+            evt_str = None
+            if types.startswith("evt-cnt "):
+                evt_str = types[8:].strip()
+
            current_op = {
                'name':         op_name,
+                'names':        names,
                'dims':         dims,
                'types':        types,
                'strides':      strides,
                'params':       params,
+                'evt':          evt_str,
                'op_text':      op_text,
                'usec':         int(op_match.group('usec')),
                'cycles':       int(op_match.group('cycles')),
@@ -127,20 +148,22 @@ def parse_log(file_path):
            continue

        trace_match = trace_pattern.search(line)
-        if trace_match and current_op:
-            if trace_match.group('op_name') == current_op['name']:
-                raw_cyc = int(trace_match.group('cycles'))
-                current_op['trace_events'].append({
-                    'thread': int(trace_match.group('thread')),
-                    'event':  trace_match.group('event'),
-                    'info':   int(trace_match.group('info')),
-                    'cycles': raw_cyc,
-                    'unwrapped_cycles': unwrapper.unwrap(raw_cyc),
-                    'state':  trace_match.group('state')
-                })
+        if trace_match:
+            raw_cyc = int(trace_match.group('cycles'))
+            unwrapped_cyc = None
+            if trace_unwrapper is not None:
+                unwrapped_cyc = trace_unwrapper.unwrap(raw_cyc)
+            all_traces.append({
+                'thread': int(trace_match.group('thread')),
+                'event':  trace_match.group('event'),
+                'info':   int(trace_match.group('info')),
+                'cycles': raw_cyc,
+                'unwrapped_cycles': unwrapped_cyc,
+                'state':  trace_match.group('state')
+            })

    f.close()
-    return all_ops
+    return all_ops, all_traces

 # --- Simple protobuf encoder ---

@@ -246,7 +269,7 @@ def write_trace_packet_to_file(f, packet_bytes):
 # --- End Protobuf Encoder ---


-def generate_perfetto_trace(filtered_ops, output_path):
+def generate_perfetto_trace(filtered_ops, trace_events, output_path):
    if not filtered_ops:
        logger.warning("No operators found after filtering.")
        return
@@ -269,14 +292,12 @@ def generate_perfetto_trace(filtered_ops, output_path):

    # Process events
    completed_events = []
-    for op in filtered_ops:
-        events = op['trace_events']
-        if not events:
-            continue
-        events = sorted(events, key=lambda e: e['unwrapped_cycles'])
+    if trace_events:
+        trace_events = sorted(trace_events, key=lambda e: e['unwrapped_cycles'])
+        one_usec_cycles = max(avg_freq_mhz, 1.0)

        active_starts = {}
-        for e in events:
+        for e in trace_events:
            t = e['thread']
            evt = e['event']
            info = e['info']
@@ -285,6 +306,17 @@ def generate_perfetto_trace(filtered_ops, output_path):

            key = (t, evt, info)
            if state == 'start':
+                # Handle missing stop (start followed by another start)
+                if key in active_starts:
+                    prev_start = active_starts[key]
+                    completed_events.append({
+                        'thread': t,
+                        'event': evt,
+                        'info': info,
+                        'start_cyc': prev_start,
+                        'end_cyc': prev_start + one_usec_cycles,
+                        'missing_stop': True,
+                    })
                active_starts[key] = cyc
            elif state == 'stop':
                if key in active_starts:
@@ -296,8 +328,29 @@ def generate_perfetto_trace(filtered_ops, output_path):
                        'info': info,
                        'start_cyc': start_cyc,
                        'end_cyc': cyc,
-                        'op_name': op['name']
                    })
+                else:
+                    # Handle missing start (stop without start)
+                    completed_events.append({
+                        'thread': t,
+                        'event': evt,
+                        'info': info,
+                        'start_cyc': cyc - one_usec_cycles,
+                        'end_cyc': cyc,
+                        'missing_start': True,
+                    })
+
+        # Clear remaining unmatched starts
+        for key, start_cyc in active_starts.items():
+            t, evt, info = key
+            completed_events.append({
+                'thread': t,
+                'event': evt,
+                'info': info,
+                'start_cyc': start_cyc,
+                'end_cyc': start_cyc + one_usec_cycles,
+                'missing_stop': True,
+            })

    completed_events.sort(key=lambda e: e['start_cyc'])

@@ -316,7 +369,7 @@ def generate_perfetto_trace(filtered_ops, output_path):
        ts = e['ts_ns']
        dur = e['dur_ns']

-        norm_evt = normalize_event_name(evt)
+        norm_evt = normalize_event_name(evt, e['info'])
        if norm_evt == "DMA":
            track_key = (t, "DMA")
        elif t == 10:
@@ -343,7 +396,7 @@ def generate_perfetto_trace(filtered_ops, output_path):
        evt = e['event']
        slot = e['slot']

-        norm_evt = normalize_event_name(evt)
+        norm_evt = normalize_event_name(evt, e['info'])
        if norm_evt == "DMA":
            track_evt = "DMA"
            evt_id = 1
@@ -421,18 +474,26 @@ def generate_perfetto_trace(filtered_ops, output_path):
        for op in filtered_ops:
            op_start_ns = int(round(((op['start_cycles'] - global_min_cyc) / avg_freq_mhz) * 1000))
            op_dur_ns = int(round((op['cycles'] / avg_freq_mhz) * 1000))
-            if op_start_ns < last_op_end_ns:
-                op_start_ns = last_op_end_ns
-            clamped_dur = max(op_dur_ns, 100) # Clamp to 100ns (0.1us)
+            if op['name'] != "OPBATCH":
+                if op_start_ns < last_op_end_ns:
+                    op_start_ns = last_op_end_ns
+                clamped_dur = max(op_dur_ns, 100) # Clamp to 100ns (0.1us)
+                last_op_end_ns = op_start_ns + clamped_dur
+            else:
+                clamped_dur = max(op_dur_ns, 100)

            # Debug annotations for Ops
            debug_annots = []
            if 'line_num' in op:
                debug_annots.append(make_debug_annotation("line", int_val=op['line_num']))
-            if 'strides' in op and op['strides']:
+            if 'names' in op and op['names'] and op['names'] != '----':
+                debug_annots.append(make_debug_annotation("names", string_val=op['names']))
+            if 'strides' in op and op['strides'] and op['strides'] != '----':
                debug_annots.append(make_debug_annotation("strides", string_val=op['strides']))
            if 'params' in op and op['params'] and op['params'] != '----':
                debug_annots.append(make_debug_annotation("params", string_val=op['params']))
+            if 'evt' in op and op['evt']:
+                debug_annots.append(make_debug_annotation("evt", string_val=op['evt']))

            # Slice Begin
            evt_begin = make_track_event(1, 2, name=f"{op['name']} ({op['dims']})", category="operator", debug_annotations=debug_annots)
@@ -444,15 +505,21 @@ def generate_perfetto_trace(filtered_ops, output_path):
            packet_end = make_trace_packet(op_start_ns + clamped_dur, track_event=evt_end)
            write_trace_packet_to_file(f, packet_end)

-            last_op_end_ns = op_start_ns + clamped_dur
-
        # Emit Thread Trace Events
        for e in completed_events:
-            norm_name = normalize_event_name(e['event'])
+            norm_name = normalize_event_name(e['event'], e['info'])
            name = f"DMA {e['info']}" if norm_name == "DMA" else norm_name
+            if e.get('missing_start') or e.get('missing_stop'):
+                name += "!"
+
+            debug_annots = []
+            if e.get('missing_start'):
+                debug_annots.append(make_debug_annotation("missing_start", string_val="true"))
+            if e.get('missing_stop'):
+                debug_annots.append(make_debug_annotation("missing_stop", string_val="true"))

            # Slice Begin
-            evt_begin = make_track_event(1, e['uuid'], name=name, category="trace")
+            evt_begin = make_track_event(1, e['uuid'], name=name, category="trace", debug_annotations=debug_annots if debug_annots else None)
            packet_begin = make_trace_packet(e['ts_ns'], track_event=evt_begin)
            write_trace_packet_to_file(f, packet_begin)

@@ -477,7 +544,7 @@ def main():
    args = parser.parse_args()
    logging.basicConfig(level=logging.INFO, format='%(message)s')

-    ops = parse_log(args.logfile)
+    ops, traces = parse_log(args.logfile)

    if args.filter:
        try:
@@ -492,7 +559,30 @@ def main():
    elif args.tail is not None:
        ops = ops[-args.tail:]

-    generate_perfetto_trace(ops, args.output)
+    if args.filter or args.head is not None or args.tail is not None:
+        valid_ranges = []
+        for op in ops:
+            start_cyc = op['unwrapped_cycles_start']
+            end_cyc = start_cyc + op['cycles'] if start_cyc is not None else None
+            if start_cyc is not None and end_cyc is not None:
+                valid_ranges.append((start_cyc, end_cyc))
+
+        valid_ranges.sort(key=lambda r: r[0])
+        range_starts = [r[0] for r in valid_ranges]
+
+        filtered_traces = []
+        for e in traces:
+            cyc = e['unwrapped_cycles']
+            if cyc is None:
+                continue
+            idx = bisect.bisect_right(range_starts, cyc) - 1
+            if idx >= 0:
+                start, end = valid_ranges[idx]
+                if start <= cyc <= end:
+                    filtered_traces.append(e)
+        traces = filtered_traces
+
+    generate_perfetto_trace(ops, traces, args.output)


 if __name__ == "__main__":
@@ -21,7 +21,7 @@ vendor = {
    f"https://raw.githubusercontent.com/yhirose/cpp-httplib/{HTTPLIB_VERSION}/split.py":  "split.py",
    f"https://raw.githubusercontent.com/yhirose/cpp-httplib/{HTTPLIB_VERSION}/LICENSE":   "vendor/cpp-httplib/LICENSE",

-    "https://raw.githubusercontent.com/sheredom/subprocess.h/b49c56e9fe214488493021017bf3954b91c7c1f5/subprocess.h": "vendor/sheredom/subprocess.h",
+    "https://raw.githubusercontent.com/sheredom/subprocess.h/8671cee1fc09f11a70ce3782a0ee13177c3aa387/subprocess.h": "vendor/sheredom/subprocess.h",
 }

 for url, filename in vendor.items():
@@ -108,6 +108,7 @@ static const std::map<llm_arch, const char *> LLM_ARCH_NAMES = {
    { LLM_ARCH_DOTS1,            "dots1"            },
    { LLM_ARCH_ARCEE,            "arcee"            },
    { LLM_ARCH_AFMOE,            "afmoe"            },
+    { LLM_ARCH_LAGUNA,           "laguna"           },
    { LLM_ARCH_ERNIE4_5,         "ernie4_5"         },
    { LLM_ARCH_ERNIE4_5_MOE,     "ernie4_5-moe"     },
    { LLM_ARCH_HUNYUAN_MOE,      "hunyuan-moe"      },
@@ -665,7 +666,7 @@ static const std::map<llm_tensor, llm_tensor_info> LLM_TENSOR_INFOS = {
    {LLM_TENSOR_HC_FFN_SCALE,               {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}},
    {LLM_TENSOR_ATTN_COMPRESSOR_WKV,        {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
    {LLM_TENSOR_ATTN_COMPRESSOR_WGATE,      {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
-    {LLM_TENSOR_ATTN_COMPRESSOR_APE,        {LLM_TENSOR_LAYER_REPEATING, GGML_OP_ADD}},
+    {LLM_TENSOR_ATTN_COMPRESSOR_APE,        {LLM_TENSOR_LAYER_REPEATING, GGML_OP_GET_ROWS}},
    {LLM_TENSOR_ATTN_COMPRESSOR_NORM,       {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}},
    {LLM_TENSOR_ATTN_K_B,                   {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
    {LLM_TENSOR_ATTN_V_B,                   {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
@@ -832,7 +833,7 @@ static const std::map<llm_tensor, llm_tensor_info> LLM_TENSOR_INFOS = {
    {LLM_TENSOR_INDEXER_ATTN_Q_B,           {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
    {LLM_TENSOR_INDEXER_COMPRESSOR_WKV,     {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
    {LLM_TENSOR_INDEXER_COMPRESSOR_WGATE,   {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
-    {LLM_TENSOR_INDEXER_COMPRESSOR_APE,     {LLM_TENSOR_LAYER_REPEATING, GGML_OP_ADD}},
+    {LLM_TENSOR_INDEXER_COMPRESSOR_APE,     {LLM_TENSOR_LAYER_REPEATING, GGML_OP_GET_ROWS}},
    {LLM_TENSOR_INDEXER_COMPRESSOR_NORM,    {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}},
    {LLM_TENSOR_FFN_GATE_TID2EID,           {LLM_TENSOR_LAYER_REPEATING, GGML_OP_GET_ROWS}},
    {LLM_TENSOR_NEXTN_PROJ_PRE,             {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
@@ -113,6 +113,7 @@ enum llm_arch {
    LLM_ARCH_DOTS1,
    LLM_ARCH_ARCEE,
    LLM_ARCH_AFMOE,
+    LLM_ARCH_LAGUNA,
    LLM_ARCH_ERNIE4_5,
    LLM_ARCH_ERNIE4_5_MOE,
    LLM_ARCH_HUNYUAN_MOE,
--- a/Show More
+++ b/Show More