sycl: fix CONV_2D wrong results for non-contiguous input (#30363)

* sycl: fix CONV_2D wrong results for non-contiguous input

The kernel indexes the `input` as a contiguous [W, H, C, N] tensor, but
supports_op claimed the op for any layout. A non-contiguous (e.g. `CWHN`
permuted) input was read with the wrong element offsets, silently
producing wrong results.

Gate the CONV_2D op on `input` and `kernel` contiguity, matching the
CUDA backend contract, and assert input contiguity in the op.

`test-backend-ops -b SYCL0 -o CONV_2D`: goes from 16 failures to 0,
none.

```text
  2006/2006 tests passed
  Backend SYCL0: OK
Backend 2/2: CPU
  Skipping
2/2 backends passed
OK
```

Assisted-by: pi (unsloth/Qwen3.8-27B-GGUF:Q4_K_M via llama.cpp)

* sycl : CONV_2D cwhn=1 support

Make the SYCL backend compute `GGML_OP_CONV_2D` correctly for `cwhn=1` (CWHN
channel-order) inputs and kernels, so the previously failing test-
backend-ops CONV_2D cases pass on SYCL instead of being skipped or
miscomputed.

- Implement CWHN layout support in `ggml/src/ggml-sycl/conv2d.cpp`
  (`ggml_sycl_op_conv2d`).
- Update SYCL `supports_op` for `GGML_OP_CONV_2D` in
  `ggml/src/ggml-sycl/ggml-sycl.cpp` to accept both layouts.

Assisted-by: pi (Greatjedi/Qwen3.8-27B-Intel-Arc-Tuned-GGUF via llama.cpp)

* sycl docs : update ops for SYCL

common/jinja/value.cpp

@@ -425,7 +425,12 @@ const func_builtins & global_builtins() {
             throw raised_exception("Jinja Exception: " + msg);
         }},
         {"dict", toobject},
-        {"namespace", toobject},
+        {"namespace", [](const func_args & args) -> value {
+            // unlike a dict, a namespace has no methods: ns.items is the attribute, as in jinja2
+            value out = toobject(args);
+            cast_val<value_object>(out)->has_builtins = false;
+            return out;
+        }},
         {"strftime_now", [](const func_args & args) -> value {
             args.ensure_vals<value_string>();
             std::string format = args.get_pos(0)->as_string().str();

conversion/__init__.py

@@ -17,8 +17,10 @@ __all__ = [
 
 
 TEXT_MODEL_MAP: dict[str, str] = {
+    "NEOChatModel": "sensenova_u1",
     "AfmoeForCausalLM": "afmoe",
     "LagunaForCausalLM": "laguna",
+    "FlintForCausalLM": "flint",
     "ApertusForCausalLM": "llama",
     "ArceeForCausalLM": "llama",
     "ArcticForCausalLM": "arctic",
@@ -367,6 +369,7 @@ MMPROJ_MODEL_MAP: dict[str, str] = {
     "MuseSparkForConditionalGeneration": "muse_spark",
     "PaddleOCRVisionModel": "ernie",
     "Phi4ForCausalLMV": "phi",
+    "NEOVisionModel": "sensenova_u1",
     "Qwen2AudioForConditionalGeneration": "ultravox",
     "Qwen2VLForConditionalGeneration": "qwenvl",
     "Qwen2VLModel": "qwenvl",

conversion/base.py

@@ -1701,7 +1701,7 @@ class TextModel(ModelBase):
             self.gguf_writer.add_expert_group_used_count(n_group_used)
             logger.info(f"gguf: expert groups used count = {n_group_used}")
 
-        if (score_func := self.find_hparam(["score_function", "scoring_func", "score_func", "moe_router_activation", "moe_router_activation_func", "expert_selection_fn", "router_score_func"], optional=True)) is not None:
+        if (score_func := self.find_hparam(["score_function", "scoring_func", "score_func", "moe_router_activation", "moe_router_activation_func", "expert_selection_fn", "router_score_func", "router_score_function"], optional=True)) is not None:
             if score_func == "sigmoid":
                 self.gguf_writer.add_expert_gating_func(gguf.ExpertGatingFuncType.SIGMOID)
             elif score_func == "softmax":
@@ -1939,6 +1939,9 @@ class TextModel(ModelBase):
         if chkhsh == "52df12b4c8d4176e7481aab4b6e8454d1fd0a210a04a574f6d4e067d10e23c3e":
             # ref: https://huggingface.co/CohereLabs/North-Mini-Code-1.0
             res = "cohere2moe"
+        if chkhsh == "d5420e601e979078c670bc503d9ac411c311fa76b76a76267c31e0d6ba816637":
+            # ref: https://huggingface.co/RecursiveSI/flint
+            res = "flint"
         if chkhsh == "e636dc30a262dcc0d8c323492e32ae2b70728f4df7dfe9737d9f920a282b8aea":
             # ref: https://huggingface.co/Qwen/Qwen1.5-7B
             res = "qwen2"
@@ -3143,7 +3146,7 @@ def get_model_architecture(hparams: dict[str, Any], model_type: ModelType) -> st
     # Kimi-K3's text_config reports "KimiLinearForCausalLM", which is the older
     # Kimi-Linear-48B architecture and cannot load K3 (no attention residuals,
     # latent MoE, situ, ...). Route on the top-level architecture instead.
-    if model_type == ModelType.TEXT and arch in ("StepVLForConditionalGeneration", "Sarashina2VisionForCausalLM", "Exaone4_5_ForConditionalGeneration", "Step3p7ForConditionalGeneration", "KimiK3ForConditionalGeneration"):
+    if model_type == ModelType.TEXT and arch in ("NEOChatModel", "StepVLForConditionalGeneration", "Sarashina2VisionForCausalLM", "Exaone4_5_ForConditionalGeneration", "Step3p7ForConditionalGeneration", "KimiK3ForConditionalGeneration"):
         return arch
 
     # if "architectures" is found in the sub-config, use that instead

conversion/flint.py

@@ -0,0 +1,38 @@
+from __future__ import annotations
+
+from .base import ModelBase, gguf
+from .qwen import Qwen2MoeModel
+
+
+@ModelBase.register("FlintForCausalLM")
+@ModelBase.example("RecursiveSI/flint")
+class FlintModel(Qwen2MoeModel):
+    model_arch = gguf.MODEL_ARCH.FLINT
+
+    def set_gguf_parameters(self):
+        super().set_gguf_parameters()
+        dense = self.hparams["mlp_only_layers"]
+        if dense != list(range(len(dense))):
+            raise ValueError("Flint dense layers must precede the MoE layers")
+        self.gguf_writer.add_leading_dense_block_count(len(dense))
+        self.gguf_writer.add_expert_shared_count(
+            self.hparams["shared_expert_intermediate_size"] // self.hparams["moe_intermediate_size"]
+        )
+        self.gguf_writer.add_expert_weights_norm(self.hparams["norm_topk_prob"])
+        self.gguf_writer.add_expert_weights_scale(1.0)
+        self.gguf_writer.add_sliding_window(self.hparams["sliding_window"])
+        self.gguf_writer.add_sliding_window_pattern([
+            kind == "sliding_attention" for kind in self.hparams["layer_types"]
+        ])
+        self.gguf_writer.add_rope_pattern([
+            i not in self.hparams["no_rope_layer_indices"] for i in range(self.block_count)
+        ])
+
+    def modify_tensors(self, data_torch, name, bid):
+        if name.endswith(".softmax_offset"):
+            name += ".weight"
+        if name.endswith(".gate.expert_bias"):
+            if not self.hparams["router_expert_bias"]:
+                return
+            name += ".bias"
+        yield from super().modify_tensors(data_torch, name, bid)

conversion/sensenova_u1.py

@@ -0,0 +1,114 @@
+from __future__ import annotations
+
+from pathlib import Path
+
+import torch
+
+from .base import MmprojModel, ModelBase, gguf
+from .qwen import Qwen2Model
+
+
+@ModelBase.example("sensenova/SenseNova-U1.5-8B-MoT")
+@ModelBase.register("NEOChatModel")
+class SenseNovaU1Model(Qwen2Model):
+    model_arch = gguf.MODEL_ARCH.SENSENOVA_U1
+
+    @classmethod
+    def filter_tensors(cls, item):
+        name, _ = item
+        if not name.startswith("language_model.") or "_mot_gen" in name:
+            return None
+        return super().filter_tensors(item)
+
+    def set_gguf_parameters(self):
+        if int(self.hparams.get("num_experts") or 0) > 1:
+            raise ValueError("SenseNova U1 conversion currently supports the dense understanding branch")
+        if self.rope_parameters.get("rope_type", "default") not in ("default", None):
+            raise ValueError("SenseNova U1 scaled RoPE is not supported")
+        if self.hparams.get("attention_bias", False):
+            raise ValueError("SenseNova U1 attention bias is not supported")
+        if self.hparams.get("use_sliding_window", False):
+            raise ValueError("SenseNova U1 sliding attention is not supported")
+        super().set_gguf_parameters()
+        self.gguf_writer.add_rope_freq_base_spatial(self.hparams["rope_theta_hw"])
+
+    def get_tensors(self):
+        for name, gen in self.model_tensors.items():
+            if name.endswith((".q_norm_hw.weight", ".k_norm_hw.weight")):
+                continue
+            data = gen()
+            if name.endswith((".q_norm.weight", ".k_norm.weight")):
+                spatial = self.model_tensors[name.replace("_norm.weight", "_norm_hw.weight")]()
+                # Keep one weight vector; the graph normalizes its two halves independently.
+                data = torch.cat((data, spatial), dim=0)
+            yield name, data
+
+
+@ModelBase.register("NEOVisionModel")
+@ModelBase.example("sensenova/SenseNova-U1.5-8B-MoT")
+class SenseNovaU1VisionModel(MmprojModel):
+    def __init__(self, dir_model: Path, *args, **kwargs):
+        hparams = kwargs.pop("hparams", None)
+        if hparams is None:
+            hparams = ModelBase.load_hparams(dir_model, is_mistral_format=False)
+        else:
+            hparams = dict(hparams)
+
+        vision_config = dict(hparams["vision_config"])
+        # NEOVisionModel has no transformer blocks
+        # these defaults satisfy the generic mmproj metadata contract without replacing checkpoint values
+        vision_config.setdefault("image_size", 0)
+        vision_config.setdefault("intermediate_size", 0)
+        vision_config.setdefault("num_hidden_layers", 0)
+        vision_config.setdefault("num_attention_heads", 1)
+        hparams["vision_config"] = vision_config
+        super().__init__(dir_model, *args, hparams=hparams, **kwargs)
+
+        self.preprocessor_config.setdefault("image_mean", [0.485, 0.456, 0.406])
+        self.preprocessor_config.setdefault("image_std", [0.229, 0.224, 0.225])
+
+    @classmethod
+    def filter_tensors(cls, item):
+        name, gen = item
+        if not name.startswith("vision_model."):
+            return None
+        return super().filter_tensors((name, gen))
+
+    def set_gguf_parameters(self):
+        super().set_gguf_parameters()
+        self.gguf_writer.add_clip_projector_type("sensenova_u1")
+        vision_config = self.hparams["vision_config"]
+        # These defaults match the known U1.5 checkpoint when older configs omit fields.
+        hidden_size = int(vision_config.get("hidden_size", 1024))
+        min_pixels = int(vision_config.get("min_pixels", 65536))
+        max_pixels = int(vision_config.get("max_pixels", 16777216))
+        downsample_ratio = float(vision_config.get("downsample_ratio", 0.5))
+        if downsample_ratio <= 0:
+            raise ValueError(f"SenseNova U1 downsample_ratio must be positive, got {downsample_ratio}")
+        merge_size = round(1.0 / downsample_ratio)
+        if abs((1.0 / downsample_ratio) - merge_size) > 1e-6:
+            raise ValueError(f"SenseNova U1 downsample_ratio must have an integer reciprocal, got {downsample_ratio}")
+        norm_eps = float(vision_config.get("layer_norm_eps", vision_config.get("rms_norm_eps", 1e-6)))
+        self.gguf_writer.add_vision_head_dim(hidden_size)
+        self.gguf_writer.add_vision_attention_layernorm_eps(norm_eps)
+        self.gguf_writer.add_vision_min_pixels(min_pixels)
+        self.gguf_writer.add_vision_max_pixels(max_pixels)
+        self.gguf_writer.add_vision_spatial_merge_size(merge_size)
+
+    def modify_tensors(self, data_torch, name, bid):
+        tensor_names = {
+            "vision_model.embeddings.patch_embedding.weight": "v.patch_embd.weight",
+            "vision_model.embeddings.patch_embedding.bias":   "v.patch_embd.bias",
+            "vision_model.embeddings.dense_embedding.weight": self.format_tensor_name(gguf.MODEL_TENSOR.V_MMPROJ, 0),
+            "vision_model.embeddings.dense_embedding.bias": self.format_tensor_name(
+                gguf.MODEL_TENSOR.V_MMPROJ, 0, suffix=".bias"
+            ),
+        }
+        if name not in tensor_names:
+            raise ValueError(f"Unexpected SenseNova U1 vision tensor: {name!r}")
+        return [(tensor_names[name], data_torch)]
+
+    def tensor_force_quant(self, name, new_name, bid, n_dims):
+        if new_name == self.format_tensor_name(gguf.MODEL_TENSOR.V_MMPROJ, 0):
+            return gguf.GGMLQuantizationType.F16 if self.ftype == gguf.LlamaFileType.MOSTLY_F16 else gguf.GGMLQuantizationType.F32
+        return super().tensor_force_quant(name, new_name, bid, n_dims)

convert_hf_to_gguf_update.py

@@ -101,6 +101,7 @@ models = [
     {"name": "command-r",        "tokt": TOKENIZER_TYPE.BPE, "repo": "https://huggingface.co/CohereForAI/c4ai-command-r-v01", },
     {"name": "tiny_aya",         "tokt": TOKENIZER_TYPE.BPE, "repo": "https://huggingface.co/CohereLabs/tiny-aya-base", },
     {"name": "cohere2moe",       "tokt": TOKENIZER_TYPE.BPE, "repo": "https://huggingface.co/CohereLabs/North-Mini-Code-1.0", },
+    {"name": "flint",            "tokt": TOKENIZER_TYPE.BPE, "repo": "https://huggingface.co/RecursiveSI/flint", },
     {"name": "qwen2",            "tokt": TOKENIZER_TYPE.BPE, "repo": "https://huggingface.co/Qwen/Qwen1.5-7B", },
     {"name": "olmo",             "tokt": TOKENIZER_TYPE.BPE, "repo": "https://huggingface.co/allenai/OLMo-1.7-7B-hf", },
     {"name": "dbrx",             "tokt": TOKENIZER_TYPE.BPE, "repo": "https://huggingface.co/databricks/dbrx-base", },

docs/ops.md

@@ -27,7 +27,7 @@ Legend:
 |                        COL2IM_1D | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
 |                           CONCAT | ❌ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | 🟡 | ✅ | ✅ | 🟡 | ❌ | ❌ |
 |                             CONT | ❌ | 🟡 | ✅ | ✅ | 🟡 | 🟡 | ✅ | 🟡 | ✅ | ✅ | 🟡 | ❌ | ❌ |
-|                       CONV_1D_DW | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ✅ |
+|                       CONV_1D_DW | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ✅ |
 |                          CONV_2D | ❌ | ❌ | 🟡 | ✅ | ❌ | ❌ | ✅ | ✅ | ✅ | 🟡 | ✅ | ❌ | ❌ |
 |                       CONV_2D_DW | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                          CONV_3D | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
@@ -76,10 +76,10 @@ Legend:
 |                          MUL_MAT | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 |
 |                 MUL_MAT_HADAMARD | ❌ | ❌ | ✅ | ❌ | ✅ | 🟡 | ❌ | ❌ | ✅ | ✅ | ✅ | ❌ | 🟡 |
 |                       MUL_MAT_ID | ❌ | 🟡 | ✅ | ✅ | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | ✅ | 🟡 | 🟡 | ❌ |
-|                  MUL_MAT_ID_W4A4 | ❌ | ❌ | ✅ | ❌ | ❌ | 🟡 | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
-|                  MUL_MAT_ID_W4A8 | ❌ | ❌ | ✅ | ❌ | ❌ | 🟡 | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
-|                     MUL_MAT_W4A4 | ❌ | ❌ | ✅ | ❌ | ❌ | 🟡 | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
-|                     MUL_MAT_W4A8 | ❌ | ❌ | ✅ | ❌ | ❌ | 🟡 | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
+|                  MUL_MAT_ID_W4A4 | ❌ | ❌ | ✅ | ❌ | ❌ | 🟡 | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ |
+|                  MUL_MAT_ID_W4A8 | ❌ | ❌ | ✅ | ❌ | ❌ | 🟡 | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ |
+|                     MUL_MAT_W4A4 | ❌ | ❌ | ✅ | ❌ | ❌ | 🟡 | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ |
+|                     MUL_MAT_W4A8 | ❌ | ❌ | ✅ | ❌ | ❌ | 🟡 | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ |
 |                              NEG | ❌ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                             NORM | ❌ | ✅ | ✅ | ✅ | 🟡 | 🟡 | ✅ | ✅ | ✅ | ✅ | 🟡 | ❌ | ❌ |
 |                   OPT_STEP_ADAMW | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |

ggml/include/ggml-rpc.h

@@ -8,10 +8,10 @@ extern "C" {
 
 #define RPC_PROTO_MAJOR_VERSION    8
 #define RPC_PROTO_MINOR_VERSION    0
-#define RPC_PROTO_PATCH_VERSION    0
+#define RPC_PROTO_PATCH_VERSION    1
 
 #ifdef  __cplusplus
-static_assert(GGML_OP_COUNT == 101, "GGML_OP_COUNT has changed - update RPC_PROTO_PATCH_VERSION");
+static_assert(GGML_OP_COUNT == 102, "GGML_OP_COUNT has changed - update RPC_PROTO_PATCH_VERSION");
 #endif
 
 #define GGML_RPC_MAX_SERVERS       16

ggml/include/ggml.h

@@ -585,6 +585,7 @@ extern "C" {
         GGML_OP_DSV4_HC_COMB,
         GGML_OP_DSV4_HC_PRE,
         GGML_OP_DSV4_HC_POST,
+        GGML_OP_TOP_K_MASK,
 
         GGML_OP_UNARY,
 
@@ -2732,6 +2733,26 @@ extern "C" {
             struct ggml_tensor  * post,
             struct ggml_tensor  * comb);
 
+    // top-k sparse attention mask: keep the mask cells of the live picks and the extra cells, set the rest to -inf
+    //
+    // mask:   [n_kv,    m1,  m2,  m3]  f16 or f32, broadcast to the rows of top_k
+    // top_k:  [n_top,   ne1, ne2, ne3] i32, picks of each row
+    // score:  [n_pick,  ne1, ne2, ne3] f32, can be NULL: a pick is live when its score is above -inf, else every pick is live
+    // blocks: [n_blk,   n_pick]        i32, can be NULL: the member cells of each pick, else a pick is a cell
+    // extra:  [n_extra, ne1, ne2, ne3] i32, can be NULL: cells of each row kept without a pick
+    // res:    [n_kv,    ne1, ne2, ne3] same type as mask
+    //
+    // res[c, i] = mask[c, i % m] when c is a cell of a live pick of row i or an extra cell of row i, else -inf
+    // picks and cells out of range are skipped
+    //
+    GGML_API struct ggml_tensor * ggml_top_k_mask(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * mask,
+            struct ggml_tensor  * top_k,
+            struct ggml_tensor  * score,
+            struct ggml_tensor  * blocks,
+            struct ggml_tensor  * extra);
+
     // custom operators
 
     typedef void (*ggml_custom1_op_t)(struct ggml_tensor * dst , const struct ggml_tensor * a, int ith, int nth, void * userdata);

ggml/src/ggml-backend-meta.cpp

@@ -1055,7 +1055,8 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state(
             } break;
             case GGML_OP_DSV4_HC_COMB:
             case GGML_OP_DSV4_HC_PRE:
-            case GGML_OP_DSV4_HC_POST: {
+            case GGML_OP_DSV4_HC_POST:
+            case GGML_OP_TOP_K_MASK: {
                 split_state = handle_generic(src_ss, /*scalar_only =*/ true);
             } break;
             case GGML_OP_UNARY: {

ggml/src/ggml-cpu/CMakeLists.txt

@@ -650,9 +650,9 @@ function(ggml_add_cpu_backend_variant_impl tag_name)
 
         # Fetch KleidiAI sources:
         include(FetchContent)
-        set(KLEIDIAI_COMMIT_TAG "v1.31.0")
+        set(KLEIDIAI_COMMIT_TAG "v1.32.0")
         set(KLEIDIAI_DOWNLOAD_URL "https://github.com/ARM-software/kleidiai/releases/download/${KLEIDIAI_COMMIT_TAG}/kleidiai-${KLEIDIAI_COMMIT_TAG}-src.tar.gz")
-        set(KLEIDIAI_RELEASE_ARCHIVE_MD5  "930495f3aeb16b25dc18a2bde45924f4")
+        set(KLEIDIAI_RELEASE_ARCHIVE_MD5  "05f1953428ffe8e8c7f585dc6d508887")
 
         set(KLEIDIAI_FETCH_ARGS
             URL ${KLEIDIAI_DOWNLOAD_URL}

ggml/src/ggml-cpu/ggml-cpu.c

@@ -2105,6 +2105,10 @@ static void ggml_compute_forward(struct ggml_compute_params * params, struct ggm
             {
                 ggml_compute_forward_dsv4_hc_post(params, tensor);
             } break;
+        case GGML_OP_TOP_K_MASK:
+            {
+                ggml_compute_forward_top_k_mask(params, tensor);
+            } break;
         case GGML_OP_MAP_CUSTOM1:
             {
                 ggml_compute_forward_map_custom1(params, tensor);
@@ -2288,6 +2292,7 @@ static int ggml_get_n_tasks(struct ggml_tensor * node, int n_threads) {
         case GGML_OP_DSV4_HC_COMB:
         case GGML_OP_DSV4_HC_PRE:
         case GGML_OP_DSV4_HC_POST:
+        case GGML_OP_TOP_K_MASK:
             {
                 n_tasks = n_threads;
             } break;

ggml/src/ggml-cpu/kleidiai/kernels.cpp

@@ -201,85 +201,117 @@ static struct kai_matmul_uker_config q4_k_matmul_config() {
     return config;
 }
 
-template<struct kai_matmul_uker_api (*GetApi)(void)>
-static size_t q4_k_get_m_step() {
-    const struct kai_matmul_uker_api api = GetApi();
-    const struct kai_matmul_uker_config config = q4_k_matmul_config();
-    return api.get_step(&config).m;
+static struct kai_matmul_uker_config default_matmul_config() {
+    return {};
 }
 
-template<struct kai_matmul_uker_api (*GetApi)(void)>
-static size_t q4_k_get_n_step() {
-    const struct kai_matmul_uker_api api = GetApi();
-    const struct kai_matmul_uker_config config = q4_k_matmul_config();
-    return api.get_step(&config).n;
-}
+template<
+        struct kai_matmul_uker_api (*GetApi)(void),
+        struct kai_matmul_uker_config (*GetConfig)(void),
+        size_t MR, size_t NR, size_t KR, size_t SR>
+struct matmul_uker_adapter {
+    static size_t get_m_step() {
+        const struct kai_matmul_uker_api api = GetApi();
+        const struct kai_matmul_uker_config config = GetConfig();
+        return api.get_step(&config).m;
+    }
 
-template<size_t Value>
-static size_t q4_k_get_constant() {
-    return Value;
-}
+    static size_t get_n_step() {
+        const struct kai_matmul_uker_api api = GetApi();
+        const struct kai_matmul_uker_config config = GetConfig();
+        return api.get_step(&config).n;
+    }
 
-template<struct kai_matmul_uker_api (*GetApi)(void)>
-static size_t q4_k_get_dst_offset(size_t m_idx, size_t n_idx, size_t dst_stride) {
-    const struct kai_matmul_uker_api api = GetApi();
-    const struct kai_matmul_uker_config config = q4_k_matmul_config();
-    const struct kai_matmul_uker_dst_dim_args index = { m_idx, n_idx };
-    const struct kai_matmul_uker_dst_stride_args stride = { dst_stride };
-    return api.get_dst_offset(&config, &index, &stride);
-}
+    static size_t get_mr() {
+        if constexpr (MR == 0) {
+            return get_m_step();
+        }
+        return MR;
+    }
 
-template<struct kai_matmul_uker_api (*GetApi)(void)>
-static size_t q4_k_get_lhs_packed_offset(size_t m_idx, size_t k, size_t bl) {
-    GGML_ASSERT(bl == KLEIDIAI_Q4_K_BLOCK_LENGTH);
+    static size_t get_nr() {
+        if constexpr (NR == 0) {
+            return get_n_step();
+        }
+        return NR;
+    }
 
-    const struct kai_matmul_uker_api api = GetApi();
-    const struct kai_matmul_uker_config config = q4_k_matmul_config();
-    const struct kai_matmul_uker_lhs_dim_args shape = { m_idx, k };
-    const struct kai_matmul_uker_lhs_dim_args index = { m_idx, 0 };
-    const struct kai_matmul_uker_lhs_stride_args stride = api.get_lhs_stride(&config, &shape);
-    return api.get_lhs_offset(&config, &index, &stride);
-}
+    static size_t get_kr() {
+        return KR;
+    }
 
-template<struct kai_matmul_uker_api (*GetApi)(void)>
-static size_t q4_k_get_rhs_packed_offset(size_t n_idx, size_t k, size_t bl) {
-    GGML_ASSERT(bl == KLEIDIAI_Q4_K_BLOCK_LENGTH);
+    static size_t get_sr() {
+        return SR;
+    }
 
-    const struct kai_matmul_uker_api api = GetApi();
-    const struct kai_matmul_uker_config config = q4_k_matmul_config();
-    const struct kai_matmul_uker_rhs_dim_args shape = { n_idx, k };
-    const struct kai_matmul_uker_rhs_dim_args index = { n_idx, 0 };
-    const struct kai_matmul_uker_rhs_stride_args stride = api.get_rhs_stride(&config, &shape);
-    return api.get_rhs_offset(&config, &index, &stride);
-}
+    static size_t get_dst_offset(size_t m_idx, size_t n_idx, size_t dst_stride) {
+        const struct kai_matmul_uker_api api = GetApi();
+        const struct kai_matmul_uker_config config = GetConfig();
+        const struct kai_matmul_uker_dst_dim_args index = { m_idx, n_idx };
+        const struct kai_matmul_uker_dst_stride_args stride = { dst_stride };
+        return api.get_dst_offset(&config, &index, &stride);
+    }
 
-template<struct kai_matmul_uker_api (*GetApi)(void)>
-static void q4_k_run_matmul(
-        size_t m, size_t n, size_t k, size_t bl,
-        const void * lhs_packed, const void * rhs_packed,
-        void * dst, size_t dst_stride_row, size_t dst_stride_col,
-        float clamp_min, float clamp_max) {
-    GGML_ASSERT(bl == KLEIDIAI_Q4_K_BLOCK_LENGTH);
-    GGML_ASSERT(dst_stride_col == sizeof(float));
+    static size_t get_lhs_packed_offset(size_t m_idx, size_t k, size_t bl) {
+        const struct kai_matmul_uker_api api = GetApi();
+        const struct kai_matmul_uker_config config = GetConfig();
+        GGML_ASSERT(bl == config.format.bl);
+        const struct kai_matmul_uker_lhs_dim_args shape = { m_idx, k };
+        const struct kai_matmul_uker_lhs_dim_args index = { m_idx, 0 };
+        const struct kai_matmul_uker_lhs_stride_args stride = api.get_lhs_stride(&config, &shape);
+        return api.get_lhs_offset(&config, &index, &stride);
+    }
 
-    const struct kai_matmul_uker_api api = GetApi();
-    const struct kai_matmul_uker_config config = q4_k_matmul_config();
-    const struct kai_matmul_uker_lhs_dim_args lhs_shape = { m, k };
-    const struct kai_matmul_uker_rhs_dim_args rhs_shape = { n, k };
+    static size_t get_rhs_packed_offset(size_t n_idx, size_t k, size_t bl) {
+        const struct kai_matmul_uker_api api = GetApi();
+        const struct kai_matmul_uker_config config = GetConfig();
+        GGML_ASSERT(bl == config.format.bl);
+        const struct kai_matmul_uker_rhs_dim_args shape = { n_idx, k };
+        const struct kai_matmul_uker_rhs_dim_args index = { n_idx, 0 };
+        const struct kai_matmul_uker_rhs_stride_args stride = api.get_rhs_stride(&config, &shape);
+        return api.get_rhs_offset(&config, &index, &stride);
+    }
 
-    struct kai_matmul_uker_args args = {};
-    args.flags = KAI_MATMUL_UKER_FLAGS_ARGS_CLAMP;
-    args.shape = { m, n, k };
-    args.operand.lhs.ptr = lhs_packed;
-    args.operand.lhs.stride = api.get_lhs_stride(&config, &lhs_shape);
-    args.operand.rhs.ptr = rhs_packed;
-    args.operand.rhs.stride = api.get_rhs_stride(&config, &rhs_shape);
-    args.operand.dst.ptr = dst;
-    args.operand.dst.stride.m = dst_stride_row;
-    args.activation.clamp.min_ptr = &clamp_min;
-    args.activation.clamp.max_ptr = &clamp_max;
-    api.run(&config, &args);
-}
+    static void run(
+            size_t m, size_t n, size_t k, size_t bl,
+            const void * lhs_packed, const void * rhs_packed,
+            void * dst, size_t dst_stride_row, size_t dst_stride_col,
+            float clamp_min, float clamp_max) {
+        const struct kai_matmul_uker_api api = GetApi();
+        const struct kai_matmul_uker_config config = GetConfig();
+        GGML_ASSERT(bl == config.format.bl);
+        GGML_ASSERT(dst_stride_col == sizeof(float));
+        const struct kai_matmul_uker_lhs_dim_args lhs_shape = { m, k };
+        const struct kai_matmul_uker_rhs_dim_args rhs_shape = { n, k };
+
+        struct kai_matmul_uker_args args = {};
+        args.flags = KAI_MATMUL_UKER_FLAGS_ARGS_CLAMP;
+        args.shape = { m, n, k };
+        args.operand.lhs.ptr = lhs_packed;
+        args.operand.lhs.stride = api.get_lhs_stride(&config, &lhs_shape);
+        args.operand.rhs.ptr = rhs_packed;
+        args.operand.rhs.stride = api.get_rhs_stride(&config, &rhs_shape);
+        args.operand.dst.ptr = dst;
+        args.operand.dst.stride.m = dst_stride_row;
+        args.activation.clamp.min_ptr = &clamp_min;
+        args.activation.clamp.max_ptr = &clamp_max;
+        api.run(&config, &args);
+    }
+};
+
+using q4_k_gemm_adapter = matmul_uker_adapter<
+    kai_matmul_clamp_f32_f16p4vsx2_qai4c32p16vsx4s1s0sf16_4vsx16vs_sme2_mopa,
+    q4_k_matmul_config, 0, 0, 4, 2>;
+using q4_k_gemv_adapter = matmul_uker_adapter<
+    kai_matmul_clamp_f32_qsi8d32p1x4_qai4c32p16vsx4s1s0sf16_1x16vs_sme2_dot,
+    q4_k_matmul_config, 0, 0, 4, 2>;
+// GEMM and GEMV share the RHS layout only at SVE256 (NR=8, KR=8, SR=1).
+using sve_q8_gemm_adapter = matmul_uker_adapter<
+    kai_matmul_clamp_f32_qai8dxp4x8sf32_qsi8cxp4vsx8sf32bf32_16x4vs_sve_i8mm,
+    default_matmul_config, 4, 0, 8, 1>;
+using sve_q8_gemv_adapter = matmul_uker_adapter<
+    kai_matmul_clamp_f32_qai8dxp1x8_qsi8cxp8x8_1x8_sve_dot,
+    default_matmul_config, 1, 8, 8, 1>;
 
 static struct kai_matmul_pack_lhs_uker_config q4_k_lhs_pack_config() {
     struct kai_matmul_pack_lhs_uker_config config = {};
@@ -628,7 +660,7 @@ static ggml_kleidiai_kernels gemm_gemv_kernels[] = {
             /* .packed_size_ex        = */ &rhs_ps_fn5<kai_get_rhs_packed_size_rhs_pack_nxk_qsi4c32pscalef16_qsu4c32s16s0>,
             /* .pack_func_ex          = */ &rhs_pack_fn12<kai_run_rhs_pack_nxk_qsi4c32pscalef16_qsu4c32s16s0>,
         },
-        /* .required_cpu       = */ CPU_FEATURE_SVE | CPU_FEATURE_I8MM | CPU_FEATURE_DOTPROD,
+        /* .required_cpu       = */ CPU_FEATURE_SVE | CPU_FEATURE_SVEI8MM | CPU_FEATURE_DOTPROD,
         /* .lhs_type           = */ GGML_TYPE_F32,
         /* .rhs_type           = */ GGML_TYPE_Q4_0,
         /* .op_type            = */ GGML_TYPE_F32,
@@ -737,16 +769,16 @@ static ggml_kleidiai_kernels gemm_gemv_kernels_q4_k[] = {
     {
         /* SME2 GEMM */
         {
-            /* .get_m_step            = */ &q4_k_get_m_step<kai_matmul_clamp_f32_f16p4vsx2_qai4c32p16vsx4s1s0sf16_4vsx16vs_sme2_mopa>,
-            /* .get_n_step            = */ &q4_k_get_n_step<kai_matmul_clamp_f32_f16p4vsx2_qai4c32p16vsx4s1s0sf16_4vsx16vs_sme2_mopa>,
-            /* .get_mr                = */ &q4_k_get_m_step<kai_matmul_clamp_f32_f16p4vsx2_qai4c32p16vsx4s1s0sf16_4vsx16vs_sme2_mopa>,
-            /* .get_nr                = */ &q4_k_get_n_step<kai_matmul_clamp_f32_f16p4vsx2_qai4c32p16vsx4s1s0sf16_4vsx16vs_sme2_mopa>,
-            /* .get_kr                = */ &q4_k_get_constant<4>,
-            /* .get_sr                = */ &q4_k_get_constant<2>,
-            /* .get_dst_offset        = */ &q4_k_get_dst_offset<kai_matmul_clamp_f32_f16p4vsx2_qai4c32p16vsx4s1s0sf16_4vsx16vs_sme2_mopa>,
-            /* .get_lhs_offset_ex     = */ &q4_k_get_lhs_packed_offset<kai_matmul_clamp_f32_f16p4vsx2_qai4c32p16vsx4s1s0sf16_4vsx16vs_sme2_mopa>,
-            /* .get_rhs_packed_offset_ex = */ &q4_k_get_rhs_packed_offset<kai_matmul_clamp_f32_f16p4vsx2_qai4c32p16vsx4s1s0sf16_4vsx16vs_sme2_mopa>,
-            /* .run_kernel_ex         = */ &q4_k_run_matmul<kai_matmul_clamp_f32_f16p4vsx2_qai4c32p16vsx4s1s0sf16_4vsx16vs_sme2_mopa>,
+            /* .get_m_step            = */ q4_k_gemm_adapter::get_m_step,
+            /* .get_n_step            = */ q4_k_gemm_adapter::get_n_step,
+            /* .get_mr                = */ q4_k_gemm_adapter::get_mr,
+            /* .get_nr                = */ q4_k_gemm_adapter::get_nr,
+            /* .get_kr                = */ q4_k_gemm_adapter::get_kr,
+            /* .get_sr                = */ q4_k_gemm_adapter::get_sr,
+            /* .get_dst_offset        = */ q4_k_gemm_adapter::get_dst_offset,
+            /* .get_lhs_offset_ex     = */ q4_k_gemm_adapter::get_lhs_packed_offset,
+            /* .get_rhs_packed_offset_ex = */ q4_k_gemm_adapter::get_rhs_packed_offset,
+            /* .run_kernel_ex         = */ q4_k_gemm_adapter::run,
         },
         {
             /* .get_offset            = */ kai_get_lhs_offset_lhs_pack_f16pmrx2_f32_neon,
@@ -756,16 +788,16 @@ static ggml_kleidiai_kernels gemm_gemv_kernels_q4_k[] = {
         },
         /* SME2 GEMV */
         {
-            /* .get_m_step            = */ &q4_k_get_m_step<kai_matmul_clamp_f32_qsi8d32p1x4_qai4c32p16vsx4s1s0sf16_1x16vs_sme2_dot>,
-            /* .get_n_step            = */ &q4_k_get_n_step<kai_matmul_clamp_f32_qsi8d32p1x4_qai4c32p16vsx4s1s0sf16_1x16vs_sme2_dot>,
-            /* .get_mr                = */ &q4_k_get_m_step<kai_matmul_clamp_f32_qsi8d32p1x4_qai4c32p16vsx4s1s0sf16_1x16vs_sme2_dot>,
-            /* .get_nr                = */ &q4_k_get_n_step<kai_matmul_clamp_f32_qsi8d32p1x4_qai4c32p16vsx4s1s0sf16_1x16vs_sme2_dot>,
-            /* .get_kr                = */ &q4_k_get_constant<4>,
-            /* .get_sr                = */ &q4_k_get_constant<2>,
-            /* .get_dst_offset        = */ &q4_k_get_dst_offset<kai_matmul_clamp_f32_qsi8d32p1x4_qai4c32p16vsx4s1s0sf16_1x16vs_sme2_dot>,
-            /* .get_lhs_offset_ex     = */ &q4_k_get_lhs_packed_offset<kai_matmul_clamp_f32_qsi8d32p1x4_qai4c32p16vsx4s1s0sf16_1x16vs_sme2_dot>,
-            /* .get_rhs_packed_offset_ex = */ &q4_k_get_rhs_packed_offset<kai_matmul_clamp_f32_qsi8d32p1x4_qai4c32p16vsx4s1s0sf16_1x16vs_sme2_dot>,
-            /* .run_kernel_ex         = */ &q4_k_run_matmul<kai_matmul_clamp_f32_qsi8d32p1x4_qai4c32p16vsx4s1s0sf16_1x16vs_sme2_dot>,
+            /* .get_m_step            = */ q4_k_gemv_adapter::get_m_step,
+            /* .get_n_step            = */ q4_k_gemv_adapter::get_n_step,
+            /* .get_mr                = */ q4_k_gemv_adapter::get_mr,
+            /* .get_nr                = */ q4_k_gemv_adapter::get_nr,
+            /* .get_kr                = */ q4_k_gemv_adapter::get_kr,
+            /* .get_sr                = */ q4_k_gemv_adapter::get_sr,
+            /* .get_dst_offset        = */ q4_k_gemv_adapter::get_dst_offset,
+            /* .get_lhs_offset_ex     = */ q4_k_gemv_adapter::get_lhs_packed_offset,
+            /* .get_rhs_packed_offset_ex = */ q4_k_gemv_adapter::get_rhs_packed_offset,
+            /* .run_kernel_ex         = */ q4_k_gemv_adapter::run,
         },
         {
             /* .get_offset            = */ q4_k_lhs_get_offset,
@@ -883,6 +915,54 @@ static ggml_kleidiai_kernels gemm_gemv_kernels_q8[] = {
         /* .rhs_type           = */ GGML_TYPE_Q8_0,
         /* .op_type            = */ GGML_TYPE_F32,
     },
+    {
+        /* SVE I8MM GEMM */
+        {
+            /* .get_m_step            = */ sve_q8_gemm_adapter::get_m_step,
+            /* .get_n_step            = */ sve_q8_gemm_adapter::get_n_step,
+            /* .get_mr                = */ sve_q8_gemm_adapter::get_mr,
+            /* .get_nr                = */ sve_q8_gemm_adapter::get_nr,
+            /* .get_kr                = */ sve_q8_gemm_adapter::get_kr,
+            /* .get_sr                = */ sve_q8_gemm_adapter::get_sr,
+            /* .get_dst_offset        = */ sve_q8_gemm_adapter::get_dst_offset,
+            /* .get_lhs_offset_ex     = */ sve_q8_gemm_adapter::get_lhs_packed_offset,
+            /* .get_rhs_packed_offset_ex = */ sve_q8_gemm_adapter::get_rhs_packed_offset,
+            /* .run_kernel_ex         = */ sve_q8_gemm_adapter::run,
+        },
+        /* .gemm_lhs_info = */ {
+            /* .get_offset            = */ kai_get_lhs_offset_lhs_quant_pack_qai8dxp_f32,
+            /* .get_packed_offset_ex  = */ &lhs_offs_fn5<kai_get_lhs_packed_offset_lhs_quant_pack_qai8dxp_f32>,
+            /* .packed_size_ex        = */ &lhs_ps_fn5<kai_get_lhs_packed_size_lhs_quant_pack_qai8dxp_f32>,
+            /* .pack_func_ex          = */ &lhs_pack_float_fn9_no_bl<kai_run_lhs_quant_pack_qai8dxp_f32>,
+        },
+        /* SVE DOT GEMV */
+        {
+            /* .get_m_step            = */ sve_q8_gemv_adapter::get_m_step,
+            /* .get_n_step            = */ sve_q8_gemv_adapter::get_n_step,
+            /* .get_mr                = */ sve_q8_gemv_adapter::get_mr,
+            /* .get_nr                = */ sve_q8_gemv_adapter::get_nr,
+            /* .get_kr                = */ sve_q8_gemv_adapter::get_kr,
+            /* .get_sr                = */ sve_q8_gemv_adapter::get_sr,
+            /* .get_dst_offset        = */ sve_q8_gemv_adapter::get_dst_offset,
+            /* .get_lhs_offset_ex     = */ sve_q8_gemv_adapter::get_lhs_packed_offset,
+            /* .get_rhs_packed_offset_ex = */ sve_q8_gemv_adapter::get_rhs_packed_offset,
+            /* .run_kernel_ex         = */ sve_q8_gemv_adapter::run,
+        },
+        /* .gemv_lhs_info = */ {
+            /* .get_offset            = */ kai_get_lhs_offset_lhs_quant_pack_qai8dxp_f32,
+            /* .get_packed_offset_ex  = */ &lhs_offs_fn5<kai_get_lhs_packed_offset_lhs_quant_pack_qai8dxp_f32>,
+            /* .packed_size_ex        = */ &lhs_ps_fn5<kai_get_lhs_packed_size_lhs_quant_pack_qai8dxp_f32>,
+            /* .pack_func_ex          = */ &lhs_pack_float_fn9_no_bl<kai_run_lhs_quant_pack_qai8dxp_f32>,
+        },
+        /* .rhs_info = */ {
+            /* .packed_size_ex        = */ &rhs_ps_fn5<kai_get_rhs_packed_size_rhs_pack_nxk_qsi8cxp_qsi8cx_neon>,
+            /* .pack_func_ex          = */ &rhs_pack_scale_fn12<kai_run_rhs_pack_nxk_qsi8cxp_qsi8cx_neon>,
+        },
+        /* .required_cpu       = */ CPU_FEATURE_SVE | CPU_FEATURE_SVEI8MM | CPU_FEATURE_DOTPROD,
+        /* .lhs_type           = */ GGML_TYPE_F32,
+        /* .rhs_type           = */ GGML_TYPE_Q8_0,
+        /* .op_type            = */ GGML_TYPE_F32,
+    },
     {
         /* I8MM GEMM */
         {

ggml/src/ggml-cpu/kleidiai/kernels.h

@@ -13,7 +13,8 @@ enum cpu_feature {
     CPU_FEATURE_SVE     = 4,
     CPU_FEATURE_SME     = 8,
     CPU_FEATURE_SME2    = 16,
-    CPU_FEATURE_FP16    = 32
+    CPU_FEATURE_FP16    = 32,
+    CPU_FEATURE_SVEI8MM = 64
 };
 
 inline cpu_feature& operator|=(cpu_feature& lhs, cpu_feature rhs) {

ggml/src/ggml-cpu/kleidiai/kleidiai.cpp

@@ -319,8 +319,10 @@ static void init_kleidiai_context(void) {
 
         size_t detected_smcus = 0;
 
+        // The fixed-width SVE kernels require 32-byte vectors, including the Q8 shared RHS layout.
         ctx.features  = (runtime_feat.has_dotprod  ? CPU_FEATURE_DOTPROD : CPU_FEATURE_NONE) |
                         (runtime_feat.has_i8mm     ? CPU_FEATURE_I8MM    : CPU_FEATURE_NONE) |
+                        (runtime_feat.has_svei8mm  ? CPU_FEATURE_SVEI8MM : CPU_FEATURE_NONE) |
                         (runtime_feat.has_fp16     ? CPU_FEATURE_FP16    : CPU_FEATURE_NONE) |
                         (runtime_feat.sve_cnt == QK8_0 ? CPU_FEATURE_SVE : CPU_FEATURE_NONE);
 
@@ -1616,6 +1618,11 @@ public:
             const size_t kr = kernel->get_kr();
             const size_t sr = kernel->get_sr();
             const ggml_type rhs_type = kernels->rhs_type;
+            if (rhs_type == GGML_TYPE_Q8_0 && (kernels->required_cpu & CPU_FEATURE_SVE)) {
+                GGML_ASSERT(nr == kernels->gemv.get_nr());
+                GGML_ASSERT(kr == kernels->gemv.get_kr());
+                GGML_ASSERT(sr == kernels->gemv.get_sr());
+            }
             const size_t block_len = rhs_type == GGML_TYPE_Q8_0 ? QK8_0 :
                                      rhs_type == GGML_TYPE_Q4_0 ? QK4_0 :
                                      rhs_type == GGML_TYPE_Q4_K ? GGML_KLEIDIAI_Q4_K_BL :

ggml/src/ggml-cpu/ops.cpp

@@ -11736,6 +11736,92 @@ void ggml_compute_forward_dsv4_hc_post(
     }
 }
 
+// ggml_compute_forward_top_k_mask
+
+void ggml_compute_forward_top_k_mask(
+        const ggml_compute_params * params,
+        ggml_tensor * dst) {
+    const ggml_tensor * mask   = dst->src[0];
+    const ggml_tensor * top_k  = dst->src[1];
+    const ggml_tensor * score  = dst->src[2];
+    const ggml_tensor * blocks = dst->src[3];
+    const ggml_tensor * extra  = dst->src[4];
+
+    GGML_ASSERT(dst->type == mask->type);
+    GGML_ASSERT(dst->nb[0] == ggml_type_size(dst->type));
+
+    const int64_t n_kv    = mask->ne[0];
+    const int64_t ne1     = dst->ne[1];
+    const int64_t ne2     = dst->ne[2];
+    const int64_t n_rows  = ggml_nrows(dst);
+    const int64_t n_top   = top_k->ne[0];
+    const int64_t n_blk   = blocks ? blocks->ne[0] : 1;
+    const int64_t n_pick  = blocks ? blocks->ne[1] : score ? score->ne[0] : n_kv;
+    const int64_t n_extra = extra ? extra->ne[0] : 0;
+
+    const size_t ts = ggml_type_size(dst->type);
+
+    const int ith = params->ith;
+    const int nth = params->nth;
+
+    for (int64_t ir = ith; ir < n_rows; ir += nth) {
+        const int64_t i1 = ir % ne1;
+        const int64_t i2 = (ir / ne1) % ne2;
+        const int64_t i3 = ir / (ne1*ne2);
+
+        const char * mask_row = (const char *) mask->data
+            + (i1 % mask->ne[1])*mask->nb[1] + (i2 % mask->ne[2])*mask->nb[2] + (i3 % mask->ne[3])*mask->nb[3];
+        char * dst_row = (char *) dst->data + i1*dst->nb[1] + i2*dst->nb[2] + i3*dst->nb[3];
+
+        if (dst->type == GGML_TYPE_F16) {
+            const ggml_fp16_t ninf = GGML_CPU_FP32_TO_FP16(-INFINITY);
+            for (int64_t c = 0; c < n_kv; ++c) {
+                ((ggml_fp16_t *) dst_row)[c] = ninf;
+            }
+        } else {
+            for (int64_t c = 0; c < n_kv; ++c) {
+                ((float *) dst_row)[c] = -INFINITY;
+            }
+        }
+
+        auto keep = [&](int32_t c) {
+            if (c >= 0 && c < n_kv) {
+                memcpy(dst_row + c*ts, mask_row + c*ts, ts);
+            }
+        };
+
+        const char * top_k_row = (const char *) top_k->data + i1*top_k->nb[1] + i2*top_k->nb[2] + i3*top_k->nb[3];
+
+        for (int64_t j = 0; j < n_top; ++j) {
+            const int32_t p = *(const int32_t *) (top_k_row + j*top_k->nb[0]);
+            if (p < 0 || p >= n_pick) {
+                continue;
+            }
+            if (score) {
+                const float s = *(const float *) ((const char *) score->data
+                    + p*score->nb[0] + i1*score->nb[1] + i2*score->nb[2] + i3*score->nb[3]);
+                if (!(s > -INFINITY)) {
+                    continue;
+                }
+            }
+            if (!blocks) {
+                keep(p);
+                continue;
+            }
+            for (int64_t k = 0; k < n_blk; ++k) {
+                keep(*(const int32_t *) ((const char *) blocks->data + k*blocks->nb[0] + p*blocks->nb[1]));
+            }
+        }
+
+        if (extra) {
+            const char * extra_row = (const char *) extra->data + i1*extra->nb[1] + i2*extra->nb[2] + i3*extra->nb[3];
+            for (int64_t j = 0; j < n_extra; ++j) {
+                keep(*(const int32_t *) (extra_row + j*extra->nb[0]));
+            }
+        }
+    }
+}
+
 // ggml_compute_forward_rwkv_wkv7
 
 static void ggml_compute_forward_rwkv_wkv7_f32(

ggml/src/ggml-cpu/ops.h

@@ -108,6 +108,7 @@ void ggml_compute_forward_lightning_indexer(const struct ggml_compute_params * p
 void ggml_compute_forward_dsv4_hc_comb(const struct ggml_compute_params * params, struct ggml_tensor * dst);
 void ggml_compute_forward_dsv4_hc_pre(const struct ggml_compute_params * params, struct ggml_tensor * dst);
 void ggml_compute_forward_dsv4_hc_post(const struct ggml_compute_params * params, struct ggml_tensor * dst);
+void ggml_compute_forward_top_k_mask(const struct ggml_compute_params * params, struct ggml_tensor * dst);
 void ggml_compute_forward_map_custom1(const struct ggml_compute_params * params, struct ggml_tensor * dst);
 void ggml_compute_forward_map_custom2(const struct ggml_compute_params * params, struct ggml_tensor * dst);
 void ggml_compute_forward_map_custom3(const struct ggml_compute_params * params, struct ggml_tensor * dst);

ggml/src/ggml-cuda/ggml-cuda.cu

@@ -70,6 +70,7 @@
 #include "ggml-cuda/cumsum.cuh"
 #include "ggml-cuda/fill.cuh"
 #include "ggml-cuda/lightning-indexer.cuh"
+#include "ggml-cuda/top-k-mask.cuh"
 #include "ggml.h"
 
 #include <algorithm>
@@ -2570,6 +2571,9 @@ static bool ggml_cuda_compute_forward(ggml_backend_cuda_context & ctx, struct gg
         case GGML_OP_LIGHTNING_INDEXER:
             ggml_cuda_lightning_indexer(ctx, dst);
             break;
+        case GGML_OP_TOP_K_MASK:
+            ggml_cuda_op_top_k_mask(ctx, dst);
+            break;
         default:
             return false;
     }
@@ -6003,6 +6007,8 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g
             return true;
         case GGML_OP_LIGHTNING_INDEXER:
             return ggml_cuda_lightning_indexer_supported(dev_ctx->device, op);
+        case GGML_OP_TOP_K_MASK:
+            return op->type == GGML_TYPE_F16 || op->type == GGML_TYPE_F32;
 
         default:
             return false;

ggml/src/ggml-cuda/top-k-mask.cu

@@ -0,0 +1,141 @@
+#include "top-k-mask.cuh"
+#include "convert.cuh"
+
+#define CUDA_TOP_K_MASK_BLOCK_SIZE 512
+
+// one block per row: fill the row with -inf, then copy the mask cells of the live picks and the extra cells
+template <typename T>
+static __global__ void top_k_mask_kernel(
+        const T * mask, const int32_t * top_k, const float * score, const int32_t * blocks, const int32_t * extra, T * dst,
+        const int64_t n_kv, const int64_t ne1, const int64_t ne2,
+        const int64_t m1, const int64_t m2, const int64_t m3,
+        const int64_t n_top, const int64_t n_pick, const int64_t n_blk, const int64_t n_extra,
+        const int64_t sm1, const int64_t sm2, const int64_t sm3,
+        const int64_t sd1, const int64_t sd2, const int64_t sd3,
+        const int64_t sk0, const int64_t sk1, const int64_t sk2, const int64_t sk3,
+        const int64_t ss0, const int64_t ss1, const int64_t ss2, const int64_t ss3,
+        const int64_t sb0, const int64_t sb1,
+        const int64_t se0, const int64_t se1, const int64_t se2, const int64_t se3,
+        const bool vec) {
+    ggml_cuda_pdl_lc();
+
+    const int64_t ir = blockIdx.x;
+    const int64_t i1 = ir % ne1;
+    const int64_t i2 = (ir / ne1) % ne2;
+    const int64_t i3 = ir / (ne1*ne2);
+
+    const T       * mask_row  = mask  + (i1 % m1)*sm1 + (i2 % m2)*sm2 + (i3 % m3)*sm3;
+    T             * dst_row   = dst   + i1*sd1 + i2*sd2 + i3*sd3;
+    const int32_t * top_k_row = top_k + i1*sk1 + i2*sk2 + i3*sk3;
+
+    ggml_cuda_pdl_sync();
+
+    const T ninf = ggml_cuda_cast<T>(-INFINITY);
+    if (vec) {
+        constexpr int nv = sizeof(int4)/sizeof(T);
+        T fill[nv];
+#pragma unroll
+        for (int j = 0; j < nv; ++j) {
+            fill[j] = ninf;
+        }
+        const int4 fill4 = *(const int4 *) fill;
+        for (int64_t i = threadIdx.x; i < n_kv/nv; i += blockDim.x) {
+            ((int4 *) dst_row)[i] = fill4;
+        }
+    } else {
+        for (int64_t i = threadIdx.x; i < n_kv; i += blockDim.x) {
+            dst_row[i] = ninf;
+        }
+    }
+
+    __syncthreads();
+
+    // one thread per pick or extra cell; cells kept twice get the same value
+    for (int64_t i = threadIdx.x; i < n_top + n_extra; i += blockDim.x) {
+        if (i >= n_top) {
+            const int32_t c = extra[(i - n_top)*se0 + i1*se1 + i2*se2 + i3*se3];
+            if (c >= 0 && c < n_kv) {
+                dst_row[c] = mask_row[c];
+            }
+            continue;
+        }
+
+        const int32_t p = top_k_row[i*sk0];
+        if (p < 0 || p >= n_pick) {
+            continue;
+        }
+        if (score && !(score[p*ss0 + i1*ss1 + i2*ss2 + i3*ss3] > -INFINITY)) {
+            continue;
+        }
+        if (!blocks) {
+            if (p < n_kv) {
+                dst_row[p] = mask_row[p];
+            }
+            continue;
+        }
+        for (int64_t k = 0; k < n_blk; ++k) {
+            const int32_t c = blocks[k*sb0 + p*sb1];
+            if (c >= 0 && c < n_kv) {
+                dst_row[c] = mask_row[c];
+            }
+        }
+    }
+}
+
+template <typename T>
+static void top_k_mask_cuda(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
+    const ggml_tensor * mask   = dst->src[0];
+    const ggml_tensor * top_k  = dst->src[1];
+    const ggml_tensor * score  = dst->src[2];
+    const ggml_tensor * blocks = dst->src[3];
+    const ggml_tensor * extra  = dst->src[4];
+
+    const int64_t n_kv = mask->ne[0];
+
+    GGML_ASSERT(dst->nb[0] == sizeof(T) && mask->nb[0] == sizeof(T));
+    GGML_ASSERT(dst->nb[1] % sizeof(T) == 0 && mask->nb[1] % sizeof(T) == 0);
+
+    // 16-byte stores need every dst row aligned
+    const bool vec = (n_kv*sizeof(T)) % sizeof(int4) == 0 &&
+        dst->nb[1] % sizeof(int4) == 0 && dst->nb[2] % sizeof(int4) == 0 && dst->nb[3] % sizeof(int4) == 0 &&
+        (uintptr_t) dst->data % sizeof(int4) == 0;
+
+    const int64_t n_pick = blocks ? blocks->ne[1] : score ? score->ne[0] : n_kv;
+
+    const size_t si = sizeof(int32_t);
+    const size_t sf = sizeof(float);
+
+    const dim3 block_dims(CUDA_TOP_K_MASK_BLOCK_SIZE, 1, 1);
+    const dim3 grid_dims(ggml_nrows(dst), 1, 1);
+    const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params(grid_dims, block_dims, 0, ctx.stream());
+
+    ggml_cuda_kernel_launch(top_k_mask_kernel<T>, launch_params,
+            (const T *) mask->data, (const int32_t *) top_k->data,
+            score  ? (const float   *) score->data  : nullptr,
+            blocks ? (const int32_t *) blocks->data : nullptr,
+            extra  ? (const int32_t *) extra->data  : nullptr,
+            (T *) dst->data,
+            n_kv, dst->ne[1], dst->ne[2],
+            mask->ne[1], mask->ne[2], mask->ne[3],
+            top_k->ne[0], n_pick, blocks ? blocks->ne[0] : 1, extra ? extra->ne[0] : 0,
+            mask->nb[1]/sizeof(T), mask->nb[2]/sizeof(T), mask->nb[3]/sizeof(T),
+            dst->nb[1]/sizeof(T),  dst->nb[2]/sizeof(T),  dst->nb[3]/sizeof(T),
+            top_k->nb[0]/si, top_k->nb[1]/si, top_k->nb[2]/si, top_k->nb[3]/si,
+            score  ? score->nb[0]/sf  : 0, score ? score->nb[1]/sf : 0, score ? score->nb[2]/sf : 0, score ? score->nb[3]/sf : 0,
+            blocks ? blocks->nb[0]/si : 0, blocks ? blocks->nb[1]/si : 0,
+            extra  ? extra->nb[0]/si  : 0, extra ? extra->nb[1]/si : 0, extra ? extra->nb[2]/si : 0, extra ? extra->nb[3]/si : 0,
+            vec);
+}
+
+void ggml_cuda_op_top_k_mask(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
+    switch (dst->type) {
+        case GGML_TYPE_F16:
+            top_k_mask_cuda<half>(ctx, dst);
+            break;
+        case GGML_TYPE_F32:
+            top_k_mask_cuda<float>(ctx, dst);
+            break;
+        default:
+            GGML_ABORT("unsupported type");
+    }
+}

ggml/src/ggml-cuda/top-k-mask.cuh

@@ -0,0 +1,3 @@
+#include "common.cuh"
+
+void ggml_cuda_op_top_k_mask(ggml_backend_cuda_context & ctx, ggml_tensor * dst);

ggml/src/ggml-feats.h

@@ -22,6 +22,10 @@
 #define HWCAP2_I8MM (1ULL << 13)
 #endif
 
+#if !defined(HWCAP2_SVEI8MM)
+#define HWCAP2_SVEI8MM (1ULL << 9)
+#endif
+
 #if !defined(HWCAP_ASIMDDP)
 #define HWCAP_ASIMDDP (1 << 20)
 #endif
@@ -63,6 +67,10 @@
 #define PF_ARM_SVE2_INSTRUCTIONS_AVAILABLE 47
 #endif
 
+#if !defined(PF_ARM_SVE_I8MM_INSTRUCTIONS_AVAILABLE)
+#define PF_ARM_SVE_I8MM_INSTRUCTIONS_AVAILABLE 57
+#endif
+
 #if !defined(PF_ARM_V82_I8MM_INSTRUCTIONS_AVAILABLE)
 #define PF_ARM_V82_I8MM_INSTRUCTIONS_AVAILABLE 66
 #endif
@@ -87,6 +95,7 @@ typedef struct ggml_feats_arch64_runtime {
     bool has_sve;
     bool has_sve2;
     bool has_i8mm;
+    bool has_svei8mm;
     bool has_sme;
     bool has_sme2;
     int sve_cnt;
@@ -104,6 +113,7 @@ static inline ggml_feats_arch64_runtime_t ggml_feats_get_arch64_runtime(void) {
     runtime_feat.has_sve     = !!(hwcap & HWCAP_SVE);
     runtime_feat.has_sve2    = !!(hwcap2 & HWCAP2_SVE2);
     runtime_feat.has_i8mm    = !!(hwcap2 & HWCAP2_I8MM);
+    runtime_feat.has_svei8mm = !!(hwcap2 & HWCAP2_SVEI8MM);
     runtime_feat.has_sme     = !!(hwcap2 & HWCAP2_SME);
     runtime_feat.has_sme2    = !!(hwcap2 & HWCAP2_SME2);
 
@@ -153,6 +163,7 @@ static inline ggml_feats_arch64_runtime_t ggml_feats_get_arch64_runtime(void) {
     runtime_feat.has_sve     = IsProcessorFeaturePresent(PF_ARM_SVE_INSTRUCTIONS_AVAILABLE) != 0;
     runtime_feat.has_sve2    = IsProcessorFeaturePresent(PF_ARM_SVE2_INSTRUCTIONS_AVAILABLE) != 0;
     runtime_feat.has_i8mm    = IsProcessorFeaturePresent(PF_ARM_V82_I8MM_INSTRUCTIONS_AVAILABLE) != 0;
+    runtime_feat.has_svei8mm = IsProcessorFeaturePresent(PF_ARM_SVE_I8MM_INSTRUCTIONS_AVAILABLE) != 0;
     runtime_feat.has_sme     = IsProcessorFeaturePresent(PF_ARM_SME_INSTRUCTIONS_AVAILABLE) != 0;
     runtime_feat.has_sme2    = IsProcessorFeaturePresent(PF_ARM_SME2_INSTRUCTIONS_AVAILABLE) != 0;
 

ggml/src/ggml-opencl/CMakeLists.txt

@@ -64,6 +64,7 @@ set(GGML_OPENCL_KERNELS
     add
     add_id
     moe_add_id_glu
+    argmax
     argsort
     tri
     fill

ggml/src/ggml-opencl/ggml-opencl.cpp

@@ -646,8 +646,8 @@ struct ggml_backend_opencl_context {
     ggml_cl_version platform_version;
     ggml_cl_version opencl_c_version;
 
-    // argsort is loaded in supports_op because its availability depends on how
-    // many workgroups are allowed, which requires kernel compilation.
+    // argmax and argsort are loaded in supports_op; argsort availability depends on the compiled kernel's workgroup limit.
+    bool kernels_loaded_argmax = false;
     bool kernels_loaded_argsort = false;
     // rest of the kernels are currently always loaded in alloc_buffer.
     bool kernels_loaded = false;
@@ -972,6 +972,7 @@ struct ggml_backend_opencl_context {
     cl_kernel kernel_solve_tri_f32;
     cl_kernel kernel_im2col_f32, kernel_im2col_f16;
     cl_kernel kernel_im2col_3d_f32, kernel_im2col_3d_f16;
+    cl_kernel kernel_argmax_f32_i32;
     cl_kernel kernel_argsort_f32_i32;
     cl_kernel kernel_sum_rows_f32, kernel_sum_rows_f32_4;
     cl_kernel kernel_cumsum_blk, kernel_cumsum_add;
@@ -1485,6 +1486,27 @@ static cl_program build_program_from_binary(cl_context ctx, cl_device_id dev, co
     return p;
 }
 
+static void load_cl_kernels_argmax(ggml_backend_opencl_context *backend_ctx) {
+    auto opencl_c_std =
+        std::string("CL") + std::to_string(backend_ctx->opencl_c_version.major) + "." + std::to_string(backend_ctx->opencl_c_version.minor);
+    std::string compile_opts = std::string("-cl-std=") + opencl_c_std;
+
+    if (!backend_ctx->kernels_loaded_argmax) {
+        cl_int err;
+#ifdef GGML_OPENCL_EMBED_KERNELS
+        const std::string kernel_src {
+            #include "argmax.cl.h"
+        };
+#else
+        const std::string kernel_src = read_file("argmax.cl");
+#endif
+        cl_program prog = build_program_from_source(backend_ctx, kernel_src.c_str(), compile_opts);
+        CL_CHECK((backend_ctx->kernel_argmax_f32_i32 = clCreateKernel(prog, "kernel_argmax_f32_i32", &err), err));
+        CL_CHECK(clReleaseProgram(prog));
+        backend_ctx->kernels_loaded_argmax = true;
+    }
+}
+
 static void load_cl_kernels_argsort(ggml_backend_opencl_context *backend_ctx) {
     // compiler options for general kernels
     auto opencl_c_std =
@@ -9693,6 +9715,11 @@ static bool ggml_opencl_supports_op(ggml_backend_dev_t dev, const struct ggml_te
             return true;
         case GGML_OP_IM2COL_3D:
             return true;
+        case GGML_OP_ARGMAX:
+            load_cl_kernels_argmax(backend_ctx);
+            return op->src[0]->type == GGML_TYPE_F32 && op->type == GGML_TYPE_I32 &&
+                   ggml_is_contiguous(op->src[0]) && ggml_is_contiguous(op) &&
+                   backend_ctx->get_kernel_workgroup_size(backend_ctx->kernel_argmax_f32_i32) >= 64;
         case GGML_OP_ARGSORT: {
             load_cl_kernels_argsort(backend_ctx);
 
@@ -30164,6 +30191,41 @@ static void ggml_cl_im2col_3d(ggml_backend_t backend, const ggml_tensor * src0,
     backend_ctx->enqueue_ndrange_kernel(kernel, 3, global_work_size, local_work_size, dst);
 }
 
+static void ggml_cl_argmax(ggml_backend_t backend, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst) {
+    GGML_ASSERT(src0);
+    GGML_ASSERT(src0->extra);
+    GGML_ASSERT(dst);
+    GGML_ASSERT(dst->extra);
+    GGML_UNUSED(src1);
+
+    GGML_ASSERT(src0->type == GGML_TYPE_F32);
+    GGML_ASSERT(dst->type == GGML_TYPE_I32);
+    GGML_ASSERT(ggml_is_contiguous(src0));
+    GGML_ASSERT(ggml_is_contiguous(dst));
+
+    ggml_backend_opencl_context * backend_ctx = (ggml_backend_opencl_context *) backend->context;
+    ggml_tensor_extra_cl * extra0 = (ggml_tensor_extra_cl *) src0->extra;
+    ggml_tensor_extra_cl * extrad = (ggml_tensor_extra_cl *) dst->extra;
+
+    const cl_ulong offset0 = extra0->offset + src0->view_offs;
+    const cl_ulong offsetd = extrad->offset + dst->view_offs;
+    const int ne00 = src0->ne[0];
+    const int nrows = ggml_nrows(src0);
+    constexpr int nth = 64;
+
+    cl_kernel kernel = backend_ctx->kernel_argmax_f32_i32;
+    CL_CHECK(clSetKernelArg(kernel, 0, sizeof(cl_mem),   &extra0->data_device));
+    CL_CHECK(clSetKernelArg(kernel, 1, sizeof(cl_ulong), &offset0));
+    CL_CHECK(clSetKernelArg(kernel, 2, sizeof(cl_mem),   &extrad->data_device));
+    CL_CHECK(clSetKernelArg(kernel, 3, sizeof(cl_ulong), &offsetd));
+    CL_CHECK(clSetKernelArg(kernel, 4, sizeof(int),      &ne00));
+    CL_CHECK(clSetKernelArg(kernel, 5, sizeof(int),      &nrows));
+
+    size_t global_work_size[] = { (size_t) nth, (size_t) CEIL_DIV(nrows, 2), 1 };
+    size_t local_work_size[] = { (size_t) nth, 1, 1 };
+    backend_ctx->enqueue_ndrange_kernel(kernel, 3, global_work_size, local_work_size, dst);
+}
+
 static void ggml_cl_argsort(ggml_backend_t backend, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst) {
     GGML_ASSERT(src0);
     GGML_ASSERT(src0->extra);
@@ -31071,6 +31133,12 @@ bool ggml_cl_compute_forward(ggml_backend_t backend, struct ggml_tensor * tensor
             }
             func = ggml_cl_im2col_3d;
             break;
+        case GGML_OP_ARGMAX:
+            if (!any_on_device) {
+                return false;
+            }
+            func = ggml_cl_argmax;
+            break;
         case GGML_OP_ARGSORT:
             if (!any_on_device) {
                 return false;

ggml/src/ggml-opencl/kernels/argmax.cl

@@ -0,0 +1,55 @@
+#define ARGMAX_LANES_PER_ROW 32
+#define ARGMAX_ROWS_PER_GROUP 2
+#define ARGMAX_BLOCK_SIZE (ARGMAX_LANES_PER_ROW * ARGMAX_ROWS_PER_GROUP)
+
+kernel void kernel_argmax_f32_i32(
+        global const char * src0,
+        ulong               offset0,
+        global char *       dst,
+        ulong               offsetd,
+        int                 ne00,
+        int                 nrows) {
+    const int local_id = get_local_id(0);
+    const int col = local_id % ARGMAX_LANES_PER_ROW;
+    const int row = get_group_id(1) * ARGMAX_ROWS_PER_GROUP + local_id / ARGMAX_LANES_PER_ROW;
+
+    float max_value = -INFINITY;
+    int max_index = 0;
+
+    if (row < nrows) {
+        global const float * src_row = (global const float *) (src0 + offset0) + row * ne00;
+        for (int i = col; i < ne00; i += ARGMAX_LANES_PER_ROW) {
+            const float value = src_row[i];
+            if (value > max_value || (value == max_value && i > max_index)) {
+                max_value = value;
+                max_index = i;
+            }
+        }
+    }
+
+    local float max_values[ARGMAX_BLOCK_SIZE];
+    local int max_indices[ARGMAX_BLOCK_SIZE];
+
+    max_values[local_id] = max_value;
+    max_indices[local_id] = max_index;
+
+    barrier(CLK_LOCAL_MEM_FENCE);
+
+    const int row_offset = local_id / ARGMAX_LANES_PER_ROW * ARGMAX_LANES_PER_ROW;
+    for (int stride = ARGMAX_LANES_PER_ROW / 2; stride > 0; stride /= 2) {
+        if (col < stride) {
+            const float value = max_values[row_offset + col + stride];
+            const int index = max_indices[row_offset + col + stride];
+            if (value > max_values[row_offset + col] || (value == max_values[row_offset + col] && index > max_indices[row_offset + col])) {
+                max_values[row_offset + col] = value;
+                max_indices[row_offset + col] = index;
+            }
+        }
+        barrier(CLK_LOCAL_MEM_FENCE);
+    }
+
+    if (col == 0 && row < nrows) {
+        global int * dst_row = (global int *) (dst + offsetd) + row;
+        dst_row[0] = max_indices[row_offset];
+    }
+}

ggml/src/ggml-sycl/conv2d.cpp

@@ -53,7 +53,16 @@ static inline int64_t whcn_output_index(int64_t n, int64_t c, int64_t y, int64_t
     return n * (P.OC * P.OW * P.OH) + c * P.OW * P.OH + y * P.OW + x;
 }
 
-template <typename T>
+// cwhn layout helpers: buffers are [N][H][W][C] input and [KH][KW][IC][OC] kernel
+static inline int64_t cwhn_input_index(int64_t n, int64_t c, int64_t y, int64_t x, const conv2d_params & P) {
+    return ((n * P.IH + y) * P.IW + x) * P.IC + c;
+}
+
+static inline int64_t cwhn_kernel_index(int64_t c_out, int64_t c_in, int64_t ky, int64_t kx, const conv2d_params & P) {
+    return ((ky * P.KW + kx) * P.IC + c_in) * P.OC + c_out;
+}
+
+template <typename T, bool CWHN>
 static void conv2d_kernel(const float * input, const T * kernel, float * output,
                           const conv2d_params P, const sycl::nd_item<3> & item_ct1) {
     const int64_t global_idx = item_ct1.get_local_id(2) +
@@ -77,8 +86,8 @@ static void conv2d_kernel(const float * input, const T * kernel, float * output,
             const int64_t in_y = calculate_input_coord(out_y, ky, P.ST_Y, P.DL_Y, P.PD_Y);
             for (int64_t kx = bounds.x_min; kx < bounds.x_max; ++kx) {
                 const int64_t in_x = calculate_input_coord(out_x, kx, P.ST_X, P.DL_X, P.PD_X);
-                const float input_val  = input[whcn_input_index(n, c_in, in_y, in_x, P)];
-                const T     kernel_val = kernel[whcn_kernel_index(c_out, c_in, ky, kx, P)];
+                const float input_val  = input[CWHN ? cwhn_input_index(n, c_in, in_y, in_x, P) : whcn_input_index(n, c_in, in_y, in_x, P)];
+                const T     kernel_val = kernel[CWHN ? cwhn_kernel_index(c_out, c_in, ky, kx, P) : whcn_kernel_index(c_out, c_in, ky, kx, P)];
                 acc += input_val * ggml_sycl_cast<float>(kernel_val);
             }
         }
@@ -87,7 +96,7 @@ static void conv2d_kernel(const float * input, const T * kernel, float * output,
     output[whcn_output_index(n, c_out, out_y, out_x, P)] = acc;
 }
 
-template <typename T>
+template <typename T, bool CWHN>
 static void conv2d_sycl(const float * X_D, const T * K_D, float * Y_D,
                         const conv2d_params P, const queue_ptr & stream) {
     const int num_blocks = (P.TOTAL + SYCL_CONV2D_BLOCK_SIZE - 1) / SYCL_CONV2D_BLOCK_SIZE;
@@ -95,7 +104,7 @@ static void conv2d_sycl(const float * X_D, const T * K_D, float * Y_D,
     const sycl::range<3> block_nums(1, 1, num_blocks);
     stream->parallel_for(sycl::nd_range<3>(block_nums * block_dims, block_dims),
         [=](sycl::nd_item<3> item_ct1) {
-            conv2d_kernel<T>(X_D, K_D, Y_D, P, item_ct1);
+            conv2d_kernel<T, CWHN>(X_D, K_D, Y_D, P, item_ct1);
         });
 }
 
@@ -108,7 +117,14 @@ void ggml_sycl_op_conv2d(ggml_backend_sycl_context & ctx, ggml_tensor * dst) {
     const float *       X_D    = (const float *) input->data;
     float *             Y_D    = (float *) dst->data;
 
-    GGML_ASSERT(ggml_is_contiguous(kernel));
+    // WHCN: contiguous [KW,KH,IC,OC] kernel and [W,H,C,N] input
+    // CWHN: channel-fastest buffer [N][H][W][C] input and [KH][KW][IC][OC] kernel (OC fastest)
+    const bool whcn = ggml_is_contiguous(kernel) && ggml_is_contiguous(input);
+    const bool cwhn = input->nb[2]  == ggml_type_size(input->type) &&
+                      kernel->nb[3] == ggml_type_size(kernel->type);
+    GGML_ASSERT(whcn || cwhn);
+    const bool cwhn_layout = !whcn && cwhn; // prefer whcn when both match (degenerate size-1 dims)
+
     GGML_ASSERT(kernel->type == GGML_TYPE_F16 || kernel->type == GGML_TYPE_F32);
     GGML_ASSERT(input->type == GGML_TYPE_F32);
     GGML_ASSERT(dst->type == GGML_TYPE_F32);
@@ -126,9 +142,6 @@ void ggml_sycl_op_conv2d(ggml_backend_sycl_context & ctx, ggml_tensor * dst) {
     const int       DL_X = p[4];
     const int       DL_Y = p[5];
 
-    // no cwhn layout support
-    GGML_ASSERT(p[6] == 0);
-
     const int IW = input->ne[0];
     const int IH = input->ne[1];
     const int OW = dst->ne[0];
@@ -142,9 +155,17 @@ void ggml_sycl_op_conv2d(ggml_backend_sycl_context & ctx, ggml_tensor * dst) {
     const int64_t     total  = (int64_t) B * OC * OH * OW;
     const conv2d_params params = { IW, IH, OW, OH, KW, KH, ST_X, ST_Y, PD_X, PD_Y, DL_X, DL_Y, IC, OC, B, total };
 
-    if (kernel->type == GGML_TYPE_F16) {
-        conv2d_sycl<sycl::half>(X_D, (const sycl::half *) K_D, Y_D, params, stream);
+    if (cwhn_layout) {
+        if (kernel->type == GGML_TYPE_F16) {
+            conv2d_sycl<sycl::half, true>(X_D, (const sycl::half *) K_D, Y_D, params, stream);
+        } else {
+            conv2d_sycl<float, true>(X_D, K_D, Y_D, params, stream);
+        }
     } else {
-        conv2d_sycl<float>(X_D, K_D, Y_D, params, stream);
+        if (kernel->type == GGML_TYPE_F16) {
+            conv2d_sycl<sycl::half, false>(X_D, (const sycl::half *) K_D, Y_D, params, stream);
+        } else {
+            conv2d_sycl<float, false>(X_D, K_D, Y_D, params, stream);
+        }
     }
 }

ggml/src/ggml-sycl/convert.cpp

@@ -657,58 +657,68 @@ static void convert_unary_sycl(const void * vx, dst_t * y, const int64_t k, dpct
 }
 
 
-to_fp16_sycl_t ggml_get_to_fp16_sycl(ggml_type type, ggml_tensor * dst) {
+static bool is_reordered_src0(const ggml_tensor * dst) {
+    return dst->src[0]->extra &&
+           ((ggml_tensor_extra_gpu *) dst->src[0]->extra)->optimized_feature.reorder;
+}
+
+template <typename dst_t>
+static to_t_sycl_t<dst_t> ggml_get_dequant_sycl(ggml_type type, ggml_tensor * dst) {
     switch (type) {
         case GGML_TYPE_Q1_0:
-            return dequantize_block_sycl<QK1_0, QR1_0, dequantize_q1_0>;
+            return dequantize_block_sycl<QK1_0, QR1_0, dequantize_q1_0, dst_t>;
         case GGML_TYPE_Q2_0:
-            return dequantize_block_sycl<QK2_0, QR2_0, dequantize_q2_0>;
+            return dequantize_block_sycl<QK2_0, QR2_0, dequantize_q2_0, dst_t>;
         case GGML_TYPE_Q4_0:
-            if (dst->src[0]->extra &&
-                ((ggml_tensor_extra_gpu*)dst->src[0]->extra)->optimized_feature.reorder) {
+            if (is_reordered_src0(dst)) {
                 return dequantize_row_q4_0_sycl_reorder;
+            } else if constexpr (std::is_same_v<dst_t, float>) {
+                return dequantize_row_q4_0_sycl;
             } else {
-                return dequantize_block_sycl<QK4_0, QR4_0, dequantize_q4_0>;
+                return dequantize_block_sycl<QK4_0, QR4_0, dequantize_q4_0, dst_t>;
             }
         case GGML_TYPE_Q4_1:
-            return dequantize_block_sycl<QK4_1, QR4_1, dequantize_q4_1>;
+            if constexpr (std::is_same_v<dst_t, float>) {
+                return dequantize_row_q4_1_sycl;
+            } else {
+                return dequantize_block_sycl<QK4_1, QR4_1, dequantize_q4_1, dst_t>;
+            }
         case GGML_TYPE_Q5_0:
-            return dequantize_block_sycl<QK5_0, QR5_0, dequantize_q5_0>;
+            return dequantize_block_sycl<QK5_0, QR5_0, dequantize_q5_0, dst_t>;
         case GGML_TYPE_Q5_1:
-            return dequantize_block_sycl<QK5_1, QR5_1, dequantize_q5_1>;
+            return dequantize_block_sycl<QK5_1, QR5_1, dequantize_q5_1, dst_t>;
         case GGML_TYPE_Q8_0:
-            if (dst->src[0]->extra &&
-                ((ggml_tensor_extra_gpu *) dst->src[0]->extra)->optimized_feature.reorder) {
+            if (is_reordered_src0(dst)) {
                 return dequantize_row_q8_0_sycl_reorder;
             } else {
-                return dequantize_block_sycl<QK8_0, QR8_0, dequantize_q8_0>;
+                return dequantize_block_sycl<QK8_0, QR8_0, dequantize_q8_0, dst_t>;
             }
         case GGML_TYPE_Q2_K:
-            if (dst->src[0]->extra && ((ggml_tensor_extra_gpu *) dst->src[0]->extra)->optimized_feature.reorder) {
+            if (is_reordered_src0(dst)) {
                 return dequantize_row_q2_K_sycl_reorder;
             } else {
                 return dequantize_row_q2_K_sycl;
             }
         case GGML_TYPE_Q3_K:
-            if (dst->src[0]->extra && ((ggml_tensor_extra_gpu *) dst->src[0]->extra)->optimized_feature.reorder) {
+            if (is_reordered_src0(dst)) {
                 return dequantize_row_q3_K_sycl_reorder;
             } else {
                 return dequantize_row_q3_K_sycl;
             }
         case GGML_TYPE_Q4_K:
-            if (dst->src[0]->extra && ((ggml_tensor_extra_gpu *) dst->src[0]->extra)->optimized_feature.reorder) {
+            if (is_reordered_src0(dst)) {
                 return dequantize_row_q4_K_sycl_reorder;
             } else {
                 return dequantize_row_q4_K_sycl;
             }
         case GGML_TYPE_Q5_K:
-            if (dst->src[0]->extra && ((ggml_tensor_extra_gpu *) dst->src[0]->extra)->optimized_feature.reorder) {
+            if (is_reordered_src0(dst)) {
                 return dequantize_row_q5_K_sycl_reorder;
             } else {
                 return dequantize_row_q5_K_sycl;
             }
         case GGML_TYPE_Q6_K:
-            if (dst->src[0]->extra && ((ggml_tensor_extra_gpu *) dst->src[0]->extra)->optimized_feature.reorder) {
+            if (is_reordered_src0(dst)) {
                 return dequantize_row_q6_K_sycl_reorder;
             } else {
                 return dequantize_row_q6_K_sycl;
@@ -724,15 +734,13 @@ to_fp16_sycl_t ggml_get_to_fp16_sycl(ggml_type type, ggml_tensor * dst) {
         case GGML_TYPE_IQ2_S:
             return dequantize_row_iq2_s_sycl;
         case GGML_TYPE_IQ3_XXS:
-            if (dst->src[0]->extra &&
-                ((ggml_tensor_extra_gpu *) dst->src[0]->extra)->optimized_feature.reorder) {
+            if (is_reordered_src0(dst)) {
                 return dequantize_row_iq3_xxs_sycl_reorder;
             } else {
                 return dequantize_row_iq3_xxs_sycl;
             }
         case GGML_TYPE_IQ3_S:
-            if (dst->src[0]->extra &&
-                ((ggml_tensor_extra_gpu *) dst->src[0]->extra)->optimized_feature.reorder) {
+            if (is_reordered_src0(dst)) {
                 return dequantize_row_iq3_s_sycl_reorder;
             } else {
                 return dequantize_row_iq3_s_sycl;
@@ -742,124 +750,38 @@ to_fp16_sycl_t ggml_get_to_fp16_sycl(ggml_type type, ggml_tensor * dst) {
         case GGML_TYPE_IQ4_NL:
             return dequantize_row_iq4_nl_sycl;
         case GGML_TYPE_MXFP4:
-            if (dst->src[0]->extra && ((ggml_tensor_extra_gpu *) dst->src[0]->extra)->optimized_feature.reorder) {
+            if (is_reordered_src0(dst)) {
                 return dequantize_row_mxfp4_sycl_reorder;
             }
             return dequantize_row_mxfp4_sycl;
         case GGML_TYPE_NVFP4:
             return dequantize_row_nvfp4_sycl;
         case GGML_TYPE_F32:
-            return convert_unary_sycl<float>;
+            if constexpr (std::is_same_v<dst_t, sycl::half>) {
+                return convert_unary_sycl<float>;
+            }
+            break;
+        case GGML_TYPE_F16:
+            if constexpr (std::is_same_v<dst_t, float>) {
+                return convert_unary_sycl<sycl::half>;
+            }
+            break;
 #ifdef GGML_SYCL_HAS_BF16
         case GGML_TYPE_BF16:
             return convert_unary_sycl<sycl::ext::oneapi::bfloat16>;
 #endif
         default:
-            GGML_ABORT("fatal error: unsupport data type=%s\n", ggml_type_name(type));
+            break;
     }
+    GGML_ABORT("fatal error: unsupport data type=%s\n", ggml_type_name(type));
+}
+
+to_fp16_sycl_t ggml_get_to_fp16_sycl(ggml_type type, ggml_tensor * dst) {
+    return ggml_get_dequant_sycl<sycl::half>(type, dst);
 }
 
 to_fp32_sycl_t ggml_get_to_fp32_sycl(ggml_type type, ggml_tensor *dst) {
-    switch (type) {
-        case GGML_TYPE_Q1_0:
-            return dequantize_block_sycl<QK1_0, QR1_0, dequantize_q1_0>;
-        case GGML_TYPE_Q2_0:
-            return dequantize_block_sycl<QK2_0, QR2_0, dequantize_q2_0>;
-        case GGML_TYPE_Q4_0:
-            if (dst->src[0]->extra &&
-                ((ggml_tensor_extra_gpu*)dst->src[0]->extra)->optimized_feature.reorder) {
-                return dequantize_row_q4_0_sycl_reorder;
-            } else {
-                return dequantize_row_q4_0_sycl;
-            }
-        case GGML_TYPE_Q4_1:
-            return dequantize_row_q4_1_sycl;
-        case GGML_TYPE_Q5_0:
-            return dequantize_block_sycl<QK5_0, QR5_0, dequantize_q5_0>;
-        case GGML_TYPE_Q5_1:
-            return dequantize_block_sycl<QK5_1, QR5_1, dequantize_q5_1>;
-        case GGML_TYPE_Q8_0:
-            if (dst->src[0]->extra &&
-                ((ggml_tensor_extra_gpu*)dst->src[0]->extra)->optimized_feature.reorder) {
-                return dequantize_row_q8_0_sycl_reorder;
-            } else {
-                return dequantize_block_sycl<QK8_0, QR8_0, dequantize_q8_0>;
-            }
-        case GGML_TYPE_Q2_K:
-            if (dst->src[0]->extra && ((ggml_tensor_extra_gpu *) dst->src[0]->extra)->optimized_feature.reorder) {
-                return dequantize_row_q2_K_sycl_reorder;
-            } else {
-                return dequantize_row_q2_K_sycl;
-            }
-        case GGML_TYPE_Q3_K:
-            if (dst->src[0]->extra && ((ggml_tensor_extra_gpu *) dst->src[0]->extra)->optimized_feature.reorder) {
-                return dequantize_row_q3_K_sycl_reorder;
-            } else {
-                return dequantize_row_q3_K_sycl;
-            }
-        case GGML_TYPE_Q4_K:
-            if (dst->src[0]->extra &&
-                ((ggml_tensor_extra_gpu*)dst->src[0]->extra)->optimized_feature.reorder) {
-                return dequantize_row_q4_K_sycl_reorder;
-            } else {
-                return dequantize_row_q4_K_sycl;
-            }
-        case GGML_TYPE_Q5_K:
-            if (dst->src[0]->extra && ((ggml_tensor_extra_gpu *) dst->src[0]->extra)->optimized_feature.reorder) {
-                return dequantize_row_q5_K_sycl_reorder;
-            } else {
-                return dequantize_row_q5_K_sycl;
-            }
-        case GGML_TYPE_Q6_K:
-            if (dst->src[0]->extra && ((ggml_tensor_extra_gpu *) dst->src[0]->extra)->optimized_feature.reorder) {
-                return dequantize_row_q6_K_sycl_reorder;
-            } else {
-                return dequantize_row_q6_K_sycl;
-            }
-        case GGML_TYPE_IQ1_S:
-            return dequantize_row_iq1_s_sycl;
-        case GGML_TYPE_IQ1_M:
-            return dequantize_row_iq1_m_sycl;
-        case GGML_TYPE_IQ2_XXS:
-            return dequantize_row_iq2_xxs_sycl;
-        case GGML_TYPE_IQ2_XS:
-            return dequantize_row_iq2_xs_sycl;
-        case GGML_TYPE_IQ2_S:
-            return dequantize_row_iq2_s_sycl;
-        case GGML_TYPE_IQ3_XXS:
-            if (dst->src[0]->extra &&
-                ((ggml_tensor_extra_gpu *) dst->src[0]->extra)->optimized_feature.reorder) {
-                return dequantize_row_iq3_xxs_sycl_reorder;
-            } else {
-                return dequantize_row_iq3_xxs_sycl;
-            }
-        case GGML_TYPE_IQ3_S:
-            if (dst->src[0]->extra &&
-                ((ggml_tensor_extra_gpu *) dst->src[0]->extra)->optimized_feature.reorder) {
-                return dequantize_row_iq3_s_sycl_reorder;
-            } else {
-                return dequantize_row_iq3_s_sycl;
-            }
-        case GGML_TYPE_IQ4_XS:
-            return dequantize_row_iq4_xs_sycl;
-        case GGML_TYPE_IQ4_NL:
-            return dequantize_row_iq4_nl_sycl;
-        case GGML_TYPE_MXFP4:
-            if (dst->src[0]->extra && ((ggml_tensor_extra_gpu *) dst->src[0]->extra)->optimized_feature.reorder) {
-                return dequantize_row_mxfp4_sycl_reorder;
-            }
-            return dequantize_row_mxfp4_sycl;
-        case GGML_TYPE_NVFP4:
-            return dequantize_row_nvfp4_sycl;
-        case GGML_TYPE_F16:
-            return convert_unary_sycl<sycl::half>;
-#ifdef GGML_SYCL_HAS_BF16
-        case GGML_TYPE_BF16:
-            return convert_unary_sycl<sycl::ext::oneapi::bfloat16>;
-#endif
-        default:
-            GGML_ABORT("fatal error: unsupport data type=%s\n", ggml_type_name(type));
-    }
+    return ggml_get_dequant_sycl<float>(type, dst);
 }
 
 

ggml/src/ggml-sycl/ggml-sycl.cpp

@@ -6912,6 +6912,11 @@ static bool do_ggml_backend_sycl_device_supports_op(ggml_backend_dev_t dev, cons
                 return false;
             }
         case GGML_OP_CONV_2D:
+            // WHCN: contiguous [KW,KH,IC,OC] kernel and [W,H,C,N] input
+            // CWHN: channel-fastest input (C fastest) and kernel with OC fastest
+            return (ggml_is_contiguous(op->src[0]) && ggml_is_contiguous(op->src[1])) ||
+                   (op->src[1]->nb[2] == ggml_type_size(op->src[1]->type) &&
+                    op->src[0]->nb[3] == ggml_type_size(op->src[0]->type));
         case GGML_OP_CONV_2D_DW:
         case GGML_OP_CONV_TRANSPOSE_2D:
             return true;

ggml/src/ggml-webgpu/ggml-webgpu-shader-lib.hpp

@@ -1655,6 +1655,7 @@ class ggml_webgpu_shader_lib {
 
                     switch (key.src_type) {
                         case GGML_TYPE_Q1_0:
+                        case GGML_TYPE_Q2_0:
                         case GGML_TYPE_Q4_0:
                         case GGML_TYPE_Q5_0:
                         case GGML_TYPE_Q8_0:
@@ -1699,7 +1700,7 @@ class ggml_webgpu_shader_lib {
                     } else if ((key.src_type >= GGML_TYPE_Q4_0 && key.src_type <= GGML_TYPE_Q8_1) ||
                                key.src_type == GGML_TYPE_IQ4_NL || key.src_type == GGML_TYPE_MXFP4) {
                         defines.push_back("BLOCK_SIZE=32u");
-                    } else if (key.src_type == GGML_TYPE_NVFP4) {
+                    } else if (key.src_type == GGML_TYPE_Q2_0 || key.src_type == GGML_TYPE_NVFP4) {
                         defines.push_back("BLOCK_SIZE=64u");
                     } else if (key.src_type >= GGML_TYPE_Q2_K) {
                         defines.push_back("BLOCK_SIZE=256u");

ggml/src/ggml-webgpu/ggml-webgpu.cpp

@@ -4395,6 +4395,7 @@ static bool ggml_backend_webgpu_device_supports_buft(ggml_backend_dev_t dev, ggm
 static bool ggml_webgpu_supported_qtype(ggml_type type) {
     switch (type) {
         case GGML_TYPE_Q1_0:
+        case GGML_TYPE_Q2_0:
         case GGML_TYPE_Q4_0:
         case GGML_TYPE_Q4_1:
         case GGML_TYPE_Q5_0:
@@ -4500,6 +4501,7 @@ static bool ggml_backend_webgpu_device_supports_op(ggml_backend_dev_t dev, const
                             case GGML_TYPE_F16:
                             case GGML_TYPE_BF16:
                             case GGML_TYPE_Q1_0:
+                            case GGML_TYPE_Q2_0:
                             case GGML_TYPE_Q4_0:
                             case GGML_TYPE_Q4_1:
                             case GGML_TYPE_Q5_0:

ggml/src/ggml-webgpu/wgsl-shaders/get_rows.wgsl

@@ -51,6 +51,20 @@ fn copy_elements(src_base: u32, dst_base: u32, offset: u32) {
 }
 #endif
 
+#ifdef Q2_0
+fn copy_elements(src_base: u32, dst_base: u32, offset: u32) {
+    let block_byte_base = (src_base + offset) * 18u;
+    let d = load_f16_as_f32_at_src(block_byte_base);
+    for (var j = 0u; j < 4u; j++) {
+        let q_packed = load_u32_at_src(block_byte_base + 2u + j * 4u);
+        for (var k = 0u; k < 16u; k++) {
+            let q = (q_packed >> (2u * k)) & 3u;
+            dst[dst_base + offset * 64u + j * 16u + k] = (f32(q) - 1.0) * d;
+        }
+    }
+}
+#endif
+
 #ifdef Q4_0
 fn copy_elements(src_base: u32, dst_base: u32, offset: u32) {
     let block_byte_base = (src_base + offset) * 18; // Block stride: 18 bytes

ggml/src/ggml-webgpu/wgsl-shaders/mul_mat_decls.tmpl

@@ -111,6 +111,29 @@ fn init_shmem_src0(thread_id: u32, batch_offset: u32, offset_m: u32, k_outer: u3
 }
 #endif // INIT_SRC0_SHMEM_Q1_0
 
+#ifdef INIT_SRC0_SHMEM_Q2_0
+const BLOCK_SIZE = 64u;
+const BLOCK_SIZE_BYTES = 18u;
+const NQ = 16u;
+
+fn init_shmem_src0(thread_id: u32, batch_offset: u32, offset_m: u32, k_outer: u32) {
+    for (var i = thread_id * NQ; i < TILE_SRC0_SHMEM; i += TOTAL_WORKGROUP_SIZE * NQ) {
+        let global_m = offset_m + i / TILE_K;
+        let global_k = k_outer + i % TILE_K;
+        if (global_m < params.m && global_k < params.k) {
+            let block_idx = batch_offset + global_m * params.stride_01 + global_k / BLOCK_SIZE;
+            let block_byte_base = block_idx * BLOCK_SIZE_BYTES;
+            let d = load_f16_at_src0(block_byte_base);
+            let q_packed = load_u32_at_src0(block_byte_base + 2u + (global_k % BLOCK_SIZE) / 4u);
+            for (var k = 0u; k < NQ; k++) {
+                let q = (q_packed >> (2u * k)) & 3u;
+                shmem[i + k] = (f16(q) - 1.0) * d;
+            }
+        }
+    }
+}
+#endif // INIT_SRC0_SHMEM_Q2_0
+
 // legacy-quants
 #if defined(INIT_SRC0_SHMEM_Q4_0) || defined(INIT_SRC0_SHMEM_Q4_1) || defined(INIT_SRC0_SHMEM_Q5_0) || defined(INIT_SRC0_SHMEM_Q5_1) || defined(INIT_SRC0_SHMEM_Q8_0) || defined(INIT_SRC0_SHMEM_Q8_1) || defined(INIT_SRC0_SHMEM_MXFP4)
 const BLOCK_SIZE = 32u;

ggml/src/ggml-webgpu/wgsl-shaders/mul_mat_vec_acc.tmpl

@@ -119,6 +119,46 @@ fn accumulate_vec_dot(thread_id: u32, row_base: u32, src0_batch_offset: u32, src
 }
 #endif
 
+#ifdef MUL_ACC_Q2_0
+#define BLOCK_SIZE 64
+#define BLOCK_SIZE_BYTES 18
+#define THREADS_PER_BLOCK 4
+#define ELEMS_PER_THREAD (BLOCK_SIZE/THREADS_PER_BLOCK)
+fn accumulate_vec_dot(thread_id: u32, row_base: u32, src0_batch_offset: u32, src1_idx_base: u32) -> array<array<f32, OUTPUTS_PER_WG>, NUM_COLS> {
+    var acc: array<array<f32, OUTPUTS_PER_WG>, NUM_COLS>;
+
+    let num_blocks = params.k / BLOCK_SIZE;
+    let thread_within_block = thread_id % THREADS_PER_BLOCK;
+    for (var block = thread_id / THREADS_PER_BLOCK; block < num_blocks; block += WG_SIZE / THREADS_PER_BLOCK) {
+        let x_base = src1_idx_base + block * BLOCK_SIZE + thread_within_block * ELEMS_PER_THREAD;
+        var x_block: array<array<f32, ELEMS_PER_THREAD>, NUM_COLS>;
+        for (var col = 0u; col < NUM_COLS; col++) {
+            for (var i = 0u; i < ELEMS_PER_THREAD; i++) {
+                x_block[col][i] = f32(SRC1[x_base + col * params.stride_11 + i]);
+            }
+        }
+        for (var row = 0u; row < OUTPUTS_PER_WG; row++) {
+            let output_row = row_base + row;
+            if (output_row < params.m) {
+                let block_byte_base = (src0_batch_offset + output_row * params.stride_01 + block) * BLOCK_SIZE_BYTES;
+                let d = f32(load_f16_at_src0(block_byte_base));
+                let q_packed = load_u32_at_src0(block_byte_base + 2u + 4u * thread_within_block);
+                for (var col = 0u; col < NUM_COLS; col++) {
+                    var row_sum = 0.0;
+                    for (var i = 0u; i < ELEMS_PER_THREAD; i++) {
+                        let q = (q_packed >> (2u * i)) & 3u;
+                        row_sum += (f32(q) - 1.0) * x_block[col][i];
+                    }
+                    acc[col][row] += d * row_sum;
+                }
+            }
+        }
+    }
+
+    return acc;
+}
+#endif
+
 #ifdef MUL_ACC_Q4_0
 #define BLOCK_SIZE 32
 #define BLOCK_SIZE_BYTES 18

ggml/src/ggml.c

@@ -1095,6 +1095,7 @@ static const char * GGML_OP_NAME[GGML_OP_COUNT] = {
     "DSV4_HC_COMB",
     "DSV4_HC_PRE",
     "DSV4_HC_POST",
+    "TOP_K_MASK",
 
     "UNARY",
 
@@ -1112,7 +1113,7 @@ static const char * GGML_OP_NAME[GGML_OP_COUNT] = {
     "GLU",
 };
 
-static_assert(GGML_OP_COUNT == 101, "GGML_OP_COUNT != 101");
+static_assert(GGML_OP_COUNT == 102, "GGML_OP_COUNT != 102");
 
 static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = {
     "none",
@@ -1210,6 +1211,7 @@ static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = {
     "dsv4_hc_comb(mixes, scale, base)",
     "dsv4_hc_pre(x, weights)",
     "dsv4_hc_post(x, residual, post, comb)",
+    "top_k_mask(mask, top_k, score, blocks, extra)",
 
     "unary(x)",
 
@@ -1227,7 +1229,7 @@ static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = {
     "glu(x)",
 };
 
-static_assert(GGML_OP_COUNT == 101, "GGML_OP_COUNT != 101");
+static_assert(GGML_OP_COUNT == 102, "GGML_OP_COUNT != 102");
 
 static_assert(GGML_OP_POOL_COUNT == 2, "GGML_OP_POOL_COUNT != 2");
 
@@ -6630,6 +6632,48 @@ struct ggml_tensor * ggml_dsv4_hc_post(
     return result;
 }
 
+// ggml_top_k_mask
+
+struct ggml_tensor * ggml_top_k_mask(
+        struct ggml_context * ctx,
+        struct ggml_tensor  * mask,
+        struct ggml_tensor  * top_k,
+        struct ggml_tensor  * score,
+        struct ggml_tensor  * blocks,
+        struct ggml_tensor  * extra) {
+    GGML_ASSERT(mask->type == GGML_TYPE_F16 || mask->type == GGML_TYPE_F32);
+    GGML_ASSERT(mask->nb[0] == ggml_type_size(mask->type));
+    GGML_ASSERT(top_k->type == GGML_TYPE_I32);
+    GGML_ASSERT(top_k->ne[1] % mask->ne[1] == 0 && top_k->ne[2] % mask->ne[2] == 0 && top_k->ne[3] % mask->ne[3] == 0);
+
+    if (score) {
+        GGML_ASSERT(score->type == GGML_TYPE_F32);
+        GGML_ASSERT(score->ne[1] == top_k->ne[1] && score->ne[2] == top_k->ne[2] && score->ne[3] == top_k->ne[3]);
+    }
+
+    if (blocks) {
+        GGML_ASSERT(blocks->type == GGML_TYPE_I32);
+        GGML_ASSERT(blocks->ne[2] == 1 && blocks->ne[3] == 1);
+        GGML_ASSERT(!score || blocks->ne[1] == score->ne[0]);
+    }
+
+    if (extra) {
+        GGML_ASSERT(extra->type == GGML_TYPE_I32);
+        GGML_ASSERT(extra->ne[1] == top_k->ne[1] && extra->ne[2] == top_k->ne[2] && extra->ne[3] == top_k->ne[3]);
+    }
+
+    struct ggml_tensor * result = ggml_new_tensor_4d(ctx, mask->type, mask->ne[0], top_k->ne[1], top_k->ne[2], top_k->ne[3]);
+
+    result->op     = GGML_OP_TOP_K_MASK;
+    result->src[0] = mask;
+    result->src[1] = top_k;
+    result->src[2] = score;
+    result->src[3] = blocks;
+    result->src[4] = extra;
+
+    return result;
+}
+
 ////////////////////////////////////////////////////////////////////////////////
 
 struct ggml_hash_set ggml_hash_set_new(size_t size) {

gguf-py/gguf/constants.py

@@ -263,6 +263,7 @@ class Keys:
         DIMENSION_SECTIONS        = "{arch}.rope.dimension_sections"
         SECTION_ORDER             = "{arch}.rope.section_order"
         FREQ_BASE                 = "{arch}.rope.freq_base"
+        FREQ_BASE_SPATIAL         = "{arch}.rope.freq_base_spatial"
         FREQ_BASE_SWA             = "{arch}.rope.freq_base_swa"
         SCALING_TYPE              = "{arch}.rope.scaling.type"
         SCALING_FACTOR            = "{arch}.rope.scaling.factor"
@@ -370,6 +371,7 @@ class Keys:
         DNA_AUTO_TAGS        = "tokenizer.ggml.dna_auto_tags"
         ADD_PREFIX           = "tokenizer.ggml.add_space_prefix"
         REMOVE_EXTRA_WS      = "tokenizer.ggml.remove_extra_whitespaces"
+        CLEAN_SPACES         = "tokenizer.ggml.clean_spaces"
         PRECOMPILED_CHARSMAP = "tokenizer.ggml.precompiled_charsmap"
         SUPPRESS_TOKENS      = "tokenizer.ggml.suppress_tokens"
         HF_JSON              = "tokenizer.huggingface.json"
@@ -573,8 +575,10 @@ class MODEL_ARCH(IntEnum):
     QWEN2            = auto()
     QWEN2MOE         = auto()
     QWEN2VL          = auto()
+    SENSENOVA_U1     = auto()
     QWEN3            = auto()
     QWEN3MOE         = auto()
+    FLINT            = auto()
     QWEN3NEXT        = auto()
     HRM_TEXT         = auto()
     BEAM             = auto()
@@ -1390,7 +1394,9 @@ MODEL_ARCH_NAMES: dict[MODEL_ARCH, str] = {
     MODEL_ARCH.QWEN2:            "qwen2",
     MODEL_ARCH.QWEN2MOE:         "qwen2moe",
     MODEL_ARCH.QWEN2VL:          "qwen2vl",
+    MODEL_ARCH.SENSENOVA_U1:     "sensenova_u1",
     MODEL_ARCH.QWEN3:            "qwen3",
+    MODEL_ARCH.FLINT:            "flint",
     MODEL_ARCH.QWEN3MOE:         "qwen3moe",
     MODEL_ARCH.QWEN3NEXT:        "qwen3next",
     MODEL_ARCH.HRM_TEXT:         "hrm_text",
@@ -2916,6 +2922,24 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = {
         MODEL_TENSOR.FFN_DOWN_SHEXP,
         MODEL_TENSOR.FFN_UP_SHEXP,
     ],
+    MODEL_ARCH.SENSENOVA_U1: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.ROPE_FREQS,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_QKV,
+        MODEL_TENSOR.ATTN_Q,
+        MODEL_TENSOR.ATTN_Q_NORM,
+        MODEL_TENSOR.ATTN_K,
+        MODEL_TENSOR.ATTN_K_NORM,
+        MODEL_TENSOR.ATTN_V,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.FFN_NORM,
+        MODEL_TENSOR.FFN_GATE,
+        MODEL_TENSOR.FFN_DOWN,
+        MODEL_TENSOR.FFN_UP,
+    ],
     MODEL_ARCH.QWEN3: [
         MODEL_TENSOR.TOKEN_EMBD,
         MODEL_TENSOR.OUTPUT_NORM,
@@ -2953,6 +2977,33 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = {
         MODEL_TENSOR.FFN_DOWN_EXP,
         MODEL_TENSOR.FFN_UP_EXP,
     ],
+
+    MODEL_ARCH.FLINT: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_QKV,
+        MODEL_TENSOR.ATTN_Q,
+        MODEL_TENSOR.ATTN_Q_NORM,
+        MODEL_TENSOR.ATTN_K,
+        MODEL_TENSOR.ATTN_K_NORM,
+        MODEL_TENSOR.ATTN_V,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.FFN_NORM,
+        MODEL_TENSOR.FFN_GATE_INP,
+        MODEL_TENSOR.FFN_GATE_EXP,
+        MODEL_TENSOR.FFN_DOWN_EXP,
+        MODEL_TENSOR.FFN_UP_EXP,
+        MODEL_TENSOR.ATTN_SINKS,
+        MODEL_TENSOR.FFN_GATE,
+        MODEL_TENSOR.FFN_UP,
+        MODEL_TENSOR.FFN_DOWN,
+        MODEL_TENSOR.FFN_GATE_SHEXP,
+        MODEL_TENSOR.FFN_UP_SHEXP,
+        MODEL_TENSOR.FFN_DOWN_SHEXP,
+        MODEL_TENSOR.FFN_EXP_PROBS_B,
+    ],
     MODEL_ARCH.QWEN3NEXT: [
         MODEL_TENSOR.TOKEN_EMBD,
         MODEL_TENSOR.OUTPUT_NORM,
@@ -6485,6 +6536,7 @@ KEY_ATTENTION_LAYERNORM_RMS_EPS = Keys.Attention.LAYERNORM_RMS_EPS
 # RoPE
 KEY_ROPE_DIMENSION_COUNT           = Keys.Rope.DIMENSION_COUNT
 KEY_ROPE_FREQ_BASE                 = Keys.Rope.FREQ_BASE
+KEY_ROPE_FREQ_BASE_SPATIAL         = Keys.Rope.FREQ_BASE_SPATIAL
 KEY_ROPE_SCALING_TYPE              = Keys.Rope.SCALING_TYPE
 KEY_ROPE_SCALING_FACTOR            = Keys.Rope.SCALING_FACTOR
 KEY_ROPE_SCALING_ORIG_CTX_LEN      = Keys.Rope.SCALING_ORIG_CTX_LEN

gguf-py/gguf/gguf_writer.py

@@ -1167,6 +1167,9 @@ class GGUFWriter:
     def add_rope_freq_base(self, value: float) -> None:
         self.add_float32(Keys.Rope.FREQ_BASE.format(arch=self.arch), value)
 
+    def add_rope_freq_base_spatial(self, value: float) -> None:
+        self.add_float32(Keys.Rope.FREQ_BASE_SPATIAL.format(arch=self.arch), value)
+
     def add_rope_scaling_type(self, value: RopeScalingType) -> None:
         self.add_string(Keys.Rope.SCALING_TYPE.format(arch=self.arch), value.value)
 
@@ -1323,6 +1326,9 @@ class GGUFWriter:
     def add_remove_extra_whitespaces(self, value: bool) -> None:
         self.add_bool(Keys.Tokenizer.REMOVE_EXTRA_WS, value)
 
+    def add_clean_spaces(self, value: bool) -> None:
+        self.add_bool(Keys.Tokenizer.CLEAN_SPACES, value)
+
     def add_precompiled_charsmap(self, charsmap: bytes) -> None:
         self.add_array(Keys.Tokenizer.PRECOMPILED_CHARSMAP, charsmap)
 

gguf-py/gguf/tensor_mapping.py

@@ -429,6 +429,7 @@ class TensorNameMap:
         ),
 
         MODEL_TENSOR.ATTN_SINKS: (
+            "model.layers.{bid}.self_attn.softmax_offset", # flint
             "model.layers.{bid}.self_attn.sinks", # openai-moe
             "model.layers.{bid}.self_attn.attention_sink_bias", # mimov2
             "model.layers.{bid}.self_attn.learnable_sink_param", # hy-v4

gguf-py/gguf/vocab.py

@@ -55,6 +55,7 @@ class SpecialVocab:
     normalizer_nfc: bool | None
     normalizer_lowercase: bool | None
     normalizer_strip_accents: bool | None
+    clean_spaces: bool | None
 
     def __init__(
         self, path: str | os.PathLike[str], load_merges: bool = False,
@@ -70,6 +71,7 @@ class SpecialVocab:
         self.normalizer_nfc = None
         self.normalizer_lowercase = None
         self.normalizer_strip_accents = None
+        self.clean_spaces = None
         if special_token_types is not None:
             self.special_token_types = special_token_types
         else:
@@ -120,6 +122,10 @@ class SpecialVocab:
             if not quiet:
                 logger.info(f'Setting normalizer_strip_accents to {self.normalizer_strip_accents}')
             gw.add_normalizer_strip_accents(self.normalizer_strip_accents)
+        if self.clean_spaces is not None:
+            if not quiet:
+                logger.info(f'Setting clean_spaces to {self.clean_spaces}')
+            gw.add_clean_spaces(self.clean_spaces)
 
     def _load(self, path: Path) -> None:
         self._try_load_from_tokenizer_json(path)
@@ -345,6 +351,10 @@ class SpecialVocab:
             self.chat_template = chat_template
         else:
             logger.warning(f'Bad type for chat_template field in {tokenizer_config_file!r} - ignoring')
+        # decode clean-up of spaces before punctuation (transformers clean_up_tokenization_spaces)
+        clean_spaces = tokenizer_config.get('clean_up_tokenization_spaces')
+        if isinstance(clean_spaces, bool):
+            self.clean_spaces = clean_spaces
         for typ in self.special_token_types:
             add_entry = tokenizer_config.get(f'add_{typ}_token')
             if isinstance(add_entry, bool):

src/llama-adapter.cpp

@@ -334,9 +334,10 @@ static void llama_adapter_lora_init_impl(llama_model & model, FILE * file, llama
         if (!model_tensor) {
             throw std::runtime_error("LoRA tensor '" + name + "' does not exist in base model (hint: maybe wrong base model?)");
         }
-        // the LoRA delta of a Hadamard-folded weight reads the transformed input, so the result is wrong
-        if (model.hdmd.weight_blocks.count(name) || model.hdmd.inverse_blocks.count(name)) {
-            throw std::runtime_error("LoRA tensor '" + name + "' targets a prism.hadamard folded weight, which is not supported");
+        // build_lora_mm feeds a folded weight's LoRA delta the primal input, so folded matmul weights are fine;
+        // a latent lookup table (token_embd, also the head when the output is tied) stores rotated rows, refuse it
+        if (model.hdmd.inverse_blocks.count(name)) {
+            throw std::runtime_error("LoRA tensor '" + name + "' targets a prism.hadamard latent lookup table, which is not supported");
         }
 
         auto * buft = ggml_backend_buffer_get_type(model_tensor->buffer);

src/llama-arch.cpp

@@ -33,7 +33,9 @@ static const std::map<llm_arch, const char *> LLM_ARCH_NAMES = {
     { LLM_ARCH_QWEN2,            "qwen2"            },
     { LLM_ARCH_QWEN2MOE,         "qwen2moe"         },
     { LLM_ARCH_QWEN2VL,          "qwen2vl"          },
+    { LLM_ARCH_SENSENOVA_U1,     "sensenova_u1"     },
     { LLM_ARCH_QWEN3,            "qwen3"            },
+    { LLM_ARCH_FLINT,            "flint"            },
     { LLM_ARCH_QWEN3MOE,         "qwen3moe"         },
     { LLM_ARCH_QWEN3NEXT,        "qwen3next"        },
     { LLM_ARCH_QWEN3VL,          "qwen3vl"          },
@@ -339,6 +341,7 @@ static const std::map<llm_kv, const char *> LLM_KV_NAMES = {
     { LLM_KV_ROPE_DIMENSION_SECTIONS,        "%s.rope.dimension_sections"              },
     { LLM_KV_ROPE_SECTION_ORDER,             "%s.rope.section_order"                   },
     { LLM_KV_ROPE_FREQ_BASE,                 "%s.rope.freq_base"                       },
+    { LLM_KV_ROPE_FREQ_BASE_SPATIAL,         "%s.rope.freq_base_spatial"               },
     { LLM_KV_ROPE_FREQ_BASE_SWA,             "%s.rope.freq_base_swa"                   },
     { LLM_KV_ROPE_SCALE_LINEAR,              "%s.rope.scale_linear"                    },
     { LLM_KV_ROPE_SCALING_TYPE,              "%s.rope.scaling.type"                    },
@@ -434,6 +437,7 @@ static const std::map<llm_kv, const char *> LLM_KV_NAMES = {
     { LLM_KV_TOKENIZER_DNA_AUTO_TAGS,            "tokenizer.ggml.dna_auto_tags"            },
     { LLM_KV_TOKENIZER_ADD_PREFIX,               "tokenizer.ggml.add_space_prefix"         },
     { LLM_KV_TOKENIZER_REMOVE_EXTRA_WS,          "tokenizer.ggml.remove_extra_whitespaces" },
+    { LLM_KV_TOKENIZER_CLEAN_SPACES,             "tokenizer.ggml.clean_spaces"             },
     { LLM_KV_TOKENIZER_PRECOMPILED_CHARSMAP,     "tokenizer.ggml.precompiled_charsmap"     },
     { LLM_KV_TOKENIZER_HF_JSON,                  "tokenizer.huggingface.json"              },
     { LLM_KV_TOKENIZER_RWKV,                     "tokenizer.rwkv.world"                    },

src/llama-arch.h

@@ -38,8 +38,10 @@ enum llm_arch {
     LLM_ARCH_QWEN2,
     LLM_ARCH_QWEN2MOE,
     LLM_ARCH_QWEN2VL,
+    LLM_ARCH_SENSENOVA_U1,
     LLM_ARCH_QWEN3,
     LLM_ARCH_QWEN3MOE,
+    LLM_ARCH_FLINT,
     LLM_ARCH_QWEN3NEXT,
     LLM_ARCH_QWEN3VL,
     LLM_ARCH_QWEN3VLMOE,
@@ -344,6 +346,7 @@ enum llm_kv {
     LLM_KV_ROPE_DIMENSION_SECTIONS,
     LLM_KV_ROPE_SECTION_ORDER,
     LLM_KV_ROPE_FREQ_BASE,
+    LLM_KV_ROPE_FREQ_BASE_SPATIAL,
     LLM_KV_ROPE_FREQ_BASE_SWA,
     LLM_KV_ROPE_SCALE_LINEAR,
     LLM_KV_ROPE_SCALING_TYPE,
@@ -397,6 +400,7 @@ enum llm_kv {
     LLM_KV_TOKENIZER_DNA_AUTO_TAGS,
     LLM_KV_TOKENIZER_ADD_PREFIX,
     LLM_KV_TOKENIZER_REMOVE_EXTRA_WS,
+    LLM_KV_TOKENIZER_CLEAN_SPACES,
     LLM_KV_TOKENIZER_PRECOMPILED_CHARSMAP,
     LLM_KV_TOKENIZER_HF_JSON,
     LLM_KV_TOKENIZER_RWKV,

src/llama-context.cpp

@@ -165,6 +165,12 @@ static const llm_fused_op_probe llm_fused_op_dsv4_hc_post_probe = {
     /*.n_tokens_per_seq =*/ 1,
 };
 
+static const llm_fused_op_probe llm_fused_op_top_k_mask_probe = {
+    /*.op               =*/ LLM_FUSED_OP_TOP_K_MASK,
+    /*.name             =*/ "fused top-k mask",
+    /*.n_tokens_per_seq =*/ 1,
+};
+
 llama_context::llama_context(
         const llama_model & model,
               llama_context_params params) :
@@ -328,6 +334,10 @@ llama_context::llama_context(
     cparams.fused_dsv4_hc_post = true;
     cparams.auto_fhc           = true;
 
+    // only the indexer models build this mask, skip the probe for the others
+    cparams.fused_top_k_mask = true;
+    cparams.auto_ftkm        = hparams.indexer_top_k > 0;
+
     // with causal attention, the batch size is limited by the context size
     cparams.n_batch = cparams.causal_attn ? std::min(cparams.n_ctx, params.n_batch) : params.n_batch;
 
@@ -671,6 +681,12 @@ void llama_context::resolve_fused_ops(const llama_memory_context_i * mctx, uint3
         resolve(llm_fused_op_dsv4_hc_post_probe, cparams.fused_dsv4_hc_post);
         cparams.auto_fhc = false;
     }
+
+    if (cparams.auto_ftkm) {
+        LLAMA_LOG_INFO("%s: resolving fused top-k mask support:\n", func);
+        resolve(llm_fused_op_top_k_mask_probe, cparams.fused_top_k_mask);
+        cparams.auto_ftkm = false;
+    }
 }
 
 static int llama_graph_n_input_tensors(ggml_cgraph * gf) {

src/llama-cparams.h

@@ -48,6 +48,8 @@ struct llama_cparams {
     bool fused_dsv4_hc_comb;
     bool fused_dsv4_hc_post;
     bool auto_fhc;
+    bool fused_top_k_mask;   // use fused top-k sparse attention mask
+    bool auto_ftkm;
     bool no_perf;
     bool warmup;             // TODO: remove [TAG_LLAMA_GRAPH_NO_WARMUP]
     bool op_offload;

src/llama-graph.cpp

@@ -1630,8 +1630,9 @@ ggml_tensor * llm_graph_context::build_lora_mm(
           ggml_tensor * w,
           ggml_tensor * cur,
           ggml_tensor * w_s) const {
-    cur = build_hadamard_input(w, cur);
-    ggml_tensor * res = ggml_mul_mat(ctx0, w, cur);
+    // a Hadamard-folded weight reads the transformed input; a LoRA delta lives in the primal basis and reads `cur`
+    ggml_tensor * cur_mm = build_hadamard_input(w, cur);
+    ggml_tensor * res = ggml_mul_mat(ctx0, w, cur_mm);
 
     if (prec_policy) {
         prec_policy->apply(res);
@@ -1684,7 +1685,10 @@ llm_graph_amax_src1 llm_graph_context::build_amax_src1_f16(
     constexpr float k_ln_2   = 0.6931471805599453f;
 
     ggml_tensor * e;
-    e = ggml_scale(ctx0, amax, 1.0f/32768.0f);
+    // note: this can technically be:
+    // e = ggml_scale(ctx0, amax, 1.0f/32768.0f);
+    // but we leave x32 headroom for the F16 sums of 32 values in Q8 activation blocks
+    e = ggml_scale(ctx0, amax, 1.0f/1024.0f);
     e = ggml_log  (ctx0, e);
     e = ggml_scale(ctx0, e, k_log2_e);
     e = ggml_ceil (ctx0, e);
@@ -1706,11 +1710,12 @@ ggml_tensor * llm_graph_context::build_lora_mm_id(
           ggml_tensor * ids,
           ggml_tensor * w_s,
           ggml_tensor * slots) const {
-    cur = build_hadamard_input(w, cur);
+    // as in build_lora_mm: the folded weight reads the transformed input, the LoRA delta reads `cur`
+    ggml_tensor * cur_mm = build_hadamard_input(w, cur);
     // the experts in the MoE cache are selected by their slots
     ggml_tensor * res = slots == nullptr ?
-        ggml_mul_mat_id(ctx0, w, cur, ids) :
-        ggml_mul_mat_id(ctx0, moe_cache->get_experts(w), cur, slots);
+        ggml_mul_mat_id(ctx0, w, cur_mm, ids) :
+        ggml_mul_mat_id(ctx0, moe_cache->get_experts(w), cur_mm, slots);
 
     if (prec_policy) {
         prec_policy->apply(res);
@@ -2505,7 +2510,7 @@ ggml_tensor * llm_graph_context::build_moe_ffn(
         //       ref: https://github.com/ggml-org/llama.cpp/pull/30656#issuecomment-6269597922
         llm_graph_amax_src1 amax_src1 = { nullptr, nullptr };
 
-        if (arch == LLM_ARCH_MISTRAL4 || arch == LLM_ARCH_MELLUM) {
+        if (arch == LLM_ARCH_MISTRAL4 || arch == LLM_ARCH_MELLUM || arch == LLM_ARCH_FLINT) {
             amax_src1 = build_amax_src1_f16(cur);
             cur = amax_src1.res;
         }
@@ -3375,6 +3380,27 @@ ggml_tensor * llm_graph_context::build_attn(
     return cur;
 }
 
+ggml_tensor * llm_graph_context::build_fused_top_k_mask(
+        ggml_tensor * kq_mask,
+        ggml_tensor * top_k,
+        ggml_tensor * score,
+        ggml_tensor * blocks,
+        ggml_tensor * extra,
+                int   il) const {
+    const int64_t m1 = kq_mask->ne[1];
+    const int64_t m2 = kq_mask->ne[2];
+    const int64_t m3 = kq_mask->ne[3];
+
+    auto rows = [&](ggml_tensor * t) {
+        return t ? ggml_reshape_4d(ctx0, t, t->ne[0], m1, m2, m3) : nullptr;
+    };
+
+    ggml_tensor * cur = ggml_top_k_mask(ctx0, kq_mask, rows(top_k), rows(score), blocks, rows(extra));
+    res->add_fused_node({LLM_FUSED_OP_TOP_K_MASK, cur, il});
+
+    return cur;
+}
+
 ggml_tensor * llm_graph_context::build_attn(
         llm_graph_input_attn_k_dsa * inp,
         ggml_tensor * wo,
@@ -3407,31 +3433,36 @@ ggml_tensor * llm_graph_context::build_attn(
 
     const auto & kq_mask = inp->get_kq_mask_mla();
 
-    // prepare new kq mask - starts filled with -INFINITY
-    ggml_tensor * kq_mask_all = ggml_fill(ctx0, kq_mask, -INFINITY);
+    ggml_tensor * kq_mask_top_k = nullptr;
+    if (cparams.fused_top_k_mask) {
+        kq_mask_top_k = build_fused_top_k_mask(kq_mask, top_k, nullptr, nullptr, nullptr, il);
+    } else {
+        // prepare new kq mask - starts filled with -INFINITY
+        ggml_tensor * kq_mask_all = ggml_fill(ctx0, kq_mask, -INFINITY);
 
-    // reshape KQ mask into tensor with rows of size 1:
-    // [n_kv, n_batch, 1, n_stream] -> [1, n_kv, n_batch, n_stream]
-    kq_mask_all = ggml_view_4d(ctx0, kq_mask_all, 1, kq_mask_all->ne[0], kq_mask_all->ne[1], kq_mask_all->ne[3], kq_mask_all->nb[0], kq_mask_all->nb[1], kq_mask_all->nb[2], 0);
+        // reshape KQ mask into tensor with rows of size 1:
+        // [n_kv, n_batch, 1, n_stream] -> [1, n_kv, n_batch, n_stream]
+        kq_mask_all = ggml_view_4d(ctx0, kq_mask_all, 1, kq_mask_all->ne[0], kq_mask_all->ne[1], kq_mask_all->ne[3], kq_mask_all->nb[0], kq_mask_all->nb[1], kq_mask_all->nb[2], 0);
 
-    // reshape top_k indices: [n_top_k, n_batch, 1, n_stream] -> [n_top_k, n_batch, n_stream, 1]
-    ggml_tensor * top_k_3d = ggml_view_4d(ctx0, top_k, top_k->ne[0], top_k->ne[1], top_k->ne[3], 1, top_k->nb[1], top_k->nb[2], top_k->ne[3]*top_k->nb[3], 0);
+        // reshape top_k indices: [n_top_k, n_batch, 1, n_stream] -> [n_top_k, n_batch, n_stream, 1]
+        ggml_tensor * top_k_3d = ggml_view_4d(ctx0, top_k, top_k->ne[0], top_k->ne[1], top_k->ne[3], 1, top_k->nb[1], top_k->nb[2], top_k->ne[3]*top_k->nb[3], 0);
 
-    // prepare zero-filled tensor with rows of size 1: [1, n_top_k, n_batch, n_stream]
-    // this will be our source of zero values for unmasking top k mask elements
-    ggml_tensor * zeros = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, 1, top_k_3d->ne[0], top_k_3d->ne[1], top_k_3d->ne[2]);
-    zeros = ggml_fill(ctx0, zeros, 0.0f);
+        // prepare zero-filled tensor with rows of size 1: [1, n_top_k, n_batch, n_stream]
+        // this will be our source of zero values for unmasking top k mask elements
+        ggml_tensor * zeros = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, 1, top_k_3d->ne[0], top_k_3d->ne[1], top_k_3d->ne[2]);
+        zeros = ggml_fill(ctx0, zeros, 0.0f);
 
-    // modify KQ mask by unmasking elements that are in top_k indices
-    // ggml_set_rows([1, n_kv, n_batch, n_stream], [1, n_top_k, n_batch, n_stream], [n_top_k, n_batch, n_stream, 1])
-    ggml_tensor * kq_mask_top_k = ggml_set_rows(ctx0, kq_mask_all, zeros, top_k_3d);
+        // modify KQ mask by unmasking elements that are in top_k indices
+        // ggml_set_rows([1, n_kv, n_batch, n_stream], [1, n_top_k, n_batch, n_stream], [n_top_k, n_batch, n_stream, 1])
+        kq_mask_top_k = ggml_set_rows(ctx0, kq_mask_all, zeros, top_k_3d);
 
-    // reshape to restore the original shape of KQ mask:
-    // [1, n_kv, n_batch, n_stream] -> [n_kv, n_batch, 1, n_stream]
-    kq_mask_top_k = ggml_view_4d(ctx0, kq_mask_top_k, kq_mask_top_k->ne[1], kq_mask_top_k->ne[2], 1, kq_mask_top_k->ne[3], kq_mask_top_k->nb[2], kq_mask_top_k->nb[3], kq_mask_top_k->nb[3], 0);
+        // reshape to restore the original shape of KQ mask:
+        // [1, n_kv, n_batch, n_stream] -> [n_kv, n_batch, 1, n_stream]
+        kq_mask_top_k = ggml_view_4d(ctx0, kq_mask_top_k, kq_mask_top_k->ne[1], kq_mask_top_k->ne[2], 1, kq_mask_top_k->ne[3], kq_mask_top_k->nb[2], kq_mask_top_k->nb[3], kq_mask_top_k->nb[3], 0);
 
-    // combine with the original kq mask
-    kq_mask_top_k = ggml_add(ctx0, kq_mask_top_k, kq_mask);
+        // combine with the original kq mask
+        kq_mask_top_k = ggml_add(ctx0, kq_mask_top_k, kq_mask);
+    }
 
     ggml_tensor * q = q_cur;
     ggml_tensor * k = mctx_cur->get_k(ctx0, il);

src/llama-graph.h

@@ -53,6 +53,7 @@ enum llm_fused_op {
     LLM_FUSED_OP_DSV4_HC_PRE,
     LLM_FUSED_OP_DSV4_HC_COMB,
     LLM_FUSED_OP_DSV4_HC_POST,
+    LLM_FUSED_OP_TOP_K_MASK,
 };
 
 enum llm_ffn_op_type : int {
@@ -1121,7 +1122,7 @@ struct llm_graph_context {
               ggml_tensor * slots = nullptr) const;
 
     // rescale `cur` (the src1 of a matmul) by a power of two, computed from its maximum absolute
-    // when the values already fit in F16, the factors are exactly 1.0f
+    // keep both individual values and Q8 activation block sums within F16 range
     llm_graph_amax_src1 build_amax_src1_f16(
               ggml_tensor * cur) const;
 
@@ -1335,6 +1336,15 @@ struct llm_graph_context {
                   float   kq_scale,
                     int   il) const;
 
+    // ggml_top_k_mask with top_k, score and extra viewed in the row layout of kq_mask, see ggml.h
+    ggml_tensor * build_fused_top_k_mask(
+            ggml_tensor * kq_mask, // [n_kv, n_tokens/n_stream, 1, n_stream]
+            ggml_tensor * top_k,   // [n_top, n_tokens]
+            ggml_tensor * score,   // [n_pick, n_tokens], can be null
+            ggml_tensor * blocks,  // [n_blk, n_pick], can be null
+            ggml_tensor * extra,   // [n_extra, n_tokens], can be null
+                    int   il) const;
+
     llm_graph_input_attn_kv_iswa * build_attn_inp_kv_iswa() const;
 
     llm_graph_input_dsv4 * build_inp_dsv4() const;

src/llama-kv-cache.cpp

@@ -882,7 +882,10 @@ bool llama_kv_cache::update(llama_context * lctx, bool do_shift, const stream_co
         LLAMA_LOG_DEBUG("%s: applying K-shift\n", __func__);
 
         // apply K-shift if needed
-        if (hparams.rope_type != LLAMA_ROPE_TYPE_NONE) {
+        if (hparams.rope_type != LLAMA_ROPE_TYPE_NONE &&
+                std::any_of(layers.begin(), layers.end(), [&](const auto & layer) {
+                    return hparams.has_rope(layer.il);
+                })) {
             ggml_backend_sched_reset(sched);
 
             auto * res = lctx->get_gf_res_reserve();

src/llama-model-saver.cpp

@@ -15,6 +15,7 @@
 
 bool llama_model_saver_supports_arch(llm_arch arch) {
     switch (arch) {
+        case LLM_ARCH_SENSENOVA_U1:
         case LLM_ARCH_GEMMA3N:
         case LLM_ARCH_BITNET:
         case LLM_ARCH_T5:
@@ -423,6 +424,7 @@ void llama_model_saver::add_kv_from_model() {
     add_kv(LLM_KV_TOKENIZER_ADD_SEP,                 vocab.get_add_sep());
     add_kv(LLM_KV_TOKENIZER_ADD_PREFIX,              vocab.get_add_space_prefix());
     add_kv(LLM_KV_TOKENIZER_REMOVE_EXTRA_WS,         vocab.get_remove_extra_whitespaces());
+    add_kv(LLM_KV_TOKENIZER_CLEAN_SPACES,            vocab.get_clean_spaces());
     add_kv(LLM_KV_TOKENIZER_PRECOMPILED_CHARSMAP,    vocab.get_precompiled_charsmap());
     // add_kv(LLM_KV_TOKENIZER_HF_JSON,                 ???);
     // add_kv(LLM_KV_TOKENIZER_RWKV,                    ???);

src/llama-model.cpp

@@ -110,8 +110,12 @@ static llama_model * llama_model_mapping(llm_arch arch, const llama_model_params
             return new llama_model_qwen2vl(params);
         case LLM_ARCH_QWEN2MOE:
             return new llama_model_qwen2moe(params);
+        case LLM_ARCH_SENSENOVA_U1:
+            return new llama_model_sensenova_u1(params);
         case LLM_ARCH_QWEN3:
             return new llama_model_qwen3(params);
+        case LLM_ARCH_FLINT:
+            return new llama_model_flint(params);
         case LLM_ARCH_QWEN3MOE:
             return new llama_model_qwen3moe(params);
         case LLM_ARCH_QWEN3VL:
@@ -3569,6 +3573,7 @@ llama_rope_type llama_model_rope_type(const llama_model * model) {
         case LLM_ARCH_QWEN2MOE:
         case LLM_ARCH_QWEN3:
         case LLM_ARCH_VLLM_DECISION2:
+        case LLM_ARCH_FLINT:
         case LLM_ARCH_QWEN3MOE:
         case LLM_ARCH_LLADA_MOE:
         case LLM_ARCH_RND1:
@@ -3640,6 +3645,7 @@ llama_rope_type llama_model_rope_type(const llama_model * model) {
             // DSV4 DSpark drafters use DeepSeek-V4's normal RoPE; legacy DFlash backbones are NeoX
             return model->hparams.dsv4_hc_mult > 0 ? LLAMA_ROPE_TYPE_NORM : LLAMA_ROPE_TYPE_NEOX;
 
+        case LLM_ARCH_SENSENOVA_U1:
         case LLM_ARCH_QWEN2VL:
         case LLM_ARCH_PADDLEOCR:
             return LLAMA_ROPE_TYPE_MROPE;

src/llama-model.h

@@ -779,8 +779,9 @@ struct llama_model {
     struct ggml_tensor * per_layer_proj_norm  = nullptr;
 
     // eagle3 / dflash feature fusion layer
-    struct ggml_tensor * fc   = nullptr;
-    struct ggml_tensor * fc_s = nullptr;
+    struct ggml_tensor * fc      = nullptr;
+    struct ggml_tensor * fc_s    = nullptr;
+    struct ggml_tensor * fc_in_s = nullptr;
     struct ggml_tensor * d2t = nullptr;  // draft to target vocabulary mapping
 
     // dspark

src/llama-vocab.cpp

@@ -281,6 +281,13 @@ struct llm_tokenizer_bpe : llm_tokenizer {
     llm_tokenizer_bpe(const llama_vocab & vocab) {
         GGML_ASSERT(vocab.get_type() == LLAMA_VOCAB_TYPE_BPE);
         switch (vocab.get_pre_type()) {
+            case LLAMA_VOCAB_PRE_TYPE_FLINT:
+                regex_exprs = {
+                    "\\p{N}{1,3}",
+                    "[一-龥぀-ゟ゠-ヿ]+",
+                    "[^\\r\\n\\p{L}\\p{N}]?[\\p{Lu}\\p{Lt}\\p{Lm}\\p{Lo}\\p{M}]*[\\p{Ll}\\p{Lm}\\p{Lo}\\p{M}]+(?:'[sS]|'[tT]|'[rR][eE]|'[vV][eE]|'[mM]|'[lL][lL]|'[dD])?|[^\\r\\n\\p{L}\\p{N}]?[\\p{Lu}\\p{Lt}\\p{Lm}\\p{Lo}\\p{M}]+[\\p{Ll}\\p{Lm}\\p{Lo}\\p{M}]*(?:'[sS]|'[tT]|'[rR][eE]|'[vV][eE]|'[mM]|'[lL][lL]|'[dD])?| ?[\\p{P}\\p{S}]+[\\r\\n]*|\\s*[\\r\\n]+|\\s+(?!\\S)|\\s+",
+                };
+                break;
             case LLAMA_VOCAB_PRE_TYPE_LLAMA3:
                 regex_exprs = {
                     // original regex from tokenizer.json
@@ -2417,24 +2424,32 @@ void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) {
                 LLAMA_LOG_WARN("%s: ************************************        \n", __func__);
                 LLAMA_LOG_WARN("%s:                                             \n", __func__);
                 pre_type = LLAMA_VOCAB_PRE_TYPE_DEFAULT;
+            } else if (tokenizer_pre == "flint") {
+                pre_type = LLAMA_VOCAB_PRE_TYPE_FLINT;
+                clean_spaces = false;
             } else if (tokenizer_pre == "default") {
                 pre_type = LLAMA_VOCAB_PRE_TYPE_DEFAULT;
             } else if (tokenizer_pre == "minicpm5") {
                 pre_type = LLAMA_VOCAB_PRE_TYPE_MINICPM5;
                 ignore_merges = true;
+                clean_spaces = false;
             } else if (
                     tokenizer_pre == "llama3"   ||
                     tokenizer_pre == "llama-v3" ||
                     tokenizer_pre == "llama-bpe"||
                     tokenizer_pre == "falcon3"  ||
                     tokenizer_pre == "falcon-h1" ||
-                    tokenizer_pre == "pixtral"  ||
                     tokenizer_pre == "midm-2.0" ||
-                    tokenizer_pre == "lfm2"     ||
                     tokenizer_pre == "jina-v5-nano") {
                 pre_type = LLAMA_VOCAB_PRE_TYPE_LLAMA3;
                 ignore_merges = true;
                 add_bos = true;
+            } else if (
+                    tokenizer_pre == "lfm2") {
+                pre_type = LLAMA_VOCAB_PRE_TYPE_LLAMA3;
+                ignore_merges = true;
+                add_bos = true;
+                clean_spaces = false;
             } else if (
                     tokenizer_pre == "deepseek-llm") {
                 pre_type = LLAMA_VOCAB_PRE_TYPE_DEEPSEEK_LLM;
@@ -2485,6 +2500,7 @@ void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) {
                     tokenizer_pre == "granite-embed-multi-311m") {
                 pre_type = LLAMA_VOCAB_PRE_TYPE_GEMMA4;
                 escape_whitespaces = true;
+                clean_spaces = false;
             } else if (
                     tokenizer_pre == "mmbert") {
                 pre_type = LLAMA_VOCAB_PRE_TYPE_MMBERT;
@@ -2531,6 +2547,7 @@ void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) {
             } else if (
                 tokenizer_pre == "dbrx") {
                 pre_type = LLAMA_VOCAB_PRE_TYPE_DBRX;
+                clean_spaces = false;
             } else if (
                 tokenizer_pre == "smaug-bpe") {
                 pre_type = LLAMA_VOCAB_PRE_TYPE_SMAUG;
@@ -2542,12 +2559,14 @@ void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) {
                 tokenizer_pre == "chatglm-bpe") {
                 pre_type = LLAMA_VOCAB_PRE_TYPE_CHATGLM4;
                 special_bos_id = LLAMA_TOKEN_NULL;
+                clean_spaces = false;
             } else if (
                 tokenizer_pre == "glm4" ||
                 tokenizer_pre == "glm5") {
                 pre_type = LLAMA_VOCAB_PRE_TYPE_CHATGLM4;
                 special_bos_id = LLAMA_TOKEN_NULL;
                 ignore_merges = true;
+                clean_spaces = false;
             } else if (
                 tokenizer_pre == "viking") {
                 pre_type = LLAMA_VOCAB_PRE_TYPE_VIKING;
@@ -2556,7 +2575,8 @@ void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) {
                 tokenizer_pre == "jais") {
                 pre_type = LLAMA_VOCAB_PRE_TYPE_JAIS;
             } else if (
-                tokenizer_pre == "tekken") {
+                tokenizer_pre == "tekken" ||
+                tokenizer_pre == "pixtral") {
                 pre_type = LLAMA_VOCAB_PRE_TYPE_TEKKEN;
                 clean_spaces = false;
                 ignore_merges = true;
@@ -2580,6 +2600,7 @@ void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) {
             } else if (
                 tokenizer_pre == "exaone4") {
                 pre_type = LLAMA_VOCAB_PRE_TYPE_GPT2;
+                clean_spaces = false;
             } else if (
                 tokenizer_pre == "exaone-moe") {
                 pre_type = LLAMA_VOCAB_PRE_TYPE_EXAONE_MOE;
@@ -2680,6 +2701,7 @@ void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) {
             } else if (
                 tokenizer_pre == "mellum2") {
                 pre_type = LLAMA_VOCAB_PRE_TYPE_MELLUM2;
+                clean_spaces = false;
             } else if (
                 tokenizer_pre == "k2-horizon") {
                 pre_type = LLAMA_VOCAB_PRE_TYPE_K2_HORIZON;
@@ -2716,6 +2738,7 @@ void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) {
 
         ml.get_key(LLM_KV_TOKENIZER_ADD_PREFIX,      add_space_prefix,         false);
         ml.get_key(LLM_KV_TOKENIZER_REMOVE_EXTRA_WS, remove_extra_whitespaces, false);
+        ml.get_key(LLM_KV_TOKENIZER_CLEAN_SPACES,    clean_spaces,             false);
     }
 
     const int token_idx = gguf_find_key(ctx, kv(LLM_KV_TOKENIZER_LIST).c_str());
@@ -3147,7 +3170,6 @@ void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) {
         // this is currently determined based on the token text, which is obviously not ideal
         // ref: https://github.com/ggml-org/llama.cpp/issues/9606
         special_eog_ids.clear();
-
         if (special_fim_pad_id != LLAMA_TOKEN_NULL && special_eog_ids.count(special_fim_pad_id) == 0) {
             special_eog_ids.insert(special_fim_pad_id);
         }
@@ -3181,6 +3203,7 @@ void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) {
                     || t.first == "[EOT]" // Kimi-K2
                     || t.first == "[EOS]" // Kimi-K2
                     || t.first == "<|end_of_text|>"
+                    || (pre_type == LLAMA_VOCAB_PRE_TYPE_FLINT && (t.first == "<|user|>" || t.first == "<|observation|>"))
                     || t.first == "<end_of_utterance>" // smoldocling
                     || t.first == "<eos>"            // gemma4
                     || t.first == "<turn|>"          // gemma4

src/llama-vocab.h

@@ -70,6 +70,7 @@ enum llama_vocab_pre_type {
     LLAMA_VOCAB_PRE_TYPE_UFAKZEKA          = 59,
     LLAMA_VOCAB_PRE_TYPE_MMBERT            = 60,
     LLAMA_VOCAB_PRE_TYPE_K2_HORIZON        = 61,
+    LLAMA_VOCAB_PRE_TYPE_FLINT             = 62,
 };
 
 struct LLM_KV;

src/models/deepseek4.cpp

@@ -683,6 +683,13 @@ ggml_tensor * llama_model_deepseek4::graph::build_top_k_mask(
     GGML_ASSERT(kq_mask);
     GGML_ASSERT(top_k);
 
+    if (cparams.fused_top_k_mask) {
+        ggml_tensor * kq_mask_top_k = build_fused_top_k_mask(kq_mask, top_k, nullptr, nullptr, nullptr, il);
+        cb(kq_mask_top_k, name, il);
+
+        return kq_mask_top_k;
+    }
+
     ggml_tensor * kq_mask_all = ggml_fill(ctx0, kq_mask, -INFINITY);
     kq_mask_all = ggml_view_4d(ctx0, kq_mask_all, 1, kq_mask_all->ne[0], kq_mask_all->ne[1], kq_mask_all->ne[3],
             kq_mask_all->nb[0], kq_mask_all->nb[1], kq_mask_all->nb[2], 0);

src/models/dflash.cpp

@@ -238,6 +238,7 @@ void llama_model_dflash::load_arch_tensors(llama_model_loader &) {
 
     fc              = create_tensor(tn(LLM_TENSOR_FC,              "weight"), { n_embd_inp, n_embd }, 0);
     fc_s            = create_tensor(tn(LLM_TENSOR_FC,              "scale"),  { 1 }, TENSOR_NOT_REQUIRED);
+    fc_in_s         = create_tensor(tn(LLM_TENSOR_FC,              "input_scale"), { 1 }, TENSOR_NOT_REQUIRED);
     output_norm_enc = create_tensor(tn(LLM_TENSOR_ENC_OUTPUT_NORM, "weight"), { n_embd }, 0); // encoder hidden_norm (after fc)
     output_norm     = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM,    "weight"), { n_embd }, 0); // decoder final norm
 

src/models/flint.cpp

@@ -0,0 +1,194 @@
+#include "models.h"
+
+void llama_model_flint::load_arch_hparams(llama_model_loader & ml) {
+    ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps);
+    ml.get_key(LLM_KV_LEADING_DENSE_BLOCK_COUNT,   hparams.n_layer_dense_lead);
+    ml.get_key_or_arr(LLM_KV_EXPERT_FEED_FORWARD_LENGTH, hparams.n_ff_exp_arr, hparams.n_layer_all);
+    ml.get_key(LLM_KV_EXPERT_SHARED_COUNT,         hparams.n_expert_shared);
+    ml.get_key(LLM_KV_EXPERT_WEIGHTS_SCALE,        hparams.expert_weights_scale, false);
+    ml.get_key(LLM_KV_EXPERT_WEIGHTS_NORM,         hparams.expert_weights_norm);
+    ml.get_key(LLM_KV_EXPERT_GATING_FUNC,          hparams.expert_gating_func);
+
+    ml.get_key(LLM_KV_EXPERT_SHARED_FEED_FORWARD_LENGTH, hparams.n_ff_shexp);
+    ml.get_arr(LLM_KV_ATTENTION_SLIDING_WINDOW_PATTERN, hparams.is_swa_impl);
+    ml.get_arr(LLM_KV_ATTENTION_ROPE_PATTERN, hparams.rope_pattern);
+    ml.get_key(LLM_KV_ATTENTION_SLIDING_WINDOW, hparams.n_swa);
+    hparams.swa_type = LLAMA_SWA_TYPE_STANDARD;
+    type = LLM_TYPE_UNKNOWN;
+}
+
+void llama_model_flint::load_arch_tensors(llama_model_loader &) {
+    LLAMA_LOAD_LOCALS;
+
+    const int64_t n_ff_exp        = hparams.n_ff_exp();
+
+    tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0);
+
+    output_norm = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM, "weight"), {n_embd}, 0);
+    output      = create_tensor(tn(LLM_TENSOR_OUTPUT,      "weight"), {n_embd, n_vocab}, 0);
+
+    for (int i = 0; i < n_layer; ++i) {
+        auto & layer = layers[i];
+
+        layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, 0);
+
+        create_tensor_qkv(layer, i, n_embd, n_embd_head_k * n_head, n_embd_k_gqa, n_embd_v_gqa, 0);
+        layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), {n_embd_head_k * n_head, n_embd}, 0);
+
+        layer.attn_k_norm = create_tensor(tn(LLM_TENSOR_ATTN_K_NORM, "weight", i), {n_embd_head_k}, 0);
+        layer.attn_q_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_NORM, "weight", i), {n_embd_head_k}, 0);
+
+        layer.attn_sinks = create_tensor(tn(LLM_TENSOR_ATTN_SINKS, "weight", i), {n_head}, 0);
+
+        layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, 0);
+
+        if (i < (int) hparams.n_layer_dense_lead) {
+            layer.ffn_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "weight", i), {n_embd,   n_ff}, 0);
+            layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", i), {  n_ff, n_embd}, 0);
+            layer.ffn_up   = create_tensor(tn(LLM_TENSOR_FFN_UP,   "weight", i), {n_embd,   n_ff}, 0);
+        } else {
+            layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", i), {n_embd, n_expert}, 0);
+            layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, "bias", i), {n_expert}, TENSOR_NOT_REQUIRED);
+
+            if (n_expert == 0) {
+                throw std::runtime_error("n_expert must be > 0");
+            }
+            if (n_expert_used == 0) {
+                throw std::runtime_error("n_expert_used must be > 0");
+            }
+
+            // MoE branch
+            layer.ffn_gate_exps = create_tensor(tn(LLM_TENSOR_FFN_GATE_EXPS, "weight", i), {  n_embd, n_ff_exp, n_expert}, 0);
+            layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", i), {n_ff_exp,   n_embd, n_expert}, 0);
+            layer.ffn_up_exps   = create_tensor(tn(LLM_TENSOR_FFN_UP_EXPS,   "weight", i), {  n_embd, n_ff_exp, n_expert}, 0);
+
+            // Shared expert branch
+            layer.ffn_gate_shexp = create_tensor(tn(LLM_TENSOR_FFN_GATE_SHEXP, "weight", i), {n_embd, hparams.n_ff_shexp}, 0);
+            layer.ffn_down_shexp = create_tensor(tn(LLM_TENSOR_FFN_DOWN_SHEXP, "weight", i), {        hparams.n_ff_shexp, n_embd}, 0);
+            layer.ffn_up_shexp   = create_tensor(tn(LLM_TENSOR_FFN_UP_SHEXP,   "weight", i), {n_embd, hparams.n_ff_shexp}, 0);
+        }
+    }
+}
+
+std::unique_ptr<llm_graph_context> llama_model_flint::build_arch_graph(const llm_graph_params & params) const {
+    return std::make_unique<graph>(*this, params);
+}
+
+llama_model_flint::graph::graph(const llama_model & model, const llm_graph_params & params) :
+    llm_graph_context(params) {
+    const int64_t n_embd_head = hparams.n_embd_head_v();
+
+    GGML_ASSERT(n_embd_head == hparams.n_embd_head_k());
+    GGML_ASSERT(n_embd_head == n_rot);
+
+    ggml_tensor * cur;
+    ggml_tensor * inpL;
+
+    inpL = build_inp_embd(model.tok_embd);
+
+    // inp_pos - contains the positions
+    ggml_tensor * inp_pos = build_inp_pos();
+
+    auto * inp_attn = build_attn_inp_kv_iswa();
+
+    ggml_tensor * inp_out_ids = build_inp_out_ids();
+
+    for (int il = 0; il < n_layer; ++il) {
+        res->t_layer_inp[il] = inpL;
+        ggml_tensor * inpSA = inpL;
+
+        // norm
+        cur = build_norm(inpL, model.layers[il].attn_norm, NULL, LLM_NORM_RMS, il);
+        cb(cur, "attn_norm", il);
+
+        // self_attention
+        {
+            // compute Q and K and RoPE them
+            auto [Qcur, Kcur, Vcur] = build_qkv(model.layers[il], cur,
+                    n_embd_head, n_head, n_head_kv, il);
+
+            Qcur = build_norm(Qcur, model.layers[il].attn_q_norm, nullptr, LLM_NORM_RMS, il);
+            Kcur = build_norm(Kcur, model.layers[il].attn_k_norm, nullptr, LLM_NORM_RMS, il);
+            if (hparams.has_rope(il)) {
+                Qcur = ggml_rope_ext(ctx0, Qcur, inp_pos, nullptr, n_rot, rope_type, n_ctx_orig,
+                        freq_base, freq_scale, ext_factor, attn_factor, beta_fast, beta_slow);
+                Kcur = ggml_rope_ext(ctx0, Kcur, inp_pos, nullptr, n_rot, rope_type, n_ctx_orig,
+                        freq_base, freq_scale, ext_factor, attn_factor, beta_fast, beta_slow);
+            }
+            cb(Qcur, "Qcur", il);
+            cb(Kcur, "Kcur", il);
+            cb(Vcur, "Vcur", il);
+
+            cur = build_attn(inp_attn,
+                    model.layers[il].wo, model.layers[il].wo_b, model.layers[il].wo_s,
+                    Qcur, Kcur, Vcur, nullptr, model.layers[il].attn_sinks, nullptr, 1.0f / sqrtf(float(n_embd_head)), il);
+        }
+        if (il == n_layer - 1 && inp_out_ids) {
+            cur   = ggml_get_rows(ctx0, cur, inp_out_ids);
+            inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids);
+        }
+        ggml_tensor * ffn_inp = ggml_add(ctx0, cur, inpSA);
+        cb(ffn_inp, "ffn_inp", il);
+
+        // MoE branch
+        cur = build_norm(ffn_inp, model.layers[il].ffn_norm, NULL, LLM_NORM_RMS, il);
+        cb(cur, "ffn_norm", il);
+
+        if ((uint32_t) il < hparams.n_layer_dense_lead) {
+            cur = build_ffn(cur,
+                    model.layers[il].ffn_up, NULL, model.layers[il].ffn_up_s,
+                    model.layers[il].ffn_gate, NULL, model.layers[il].ffn_gate_s,
+                    model.layers[il].ffn_down, NULL, model.layers[il].ffn_down_s,
+                    NULL, LLM_FFN_SILU, LLM_FFN_PAR, il);
+            cb(cur, "ffn_out", il);
+        } else {
+            ggml_tensor * moe_out = build_moe_ffn(cur,
+                model.layers[il].ffn_gate_inp,
+                model.layers[il].ffn_up_exps,
+                model.layers[il].ffn_gate_exps,
+                model.layers[il].ffn_down_exps,
+                model.layers[il].ffn_exp_probs_b,
+                n_expert, n_expert_used,
+                LLM_FFN_SILU, hparams.expert_weights_norm,
+                hparams.expert_weights_scale,
+                (llama_expert_gating_func_type) hparams.expert_gating_func,
+                il, nullptr, nullptr,
+                model.layers[il].ffn_up_exps_s,
+                model.layers[il].ffn_gate_exps_s,
+                model.layers[il].ffn_down_exps_s);
+            cb(moe_out, "ffn_moe_out", il);
+
+            {
+                ggml_tensor * ffn_shexp =
+                    build_ffn(cur,
+                        model.layers[il].ffn_up_shexp, NULL, model.layers[il].ffn_up_shexp_s,
+                        model.layers[il].ffn_gate_shexp, NULL, model.layers[il].ffn_gate_shexp_s,
+                        model.layers[il].ffn_down_shexp, NULL, model.layers[il].ffn_down_shexp_s,
+                        NULL, LLM_FFN_SILU, LLM_FFN_PAR, il);
+                cb(ffn_shexp, "ffn_shexp", il);
+
+                cur = ggml_add(ctx0, moe_out, ffn_shexp);
+                cb(cur, "ffn_out", il);
+            }
+        }
+        cur = ggml_add(ctx0, cur, ffn_inp);
+
+        cur = build_cvec(cur, il);
+        cb(cur, "l_out", il);
+
+        // input for next layer
+        inpL = cur;
+    }
+    cur = inpL;
+
+    cur = build_norm(cur, model.output_norm, NULL, LLM_NORM_RMS, -1);
+
+    cb(cur, "result_norm", -1);
+    res->t_embd = cur;
+
+    // lm_head
+    cur = build_lora_mm(model.output, cur, model.output_s);
+
+    cb(cur, "result_output", -1);
+    res->t_logits = cur;
+}

src/models/glm5-next.cpp

@@ -323,13 +323,15 @@ llama_model_glm5_next::llm_graph_input_kpool * llama_model_glm5_next::graph::bui
         const int64_t n_sel      = kpool*n_top_pool + (hparams.indexer_kpool_select_tail ? kpool - 1 : 0);
         inp->n_sel = (uint32_t) n_sel;
 
-        // the scatter maps the dead slots of this mask to dump rows
-        inp->sel_mask = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, n_sel, 1, 1, n_tokens);
-        ggml_set_input(inp->sel_mask);
-        // set_input_kpool always fills the mask, so it stays allocated in every graph
-        ggml_build_forward_expand(gf, inp->sel_mask);
-        // one view for all layers: each view of a host input is a separate upload
-        inp->sel_mask_2d = ggml_reshape_2d(ctx0, inp->sel_mask, n_sel, n_tokens);
+        // the scatter maps the dead slots of this mask to dump rows, the fused mask op reads liveness from the scores instead
+        if (!cparams.fused_top_k_mask) {
+            inp->sel_mask = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, n_sel, 1, 1, n_tokens);
+            ggml_set_input(inp->sel_mask);
+            // set_input_kpool always fills the mask, so it stays allocated in every graph
+            ggml_build_forward_expand(gf, inp->sel_mask);
+            // one view for all layers: each view of a host input is a separate upload
+            inp->sel_mask_2d = ggml_reshape_2d(ctx0, inp->sel_mask, n_sel, n_tokens);
+        }
     }
 
     inp->n_new = n_new;
@@ -973,6 +975,17 @@ ggml_tensor * llama_model_glm5_next::graph::build_kpool_select(
         const int64_t n_top_pool = std::min<int64_t>(n_pool, hparams.indexer_top_k / kpool);
         ggml_tensor * top_k = ggml_top_k(ctx0, score, n_top_pool); // [n_top_pool, n_tokens], UNORDERED
 
+        // keep the kq_mask cells of the live picked pools and the tail, -inf elsewhere
+        if (cparams.fused_top_k_mask) {
+            ggml_tensor * tail = hparams.indexer_kpool_select_tail ? inp_kpool->tail_idxs : nullptr;
+            GGML_ASSERT(kpool*n_top_pool + (tail ? tail->ne[0] : 0) == (int64_t) inp_kpool->n_sel);
+
+            ggml_tensor * sel = build_fused_top_k_mask(kq_mask, top_k, score, inp_kpool->pool_idxs, tail, il);
+            cb(sel, "indexer_sel", il);
+
+            return sel;
+        }
+
         // The selection mask marks the first min(nv, n_top_pool) slots as the visible pools, so order the set by descending score.
         ggml_tensor * sel_score = ggml_get_rows(ctx0,
                 ggml_reshape_3d(ctx0, score, 1, n_pool, n_tokens), top_k); // [1, n_top_pool, n_tokens]

src/models/granite-hybrid.cpp

@@ -153,6 +153,8 @@ llama_model_granite_hybrid::graph::graph(const llama_model & model, const llm_gr
     }
 
     for (int il = 0; il < n_layer; ++il) {
+        res->t_layer_inp[il] = inpL;
+
         struct ggml_tensor * inpSA = inpL;
 
         // norm

src/models/granite-swa.cpp

@@ -154,6 +154,8 @@ llama_model_granite_swa::graph::graph(
             cb(inpL, "deepstack_in", il);
         }
 
+        res->t_layer_inp[il] = inpL;
+
         ggml_tensor * inpSA = inpL;
 
         // norm

src/models/granite-switch.cpp

@@ -293,6 +293,8 @@ llama_model_granite_switch::graph::graph(
     ggml_tensor * cur;
 
     for (int il = 0; il < n_layer; ++il) {
+        res->t_layer_inp[il] = inpL;
+
         ggml_tensor * inpSA = inpL;
 
         cur = build_norm(inpL, model.layers[il].attn_norm, NULL, LLM_NORM_RMS, il);

src/models/granite.cpp

@@ -159,6 +159,8 @@ llama_model_granite::graph::graph(
             cb(inpL, "deepstack_in", il);
         }
 
+        res->t_layer_inp[il] = inpL;
+
         ggml_tensor * inpSA = inpL;
 
         // norm

src/models/models.h

@@ -614,6 +614,21 @@ struct llama_model_qwen2moe : public llama_model_base {
 };
 
 
+struct llama_model_sensenova_u1 : public llama_model_base {
+    llama_model_sensenova_u1(const struct llama_model_params & params) : llama_model_base(params) {}
+    void load_arch_hparams(llama_model_loader & ml) override;
+    void load_arch_tensors(llama_model_loader & ml) override;
+
+    float rope_freq_base_spatial = 10000.0f;
+
+    struct graph : public llm_graph_context {
+        graph(const llama_model_sensenova_u1 & model, const llm_graph_params & params);
+    };
+
+    std::unique_ptr<llm_graph_context> build_arch_graph(const llm_graph_params & params) const override;
+};
+
+
 struct llama_model_qwen3 : public llama_model_base {
     llama_model_qwen3(const struct llama_model_params & params) : llama_model_base(params) {}
     void load_arch_hparams(llama_model_loader & ml) override;
@@ -678,6 +693,19 @@ struct llama_model_qwen3moe : public llama_model_base {
 };
 
 
+struct llama_model_flint : public llama_model_base {
+    llama_model_flint(const struct llama_model_params & params) : llama_model_base(params) {}
+    void load_arch_hparams(llama_model_loader & ml) override;
+    void load_arch_tensors(llama_model_loader & ml) override;
+
+    struct graph : public llm_graph_context {
+        graph(const llama_model & model, const llm_graph_params & params);
+    };
+
+    std::unique_ptr<llm_graph_context> build_arch_graph(const llm_graph_params & params) const override;
+};
+
+
 struct llama_model_qwen3vl : public llama_model_base {
     llama_model_qwen3vl(const struct llama_model_params & params) : llama_model_base(params) {}
     void load_arch_hparams(llama_model_loader & ml) override;

src/models/qwen4exp.cpp

@@ -813,6 +813,16 @@ ggml_tensor * llama_model_qwen4exp::graph::build_qsa_sel(
     ggml_tensor * top_k = ggml_top_k(ctx0, score, n_top_pool); // [n_top_pool, n_tokens], unordered
     cb(top_k, "indexer_top_k", il);
 
+    // keep the kq_mask cells of the live picked pools and the tail, -inf elsewhere
+    if (cparams.fused_top_k_mask) {
+        GGML_ASSERT(kpool*n_top_pool + inp_kpool->tail_idxs->ne[0] == inp_kpool->n_sel);
+
+        ggml_tensor * sel = build_fused_top_k_mask(kq_mask, top_k, score, inp_kpool->pool_idxs, inp_kpool->tail_idxs, il);
+        cb(sel, "indexer_sel", il);
+
+        return sel;
+    }
+
     // the top blocks, then the incomplete tail with n_kv for missing cells
     ggml_tensor * sel_idx = ggml_get_rows(ctx0, inp_kpool->pool_idxs,
             ggml_reshape_1d(ctx0, top_k, n_top_pool*n_tokens)); // [kpool, n_top_pool*n_tokens]

src/models/sensenova-u1.cpp

@@ -0,0 +1,190 @@
+#include "models.h"
+
+namespace {
+class llm_graph_input_pos_sensenova_u1 : public llm_graph_input_pos {
+public:
+    llm_graph_input_pos_sensenova_u1() : llm_graph_input_pos(4) {}
+
+    void set_input(const llama_ubatch * ubatch) override {
+        if (ubatch->pos && pos) {
+            const int64_t n_tokens = ubatch->n_tokens;
+
+            // U1 expects text-token positions in the temporal plane only: [t, 0, 0, 0].
+            // The generic M-RoPE input uses [t, t, t, 0], so U1 needs a custom input class.
+            if (ubatch->token) {
+                std::vector<llama_pos> pos_data(n_tokens*n_pos_per_embd, 0);
+                for (int i = 0; i < n_tokens; ++i) {
+                    pos_data[i] = ubatch->pos[i];
+                }
+                ggml_backend_tensor_set(pos, pos_data.data(), 0, pos_data.size()*ggml_element_size(pos));
+            } else {
+                ggml_backend_tensor_set(pos, ubatch->pos, 0, n_tokens*n_pos_per_embd*ggml_element_size(pos));
+            }
+        }
+    }
+};
+}
+
+void llama_model_sensenova_u1::load_arch_hparams(llama_model_loader & ml) {
+    ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps);
+
+    ml.get_key(LLM_KV_ROPE_FREQ_BASE_SPATIAL, rope_freq_base_spatial);
+    if (hparams.n_embd_head_k() % 4 != 0 || hparams.n_embd_head_v() != hparams.n_embd_head_k()) {
+        throw std::runtime_error("SenseNova U1 requires equal Q/K/V head sizes divisible by four");
+    }
+    switch (hparams.n_layer()) {
+        case 42: type = LLM_TYPE_8B; break;
+        default: type = LLM_TYPE_UNKNOWN;
+    }
+}
+
+void llama_model_sensenova_u1::load_arch_tensors(llama_model_loader &) {
+    LLAMA_LOAD_LOCALS;
+
+    tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0);
+
+    // output
+    output_norm = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM, "weight"), {n_embd}, 0);
+    output      = create_tensor(tn(LLM_TENSOR_OUTPUT,      "weight"), {n_embd, n_vocab}, TENSOR_NOT_REQUIRED);
+    // if output is NULL, init from the input tok embed
+    if (output == NULL) {
+        output = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, TENSOR_DUPLICATED);
+    }
+
+    for (int i = 0; i < n_layer; ++i) {
+        auto & layer = layers[i];
+
+        layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, 0);
+
+        create_tensor_qkv(layer, i, n_embd, n_embd_head_k * n_head, n_embd_gqa, n_embd_gqa, 0);
+        layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), {n_embd_head_k * n_head, n_embd}, 0);
+
+        layer.attn_k_norm = create_tensor(tn(LLM_TENSOR_ATTN_K_NORM, "weight", i), {n_embd_head_k}, 0);
+        layer.attn_q_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_NORM, "weight", i), {n_embd_head_k}, 0);
+
+        layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, 0);
+        layer.ffn_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "weight", i), {n_embd,   n_ff}, 0);
+        layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", i), {  n_ff, n_embd}, 0);
+        layer.ffn_up   = create_tensor(tn(LLM_TENSOR_FFN_UP,   "weight", i), {n_embd,   n_ff}, 0);
+    }
+}
+
+std::unique_ptr<llm_graph_context> llama_model_sensenova_u1::build_arch_graph(const llm_graph_params & params) const {
+    return std::make_unique<graph>(*this, params);
+}
+
+llama_model_sensenova_u1::graph::graph(const llama_model_sensenova_u1 & model, const llm_graph_params & params) : llm_graph_context(params) {
+    const int64_t n_embd_head = hparams.n_embd_head_v();
+
+    GGML_ASSERT(n_embd_head == hparams.n_embd_head_k());
+
+    ggml_tensor * cur;
+    ggml_tensor * inpL;
+
+    auto normalize = [&](ggml_tensor * x, ggml_tensor * weight) {
+        x = ggml_rms_norm(ctx0, x, hparams.f_norm_rms_eps);
+
+        return ggml_mul(ctx0, x, weight);
+    };
+
+    inpL = build_inp_embd(model.tok_embd);
+
+    ggml_tensor * pos[3]{};
+    auto inp = std::make_unique<llm_graph_input_pos_sensenova_u1>();
+    auto & inp_pos = inp->pos;
+    inp_pos = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, (int64_t)n_tokens*hparams.n_pos_per_embd());
+    ggml_set_input(inp_pos);
+    res->add_input(std::move(inp));
+
+    for (int axis = 0; axis < 3; ++axis) {
+        pos[axis] = ggml_view_1d(ctx0, inp_pos, n_tokens, axis*n_tokens*sizeof(int32_t));
+    }
+
+    auto normalize_and_rotate = [&](ggml_tensor * x, ggml_tensor * weight) {
+        const int64_t heads = x->ne[1];
+        x = ggml_reshape_4d(ctx0, x, n_embd_head/2, 2, heads, n_tokens);
+        x = normalize(x, ggml_reshape_2d(ctx0, weight, n_embd_head/2, 2));
+        x = ggml_reshape_3d(ctx0, x, n_embd_head, heads, n_tokens);
+
+        // Normalize temporal and spatial halves separately; rotate time, height, then width.
+        for (int axis = 0; axis < 3; ++axis) {
+            const int dims = axis == 0 ? n_embd_head/2 : n_embd_head/4;
+            const int offset = axis == 0 ? 0 : n_embd_head/2 + (axis - 1)*n_embd_head/4;
+            const float base = axis == 0 ? freq_base : model.rope_freq_base_spatial;
+            x = ggml_rope_ext(ctx0, x, pos[axis], nullptr, dims, GGML_ROPE_TYPE_NEOX,
+                    n_ctx_orig, base, 1.0f, 0.0f, 1.0f, 32.0f, 1.0f);
+            ggml_rope_set_offset(x, offset);
+        }
+        return x;
+    };
+
+    auto * inp_attn = build_attn_inp_kv();
+
+    ggml_tensor * inp_out_ids = build_inp_out_ids();
+
+    for (int il = 0; il < n_layer; ++il) {
+        res->t_layer_inp[il] = inpL;
+
+        ggml_tensor * inpSA = inpL;
+
+        // norm
+        cur = normalize(inpL, model.layers[il].attn_norm);
+        cb(cur, "attn_norm", il);
+
+        // self-attention
+        {
+            // compute Q and K and RoPE them
+            auto [Qcur, Kcur, Vcur] = build_qkv(model.layers[il], cur,
+                    n_embd_head, n_head, n_head_kv, il);
+
+            Qcur = normalize_and_rotate(Qcur, model.layers[il].attn_q_norm);
+            Kcur = normalize_and_rotate(Kcur, model.layers[il].attn_k_norm);
+
+            cb(Qcur, "Qcur", il);
+            cb(Kcur, "Kcur", il);
+            cb(Vcur, "Vcur", il);
+
+            cur = build_attn(inp_attn,
+                    model.layers[il].wo, model.layers[il].wo_b, model.layers[il].wo_s,
+                    Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, 1.0f/sqrtf(float(n_embd_head)), il);
+        }
+        if (il == n_layer - 1 && inp_out_ids) {
+            cur   = ggml_get_rows(ctx0,   cur, inp_out_ids);
+            inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids);
+        }
+        ggml_tensor * ffn_inp = ggml_add(ctx0, cur, inpSA);
+        cb(ffn_inp, "ffn_inp", il);
+
+        // feed-forward network
+        cur = normalize(ffn_inp, model.layers[il].ffn_norm);
+        cb(cur, "ffn_norm", il);
+
+        cur = build_ffn(cur,
+            model.layers[il].ffn_up,   NULL, model.layers[il].ffn_up_s,
+            model.layers[il].ffn_gate, NULL, model.layers[il].ffn_gate_s,
+            model.layers[il].ffn_down, NULL, model.layers[il].ffn_down_s,
+            NULL,
+            LLM_FFN_SILU, LLM_FFN_PAR, il);
+        cb(cur, "ffn_out", il);
+
+        cur = ggml_add(ctx0, cur, ffn_inp);
+
+        cur = build_cvec(cur, il);
+        cb(cur, "l_out", il);
+
+        // input for next layer
+        inpL = cur;
+    }
+    cur = inpL;
+
+    cur = normalize(cur, model.output_norm);
+
+    cb(cur, "result_norm", -1);
+    res->t_embd = cur;
+
+    // lm_head
+    cur = build_lora_mm(model.output, cur, model.output_s);
+
+    cb(cur, "result_output", -1);
+    res->t_logits = cur;
+}

tests/CMakeLists.txt

@@ -142,7 +142,7 @@ llama_test(test-tokenizer-0 NAME test-tokenizer-0-command-r         ARGS ${PROJE
 llama_test(test-tokenizer-0 NAME test-tokenizer-0-deepseek-coder    ARGS ${PROJECT_SOURCE_DIR}/models/ggml-vocab-deepseek-coder.gguf)
 llama_test(test-tokenizer-0 NAME test-tokenizer-0-deepseek-llm      ARGS ${PROJECT_SOURCE_DIR}/models/ggml-vocab-deepseek-llm.gguf)
 llama_test(test-tokenizer-0 NAME test-tokenizer-0-falcon            ARGS ${PROJECT_SOURCE_DIR}/models/ggml-vocab-falcon.gguf)
-llama_test(test-tokenizer-0 NAME test-tokenizer-0-gemma-4           ARGS ${PROJECT_SOURCE_DIR}/models/ggml-vocab-gemma-4.gguf)
+llama_test(test-tokenizer-0 NAME test-tokenizer-0-gemma-4           ARGS ${PROJECT_SOURCE_DIR}/models/ggml-vocab-gemma-4.gguf --check-decode)
 llama_test(test-tokenizer-0 NAME test-tokenizer-0-gpt-2             ARGS ${PROJECT_SOURCE_DIR}/models/ggml-vocab-gpt-2.gguf)
 llama_test(test-tokenizer-0 NAME test-tokenizer-0-llama-bpe         ARGS ${PROJECT_SOURCE_DIR}/models/ggml-vocab-llama-bpe.gguf)
 llama_test(test-tokenizer-0 NAME test-tokenizer-0-llama-spm         ARGS ${PROJECT_SOURCE_DIR}/models/ggml-vocab-llama-spm.gguf)

tests/fusion/MTL.csv

@@ -71,6 +71,9 @@ falcon              ,0   ,any     ,ADD+ADD                     ,      2
 falcon              ,0   ,any     ,NORM+MUL+ADD                ,      5
 falcon-h1           ,0   ,any     ,ADD+ADD                     ,      2
 falcon-h1           ,0   ,any     ,RMS_NORM+MUL                ,      9
+flint               ,0   ,any     ,ADD+ADD                     ,      1
+flint               ,0   ,any     ,MUL+ADD                     ,      2
+flint               ,0   ,any     ,RMS_NORM+MUL                ,      9
 gemma               ,0   ,any     ,RMS_NORM+MUL                ,      5
 gemma2              ,0   ,any     ,RMS_NORM+MUL                ,      5
 gemma2              ,0   ,any     ,RMS_NORM+MUL+ADD            ,      4

tests/test-backend-ops.cpp

@@ -9128,6 +9128,148 @@ struct test_lightning_indexer : public test_case {
     }
 };
 
+// GGML_OP_TOP_K_MASK
+struct test_top_k_mask : public test_case {
+    const ggml_type type;
+    const int64_t n_kv;
+    const int64_t n_blk;   // 0: no blocks, a pick is a cell
+    const int64_t n_pick;  // number of blocks, ignored without blocks
+    const int64_t n_top;
+    const int64_t n_extra; // 0: no extra cells
+    const bool    score;   // false: every pick is live
+    const int64_t nb;      // rows per stream
+    const int64_t nh;      // rows of top_k per mask row, broadcast over ne2
+    const int64_t ns;      // streams
+    const bool    real;    // blocks of consecutive cells, live picks, as in a real cache
+
+    std::string vars() override {
+        return VARS_TO_STR11(type, n_kv, n_blk, n_pick, n_top, n_extra, score, nb, nh, ns, real);
+    }
+
+    // the result is a copy of the mask or -inf, the harness checks the -inf pattern, the finite values must match exactly
+    double max_err() override {
+        return 0.0;
+    }
+
+    double err(const float * a, const float * b, size_t n) override {
+        double result = 0.0;
+        for (size_t i = 0; i < n; ++i) {
+            if (std::isfinite(a[i]) && std::isfinite(b[i])) {
+                result = std::max(result, (double) fabsf(a[i] - b[i]));
+            }
+        }
+        return result;
+    }
+
+    size_t op_size(ggml_tensor * t) override {
+        const size_t ts = ggml_type_size(type);
+        const size_t nc = std::max<int64_t>(n_blk, 1);
+        return ggml_nbytes(t) + nb*nh*ns*(n_top*(sizeof(int32_t) + (score ? sizeof(float) : 0) + nc*(sizeof(int32_t) + 2*ts)) + n_extra*(sizeof(int32_t) + 2*ts));
+    }
+
+    test_top_k_mask(ggml_type type = GGML_TYPE_F16, int64_t n_kv = 256, int64_t n_blk = 4, int64_t n_pick = 64,
+            int64_t n_top = 16, int64_t n_extra = 3, bool score = true, int64_t nb = 32, int64_t nh = 1, int64_t ns = 1, bool real = false)
+        : type(type), n_kv(n_kv), n_blk(n_blk), n_pick(n_blk ? n_pick : n_kv), n_top(n_top), n_extra(n_extra), score(score),
+          nb(nb), nh(nh), ns(ns), real(real) {}
+
+    ggml_tensor * build_graph(ggml_context * ctx) override {
+        ggml_tensor * m = ggml_new_tensor_4d(ctx, type, n_kv, nb, 1, ns);
+        ggml_set_name(m, "m");
+
+        ggml_tensor * top_k = ggml_new_tensor_4d(ctx, GGML_TYPE_I32, n_top, nb, nh, ns);
+        ggml_set_name(top_k, "top_k");
+
+        ggml_tensor * s = nullptr;
+        if (score) {
+            s = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, n_pick, nb, nh, ns);
+            ggml_set_name(s, "score");
+        }
+
+        ggml_tensor * blocks = nullptr;
+        if (n_blk > 0) {
+            blocks = ggml_new_tensor_2d(ctx, GGML_TYPE_I32, n_blk, n_pick);
+            ggml_set_name(blocks, "blocks");
+        }
+
+        ggml_tensor * extra = nullptr;
+        if (n_extra > 0) {
+            extra = ggml_new_tensor_4d(ctx, GGML_TYPE_I32, n_extra, nb, nh, ns);
+            ggml_set_name(extra, "extra");
+        }
+
+        ggml_tensor * out = ggml_top_k_mask(ctx, m, top_k, s, blocks, extra);
+        ggml_set_name(out, "out");
+
+        return out;
+    }
+
+    void initialize_tensors(ggml_context * ctx) override {
+        std::mt19937 gen(0x4B50);
+
+        // cells and picks sometimes fall out of range, the op skips them
+        auto init_cells = [&](ggml_tensor * t) {
+            std::vector<int32_t> data(ggml_nelements(t));
+            if (real) {
+                // block p holds cells p*n_blk..p*n_blk + n_blk - 1, the extra cells are the last cells
+                const bool is_blocks = strcmp(t->name, "blocks") == 0;
+                for (int64_t i = 0; i < (int64_t) data.size(); ++i) {
+                    const int64_t c = is_blocks ? i : n_kv - 1 - i % t->ne[0];
+                    data[i] = c < n_kv ? (int32_t) c : (int32_t) n_kv;
+                }
+            } else {
+                std::uniform_int_distribution<int32_t> cell(0, n_kv - 1);
+                for (auto & c : data) {
+                    const uint32_t r = gen() % 16;
+                    c = r == 0 ? (int32_t) n_kv : r == 1 ? -1 : cell(gen);
+                }
+            }
+            ggml_backend_tensor_set(t, data.data(), 0, data.size()*sizeof(int32_t));
+        };
+
+        for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != NULL; t = ggml_get_next_tensor(ctx, t)) {
+            if (strcmp(t->name, "m") == 0) {
+                if (type == GGML_TYPE_F16) {
+                    init_tensor_kq_mask(t);
+                } else {
+                    std::vector<float> data(ggml_nelements(t));
+                    std::uniform_real_distribution<float> dis(-1.0f, 1.0f);
+                    for (auto & v : data) {
+                        v = gen() % 4 == 0 ? -INFINITY : dis(gen);
+                    }
+                    ggml_backend_tensor_set(t, data.data(), 0, data.size()*sizeof(float));
+                }
+            } else if (strcmp(t->name, "score") == 0) {
+                // -inf marks the picks a row cannot see
+                std::vector<float> data(ggml_nelements(t));
+                std::uniform_real_distribution<float> dis(-1.0f, 1.0f);
+                for (auto & v : data) {
+                    v = !real && gen() % 4 == 0 ? -INFINITY : dis(gen);
+                }
+                ggml_backend_tensor_set(t, data.data(), 0, data.size()*sizeof(float));
+            } else if (strcmp(t->name, "top_k") == 0) {
+                std::vector<int32_t> data(ggml_nelements(t));
+                std::vector<int32_t> order(n_pick);
+                for (int64_t r = 0; r < ggml_nrows(t); ++r) {
+                    for (int64_t i = 0; i < n_pick; ++i) {
+                        order[i] = i;
+                    }
+                    std::shuffle(order.begin(), order.end(), gen);
+                    for (int64_t j = 0; j < n_top; ++j) {
+                        const bool out = !real && j == 0 && r % 5 == 0;
+                        data[r*n_top + j] = out ? (r % 2 ? -1 : (int32_t) n_pick) : order[j % n_pick];
+                    }
+                }
+                ggml_backend_tensor_set(t, data.data(), 0, data.size()*sizeof(int32_t));
+            } else if (strcmp(t->name, "blocks") == 0 || strcmp(t->name, "extra") == 0) {
+                init_cells(t);
+            } else {
+                // the harness sentinels and the output, the callback checks the sentinels for NaN
+                init_tensor_uniform(t);
+            }
+        }
+    }
+};
+
 // Deserializable generic test case
 struct input_tensor {
     ggml_type type;
@@ -9912,6 +10054,25 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
     test_cases.emplace_back(new test_dsv4_hc_post(2560, 21, true, true));
     test_cases.emplace_back(new test_dsv4_hc_post(31, 17, false, true));
 
+    for (ggml_type type : {GGML_TYPE_F16, GGML_TYPE_F32}) {
+        // blocks with extra cells and scores (k-pool)
+        test_cases.emplace_back(new test_top_k_mask(type));
+        test_cases.emplace_back(new test_top_k_mask(type, 1001, 4, 70, 16, 3, true, 7));
+        test_cases.emplace_back(new test_top_k_mask(type, 256, 4, 64, 64, 0, true, 33));
+        test_cases.emplace_back(new test_top_k_mask(type, 512, 8, 16, 32, 7, true, 5, 1, 4));
+        test_cases.emplace_back(new test_top_k_mask(type, 4096, 4, 1024, 512, 3, true, 1));
+        test_cases.emplace_back(new test_top_k_mask(type, 4096, 4, 1024, 512, 3, true, 64, 1, 2));
+        // cells without scores (DSA)
+        test_cases.emplace_back(new test_top_k_mask(type, 256, 0, 0, 16, 0, false, 32));
+        test_cases.emplace_back(new test_top_k_mask(type, 1001, 0, 0, 64, 0, false, 7, 1, 3));
+        test_cases.emplace_back(new test_top_k_mask(type, 4096, 0, 0, 2048, 0, false, 5));
+        // cells with scores and extra cells
+        test_cases.emplace_back(new test_top_k_mask(type, 512, 0, 0, 32, 5, true, 9));
+        // blocks picked per head, the mask broadcast over the heads
+        test_cases.emplace_back(new test_top_k_mask(type, 512, 16, 32, 4, 0, true, 6, 4, 1));
+        test_cases.emplace_back(new test_top_k_mask(type, 1024, 8, 128, 16, 8, false, 3, 2, 2));
+    }
+
     // glu ops
     for (ggml_type type : {GGML_TYPE_F16, GGML_TYPE_F32, GGML_TYPE_BF16}) {
         for (int v : {0, 1}) {
@@ -12745,6 +12906,14 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_perf() {
         }
     }
 
+    // top_k_mask, Flash-Next QSA shapes: 512 pools of 4 + a 3-cell tail, and DSA shapes: 2048 cells
+    for (int64_t n_kv : { 18432, 131072 }) {
+        for (int64_t nb : { 1, 2048 }) {
+            test_cases.emplace_back(new test_top_k_mask(GGML_TYPE_F16, n_kv, 4, n_kv/4, 512, 3, true, nb, 1, 1, true));
+            test_cases.emplace_back(new test_top_k_mask(GGML_TYPE_F16, n_kv, 0, 0, 2048, 0, false, nb, 1, 1, true));
+        }
+    }
+
     // launch-overhead isolation: single L2_NORM launch vs batched siblings at the GDN
     // production shape (strided qkv views) -- perf-mode only, the eval list has its own
     // 2/4-wide coverage

tests/test-jinja.cpp

@@ -1299,6 +1299,18 @@ static void test_namespace(testing & t) {
         json::object(),
         "3"
     );
+
+    test_template(t, "namespace attribute named like a dict method",
+        "{% set ns = namespace(items=[]) %}{{ 'yes' if ns.items else 'no' }}{% set ns.items = ns.items + [1] %}{{ ns.items | length }}{{ 'yes' if ns.items else 'no' }}",
+        json::object(),
+        "no1yes"
+    );
+
+    test_template(t, "dict keeps its methods",
+        "{% set d = dict(items=1) %}{% for k, v in d.items() %}{{ k }}={{ v }}{% endfor %}",
+        json::object(),
+        "items=1"
+    );
 }
 
 static void test_tests(testing & t) {

tests/test-llama-archs.cpp

@@ -329,6 +329,9 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) {
         // partial rotary: n_rot must not exceed the indexer key length (64)
         ms.add_kv(LLM_KV_ROPE_DIMENSION_COUNT,       uint32_t(64));
     }
+    if (arch == LLM_ARCH_SENSENOVA_U1) {
+        ms.add_kv(LLM_KV_ROPE_FREQ_BASE_SPATIAL, 10000.0f);
+    }
     ms.add_kv(LLM_KV_ATTENTION_CLAMP_KQV,              1.0f);
     ms.add_kv(LLM_KV_ATTENTION_LAYERNORM_EPS,          1e-5f);
     ms.add_kv(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS,      1e-5f);
@@ -358,13 +361,16 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) {
         ms.add_kv(LLM_KV_ATTENTION_SLIDING_WINDOW_PATTERN, uint32_t(5));
     } else if (arch == LLM_ARCH_COHERE2MOE || arch == LLM_ARCH_MIMO2 || arch == LLM_ARCH_STEP35 || arch == LLM_ARCH_SPARK2_5 ||
             arch == LLM_ARCH_MUSE_GLIMMER || arch == LLM_ARCH_MUSE_SPARK || arch == LLM_ARCH_GRANITE_SWA || arch == LLM_ARCH_DOTS3NOTE ||
-            arch == LLM_ARCH_KOLIBRI1 || arch == LLM_ARCH_MAPLE) {
+            arch == LLM_ARCH_KOLIBRI1 || arch == LLM_ARCH_MAPLE || arch == LLM_ARCH_FLINT) {
         std::vector<uint32_t> pattern;
         pattern.reserve(n_layer);
         for (uint32_t il = 0; il < n_layer; il++) {
             pattern.push_back(il % 2);
         }
         ms.add_kv(LLM_KV_ATTENTION_SLIDING_WINDOW_PATTERN, pattern);
+        if (arch == LLM_ARCH_FLINT) {
+            ms.add_kv(LLM_KV_ATTENTION_ROPE_PATTERN, pattern);
+        }
     } else {
         ms.add_kv(LLM_KV_ATTENTION_SLIDING_WINDOW_PATTERN, uint32_t(2));
     }
@@ -522,6 +528,9 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) {
         ms.add_kv(LLM_KV_EXPERT_GATING_FUNC,         gating_func); // sigmoid : softmax : sqrtsoftplus
         ms.add_kv(LLM_KV_EXPERT_GROUP_SCALE,         1.0f);
         ms.add_kv(LLM_KV_EXPERTS_PER_GROUP,          uint32_t(1));
+        if (arch == LLM_ARCH_FLINT) {
+            ms.add_kv(LLM_KV_EXPERT_WEIGHTS_NORM, true);
+        }
         if (arch == LLM_ARCH_K2_HORIZON) {
             ms.add_kv(LLM_KV_ATTENTION_VALUE_EXPERT_COUNT,      uint32_t(2));
             ms.add_kv(LLM_KV_ATTENTION_VALUE_EXPERT_USED_COUNT, uint32_t(2));
@@ -751,6 +760,7 @@ static bool check_causal_attn_toggle(
 
 static bool moe_mandatory(const llm_arch arch) {
     switch (arch) {
+        case LLM_ARCH_FLINT:
         case LLM_ARCH_LLAMA4:
         case LLM_ARCH_MAION_CODER:
         case LLM_ARCH_COHERE2MOE:

tests/test-tokenizer-0.cpp

@@ -125,7 +125,7 @@ static llama_tests read_tests(const std::string & fname_inp, const std::string &
 
 int main(int argc, char **argv) {
     if (argc < 2) {
-        fprintf(stderr, "Usage: %s vocab-file [text-file]\n", argv[0]);
+        fprintf(stderr, "Usage: %s vocab-file [text-file] [--check-decode]\n", argv[0]);
         return 1;
     }
 
@@ -135,8 +135,13 @@ int main(int argc, char **argv) {
     const std::string fname_out = fname + ".out";
 
     std::string fname_text;
-    if (argc > 2) {
-        fname_text = argv[2];
+    bool check_decode = false; // also check that the expected tokens detokenize to the input text
+    for (int i = 2; i < argc; i++) {
+        if (std::string(argv[i]) == "--check-decode") {
+            check_decode = true;
+        } else {
+            fname_text = argv[i];
+        }
     }
 
     fprintf(stderr, "%s : reading vocab from: '%s'\n", __func__, fname.c_str());
@@ -244,6 +249,14 @@ int main(int argc, char **argv) {
 
                     success = false;
                 }
+
+                if (check_decode) {
+                    const std::string dec = common_detokenize(ctx, test_kv.second);
+                    if (dec != test_kv.first) {
+                        fprintf(stderr, "%s : failed decode:  '%s' detokenized to '%s'\n", __func__, test_kv.first.c_str(), dec.c_str());
+                        success = false;
+                    }
+                }
             }
         });
     }

tools/mtmd/CMakeLists.txt

@@ -55,6 +55,7 @@ add_library(mtmd
             models/paddleocr.cpp
             models/pixtral.cpp
             models/qwen2vl.cpp
+            models/sensenova-u1.cpp
             models/minimax-m3.cpp
             models/qwen3vl.cpp
             models/ling3vl.cpp

tools/mtmd/clip-impl.h

@@ -547,6 +547,7 @@ enum projector_type {
     PROJECTOR_TYPE_POCKETTTS_SPKENC,
     PROJECTOR_TYPE_POCKETTTS_GEN,
     PROJECTOR_TYPE_MUSE_GLIMMER,
+    PROJECTOR_TYPE_SENSENOVA_U1,
     PROJECTOR_TYPE_COHERE2V,
     PROJECTOR_TYPE_MUSE_SPARK_AUDIO,
     PROJECTOR_TYPE_UNKNOWN,
@@ -621,6 +622,7 @@ static std::map<projector_type, std::string> PROJECTOR_TYPE_NAMES = {
     { PROJECTOR_TYPE_POCKETTTS_SPKENC,  "pockettts_spkenc"},
     { PROJECTOR_TYPE_POCKETTTS_GEN,     "pockettts_gen"},
     { PROJECTOR_TYPE_MUSE_GLIMMER,      "muse-glimmer"},
+    { PROJECTOR_TYPE_SENSENOVA_U1,      "sensenova_u1"},
     { PROJECTOR_TYPE_COHERE2V,          "cohere2v"},
     { PROJECTOR_TYPE_MUSE_SPARK_AUDIO,  "muse-spark-audio"},
 };

tools/mtmd/clip.cpp

@@ -997,6 +997,10 @@ static std::unique_ptr<clip_graph> clip_get_graph_builder(clip_ctx * ctx, const
             {
                 builder = std::make_unique<clip_graph_qwen3vl>(ctx, img);
             } break;
+        case PROJECTOR_TYPE_SENSENOVA_U1:
+            {
+                builder = std::make_unique<clip_graph_sensenova_u1>(ctx, img);
+            } break;
         case PROJECTOR_TYPE_LING3VL:
             {
                 builder = std::make_unique<clip_graph_ling3vl>(ctx, img);
@@ -1795,6 +1799,20 @@ struct clip_model_loader {
                             LOG_WRN("%s: more info: https://github.com/ggml-org/llama.cpp/issues/16842\n\n", __func__);
                         }
                     } break;
+                case PROJECTOR_TYPE_SENSENOVA_U1:
+                    {
+                        hparams.n_merge = 2;
+                        hparams.image_resize_algo = RESIZE_ALGO_BICUBIC;
+                        get_u32(KEY_SPATIAL_MERGE_SIZE, hparams.n_merge, false);
+                        get_u32(KEY_IMAGE_MIN_PIXELS, hparams.image_min_pixels);
+                        get_u32(KEY_IMAGE_MAX_PIXELS, hparams.image_max_pixels);
+                        hparams.rope_theta = 10000.0f;
+                        hparams.set_limit_image_tokens(16, 4096);
+                        if (hparams.n_merge != 2) {
+                            throw std::runtime_error(string_format("%s: SenseNova U1 requires spatial_merge_size=2, got %u", __func__, hparams.n_merge));
+                        }
+                        hparams.set_warmup_n_tokens(16*16);
+                    } break;
                 case PROJECTOR_TYPE_MINIMAX_M3:
                     {
                         hparams.n_merge = 2; // spatial_merge_size
@@ -3837,6 +3855,11 @@ struct clip_model_loader {
                     }
 
                 } break;
+            case PROJECTOR_TYPE_SENSENOVA_U1:
+                {
+                    model.mm_0_w = get_tensor(string_format(TN_LLAVA_PROJ, 0, "weight"));
+                    model.mm_0_b = get_tensor(string_format(TN_LLAVA_PROJ, 0, "bias"));
+                } break;
             default:
                 GGML_ASSERT(false && "unknown projector type");
         }
@@ -4311,6 +4334,7 @@ int clip_n_output_tokens_x(const clip_ctx * ctx, const clip_image_f32 * img) {
         case PROJECTOR_TYPE_QWEN2VL:
         case PROJECTOR_TYPE_QWEN25VL:
         case PROJECTOR_TYPE_QWEN3VL:
+        case PROJECTOR_TYPE_SENSENOVA_U1:
         case PROJECTOR_TYPE_LING3VL:
         case PROJECTOR_TYPE_EXAONE4_5:
         case PROJECTOR_TYPE_MIMOVL:
@@ -4341,6 +4365,7 @@ int clip_n_output_tokens_y(const clip_ctx * ctx, const clip_image_f32 * img) {
         case PROJECTOR_TYPE_QWEN2VL:
         case PROJECTOR_TYPE_QWEN25VL:
         case PROJECTOR_TYPE_QWEN3VL:
+        case PROJECTOR_TYPE_SENSENOVA_U1:
         case PROJECTOR_TYPE_LING3VL:
         case PROJECTOR_TYPE_EXAONE4_5:
         case PROJECTOR_TYPE_MIMOVL:
@@ -4426,6 +4451,7 @@ int clip_n_output_tokens(const clip_ctx * ctx, const clip_image_f32 * img) {
         case PROJECTOR_TYPE_QWEN2VL:
         case PROJECTOR_TYPE_QWEN25VL:
         case PROJECTOR_TYPE_QWEN3VL:
+        case PROJECTOR_TYPE_SENSENOVA_U1:
         case PROJECTOR_TYPE_LING3VL:
         case PROJECTOR_TYPE_EXAONE4_5:
         case PROJECTOR_TYPE_MIMOVL:
@@ -5525,6 +5551,19 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) {
                 }
                 set_input_i32("pos_w", pos_data);
             } break;
+        case PROJECTOR_TYPE_SENSENOVA_U1:
+            {
+                const int n_patches_per_row = image_size_width / patch_size;
+                std::vector<int> pos_data(n_pos);
+                for (int i = 0; i < n_pos; i++) {
+                    pos_data[i] = i / n_patches_per_row;
+                }
+                set_input_i32("pos_y", pos_data);
+                for (int i = 0; i < n_pos; i++) {
+                    pos_data[i] = i % n_patches_per_row;
+                }
+                set_input_i32("pos_x", pos_data);
+            } break;
         case PROJECTOR_TYPE_DEEPSEEK4V:
             {
                 // set the 2D positions (mrope layout, only the first 2 channels are used)
@@ -6411,6 +6450,8 @@ int clip_n_mmproj_embd(const struct clip_ctx * ctx) {
         case PROJECTOR_TYPE_JANUS_PRO:
         case PROJECTOR_TYPE_YOUTUVL:
             return ctx->model.mm_1_b->ne[0];
+        case PROJECTOR_TYPE_SENSENOVA_U1:
+            return ctx->model.hparams.projection_dim;
         case PROJECTOR_TYPE_QWEN3VL:
             // main path + deepstack paths
             return ctx->model.mm_1_b->ne[0] * (1 + ctx->model.n_deepstack_layers);

tools/mtmd/models/models.h

@@ -50,6 +50,11 @@ struct clip_graph_qwen3vl : clip_graph_qwen2vl {
     ggml_cgraph * build() override;
 };
 
+struct clip_graph_sensenova_u1 : clip_graph {
+    clip_graph_sensenova_u1(clip_ctx * ctx, const clip_image_f32 & img) : clip_graph(ctx, img) {}
+    ggml_cgraph * build() override;
+};
+
 struct clip_graph_ling3vl : clip_graph_qwen3vl {
     clip_graph_ling3vl(clip_ctx * ctx, const clip_image_f32 & img) : clip_graph_qwen3vl(ctx, img) {}
     ggml_cgraph * build() override;

tools/mtmd/models/sensenova-u1.cpp

@@ -0,0 +1,40 @@
+#include "models.h"
+
+ggml_cgraph * clip_graph_sensenova_u1::build() {
+    GGML_ASSERT(model.patch_embeddings_0 != nullptr);
+    GGML_ASSERT(model.patch_bias != nullptr);
+    GGML_ASSERT(model.mm_0_w != nullptr);
+    GGML_ASSERT(model.mm_0_b != nullptr);
+    GGML_ASSERT(hparams.n_merge == 2);
+
+    ggml_tensor * pos_x = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_patches);
+    ggml_set_name(pos_x, "pos_x");
+    ggml_set_input(pos_x);
+
+    ggml_tensor * pos_y = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_patches);
+    ggml_set_name(pos_y, "pos_y");
+    ggml_set_input(pos_y);
+
+    ggml_tensor * inp = build_inp();
+    inp = ggml_gelu(ctx0, inp);
+
+    // Apply adjacent-pair 2D RoPE to patch embeddings before downsampling.
+    inp = ggml_reshape_4d(ctx0, inp, n_embd, 1, n_patches, n_batch);
+    inp = build_rope_2d(ctx0, inp, pos_x, pos_y, hparams.rope_theta, false);
+    inp = ggml_reshape_3d(ctx0, inp, n_embd, n_patches, n_batch);
+    cb(inp, "patch_embd_rope", -1);
+
+    inp = ggml_reshape_4d(ctx0, inp, n_embd, n_patches_x, n_patches_y, n_batch);
+    inp = ggml_cont(ctx0, ggml_permute(ctx0, inp, 2, 0, 1, 3));
+    inp = ggml_conv_2d(ctx0, model.mm_0_w, inp,
+            hparams.n_merge, hparams.n_merge, 0, 0, 1, 1);
+
+    const int n_output_patches = n_patches / (hparams.n_merge * hparams.n_merge);
+    inp = ggml_cont(ctx0, ggml_permute(ctx0, inp, 1, 2, 0, 3));
+    inp = ggml_reshape_3d(ctx0, inp, n_mmproj_embd, n_output_patches, n_batch);
+    inp = ggml_add(ctx0, inp, model.mm_0_b);
+    cb(inp, "proj_out", -1);
+
+    ggml_build_forward_expand(gf, inp);
+    return gf;
+}

tools/mtmd/mtmd.cpp

@@ -225,6 +225,7 @@ enum mtmd_pos_type {
     MTMD_POS_TYPE_NORMAL,    // number of positions equals to number of tokens
     MTMD_POS_TYPE_MROPE,     // qwen-vl mrope style, each image takes max(t,h,w) position indexes
     MTMD_POS_TYPE_HUNYUANVL, // HunyuanVL mrope + BOI/EOI/newline layout with XD-RoPE dim-3
+    MTMD_POS_TYPE_SENSENOVA_U1, // U1 uses zero-based H/W and one temporal position per image
     MTMD_POS_TYPE_CANVAS,    // MiniCPM-V 4.7: overview + slices in one chunk, sharing one 2D canvas (see mtmd_image_tokens::canvas_tile_grid)
     MTMD_POS_TYPE_COUNT,     // for validation
 };
@@ -751,6 +752,12 @@ struct mtmd_context {
                     img_end = "]<]end of image[>[";
                     image_preproc = std::make_unique<mtmd_image_preprocessor_dyn_size>(ctx_v);
                 } break;
+            case PROJECTOR_TYPE_SENSENOVA_U1:
+                {
+                    img_beg = "<img>";
+                    img_end = "</img>";
+                    image_preproc = std::make_unique<mtmd_image_preprocessor_dyn_size>(ctx_v);
+                } break;
             case PROJECTOR_TYPE_MUSE_GLIMMER:
                 {
                     img_beg = "<|image_start|>";
@@ -1615,10 +1622,15 @@ struct mtmd_tokenizer {
                     image_tokens->ny = 1;
                 }
                 image_tokens->pos = ctx->pos_type;
+                // U1 uses M-RoPE, but unlike Qwen its image H/W coordinates are
+                // zero-based and the whole image occupies one temporal position.
+                if (ctx->proj_type_v() == PROJECTOR_TYPE_SENSENOVA_U1) {
+                    image_tokens->pos = MTMD_POS_TYPE_SENSENOVA_U1;
+                }
                 // HunyuanVL wraps the image grid with BOI/EOI and adds one newline per row,
                 // and uses XD-RoPE (dim-3 = image index). Override the position type so that
                 // n_tokens() and mtmd_image_tokens_get_decoder_pos pick the HunyuanVL layout.
-                if (ctx->proj_type_v() == PROJECTOR_TYPE_HUNYUANVL) {
+                else if (ctx->proj_type_v() == PROJECTOR_TYPE_HUNYUANVL) {
                     image_tokens->pos       = MTMD_POS_TYPE_HUNYUANVL;
                     image_tokens->image_idx = n_images_added;
                     GGML_ASSERT(n_tokens == (size_t)image_tokens->n_tokens());
@@ -2398,6 +2410,7 @@ bool mtmd_decode_use_non_causal(const mtmd_context * ctx, const mtmd_input_chunk
             }
         case PROJECTOR_TYPE_GEMMA4UV:
         case PROJECTOR_TYPE_GEMMA3:
+        case PROJECTOR_TYPE_SENSENOVA_U1:
         case PROJECTOR_TYPE_DEEPSEEK4V:
             return true;
         default:
@@ -2761,6 +2774,13 @@ mtmd_decoder_pos mtmd_image_tokens_get_decoder_pos(const mtmd_image_tokens * ima
                 pos.y = pos_0 + i;
                 pos.z = pos_0 + i;
             } break;
+        case MTMD_POS_TYPE_SENSENOVA_U1:
+            {
+                pos.t = pos_0;
+                pos.x = i % image_tokens->nx;
+                pos.y = i / image_tokens->nx;
+                pos.z = 0;
+            } break;
         case MTMD_POS_TYPE_HUNYUANVL:
             {
                 // HunyuanVL layout: [BOI] [row0 tokens + newline] ... [row(ny-1) tokens + newline] [EOI]
@@ -2813,6 +2833,8 @@ llama_pos mtmd_image_tokens_get_n_pos(const mtmd_image_tokens * image_tokens) {
             return std::max(image_tokens->nx, image_tokens->ny);
         case MTMD_POS_TYPE_NORMAL:
             return image_tokens->n_tokens();
+        case MTMD_POS_TYPE_SENSENOVA_U1:
+            return 1;
         case MTMD_POS_TYPE_HUNYUANVL:
             // HunyuanVL: the sequential (dim-0) position advances by the full token count
             // (includes BOI/EOI and row newline tokens), not by max(nx, ny)

tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/ChatMessageToolCallBlock.svelte

@@ -23,13 +23,15 @@
 	interface Props {
 		section: AgenticSection;
 		attachments?: DatabaseMessageExtra[];
+		conversationCwd?: string;
 		open: boolean;
 		isStreaming: boolean;
 		isExecuting?: boolean;
 		onToggle?: () => void;
 	}
 
-	let { attachments, isExecuting, isStreaming, onToggle, open, section }: Props = $props();
+	let { attachments, conversationCwd, isExecuting, isStreaming, onToggle, open, section }: Props =
+		$props();
 
 	// Runs for every tool block on mount, before the body renders: the cheap
 	// content prefilter and the tool-name allow-list come first so blobs from
@@ -60,6 +62,7 @@
 {:else if section.toolName === BuiltInTool.SERVER_EXEC_SHELL_COMMAND}
 	<ChatMessageToolCallBlockExecShellCommand
 		{attachments}
+		{conversationCwd}
 		{isExecuting}
 		{isStreaming}
 		{onToggle}

tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/ChatMessageToolCallBlockEditFile.svelte

@@ -2,10 +2,12 @@
 	import { parseEditFileMeta, parseEditFileTitleMeta } from './parsers/edit-file';
 	import ToolCallBlock from './ToolCallBlock.svelte';
 	import { XCircle } from '@lucide/svelte';
+	import { SyntaxHighlightedCode } from '$lib/components/app';
 	import { MAX_HEIGHT_CODE_BLOCK, RESULT_STAT_SEPARATOR } from '$lib/constants';
+	import { FileTypeText } from '$lib/enums';
 	import { toolsStore } from '$lib/stores';
 	import type { AgenticSection } from '$lib/types';
-	import { abbreviateHome, computeLineDiff, prefixFor } from '$lib/utils';
+	import { abbreviateHome, computeLineDiff, formatJsonPretty, prefixFor } from '$lib/utils';
 
 	interface Props {
 		section: AgenticSection;
@@ -44,16 +46,16 @@
 		</span>
 	{/snippet}
 
-	{#snippet children(meta, _ctx)}
-		{#if meta?.errorMessage}
+	{#snippet children(_meta, ctx)}
+		{#if editFileBody?.errorMessage}
 			<div
 				class="flex items-start gap-2 rounded bg-red-500/10 p-2 text-xs text-red-600 italic dark:text-red-400"
 			>
 				<XCircle class="mt-0.5 h-3 w-3 shrink-0" />
 
-				<span>{meta.errorMessage}</span>
+				<span>{editFileBody.errorMessage}</span>
 			</div>
-		{:else if meta && editFileBody && editFileBody.edits.length > 0}
+		{:else if editFileBody?.edits.length}
 			{#each editDiffs as diffLines, ei (ei)}
 				<div class={ei === 0 ? '' : 'mt-3'}>
 					<div class="mb-1.5 text-xs text-muted-foreground/70 italic">
@@ -79,14 +81,28 @@
 			{/each}
 
 			<div class="mt-1.5 text-xs text-muted-foreground/70 italic">
-				{#if meta.resultMessage}
-					{meta.resultMessage}{meta.editsApplied != null ? RESULT_STAT_SEPARATOR : ''}{/if}
+				{#if editFileBody.resultMessage}
+					{editFileBody.resultMessage}{editFileBody.editsApplied != null
+						? RESULT_STAT_SEPARATOR
+						: ''}{/if}
 
-				{#if meta.editsApplied != null}
-					<span class="font-mono">{meta.editsApplied}</span>
-					{meta.editsApplied === 1 ? 'edit' : 'edits'}&nbsp;applied
+				{#if editFileBody.editsApplied != null}
+					<span class="font-mono">{editFileBody.editsApplied}</span>
+					{editFileBody.editsApplied === 1 ? 'edit' : 'edits'}&nbsp;applied
 				{/if}
 			</div>
+		{:else if section.toolArgs}
+			<!-- no usable edits: not streamed yet, or a same-named tool with other arg names -->
+			<SyntaxHighlightedCode
+				code={formatJsonPretty(section.toolArgs)}
+				language={FileTypeText.JSON}
+				maxHeight={MAX_HEIGHT_CODE_BLOCK}
+				streaming={ctx.isCodeStreaming}
+			/>
+		{:else if ctx.isStreaming}
+			<div class="rounded bg-muted/20 p-2 text-xs text-muted-foreground/70 italic">
+				Receiving arguments...
+			</div>
 		{:else}
 			<div class="rounded bg-muted/20 p-2 text-xs text-muted-foreground/70 italic">No edits</div>
 		{/if}

tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/ChatMessageToolCallBlockExecShellCommand.svelte

@@ -11,7 +11,7 @@
 	import { AlertTriangle, Check, Loader2, XCircle } from '@lucide/svelte';
 	import { CollapsibleTerminalBlock } from '$lib/components/app';
 	import { SETTINGS_KEYS, TOOL_RUNTIME_SCROLL_AT_BOTTOM_THRESHOLD_PX } from '$lib/constants';
-	import { AttachmentType } from '$lib/enums';
+	import { AgenticSectionType, AttachmentType } from '$lib/enums';
 	import { settingsStore, toolsStore } from '$lib/stores';
 	import type { AgenticSection, DatabaseMessageExtra, ToolResultLine } from '$lib/types';
 	import {
@@ -32,17 +32,40 @@
 		 *  tool call. Drives max-height + auto-scroll while true; releases
 		 *  them when the loop reports this call as done. */
 		isExecuting?: boolean;
+		/** Working directory of the conversation this call belongs to; stands
+		 *  in for `section.toolCwd` until the result message persists the
+		 *  real one. */
+		conversationCwd?: string;
 		attachments?: DatabaseMessageExtra[];
 		onToggle?: () => void;
 	}
 
-	let { attachments, isExecuting = false, isStreaming, onToggle, open, section }: Props = $props();
+	let {
+		attachments,
+		conversationCwd,
+		isExecuting = false,
+		isStreaming,
+		onToggle,
+		open,
+		section
+	}: Props = $props();
 
 	// `isLive` covers all in-flight phases: pre-chunk spinner and
 	// streaming itself. Frozen output (tool done while agent continues)
 	// is not live.
 	const isLive = $derived(isExecuting);
 
+	const isPreResult = $derived(
+		section.type === AgenticSectionType.TOOL_CALL_PENDING ||
+			section.type === AgenticSectionType.TOOL_CALL_STREAMING
+	);
+
+	// The model is still typing the call while the args JSON doesn't parse;
+	// the title keeps a terminal cursor until it does.
+	const commandStreaming = $derived(
+		isStreaming && isPreResult && !isParsableJson(section.toolArgs)
+	);
+
 	const execShellMeta = $derived(parseExecShellCommandMeta(section));
 	const execShellError = $derived(parseExecShellCommandError(section.toolResult));
 	const execShellExitStatus: ExecShellExitStatus | undefined = $derived(
@@ -76,9 +99,11 @@
 
 	// The working directory the command ran with, persisted per call on the
 	// tool result message (it travels via the x-tool-cwd header, not the tool
-	// args). Reading it from the section keeps it accurate even if the
-	// conversation cwd changes later.
-	const cwd = $derived(section.toolCwd);
+	// args). Until that message exists, the conversation cwd is the best known
+	// value - it is what the agentic loop passes when executing the call.
+	const cwd = $derived(
+		section.toolCwd ?? (isStreaming && isPreResult ? conversationCwd : undefined)
+	);
 	const home = $derived(toolsStore.serverHome);
 	const wdDisplay = $derived(abbreviateHome(cwd ?? '', home));
 
@@ -112,6 +137,18 @@
 		);
 	}
 
+	function isParsableJson(value: string | undefined): boolean {
+		if (!value) return false;
+
+		try {
+			JSON.parse(value);
+
+			return true;
+		} catch {
+			return false;
+		}
+	}
+
 	function scrollToBottomOnFrame() {
 		if (pendingFrame !== null || !scrollEl || userScrolledUp) return;
 
@@ -185,6 +222,10 @@
 	{:else}
 		<span class="font-mono">{execShellMeta?.command}</span>
 	{/if}
+
+	{#if commandStreaming}
+		<span class="exec-cursor"></span>
+	{/if}
 {/snippet}
 
 <ToolCallBlock
@@ -202,10 +243,10 @@
 	{/snippet}
 
 	{#snippet children(_meta, ctx)}
-		{#if ctx.isPending}
+		{#if (ctx.isPending && !commandStreaming) || (isLive && !section.toolResult)}
 			<div class="flex items-start gap-2 text-xs text-muted-foreground/70">
 				<Loader2 class="h-3 w-3 animate-spin" />
-				Running...
+				Executing command...
 			</div>
 		{:else if execShellError}
 			<div class="flex items-start gap-2 text-xs text-red-600 italic dark:text-red-400">
@@ -277,6 +318,22 @@
 		margin-right: var(--exec-wd-margin);
 	}
 
+	.exec-cursor {
+		display: inline-block;
+		width: 0.5em;
+		height: 1.05em;
+		margin-left: 0.15rem;
+		vertical-align: text-bottom;
+		background: var(--foreground);
+		animation: exec-cursor-blink 1.1s steps(1) infinite;
+	}
+
+	@keyframes exec-cursor-blink {
+		50% {
+			opacity: 0;
+		}
+	}
+
 	.terminal-output {
 		overscroll-behavior: contain;
 	}

tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/ChatMessageToolCallBlockWriteFile.svelte

@@ -4,9 +4,10 @@
 	import { XCircle } from '@lucide/svelte';
 	import { SyntaxHighlightedCode } from '$lib/components/app';
 	import { MAX_HEIGHT_CODE_BLOCK, RESULT_STAT_SEPARATOR } from '$lib/constants';
+	import { FileTypeText } from '$lib/enums';
 	import { toolsStore } from '$lib/stores';
 	import type { AgenticSection } from '$lib/types';
-	import { abbreviateHome } from '$lib/utils';
+	import { abbreviateHome, formatJsonPretty } from '$lib/utils';
 
 	interface Props {
 		section: AgenticSection;
@@ -42,32 +43,48 @@
 		</span>
 	{/snippet}
 
-	{#snippet children(meta, ctx)}
-		{#if meta?.errorMessage}
+	{#snippet children(_meta, ctx)}
+		{#if writeFileBody?.errorMessage}
 			<div
 				class="flex items-start gap-2 rounded bg-red-500/10 p-2 text-xs text-red-600 italic dark:text-red-400"
 			>
 				<XCircle class="mt-0.5 h-3 w-3 shrink-0" />
 
-				<span>{meta.errorMessage}</span>
+				<span>{writeFileBody.errorMessage}</span>
 			</div>
-		{:else if meta}
+		{:else if writeFileBody?.content !== undefined}
 			<SyntaxHighlightedCode
-				code={writeFileBody?.content ?? ''}
-				language={meta.language}
+				code={writeFileBody.content}
+				language={writeFileBody.language}
 				maxHeight={MAX_HEIGHT_CODE_BLOCK}
 				streaming={ctx.isCodeStreaming}
 			/>
 
 			<div class="mt-1.5 text-xs text-muted-foreground/70 italic">
-				{#if meta.resultMessage}
-					{meta.resultMessage}{meta.bytesWritten != null ? RESULT_STAT_SEPARATOR : ''}{/if}
+				{#if writeFileBody.resultMessage}
+					{writeFileBody.resultMessage}{writeFileBody.bytesWritten != null
+						? RESULT_STAT_SEPARATOR
+						: ''}{/if}
 
-				{#if meta.bytesWritten != null}
-					<span class="font-mono">{meta.bytesWritten}</span>
+				{#if writeFileBody.bytesWritten != null}
+					<span class="font-mono">{writeFileBody.bytesWritten}</span>
 					bytes
 				{/if}
 			</div>
+		{:else if section.toolArgs}
+			<!-- no content arg: not streamed yet, or a same-named tool with other arg names -->
+			<SyntaxHighlightedCode
+				code={formatJsonPretty(section.toolArgs)}
+				language={FileTypeText.JSON}
+				maxHeight={MAX_HEIGHT_CODE_BLOCK}
+				streaming={ctx.isCodeStreaming}
+			/>
+		{:else if ctx.isStreaming}
+			<div class="rounded bg-muted/20 p-2 text-xs text-muted-foreground/70 italic">
+				Receiving arguments...
+			</div>
+		{:else}
+			<div class="rounded bg-muted/20 p-2 text-xs text-muted-foreground/70 italic">No content</div>
 		{/if}
 	{/snippet}
 </ToolCallBlock>

tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/parsers/_shared.ts

@@ -1,10 +1,14 @@
 // Helpers shared by the per-tool meta parsers under
 // `src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/parsers/`.
 // Each tool needs the same first three steps (tool-name check,
-// args-present check, JSON parse) - keeping them here lets each parser
-// stay focused on its own format quirks.
+// args-present check, JSON parse), plus the same path aliases - keeping
+// them here lets each parser stay focused on its own format quirks.
 
-import { TOOL_ARG_STRING_FIELD_PATTERN_TEMPLATE } from '$lib/constants';
+import {
+	FILE_PATH_SEPARATOR_REGEX,
+	TOOL_ARG_PATH_KEYS,
+	TOOL_ARG_STRING_FIELD_PATTERN_TEMPLATE
+} from '$lib/constants';
 import { BuiltInTool } from '$lib/enums';
 import type { AgenticSection } from '$lib/types/agentic';
 import { parsePartialJsonArgs } from '$lib/utils/parse-partial-json-args';
@@ -68,6 +72,34 @@ export function extractToolArgString(
 	return undefined;
 }
 
+/**
+ * File name shown in a block title, from a tool-call path. Undefined when
+ * the call carries no path; the title then shows the tool label alone.
+ */
+export function fileNameFromPath(filePath: string | undefined): string | undefined {
+	if (!filePath) return undefined;
+
+	return filePath.split(FILE_PATH_SEPARATOR_REGEX).pop() || filePath;
+}
+
+/**
+ * Read the path argument out of an already parsed args object, checking
+ * the aliases in `TOOL_ARG_PATH_KEYS`. Returns undefined when the call
+ * carries no usable path - the path is display-only, so callers must
+ * still render whatever else the call streamed.
+ */
+export function extractToolArgPath(args?: Record<string, unknown> | null): string | undefined {
+	if (!args) return undefined;
+
+	for (const key of TOOL_ARG_PATH_KEYS) {
+		const value = args[key];
+
+		if (typeof value === 'string' && value) return value;
+	}
+
+	return undefined;
+}
+
 /**
  * Parse a section's toolArgs against an expected tool name. Returns
  * `null` when:

tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/parsers/edit-file.ts

@@ -1,10 +1,16 @@
 // Meta parser for `edit_file` tool calls. Reads the file path and the
 // array of edits from the streamed args (partial JSON for incremental
 // rendering), plus the result blob for `result` / `edits_applied` /
-// `error` fields.
+// `error` fields. The path is display-only (block title), so a call
+// that has not streamed one still shows its diffs.
 
-import { extractToolArgString, parseToolArgs } from './_shared';
-import { FILE_PATH_SEPARATOR_REGEX, TOOL_ARG_PATH_KEYS } from '$lib/constants';
+import {
+	extractToolArgPath,
+	extractToolArgString,
+	fileNameFromPath,
+	parseToolArgs
+} from './_shared';
+import { TOOL_ARG_PATH_KEYS } from '$lib/constants';
 import { BuiltInTool } from '$lib/enums';
 import type { AgenticSection, EditFileEdit, EditFileMeta, EditFileTitleMeta } from '$lib/types';
 import { tryParseToolResultObject } from '$lib/utils';
@@ -14,11 +20,8 @@ export function parseEditFileMeta(section: AgenticSection): EditFileMeta | null
 
 	if (!args) return null;
 
-	const rawPath = args.path ?? args.file_path ?? args.filePath;
-
-	if (typeof rawPath !== 'string' || !rawPath) return null;
-
-	const fileName = rawPath.split(FILE_PATH_SEPARATOR_REGEX).pop() || rawPath;
+	const filePath = extractToolArgPath(args);
+	const fileName = fileNameFromPath(filePath);
 	// Filter the streamed edits array strictly: each entry must be an
 	// object with a non-empty `old_text`. Edits without an old_text
 	// would diff against empty and render as a full re-write.
@@ -61,7 +64,7 @@ export function parseEditFileMeta(section: AgenticSection): EditFileMeta | null
 		editsApplied,
 		errorMessage,
 		fileName,
-		filePath: rawPath,
+		filePath,
 		resultMessage
 	};
 }
@@ -75,18 +78,15 @@ export function parseEditFileMeta(section: AgenticSection): EditFileMeta | null
 export function parseEditFileTitleMeta(section: AgenticSection): EditFileTitleMeta | null {
 	if (section.toolName !== BuiltInTool.SERVER_EDIT_FILE || !section.toolArgs) return null;
 
-	let rawPath: string | undefined = extractToolArgString(section.toolArgs, TOOL_ARG_PATH_KEYS);
+	let filePath = extractToolArgString(section.toolArgs, TOOL_ARG_PATH_KEYS);
 
-	if (!rawPath) {
-		const args = parseToolArgs(BuiltInTool.SERVER_EDIT_FILE, section, { partial: true });
-		const fallbackPath = args?.path ?? args?.file_path ?? args?.filePath;
-
-		if (typeof fallbackPath === 'string' && fallbackPath) rawPath = fallbackPath;
+	if (!filePath) {
+		filePath = extractToolArgPath(
+			parseToolArgs(BuiltInTool.SERVER_EDIT_FILE, section, { partial: true })
+		);
 	}
 
-	if (!rawPath) return null;
-
-	const fileName = rawPath.split(FILE_PATH_SEPARATOR_REGEX).pop() || rawPath;
+	const fileName = fileNameFromPath(filePath);
 	const resultObj = tryParseToolResultObject(section.toolResult);
 
 	let resultMessage: string | undefined;
@@ -105,5 +105,5 @@ export function parseEditFileTitleMeta(section: AgenticSection): EditFileTitleMe
 		}
 	}
 
-	return { editsApplied, errorMessage, fileName, filePath: rawPath, resultMessage };
+	return { editsApplied, errorMessage, fileName, filePath, resultMessage };
 }

tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/parsers/exec-shell-command.ts

@@ -1,5 +1,7 @@
 // Meta parser for `exec_shell_command` tool calls. Surfaces the
 // command text from args `command` / `cmd` / `shell_command` aliases.
+// Args are parsed partially so the title can render the command while
+// it is still streaming.
 // The exit-status and error parsing live in their own utilities
 // (`parse-exec-shell-status.ts` / `parse-exec-shell-error.ts`) - this
 // file only deals with what's strictly about *calling* the tool, since
@@ -14,7 +16,7 @@ export type ExecShellCommandMeta = {
 };
 
 export function parseExecShellCommandMeta(section: AgenticSection): ExecShellCommandMeta | null {
-	const args = parseToolArgs(BuiltInTool.SERVER_EXEC_SHELL_COMMAND, section);
+	const args = parseToolArgs(BuiltInTool.SERVER_EXEC_SHELL_COMMAND, section, { partial: true });
 
 	if (!args) return null;
 

tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/parsers/write-file.ts

@@ -1,30 +1,40 @@
 // Meta parser for `write_file` tool calls. Reads the path/content from
 // the streamed args (partial JSON so we can render before the call
 // finishes) and surfaces `bytes`, `result`, and `error` from the
-// result blob.
+// result blob. The path is display-only (title + highlight language),
+// so a call that has not streamed one still shows its content.
 
-import { extractToolArgString, parseToolArgs } from './_shared';
-import { CODE_BLOCK, FILE_PATH_SEPARATOR_REGEX, TOOL_ARG_PATH_KEYS } from '$lib/constants';
+import {
+	extractToolArgPath,
+	extractToolArgString,
+	fileNameFromPath,
+	parseToolArgs
+} from './_shared';
+import { CODE_BLOCK, TOOL_ARG_PATH_KEYS } from '$lib/constants';
 import { BuiltInTool } from '$lib/enums';
 import type { AgenticSection, WriteFileMeta, WriteFileTitleMeta } from '$lib/types';
 import { getFileTypeByExtension, tryParseToolResultObject } from '$lib/utils';
 
+// Plain text when the call carries no path, or one with no known extension.
+function languageForPath(filePath: string | undefined): string {
+	if (!filePath) return CODE_BLOCK.DEFAULT_LANGUAGE;
+
+	return (
+		getFileTypeByExtension(filePath)?.replace(CODE_BLOCK.TEXT_LANGUAGE_PREFIX_REGEX, '') ??
+		CODE_BLOCK.DEFAULT_LANGUAGE
+	);
+}
+
 export function parseWriteFileMeta(section: AgenticSection): WriteFileMeta | null {
 	const args = parseToolArgs(BuiltInTool.SERVER_WRITE_FILE, section, { partial: true });
 
 	if (!args) return null;
 
-	// Tool contracts drifted over time: some models emit `path`,
-	// others `file_path` / `filePath`. Accept all three.
-	const rawPath = args.path ?? args.file_path ?? args.filePath;
-
-	if (typeof rawPath !== 'string' || !rawPath) return null;
-
-	const fileName = rawPath.split(FILE_PATH_SEPARATOR_REGEX).pop() || rawPath;
-	const content = typeof args.content === 'string' ? args.content : '';
-	const language =
-		getFileTypeByExtension(rawPath)?.replace(CODE_BLOCK.TEXT_LANGUAGE_PREFIX_REGEX, '') ??
-		CODE_BLOCK.DEFAULT_LANGUAGE;
+	const filePath = extractToolArgPath(args);
+	const fileName = fileNameFromPath(filePath);
+	// undefined, not '', so an intentionally empty write still renders
+	const content = typeof args.content === 'string' ? args.content : undefined;
+	const language = languageForPath(filePath);
 	const resultObj = tryParseToolResultObject(section.toolResult);
 	const bytesWritten =
 		resultObj && Number.isFinite(Number(resultObj.bytes)) ? Number(resultObj.bytes) : undefined;
@@ -36,7 +46,7 @@ export function parseWriteFileMeta(section: AgenticSection): WriteFileMeta | nul
 		content,
 		errorMessage,
 		fileName,
-		filePath: rawPath,
+		filePath,
 		language,
 		resultMessage
 	};
@@ -51,21 +61,14 @@ export function parseWriteFileMeta(section: AgenticSection): WriteFileMeta | nul
 export function parseWriteFileTitleMeta(section: AgenticSection): WriteFileTitleMeta | null {
 	if (section.toolName !== BuiltInTool.SERVER_WRITE_FILE || !section.toolArgs) return null;
 
-	let rawPath: string | undefined = extractToolArgString(section.toolArgs, TOOL_ARG_PATH_KEYS);
+	let filePath = extractToolArgString(section.toolArgs, TOOL_ARG_PATH_KEYS);
 
-	if (!rawPath) {
-		const args = parseToolArgs(BuiltInTool.SERVER_WRITE_FILE, section, { partial: true });
-		const fallbackPath = args?.path ?? args?.file_path ?? args?.filePath;
-
-		if (typeof fallbackPath === 'string' && fallbackPath) rawPath = fallbackPath;
+	if (!filePath) {
+		filePath = extractToolArgPath(
+			parseToolArgs(BuiltInTool.SERVER_WRITE_FILE, section, { partial: true })
+		);
 	}
 
-	if (!rawPath) return null;
-
-	const fileName = rawPath.split(FILE_PATH_SEPARATOR_REGEX).pop() || rawPath;
-	const language =
-		getFileTypeByExtension(rawPath)?.replace(CODE_BLOCK.TEXT_LANGUAGE_PREFIX_REGEX, '') ??
-		CODE_BLOCK.DEFAULT_LANGUAGE;
 	const resultObj = tryParseToolResultObject(section.toolResult);
 	const bytesWritten =
 		resultObj && Number.isFinite(Number(resultObj.bytes)) ? Number(resultObj.bytes) : undefined;
@@ -75,9 +78,9 @@ export function parseWriteFileTitleMeta(section: AgenticSection): WriteFileTitle
 	return {
 		bytesWritten,
 		errorMessage,
-		fileName,
-		filePath: rawPath,
-		language,
+		fileName: fileNameFromPath(filePath),
+		filePath,
+		language: languageForPath(filePath),
 		resultMessage
 	};
 }

tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessageAgenticContent.svelte

@@ -8,14 +8,14 @@
 		MarkdownContent
 	} from '$lib/components/app';
 	import { AgenticSectionType, ChatMessageStatsView, ToolPermissionDecision } from '$lib/enums';
-	import { agenticStore, settingsStore } from '$lib/stores';
+	import { agenticStore, conversationsStore, settingsStore } from '$lib/stores';
 	import type {
 		AgenticSection,
 		ChatMessageAgenticTimings,
 		ChatMessageAgenticTurnStats,
 		DatabaseMessage
 	} from '$lib/types';
-	import { deriveAgenticSections } from '$lib/utils';
+	import { deriveAgenticSections, isSectionExpandedByDefault } from '$lib/utils';
 
 	interface Props {
 		message: DatabaseMessage;
@@ -34,6 +34,7 @@
 	let expandedStates: Record<number, boolean> = $state({});
 
 	const showThoughtInProgress = $derived(Boolean(settingsStore.config.showThoughtInProgress));
+	const showToolCallInProgress = $derived(Boolean(settingsStore.config.showToolCallInProgress));
 	const alwaysShowToolCallContent = $derived(
 		Boolean(settingsStore.config.alwaysShowToolCallContent)
 	);
@@ -95,6 +96,14 @@
 		isStreaming ? agenticStore.getExecutingToolCallId(message.convId) : null
 	);
 
+	// Cwd the agentic loop passes to tool calls in this conversation; the exec
+	// block shows it until the result message persists the real one.
+	const conversationCwd = $derived(
+		conversationsStore.activeConversation?.id === message.convId
+			? conversationsStore.activeConversation?.cwd
+			: undefined
+	);
+
 	type TurnGroup = {
 		sections: AgenticSection[];
 		flatIndices: number[];
@@ -132,28 +141,17 @@
 		return groups;
 	});
 
-	function getDefaultExpanded(section: AgenticSection): boolean {
-		if (
-			section.type === AgenticSectionType.TOOL_CALL ||
-			section.type === AgenticSectionType.TOOL_CALL_PENDING ||
-			section.type === AgenticSectionType.TOOL_CALL_STREAMING
-		) {
-			return alwaysShowToolCallContent;
-		}
-
-		if (section.type === AgenticSectionType.REASONING_PENDING) {
-			return showThoughtInProgress;
-		}
-
-		return false;
-	}
-
 	function isExpanded(index: number, section: AgenticSection): boolean {
 		if (expandedStates[index] !== undefined) {
 			return expandedStates[index];
 		}
 
-		return getDefaultExpanded(section);
+		return isSectionExpandedByDefault(section, {
+			alwaysShowToolCallContent,
+			executingToolCallId: currentlyExecutingToolCallId,
+			showThoughtInProgress,
+			showToolCallInProgress
+		});
 	}
 
 	function toggleExpanded(index: number, section: AgenticSection) {
@@ -190,6 +188,7 @@
 	{:else if section.type === AgenticSectionType.TOOL_CALL || section.type === AgenticSectionType.TOOL_CALL_PENDING || section.type === AgenticSectionType.TOOL_CALL_STREAMING}
 		<ChatMessageToolCallBlock
 			attachments={section.toolResultExtras}
+			{conversationCwd}
 			isExecuting={section.toolCallId !== undefined &&
 				section.toolCallId === currentlyExecutingToolCallId}
 			{isStreaming}

tools/ui/src/lib/components/app/content/CollapsibleTerminalBlock.svelte

@@ -98,9 +98,17 @@
 		<!-- See CollapsibleContentBlock: bits-ui keeps closed content mounted, which
 		     makes a collapsed tool result re-render on every streamed token. -->
 		{#if open}
-			<div class="p-3 pt-1">
+			<div class="terminal-body p-3 pt-1">
 				{@render children()}
 			</div>
 		{/if}
 	</Collapsible.Content>
 </Collapsible.Root>
+
+<style>
+	/* A body that renders no element (exec block while its args stream)
+	   must not leave its padding as a gap under the title. */
+	.terminal-body:not(:has(*)) {
+		padding: 0;
+	}
+</style>

tools/ui/src/lib/constants/settings-keys.constants.ts

@@ -65,6 +65,7 @@ export const SETTINGS_KEYS = {
 	SHOW_RAW_OUTPUT_SWITCH: 'showRawOutputSwitch',
 	SHOW_SYSTEM_MESSAGE: 'showSystemMessage',
 	SHOW_THOUGHT_IN_PROGRESS: 'showThoughtInProgress',
+	SHOW_TOOL_CALL_IN_PROGRESS: 'showToolCallInProgress',
 	SYMBOLIC_MATH_ENABLED: 'symbolicMathEnabled',
 	SYSTEM_MESSAGE: 'systemMessage',
 	// Sampling

tools/ui/src/lib/constants/settings.constants.ts

@@ -227,9 +227,16 @@ export const SETTINGS_REGISTRY: SettingsSectionEntry[] = [
 				label: 'Show thought in progress',
 				type: SettingsFieldType.CHECKBOX
 			},
+			{
+				defaultValue: true,
+				help: 'Expand tool call details while the call is being generated and executed, and collapse them again once its result is in. Mirrors "Show thought in progress". Tool calls you expand or collapse by hand keep your choice.',
+				key: SETTINGS_KEYS.SHOW_TOOL_CALL_IN_PROGRESS,
+				label: 'Show tool call in progress',
+				type: SettingsFieldType.CHECKBOX
+			},
 			{
 				defaultValue: false,
-				help: 'Automatically expand tool call details while executing and keep them expanded after completion.',
+				help: 'Keep tool call details expanded after completion. Takes precedence over "Show tool call in progress", which only covers the in-flight phases.',
 				key: SETTINGS_KEYS.ALWAYS_SHOW_TOOL_CALL_CONTENT,
 				label: 'Always show tool call content',
 				type: SettingsFieldType.CHECKBOX

tools/ui/src/lib/stores/models/index.svelte.ts

@@ -578,7 +578,6 @@ class ModelsStore implements ModelPropsHost, ModelStatusHost {
 					const displayNameSource =
 						details?.name && details.name.trim().length > 0 ? details.name : item.id;
 					const modelId = details?.model || item.id;
-					const repo = repoOf(modelId);
 
 					return {
 						aliases: item.aliases ?? [],
@@ -594,7 +593,7 @@ class ModelsStore implements ModelPropsHost, ModelStatusHost {
 								: undefined),
 						description: details?.description,
 						details: details?.details,
-						draftSidecars: this._status.getDraftSidecars(repo),
+						draftSidecars: this._status.getDraftSidecars(modelId),
 						id: item.id,
 						meta: item.meta ?? null,
 						modalities: this.props.buildArchitectureModalities(item.architecture),

tools/ui/src/lib/stores/models/status.svelte.ts

@@ -97,10 +97,11 @@ export class ModelStatusManager {
 	/** Sidecar kinds each in-flight or paused download pulls with its model. */
 	private downloadSidecars = new SvelteMap<string, ModelSidecar[]>();
 	/**
-	 * Draft sidecars the registered models pull with them, keyed by the repo of the model
-	 * that pulls them: the router records them in the model's `--model-draft` args.
+	 * Draft sidecars the registered models pull with them, keyed by the model that pulls
+	 * them: the router records them in the model's `--model-draft` args, and each quant
+	 * of a repo can pull a different draft.
 	 */
-	private draftSidecarsByRepo = $derived.by(() => {
+	private draftSidecarsByModel = $derived.by(() => {
 		const result = new SvelteMap<string, ModelSidecarBadge[]>();
 
 		for (const m of this.host.routerModels) {
@@ -108,8 +109,6 @@ export class ModelStatusManager {
 
 			if (!args) continue;
 
-			const baseRepo = repoOf(m.id);
-
 			for (let i = 0; i < args.length - 1; i++) {
 				if (args[i] !== CLI_FLAGS.MODEL_DRAFT && args[i] !== CLI_FLAGS.MODEL_DRAFT_SHORT) {
 					continue;
@@ -128,7 +127,7 @@ export class ModelStatusManager {
 					quant: meta.quant,
 					repo: parsed.repo
 				};
-				const kinds = result.get(baseRepo);
+				const kinds = result.get(m.id);
 
 				if (kinds) {
 					if (
@@ -139,7 +138,7 @@ export class ModelStatusManager {
 						kinds.push(badge);
 					}
 				} else {
-					result.set(baseRepo, [badge]);
+					result.set(m.id, [badge]);
 				}
 			}
 		}
@@ -404,17 +403,19 @@ export class ModelStatusManager {
 	}
 
 	/**
-	 * Draft sidecars a repo pulls with its model, e.g. `mtp` at `Q8_0`: the
-	 * router records loaded models' sidecars in their `--model-draft` args, and
-	 * lists the sidecars each cached model has on disk.
+	 * Draft sidecars a model pulls with it, e.g. `mtp` at `Q8_0`: the router
+	 * records each model's sidecars in its `--model-draft` args, and lists the
+	 * sidecars each cached repo has on disk.
 	 */
-	getDraftSidecars(repoId: string): ModelSidecarBadge[] {
-		// one badge per kind: the args know the quant, the listing carries kinds only,
-		// and a draft still coming down is pending rather than on disk
-		const merged = [...(this.draftSidecarsByRepo.get(repoId) ?? [])];
+	getDraftSidecars(modelId: string): ModelSidecarBadge[] {
+		const repo = repoOf(modelId);
+		// one badge per kind: the args know the quant of this model's draft, the
+		// listing carries repo-wide kinds only, and a draft still coming down is
+		// pending rather than on disk
+		const merged = [...(this.draftSidecarsByModel.get(modelId) ?? [])];
 
 		for (const source of [this.listedSidecarsByRepo, this.pendingSidecarsByRepo]) {
-			for (const badge of source.get(repoId) ?? []) {
+			for (const badge of source.get(repo) ?? []) {
 				if (!merged.some((b) => b.kind === badge.kind && b.repo === badge.repo)) {
 					merged.push(badge);
 				}

tools/ui/src/lib/types/tools.d.ts

@@ -33,10 +33,12 @@ export interface ToolGroup {
 }
 
 export interface WriteFileMeta {
-	fileName: string;
-	filePath: string;
+	/** Undefined until the call streams a path; content renders without it. */
+	fileName?: string;
+	filePath?: string;
 	language: string;
-	content: string;
+	/** Undefined until the args carry a content string; an empty write is ''. */
+	content?: string;
 	bytesWritten?: number;
 	resultMessage?: string;
 	errorMessage?: string;
@@ -46,8 +48,8 @@ export interface WriteFileMeta {
  *  ( with the embedded file content ) stays body-only so collapsed blocks
  *  never parse the content blob. */
 export interface WriteFileTitleMeta {
-	fileName: string;
-	filePath: string;
+	fileName?: string;
+	filePath?: string;
 	language: string;
 	bytesWritten?: number;
 	resultMessage?: string;
@@ -60,8 +62,9 @@ export interface EditFileEdit {
 }
 
 export interface EditFileMeta {
-	fileName: string;
-	filePath: string;
+	/** Undefined until the call streams a path; the edits render without it. */
+	fileName?: string;
+	filePath?: string;
 	edits: EditFileEdit[];
 	resultMessage?: string;
 	editsApplied?: number;
@@ -72,8 +75,8 @@ export interface EditFileMeta {
  *  ( with the embedded edit strings ) stays body-only so collapsed blocks
  *  never parse the args blob. */
 export interface EditFileTitleMeta {
-	fileName: string;
-	filePath: string;
+	fileName?: string;
+	filePath?: string;
 	resultMessage?: string;
 	editsApplied?: number;
 	errorMessage?: string;

tools/ui/src/lib/utils/agentic.ts

@@ -540,6 +540,49 @@ export function hasAgenticContent(
 	return toolMessages.length > 0;
 }
 
+/**
+ * Default expansion of a section, used until the user toggles it by hand.
+ *
+ * @param section - The section to judge
+ * @param flags - Display settings plus the tool call the agentic loop is running now.
+ *   A running call outranks its section type: `exec_shell_command` persists an empty
+ *   result before reading output, so a live call can already be TOOL_CALL.
+ */
+export function isSectionExpandedByDefault(
+	section: AgenticSection,
+	{
+		alwaysShowToolCallContent,
+		executingToolCallId,
+		showThoughtInProgress,
+		showToolCallInProgress
+	}: {
+		alwaysShowToolCallContent: boolean;
+		executingToolCallId?: string | null;
+		showThoughtInProgress: boolean;
+		showToolCallInProgress: boolean;
+	}
+): boolean {
+	const isRunning = section.toolCallId !== undefined && section.toolCallId === executingToolCallId;
+
+	if (
+		isRunning ||
+		section.type === AgenticSectionType.TOOL_CALL_PENDING ||
+		section.type === AgenticSectionType.TOOL_CALL_STREAMING
+	) {
+		return showToolCallInProgress || alwaysShowToolCallContent;
+	}
+
+	if (section.type === AgenticSectionType.TOOL_CALL) {
+		return alwaysShowToolCallContent;
+	}
+
+	if (section.type === AgenticSectionType.REASONING_PENDING) {
+		return showThoughtInProgress;
+	}
+
+	return false;
+}
+
 /**
  * Decide how a Continue click on messages[idx] should resume generation.
  * Pure function over the persisted history snapshot.

tools/ui/src/lib/utils/index.ts

@@ -282,6 +282,7 @@ export {
 	parseToolResultWithMedia,
 	splitSearchSummaryList,
 	hasAgenticContent,
+	isSectionExpandedByDefault,
 	classifyToolResult,
 	classifyContinueIntent
 } from './agentic';

tools/ui/tests/client/exec-shell-command-streaming.svelte.test.ts

@@ -0,0 +1,216 @@
+// Guards the exec_shell_command title while the call is still being
+// generated: the command must stream into the terminal title from the
+// first character (args arrive as partial JSON), with a blinking cursor
+// while the args are incomplete, and the conversation cwd standing in
+// for the x-tool-cwd value that only arrives with the result message.
+
+import ChatMessageToolCallBlockExecShellCommand from '$lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/ChatMessageToolCallBlockExecShellCommand.svelte';
+import ChatMessageAgenticContent from '$lib/components/app/chat/ChatMessages/ChatMessageAgenticContent.svelte';
+import { SETTINGS_KEYS } from '$lib/constants';
+import { AgenticSectionType, BuiltInTool, MessageRole, MessageType } from '$lib/enums';
+import { conversationsStore, settingsStore } from '$lib/stores';
+import type { AgenticSection, DatabaseMessage } from '$lib/types';
+import { tick } from 'svelte';
+import { afterEach, beforeEach, describe, expect, it } from 'vitest';
+import { render } from 'vitest-browser-svelte';
+
+const CONV_ID = 'conv-1';
+const TOOL_CALL_ID = 'call-1';
+const TOOL_NAME = BuiltInTool.SERVER_EXEC_SHELL_COMMAND;
+const CWD = '/home/user/project';
+
+function messageWithArgs(args: string): DatabaseMessage {
+	return {
+		children: [],
+		content: '',
+		convId: CONV_ID,
+		id: 'assistant-1',
+		parent: null,
+		role: MessageRole.ASSISTANT,
+		timestamp: 0,
+		toolCalls: JSON.stringify([
+			{ function: { arguments: args, name: TOOL_NAME }, id: TOOL_CALL_ID }
+		]),
+		type: MessageType.TEXT
+	};
+}
+
+function toolResultMessage(content: string, toolCwd?: string): DatabaseMessage {
+	return {
+		children: [],
+		content,
+		convId: CONV_ID,
+		id: 'tool-1',
+		parent: null,
+		role: MessageRole.TOOL,
+		timestamp: 0,
+		toolCallId: TOOL_CALL_ID,
+		toolCalls: '',
+		toolCwd,
+		type: MessageType.TEXT
+	};
+}
+
+function triggerText(container: Element): string {
+	return container.querySelector('[data-slot="collapsible-trigger"]')?.textContent ?? '';
+}
+
+function hasCursor(container: Element): boolean {
+	return container.querySelector('.exec-cursor') !== null;
+}
+
+describe('exec_shell_command streaming title', () => {
+	beforeEach(() => {
+		settingsStore.updateMultipleConfig({
+			[SETTINGS_KEYS.ALWAYS_SHOW_TOOL_CALL_CONTENT]: false,
+			[SETTINGS_KEYS.SHOW_TOOL_CALL_IN_PROGRESS]: true
+		});
+		conversationsStore.activeConversation = {
+			currNode: null,
+			cwd: CWD,
+			id: CONV_ID,
+			lastModified: 0,
+			name: 'test'
+		};
+	});
+
+	afterEach(() => {
+		conversationsStore.activeConversation = null;
+		conversationsStore.activeMessages = [];
+	});
+
+	it('streams the partial command into the title with a cursor', async () => {
+		const { container } = render(ChatMessageAgenticContent, {
+			isStreaming: true,
+			message: messageWithArgs('{"command":"git sta'),
+			toolMessages: []
+		});
+
+		await tick();
+
+		const title = triggerText(container);
+
+		expect(title).toContain('git sta');
+		expect(title).toContain('$');
+		expect(title).toContain(CWD);
+		expect(hasCursor(container)).toBe(true);
+	});
+
+	it('drops the cursor once the args JSON is complete', async () => {
+		const screen = render(ChatMessageAgenticContent, {
+			isStreaming: true,
+			message: messageWithArgs('{"command":"git sta'),
+			toolMessages: []
+		});
+
+		await tick();
+
+		expect(hasCursor(screen.container)).toBe(true);
+
+		screen.rerender({
+			isStreaming: true,
+			message: messageWithArgs('{"command":"git status"}'),
+			toolMessages: []
+		});
+		await tick();
+
+		expect(triggerText(screen.container)).toContain('git status');
+		expect(hasCursor(screen.container)).toBe(false);
+	});
+
+	it('switches the prompt from the conversation cwd to the persisted one', async () => {
+		const screen = render(ChatMessageAgenticContent, {
+			isStreaming: true,
+			message: messageWithArgs('{"command":"git status"}'),
+			toolMessages: []
+		});
+
+		await tick();
+
+		expect(triggerText(screen.container)).toContain(CWD);
+
+		// Result message lands: its x-tool-cwd takes over the prompt.
+		screen.rerender({
+			isStreaming: true,
+			message: messageWithArgs('{"command":"git status"}'),
+			toolMessages: [toolResultMessage('', '/other/dir')]
+		});
+		await tick();
+
+		expect(triggerText(screen.container)).toContain('/other/dir');
+		expect(triggerText(screen.container)).not.toContain(CWD);
+	});
+
+	it('shows the executing label only after the command finishes streaming', async () => {
+		const screen = render(ChatMessageAgenticContent, {
+			isStreaming: true,
+			message: messageWithArgs('{"command":"git sta'),
+			toolMessages: []
+		});
+
+		await tick();
+
+		// While the args stream, the title cursor alone signals activity.
+		expect(screen.container.textContent).not.toContain('Executing command...');
+
+		screen.rerender({
+			isStreaming: true,
+			message: messageWithArgs('{"command":"git status"}'),
+			toolMessages: []
+		});
+		await tick();
+
+		expect(screen.container.textContent).toContain('Executing command...');
+	});
+
+	it('shows Executing command... while executing before the first output chunk', async () => {
+		// The result message is created empty right before execution, so
+		// the block must not go blank while the process produces nothing.
+		const section: AgenticSection = {
+			content: '',
+			toolArgs: '{"command":"sleep 1"}',
+			toolCallId: TOOL_CALL_ID,
+			toolName: TOOL_NAME,
+			toolResult: '',
+			type: AgenticSectionType.TOOL_CALL
+		};
+		const { container } = render(ChatMessageToolCallBlockExecShellCommand, {
+			isExecuting: true,
+			isStreaming: true,
+			open: true,
+			section
+		});
+
+		await tick();
+
+		expect(container.textContent).toContain('Executing command...');
+	});
+
+	it('updates the title when the store mutates toolCalls in place, as the stream does', async () => {
+		// The live flow never swaps the message object: onToolCallsStreaming
+		// writes growing toolCalls strings into the same store row per delta.
+		conversationsStore.activeMessages = [messageWithArgs('')];
+
+		const screen = render(ChatMessageAgenticContent, {
+			isStreaming: true,
+			message: conversationsStore.activeMessages[0]!,
+			toolMessages: []
+		});
+
+		await tick();
+
+		const idx = conversationsStore.findMessageIndex('assistant-1');
+
+		for (const args of ['{"command', '{"command":"git', '{"command":"git sta']) {
+			conversationsStore.updateMessageAtIndex(idx, {
+				toolCalls: JSON.stringify([
+					{ function: { arguments: args, name: TOOL_NAME }, id: TOOL_CALL_ID }
+				])
+			});
+			await tick();
+
+			expect(triggerText(screen.container)).toContain(args.slice('{"command":"'.length));
+			expect(hasCursor(screen.container)).toBe(true);
+		}
+	});
+});

tools/ui/tests/client/models-sidecar-badges.svelte.test.ts

@@ -1,6 +1,7 @@
-// Guards the draft badges a repo row shows: a loaded model records its draft in
-// the `--model-draft` args, which know the file's quant, and the listing reports
-// the kinds of every cached model, which do not. The row shows each kind once.
+// Guards the draft badges a model row shows: a registered model records its draft
+// in its `--model-draft` args, which know the file's quant, and the listing reports
+// the kinds every cached model of the repo has on disk, which do not. The row shows
+// each kind once, and two quants of one repo keep their own drafts apart.
 
 import { ModelAuxSidecar, ModelDraftSidecar, ServerModelStatus } from '$lib/enums';
 import { modelsStore } from '$lib/stores';
@@ -9,6 +10,11 @@ import { beforeEach, describe, expect, it } from 'vitest';
 
 const REPO = 'ggml-org/Qwen3.8-Flash-Next-GGUF';
 
+/** Draft file of a cached model, as a preset records it in `--model-draft`. */
+function draftPath(quant: string): string {
+	return `/cache/models--ggml-org--Qwen3.8-Flash-Next-GGUF/snapshots/${'a'.repeat(40)}/mtp-Qwen3.8-Flash-Next-${quant}.gguf`;
+}
+
 /** A router entry; `args` stand in for the args of a loaded instance. */
 function entry(
 	id: string,
@@ -36,18 +42,28 @@ describe('draft sidecar badges', () => {
 	it('shows one badge for a draft the args and the listing both report', () => {
 		modelsStore.routerModels = [
 			entry(
-				`${REPO}:IQ4_NL`,
-				[],
-				[
-					'--model-draft',
-					`/cache/models--ggml-org--Qwen3.8-Flash-Next-GGUF/snapshots/${'a'.repeat(40)}/mtp-Qwen3.8-Flash-Next-Q8_0.gguf`
-				]
-			),
-			entry(`${REPO}:Q8_0`, [ModelAuxSidecar.MMPROJ, ModelDraftSidecar.MTP])
+				`${REPO}:Q8_0`,
+				[ModelAuxSidecar.MMPROJ, ModelDraftSidecar.MTP],
+				['--model-draft', draftPath('Q8_0')]
+			)
 		];
 
 		// the args badge wins: it knows the quant the listing does not
-		expect(modelsStore.status.getDraftSidecars(REPO)).toEqual([
+		expect(modelsStore.status.getDraftSidecars(`${REPO}:Q8_0`)).toEqual([
+			{ kind: ModelDraftSidecar.MTP, quant: 'Q8_0', repo: REPO }
+		]);
+	});
+
+	it('gives each quant of a repo the draft it pulls itself', () => {
+		modelsStore.routerModels = [
+			entry(`${REPO}:IQ4_NL`, [], ['--model-draft', draftPath('Q4_0')]),
+			entry(`${REPO}:Q8_0`, [], ['--model-draft', draftPath('Q8_0')])
+		];
+
+		expect(modelsStore.status.getDraftSidecars(`${REPO}:IQ4_NL`)).toEqual([
+			{ kind: ModelDraftSidecar.MTP, quant: 'Q4_0', repo: REPO }
+		]);
+		expect(modelsStore.status.getDraftSidecars(`${REPO}:Q8_0`)).toEqual([
 			{ kind: ModelDraftSidecar.MTP, quant: 'Q8_0', repo: REPO }
 		]);
 	});
@@ -58,7 +74,7 @@ describe('draft sidecar badges', () => {
 			entry(`${REPO}:Q8_0`, [ModelDraftSidecar.MTP])
 		];
 
-		expect(modelsStore.status.getDraftSidecars(REPO)).toEqual([
+		expect(modelsStore.status.getDraftSidecars(`${REPO}:IQ4_NL`)).toEqual([
 			{ kind: ModelDraftSidecar.MTP, quant: null, repo: REPO }
 		]);
 	});
@@ -74,7 +90,7 @@ describe('draft sidecar badges', () => {
 		];
 
 		// the kinds are not on disk yet, so the row reads them as in flight
-		expect(modelsStore.status.getDraftSidecars(REPO)).toEqual([
+		expect(modelsStore.status.getDraftSidecars(`${REPO}:Q4_0`)).toEqual([
 			{ kind: ModelDraftSidecar.MTP, pending: true, quant: null, repo: REPO }
 		]);
 	});
@@ -82,17 +98,10 @@ describe('draft sidecar badges', () => {
 	it('keeps the on-disk badge of a kind a download also pulls', () => {
 		modelsStore.routerModels = [
 			entry(`${REPO}:IQ4_NL`, [ModelDraftSidecar.MTP], [], ServerModelStatus.INCOMPLETE),
-			entry(
-				`${REPO}:Q8_0`,
-				[],
-				[
-					'--model-draft',
-					`/cache/models--ggml-org--Qwen3.8-Flash-Next-GGUF/snapshots/${'a'.repeat(40)}/mtp-Qwen3.8-Flash-Next-Q8_0.gguf`
-				]
-			)
+			entry(`${REPO}:Q8_0`, [], ['--model-draft', draftPath('Q8_0')])
 		];
 
-		expect(modelsStore.status.getDraftSidecars(REPO)).toEqual([
+		expect(modelsStore.status.getDraftSidecars(`${REPO}:Q8_0`)).toEqual([
 			{ kind: ModelDraftSidecar.MTP, quant: 'Q8_0', repo: REPO }
 		]);
 	});

tools/ui/tests/client/tool-call-block-streamed-body.svelte.test.ts

@@ -0,0 +1,82 @@
+// A file tool call must show what it streamed even when it carries no
+// path: the path only names the file in the block title, so a call that
+// streams its body first (or a same-named tool with other arg names) used
+// to render an empty block.
+
+import ChatMessageToolCallBlockEditFile from '$lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/ChatMessageToolCallBlockEditFile.svelte';
+import ChatMessageToolCallBlockWriteFile from '$lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/ChatMessageToolCallBlockWriteFile.svelte';
+import { AgenticSectionType, BuiltInTool } from '$lib/enums';
+import type { AgenticSection } from '$lib/types';
+import { tick } from 'svelte';
+import { describe, expect, it } from 'vitest';
+import { render } from 'vitest-browser-svelte';
+
+const WRITE_FILE_BODY = 'body streamed without a path';
+const EDIT_FILE_NEW_TEXT = 'replacement text';
+
+function makeSection(toolName: BuiltInTool, toolArgs: string): AgenticSection {
+	return {
+		content: '',
+		toolArgs,
+		toolName,
+		toolResult: undefined,
+		type: AgenticSectionType.TOOL_CALL
+	};
+}
+
+describe('file tool calls render their body without a path', () => {
+	it('write_file shows the content it streamed', async () => {
+		render(ChatMessageToolCallBlockWriteFile, {
+			isStreaming: true,
+			open: true,
+			section: makeSection(BuiltInTool.SERVER_WRITE_FILE, `{"content":"${WRITE_FILE_BODY}"`)
+		});
+
+		await tick();
+
+		expect(document.body.textContent).toContain(WRITE_FILE_BODY);
+	});
+
+	it('edit_file shows the diff it streamed', async () => {
+		render(ChatMessageToolCallBlockEditFile, {
+			isStreaming: true,
+			open: true,
+			section: makeSection(
+				BuiltInTool.SERVER_EDIT_FILE,
+				`{"edits":[{"old_text":"stale text","new_text":"${EDIT_FILE_NEW_TEXT}"}]}`
+			)
+		});
+
+		await tick();
+
+		expect(document.body.textContent).toContain(EDIT_FILE_NEW_TEXT);
+		expect(document.body.textContent).toContain('stale text');
+	});
+
+	it('write_file falls back to the raw args when no known content key is present', async () => {
+		render(ChatMessageToolCallBlockWriteFile, {
+			isStreaming: false,
+			open: true,
+			section: makeSection(BuiltInTool.SERVER_WRITE_FILE, '{"contents":"x"}')
+		});
+
+		await tick();
+
+		expect(document.body.textContent).toContain('contents');
+	});
+
+	it('write_file renders an empty code block for a finished empty write', async () => {
+		// content of '' used to fall through to the raw-args fallback
+		const section = makeSection(BuiltInTool.SERVER_WRITE_FILE, '{"content":""}');
+
+		section.toolResult = '{"result":"File written","bytes":0}';
+
+		render(ChatMessageToolCallBlockWriteFile, { isStreaming: false, open: true, section });
+
+		await tick();
+
+		expect(document.querySelector('pre code.hljs')).toBeTruthy();
+		expect(document.body.textContent).not.toContain('"content"');
+		expect(document.body.textContent).toMatch(/0\s*bytes/);
+	});
+});

tools/ui/tests/client/tool-call-in-progress.svelte.test.ts

@@ -0,0 +1,212 @@
+// Guards the two Settings > Display toggles that drive the default expansion of
+// a tool call block.
+//
+// "Show tool call in progress" mirrors "Show thought in progress": a call is
+// expanded by default only while it is in flight and folds back on its own once
+// the result lands. "Always show tool call content" is the one that keeps
+// finished calls open and wins on its own. A manual toggle always overrides the
+// default, in both directions.
+
+import ChatMessageAgenticContent from '$lib/components/app/chat/ChatMessages/ChatMessageAgenticContent.svelte';
+import { SETTINGS_KEYS } from '$lib/constants';
+import { MessageRole, MessageType } from '$lib/enums';
+import { settingsStore } from '$lib/stores';
+import type { DatabaseMessage } from '$lib/types';
+import { tick } from 'svelte';
+import { beforeEach, describe, expect, it } from 'vitest';
+import { userEvent } from 'vitest/browser';
+import { render } from 'vitest-browser-svelte';
+
+const TOOL_NAME = 'custom_tool';
+const TOOL_CALL_ID = 'call-1';
+// Both markers sit inside the collapsible body, which CollapsibleContentBlock
+// only mounts while open - so their presence is exactly "the block is expanded".
+const PENDING_BODY = 'Waiting for result...';
+const RESULT_BODY = 'marker-tool-output';
+const TRIGGER_SELECTOR = '[data-slot="collapsible-trigger"]';
+
+function assistantMessage(): DatabaseMessage {
+	return {
+		children: [],
+		content: '',
+		convId: 'conv-1',
+		id: 'assistant-1',
+		parent: null,
+		role: MessageRole.ASSISTANT,
+		timestamp: 0,
+		toolCalls: JSON.stringify([
+			{ function: { arguments: '{"path":"a.txt"}', name: TOOL_NAME }, id: TOOL_CALL_ID }
+		]),
+		type: MessageType.TEXT
+	};
+}
+
+function toolResultMessage(): DatabaseMessage {
+	return {
+		...assistantMessage(),
+		content: RESULT_BODY,
+		id: 'tool-1',
+		role: MessageRole.TOOL,
+		toolCallId: TOOL_CALL_ID
+	};
+}
+
+function setSettings({
+	alwaysShow,
+	inProgress
+}: {
+	alwaysShow: boolean;
+	inProgress: boolean;
+}): void {
+	settingsStore.updateMultipleConfig({
+		[SETTINGS_KEYS.ALWAYS_SHOW_TOOL_CALL_CONTENT]: alwaysShow,
+		[SETTINGS_KEYS.SHOW_TOOL_CALL_IN_PROGRESS]: inProgress
+	});
+}
+
+function textOf(container: Element): string {
+	return container.textContent ?? '';
+}
+
+async function collapseBlock(screen: { container: Element }): Promise<void> {
+	const trigger = screen.container.querySelector(TRIGGER_SELECTOR);
+
+	if (!(trigger instanceof HTMLButtonElement)) throw new Error('trigger not found');
+
+	await userEvent.click(trigger);
+	await tick();
+}
+
+describe('Show tool call in progress', () => {
+	beforeEach(() => {
+		setSettings({ alwaysShow: false, inProgress: false });
+	});
+
+	it('expands a call that is still executing when the setting is on', async () => {
+		setSettings({ alwaysShow: false, inProgress: true });
+
+		const { container } = render(ChatMessageAgenticContent, {
+			isStreaming: true,
+			message: assistantMessage(),
+			toolMessages: []
+		});
+
+		await tick();
+
+		expect(textOf(container)).toContain(PENDING_BODY);
+	});
+
+	it('keeps the call collapsed while executing when the setting is off', async () => {
+		const { container } = render(ChatMessageAgenticContent, {
+			isStreaming: true,
+			message: assistantMessage(),
+			toolMessages: []
+		});
+
+		await tick();
+
+		// The header is still there, only the body stays hidden.
+		expect(textOf(container)).toContain(TOOL_NAME);
+		expect(textOf(container)).not.toContain(PENDING_BODY);
+	});
+
+	it('folds the call back when its result arrives', async () => {
+		setSettings({ alwaysShow: false, inProgress: true });
+
+		const screen = render(ChatMessageAgenticContent, {
+			isStreaming: true,
+			message: assistantMessage(),
+			toolMessages: []
+		});
+
+		await tick();
+
+		// Open while it waits for the result.
+		expect(textOf(screen.container)).toContain(PENDING_BODY);
+
+		// Same block, now finished: it must close on its own, no toggle given.
+		screen.rerender({
+			isStreaming: true,
+			message: assistantMessage(),
+			toolMessages: [toolResultMessage()]
+		});
+		await tick();
+
+		expect(textOf(screen.container)).toContain(TOOL_NAME);
+		expect(textOf(screen.container)).not.toContain(PENDING_BODY);
+		expect(textOf(screen.container)).not.toContain(RESULT_BODY);
+	});
+
+	it('keeps the manual choice once the result arrives', async () => {
+		setSettings({ alwaysShow: false, inProgress: true });
+
+		const screen = render(ChatMessageAgenticContent, {
+			isStreaming: true,
+			message: assistantMessage(),
+			toolMessages: []
+		});
+
+		await tick();
+
+		// Collapsed by hand while running, so the result must not reopen it.
+		await collapseBlock(screen);
+
+		screen.rerender({
+			isStreaming: true,
+			message: assistantMessage(),
+			toolMessages: [toolResultMessage()]
+		});
+		await tick();
+
+		expect(textOf(screen.container)).toContain(TOOL_NAME);
+		expect(textOf(screen.container)).not.toContain(RESULT_BODY);
+	});
+
+	it('still expands a finished call when "always show" is on', async () => {
+		setSettings({ alwaysShow: true, inProgress: false });
+
+		const { container } = render(ChatMessageAgenticContent, {
+			isStreaming: true,
+			message: assistantMessage(),
+			toolMessages: [toolResultMessage()]
+		});
+
+		await tick();
+
+		expect(textOf(container)).toContain(RESULT_BODY);
+	});
+
+	it('lets a manual collapse win over the in-progress default', async () => {
+		setSettings({ alwaysShow: false, inProgress: true });
+
+		const screen = render(ChatMessageAgenticContent, {
+			isStreaming: true,
+			message: assistantMessage(),
+			toolMessages: []
+		});
+
+		await tick();
+
+		expect(textOf(screen.container)).toContain(PENDING_BODY);
+
+		await collapseBlock(screen);
+
+		expect(textOf(screen.container)).not.toContain(PENDING_BODY);
+	});
+
+	it('lets a manual expand win over the collapsed default', async () => {
+		const screen = render(ChatMessageAgenticContent, {
+			isStreaming: true,
+			message: assistantMessage(),
+			toolMessages: []
+		});
+
+		await tick();
+
+		expect(textOf(screen.container)).not.toContain(PENDING_BODY);
+
+		await collapseBlock(screen);
+
+		expect(textOf(screen.container)).toContain(PENDING_BODY);
+	});
+});

tools/ui/tests/unit/agentic-section-expansion.test.ts

@@ -0,0 +1,122 @@
+// Truth table for the default expansion of agentic sections - what a tool call
+// block does before the user touches it.
+//
+// The interesting row is a call the agentic loop reports as running while its
+// section is already TOOL_CALL: exec_shell_command persists an empty result
+// before any output is read, so the section type alone cannot tell a live call
+// from a finished one.
+
+import { AgenticSectionType } from '$lib/enums';
+import type { AgenticSection } from '$lib/types';
+import { isSectionExpandedByDefault } from '$lib/utils';
+import { describe, expect, it } from 'vitest';
+
+const CALL_ID = 'call-1';
+
+function section(overrides: Partial<AgenticSection> = {}): AgenticSection {
+	return { content: '', type: AgenticSectionType.TOOL_CALL, ...overrides };
+}
+
+function expanded(section: AgenticSection, flags: Record<string, unknown> = {}): boolean {
+	return isSectionExpandedByDefault(section, {
+		alwaysShowToolCallContent: false,
+		executingToolCallId: null,
+		showThoughtInProgress: false,
+		showToolCallInProgress: false,
+		...flags
+	});
+}
+
+describe('isSectionExpandedByDefault', () => {
+	it('expands a call whose arguments are streaming when the setting is on', () => {
+		const streaming = section({
+			toolCallId: CALL_ID,
+			toolName: 'read_file',
+			type: AgenticSectionType.TOOL_CALL_STREAMING
+		});
+
+		expect(expanded(streaming, { showToolCallInProgress: true })).toBe(true);
+		expect(expanded(streaming)).toBe(false);
+	});
+
+	it('expands a call waiting for its result when the setting is on', () => {
+		const pending = section({ toolCallId: CALL_ID, type: AgenticSectionType.TOOL_CALL_PENDING });
+
+		expect(expanded(pending, { showToolCallInProgress: true })).toBe(true);
+		expect(expanded(pending)).toBe(false);
+	});
+
+	it('collapses a finished call unless "always show" is on', () => {
+		const finished = section({ toolCallId: CALL_ID, toolResult: 'output' });
+
+		expect(expanded(finished, { showToolCallInProgress: true })).toBe(false);
+		expect(expanded(finished, { alwaysShowToolCallContent: true })).toBe(true);
+	});
+
+	it('keeps the running call expanded even when its result is already persisted', () => {
+		// exec_shell_command: empty result created, output still streaming
+		const live = section({ toolCallId: CALL_ID, toolResult: '' });
+
+		expect(expanded(live, { executingToolCallId: CALL_ID, showToolCallInProgress: true })).toBe(
+			true
+		);
+	});
+
+	it('does not treat another call finishing as the running one', () => {
+		const other = section({ toolCallId: 'call-2', toolResult: 'output' });
+
+		expect(expanded(other, { executingToolCallId: CALL_ID, showToolCallInProgress: true })).toBe(
+			false
+		);
+	});
+
+	it('still leaves the running call collapsed when both settings are off', () => {
+		const live = section({ toolCallId: CALL_ID, toolResult: '' });
+
+		expect(expanded(live, { executingToolCallId: CALL_ID })).toBe(false);
+	});
+
+	it('never reads a missing executing id as a running call', () => {
+		expect(expanded(section({ toolCallId: CALL_ID }), { executingToolCallId: null })).toBe(false);
+		expect(expanded(section({ type: AgenticSectionType.TEXT }), { executingToolCallId: '' })).toBe(
+			false
+		);
+	});
+
+	it('expands a running call when "always show" is on and the in-progress one is off', () => {
+		const pending = section({ toolCallId: CALL_ID, type: AgenticSectionType.TOOL_CALL_PENDING });
+
+		expect(expanded(pending, { alwaysShowToolCallContent: true })).toBe(true);
+	});
+
+	it('follows the thought setting for reasoning and never for finished thoughts', () => {
+		expect(
+			expanded(
+				{ content: 'think', type: AgenticSectionType.REASONING_PENDING },
+				{ showThoughtInProgress: true }
+			)
+		).toBe(true);
+		expect(
+			expanded(
+				{ content: 'think', type: AgenticSectionType.REASONING },
+				{
+					showThoughtInProgress: true
+				}
+			)
+		).toBe(false);
+	});
+
+	it('never expands plain text', () => {
+		expect(
+			expanded(
+				{ content: 'hello', type: AgenticSectionType.TEXT },
+				{
+					alwaysShowToolCallContent: true,
+					executingToolCallId: CALL_ID,
+					showThoughtInProgress: true,
+					showToolCallInProgress: true
+				}
+			)
+		).toBe(false);
+	});
+});

tools/ui/tests/unit/settings-tool-call-in-progress.test.ts

@@ -0,0 +1,45 @@
+// Guards the Settings > Display wiring of "Show tool call in progress": the
+// flat config default and its presence as a checkbox in the Display section
+// come from the settings registry, so the key/label/type/section pairing is the
+// contract the dialog renders.
+
+import {
+	SETTING_CONFIG_DEFAULT,
+	SETTINGS_CHAT_SECTIONS,
+	SETTINGS_KEYS,
+	SETTINGS_SECTION_SLUGS
+} from '$lib/constants';
+import { SettingsFieldType } from '$lib/enums';
+import { describe, expect, it } from 'vitest';
+
+function displayField(key: string) {
+	const section = SETTINGS_CHAT_SECTIONS.find((s) => s.slug === SETTINGS_SECTION_SLUGS.DISPLAY);
+
+	return section?.fields?.find((field) => field?.key === key);
+}
+
+describe('showToolCallInProgress setting', () => {
+	it('is a checkbox in the Display section', () => {
+		const field = displayField(SETTINGS_KEYS.SHOW_TOOL_CALL_IN_PROGRESS);
+
+		expect(field).toBeDefined();
+		expect(field?.type).toBe(SettingsFieldType.CHECKBOX);
+		expect(field?.label).toBe('Show tool call in progress');
+	});
+
+	it('ships enabled, like "Show thought in progress"', () => {
+		expect(SETTING_CONFIG_DEFAULT[SETTINGS_KEYS.SHOW_TOOL_CALL_IN_PROGRESS]).toBe(true);
+		expect(SETTING_CONFIG_DEFAULT[SETTINGS_KEYS.SHOW_THOUGHT_IN_PROGRESS]).toBe(true);
+	});
+
+	it('is listed next to its "always show" counterpart', () => {
+		const section = SETTINGS_CHAT_SECTIONS.find(
+			(s) => s.slug === SETTINGS_SECTION_SLUGS.DISPLAY
+		) as (typeof SETTINGS_CHAT_SECTIONS)[number];
+		const keys = section.fields?.map((field) => field?.key) ?? [];
+		const inProgressIndex = keys.indexOf(SETTINGS_KEYS.SHOW_TOOL_CALL_IN_PROGRESS);
+
+		expect(inProgressIndex).toBeGreaterThanOrEqual(0);
+		expect(keys[inProgressIndex + 1]).toBe(SETTINGS_KEYS.ALWAYS_SHOW_TOOL_CALL_CONTENT);
+	});
+});

tools/ui/tests/unit/tool-calls.test.ts

@@ -1,4 +1,7 @@
-import { parseToolArgs } from '$lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/parsers/_shared';
+import {
+	extractToolArgPath,
+	parseToolArgs
+} from '$lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/parsers/_shared';
 import {
 	parseEditFileMeta,
 	parseEditFileTitleMeta
@@ -12,6 +15,7 @@ import {
 	parseWriteFileMeta,
 	parseWriteFileTitleMeta
 } from '$lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/parsers/write-file';
+import { CODE_BLOCK } from '$lib/constants';
 import { AgenticSectionType, BuiltInTool } from '$lib/enums';
 import type { AgenticSection, WriteFileMeta } from '$lib/types';
 import { abbreviateHome, formatCwdMessage, lastPathSegment, parseCwdMessage } from '$lib/utils';
@@ -164,6 +168,22 @@ describe('parseToolArgs (shared)', () => {
 	});
 });
 
+describe('extractToolArgPath', () => {
+	it('reads every path alias', () => {
+		expect(extractToolArgPath({ path: '/a' })).toBe('/a');
+		expect(extractToolArgPath({ file_path: '/b' })).toBe('/b');
+		expect(extractToolArgPath({ filePath: '/c' })).toBe('/c');
+	});
+
+	it('returns undefined when no path alias is a usable string', () => {
+		expect(extractToolArgPath({ content: 'x' })).toBeUndefined();
+		expect(extractToolArgPath({ path: '' })).toBeUndefined();
+		expect(extractToolArgPath({ path: 7 })).toBeUndefined();
+		expect(extractToolArgPath(null)).toBeUndefined();
+		expect(extractToolArgPath(undefined)).toBeUndefined();
+	});
+});
+
 describe('parseWriteFileMeta', () => {
 	it('returns null for sections with a different tool name', () => {
 		expect(
@@ -176,12 +196,23 @@ describe('parseWriteFileMeta', () => {
 		).toBeNull();
 	});
 
-	it('returns null when args have no path-like field', () => {
-		expect(
-			parseWriteFileMeta(
-				makeSection({ toolArgs: '{"content":"x"}', toolName: BuiltInTool.SERVER_WRITE_FILE })
-			)
-		).toBeNull();
+	it('returns content without a path when args have no path-like field', () => {
+		const meta = parseWriteFileMeta(
+			makeSection({ toolArgs: '{"content":"x"}', toolName: BuiltInTool.SERVER_WRITE_FILE })
+		);
+
+		expect(meta?.content).toBe('x');
+		expect(meta?.filePath).toBeUndefined();
+		expect(meta?.fileName).toBeUndefined();
+		expect(meta?.language).toBe(CODE_BLOCK.DEFAULT_LANGUAGE);
+	});
+
+	it('returns content streamed before the path has arrived', () => {
+		const meta = parseWriteFileMeta(
+			makeSection({ toolArgs: '{"content":"hello ', toolName: BuiltInTool.SERVER_WRITE_FILE })
+		);
+
+		expect(meta?.content).toBe('hello ');
 	});
 
 	it('accepts partial args (renders incrementally as content streams in)', () => {
@@ -279,6 +310,16 @@ describe('parseWriteFileTitleMeta', () => {
 		expect(parseWriteFileTitleMeta(section)?.filePath).toBe('/foo.t');
 	});
 
+	it('returns meta without a path when args have no path-like field', () => {
+		const title = parseWriteFileTitleMeta(
+			makeSection({ toolArgs: '{"content":"x"}', toolName: BuiltInTool.SERVER_WRITE_FILE })
+		);
+
+		expect(title?.filePath).toBeUndefined();
+		expect(title?.fileName).toBeUndefined();
+		expect(title?.language).toBe(CODE_BLOCK.DEFAULT_LANGUAGE);
+	});
+
 	it('returns null for sections with a different tool name', () => {
 		expect(
 			parseWriteFileTitleMeta(
@@ -324,12 +365,13 @@ describe('parseEditFileTitleMeta', () => {
 		expect(parseEditFileTitleMeta(section)?.errorMessage).toBe('permission denied');
 	});
 
-	it('returns null when args have no path-like field', () => {
-		expect(
-			parseEditFileTitleMeta(
-				makeSection({ toolArgs: '{"edits":[]}', toolName: BuiltInTool.SERVER_EDIT_FILE })
-			)
-		).toBeNull();
+	it('returns meta without a path when args have no path-like field', () => {
+		const title = parseEditFileTitleMeta(
+			makeSection({ toolArgs: '{"edits":[]}', toolName: BuiltInTool.SERVER_EDIT_FILE })
+		);
+
+		expect(title?.filePath).toBeUndefined();
+		expect(title?.fileName).toBeUndefined();
 	});
 });
 
@@ -383,6 +425,33 @@ describe('parseEditFileMeta', () => {
 		expect(meta?.errorMessage).toBe('bad path');
 		expect(meta?.resultMessage).toBeUndefined();
 	});
+
+	it('still parses edits when args have no path-like field', () => {
+		const meta = parseEditFileMeta(
+			makeSection({
+				toolArgs: '{"edits":[{"old_text":"a","new_text":"b"}]}',
+				toolName: BuiltInTool.SERVER_EDIT_FILE
+			})
+		);
+
+		expect(meta?.edits).toEqual([{ newText: 'b', oldText: 'a' }]);
+		expect(meta?.filePath).toBeUndefined();
+	});
+
+	it('parses edits streamed before the path has arrived', () => {
+		const meta = parseEditFileMeta(
+			makeSection({
+				// second edit is still mid-stream: only its old_text has landed
+				toolArgs: '{"edits":[{"old_text":"a","new_text":"b"},{"old_text":"c',
+				toolName: BuiltInTool.SERVER_EDIT_FILE
+			})
+		);
+
+		expect(meta?.edits).toEqual([
+			{ newText: 'b', oldText: 'a' },
+			{ newText: '', oldText: 'c' }
+		]);
+	});
 });
 
 describe('parseReadFileMeta', () => {
@@ -647,4 +716,37 @@ describe('parseExecShellCommandMeta', () => {
 			)
 		).toBeNull();
 	});
+
+	it('surfaces a partially streamed command before the JSON closes', () => {
+		const meta = parseExecShellCommandMeta(
+			makeSection(
+				{ toolArgs: '{"command":"git sta', toolName: BuiltInTool.SERVER_EXEC_SHELL_COMMAND },
+				BuiltInTool.SERVER_EXEC_SHELL_COMMAND
+			)
+		);
+
+		expect(meta?.command).toBe('git sta');
+	});
+
+	it('keeps a dangling escape from breaking the partial parse', () => {
+		const meta = parseExecShellCommandMeta(
+			makeSection(
+				{ toolArgs: '{"command":"echo \\"he', toolName: BuiltInTool.SERVER_EXEC_SHELL_COMMAND },
+				BuiltInTool.SERVER_EXEC_SHELL_COMMAND
+			)
+		);
+
+		expect(meta?.command).toBe('echo "he');
+	});
+
+	it('returns null while the command value is still empty', () => {
+		expect(
+			parseExecShellCommandMeta(
+				makeSection(
+					{ toolArgs: '{"command":"', toolName: BuiltInTool.SERVER_EXEC_SHELL_COMMAND },
+					BuiltInTool.SERVER_EXEC_SHELL_COMMAND
+				)
+			)
+		).toBeNull();
+	});
 });