llama-cpp-python 0.2.62__tar.gz → 0.2.63__tar.gz

llama_cpp_python-0.2.63/.git/FETCH_HEAD

@@ -0,0 +1 @@
+02812148635bf6337ffc7d1abb34093f4065df88		'02812148635bf6337ffc7d1abb34093f4065df88' of https://github.com/abetlen/llama-cpp-python

llama_cpp_python-0.2.63/.git/HEAD

@@ -0,0 +1 @@
+02812148635bf6337ffc7d1abb34093f4065df88

{llama_cpp_python-0.2.62 → llama_cpp_python-0.2.63}/.git/config

@@ -9,7 +9,7 @@
 [gc]
 	auto = 0
 [http "https://github.com/"]
-	extraheader = AUTHORIZATION: basic eC1hY2Nlc3MtdG9rZW46Z2hzX2g4ejgzMTE4dlV0N3pXMXBJa3ZtUkxLSGtnZERNaTFqTzZuag==
+	extraheader = AUTHORIZATION: basic eC1hY2Nlc3MtdG9rZW46Z2hzX3pWRkFjdDdJMG5wSmI5RlY5WHFMUTAzUkl6UXFqbTN1WGJmRA==
 [submodule "vendor/llama.cpp"]
 	active = true
 	url = https://github.com/ggerganov/llama.cpp.git

llama_cpp_python-0.2.63/.git/logs/HEAD

@@ -0,0 +1 @@
+0000000000000000000000000000000000000000 02812148635bf6337ffc7d1abb34093f4065df88 runner <runner@fv-az840-923.wndg40qfsgie5kikvjficeueua.ex.internal.cloudapp.net> 1713586238 +0000	checkout: moving from master to refs/tags/v0.2.63

llama_cpp_python-0.2.63/.git/modules/vendor/llama.cpp/HEAD

@@ -0,0 +1 @@
+0e4802b2ecbaab04b4f829fde4a3096ca19c84b5

{llama_cpp_python-0.2.62 → llama_cpp_python-0.2.63}/.git/modules/vendor/llama.cpp/config

@@ -16,7 +16,7 @@
 [gc]
 	auto = 0
 [http "https://github.com/"]
-	extraheader = AUTHORIZATION: basic eC1hY2Nlc3MtdG9rZW46Z2hzX2g4ejgzMTE4dlV0N3pXMXBJa3ZtUkxLSGtnZERNaTFqTzZuag==
+	extraheader = AUTHORIZATION: basic eC1hY2Nlc3MtdG9rZW46Z2hzX3pWRkFjdDdJMG5wSmI5RlY5WHFMUTAzUkl6UXFqbTN1WGJmRA==
 [url "https://github.com/"]
 	insteadOf = git@github.com:
 	insteadOf = org-6826477@github.com:

llama_cpp_python-0.2.63/.git/modules/vendor/llama.cpp/logs/HEAD

@@ -0,0 +1,2 @@
+0000000000000000000000000000000000000000 0e4802b2ecbaab04b4f829fde4a3096ca19c84b5 runner <runner@fv-az840-923.wndg40qfsgie5kikvjficeueua.ex.internal.cloudapp.net> 1713586239 +0000	clone: from https://github.com/ggerganov/llama.cpp.git
+0e4802b2ecbaab04b4f829fde4a3096ca19c84b5 0e4802b2ecbaab04b4f829fde4a3096ca19c84b5 runner <runner@fv-az840-923.wndg40qfsgie5kikvjficeueua.ex.internal.cloudapp.net> 1713586239 +0000	checkout: moving from master to 0e4802b2ecbaab04b4f829fde4a3096ca19c84b5

llama_cpp_python-0.2.63/.git/modules/vendor/llama.cpp/logs/refs/heads/master

@@ -0,0 +1 @@
+0000000000000000000000000000000000000000 0e4802b2ecbaab04b4f829fde4a3096ca19c84b5 runner <runner@fv-az840-923.wndg40qfsgie5kikvjficeueua.ex.internal.cloudapp.net> 1713586239 +0000	clone: from https://github.com/ggerganov/llama.cpp.git

llama_cpp_python-0.2.63/.git/modules/vendor/llama.cpp/logs/refs/remotes/origin/HEAD

@@ -0,0 +1 @@
+0000000000000000000000000000000000000000 0e4802b2ecbaab04b4f829fde4a3096ca19c84b5 runner <runner@fv-az840-923.wndg40qfsgie5kikvjficeueua.ex.internal.cloudapp.net> 1713586239 +0000	clone: from https://github.com/ggerganov/llama.cpp.git

{llama_cpp_python-0.2.62 → llama_cpp_python-0.2.63}/.git/modules/vendor/llama.cpp/modules/kompute/config

@@ -13,7 +13,7 @@
 [gc]
 	auto = 0
 [http "https://github.com/"]
-	extraheader = AUTHORIZATION: basic eC1hY2Nlc3MtdG9rZW46Z2hzX2g4ejgzMTE4dlV0N3pXMXBJa3ZtUkxLSGtnZERNaTFqTzZuag==
+	extraheader = AUTHORIZATION: basic eC1hY2Nlc3MtdG9rZW46Z2hzX3pWRkFjdDdJMG5wSmI5RlY5WHFMUTAzUkl6UXFqbTN1WGJmRA==
 [url "https://github.com/"]
 	insteadOf = git@github.com:
 	insteadOf = org-6826477@github.com:

llama_cpp_python-0.2.63/.git/modules/vendor/llama.cpp/modules/kompute/logs/HEAD

@@ -0,0 +1,2 @@
+0000000000000000000000000000000000000000 d1e3b0953cf66acc94b2e29693e221427b2c1f3f runner <runner@fv-az840-923.wndg40qfsgie5kikvjficeueua.ex.internal.cloudapp.net> 1713586240 +0000	clone: from https://github.com/nomic-ai/kompute.git
+d1e3b0953cf66acc94b2e29693e221427b2c1f3f 4565194ed7c32d1d2efa32ceab4d3c6cae006306 runner <runner@fv-az840-923.wndg40qfsgie5kikvjficeueua.ex.internal.cloudapp.net> 1713586241 +0000	checkout: moving from master to 4565194ed7c32d1d2efa32ceab4d3c6cae006306

llama_cpp_python-0.2.63/.git/modules/vendor/llama.cpp/modules/kompute/logs/refs/heads/master

@@ -0,0 +1 @@
+0000000000000000000000000000000000000000 d1e3b0953cf66acc94b2e29693e221427b2c1f3f runner <runner@fv-az840-923.wndg40qfsgie5kikvjficeueua.ex.internal.cloudapp.net> 1713586240 +0000	clone: from https://github.com/nomic-ai/kompute.git

llama_cpp_python-0.2.63/.git/modules/vendor/llama.cpp/modules/kompute/logs/refs/remotes/origin/HEAD

@@ -0,0 +1 @@
+0000000000000000000000000000000000000000 d1e3b0953cf66acc94b2e29693e221427b2c1f3f runner <runner@fv-az840-923.wndg40qfsgie5kikvjficeueua.ex.internal.cloudapp.net> 1713586240 +0000	clone: from https://github.com/nomic-ai/kompute.git

llama_cpp_python-0.2.63/.git/modules/vendor/llama.cpp/packed-refs

@@ -0,0 +1,2 @@
+# pack-refs with: peeled fully-peeled sorted 
+0e4802b2ecbaab04b4f829fde4a3096ca19c84b5 refs/remotes/origin/master

llama_cpp_python-0.2.63/.git/modules/vendor/llama.cpp/refs/heads/master

@@ -0,0 +1 @@
+0e4802b2ecbaab04b4f829fde4a3096ca19c84b5

llama_cpp_python-0.2.63/.git/modules/vendor/llama.cpp/refs/tags/b2699

@@ -0,0 +1 @@
+0e4802b2ecbaab04b4f829fde4a3096ca19c84b5

llama_cpp_python-0.2.63/.git/modules/vendor/llama.cpp/shallow

@@ -0,0 +1 @@
+0e4802b2ecbaab04b4f829fde4a3096ca19c84b5

llama_cpp_python-0.2.63/.git/refs/tags/v0.2.63

@@ -0,0 +1 @@
+02812148635bf6337ffc7d1abb34093f4065df88

llama_cpp_python-0.2.63/.git/shallow

@@ -0,0 +1 @@
+02812148635bf6337ffc7d1abb34093f4065df88

{llama_cpp_python-0.2.62 → llama_cpp_python-0.2.63}/CHANGELOG.md

@@ -7,6 +7,11 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+## [0.2.63]
+
+- feat: Update llama.cpp to ggerganov/llama.cpp@0e4802b2ecbaab04b4f829fde4a3096ca19c84b5
+- feat: Add stopping_criteria to ChatFormatter, allow stopping on arbitrary token ids, fixes llama3 instruct by @abetlen in cc81afebf04d26ca1ac3cf72f23f18da6ab58588
+
 ## [0.2.62]
 
 - feat: Update llama.cpp to ggerganov/llama.cpp@3b8f1ec4b18770531d0b1d792f3edf08254e4f0c

{llama_cpp_python-0.2.62 → llama_cpp_python-0.2.63}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: llama_cpp_python
-Version: 0.2.62
+Version: 0.2.63
 Summary: Python bindings for the llama.cpp library
 Author-Email: Andrei Betlen <abetlen@gmail.com>
 License: MIT

{llama_cpp_python-0.2.62 → llama_cpp_python-0.2.63}/llama_cpp/__init__.py

@@ -1,4 +1,4 @@
 from .llama_cpp import *
 from .llama import *
 
-__version__ = "0.2.62"
+__version__ = "0.2.63"

{llama_cpp_python-0.2.62 → llama_cpp_python-0.2.63}/llama_cpp/llama.py

@@ -426,7 +426,10 @@ class Llama:
                     print(f"Using chat bos_token: {bos_token}", file=sys.stderr)
 
                 self.chat_handler = llama_chat_format.Jinja2ChatFormatter(
-                    template=template, eos_token=eos_token, bos_token=bos_token
+                    template=template,
+                    eos_token=eos_token,
+                    bos_token=bos_token,
+                    stop_token_ids=[eos_token_id],
                 ).to_chat_handler()
 
         if self.chat_format is None and self.chat_handler is None:

{llama_cpp_python-0.2.62 → llama_cpp_python-0.2.63}/llama_cpp/llama_chat_format.py

@@ -10,6 +10,9 @@ from typing import Any, Dict, Iterator, List, Literal, Optional, Tuple, Union, P
 
 import jinja2
 
+import numpy as np
+import numpy.typing as npt
+
 import llama_cpp.llama as llama
 import llama_cpp.llama_types as llama_types
 import llama_cpp.llama_grammar as llama_grammar
@@ -150,6 +153,7 @@ class ChatFormatterResponse:
 
     prompt: str
     stop: Optional[Union[str, List[str]]] = None
+    stopping_criteria: Optional[llama.StoppingCriteriaList] = None
 
 
 class ChatFormatter(Protocol):
@@ -173,12 +177,14 @@ class Jinja2ChatFormatter(ChatFormatter):
         eos_token: str,
         bos_token: str,
         add_generation_prompt: bool = True,
+        stop_token_ids: Optional[List[int]] = None,
     ):
         """A chat formatter that uses jinja2 templates to format the prompt."""
         self.template = template
         self.eos_token = eos_token
         self.bos_token = bos_token
         self.add_generation_prompt = add_generation_prompt
+        self.stop_token_ids = set(stop_token_ids) if stop_token_ids is not None else None
 
         self._environment = jinja2.Environment(
             loader=jinja2.BaseLoader(),
@@ -211,7 +217,16 @@ class Jinja2ChatFormatter(ChatFormatter):
             tool_choice=tool_choice,
         )
 
-        return ChatFormatterResponse(prompt=prompt, stop=[self.eos_token])
+        stopping_criteria = None
+        if self.stop_token_ids is not None:
+            def stop_on_last_token(
+                    tokens: npt.NDArray[np.intc],
+                    logits: npt.NDArray[np.single]
+            ) -> bool:
+                return tokens[-1] in self.stop_token_ids
+            stopping_criteria = llama.StoppingCriteriaList([stop_on_last_token])
+
+        return ChatFormatterResponse(prompt=prompt, stop=[self.eos_token], stopping_criteria=stopping_criteria)
 
     def to_chat_handler(self) -> LlamaChatCompletionHandler:
         return chat_formatter_to_chat_completion_handler(self)
@@ -533,6 +548,10 @@ def chat_formatter_to_chat_completion_handler(
             rstop = result.stop if isinstance(result.stop, list) else [result.stop]
             stop = stop + rstop
 
+        stopping_criteria = None
+        if result.stopping_criteria is not None:
+            stopping_criteria = result.stopping_criteria
+
         if response_format is not None and response_format["type"] == "json_object":
             grammar = _grammar_for_response_format(response_format, verbose=llama.verbose)
 
@@ -598,6 +617,7 @@ def chat_formatter_to_chat_completion_handler(
             mirostat_eta=mirostat_eta,
             model=model,
             logits_processor=logits_processor,
+            stopping_criteria=stopping_criteria,
             grammar=grammar,
             logit_bias=logit_bias,
         )

{llama_cpp_python-0.2.62 → llama_cpp_python-0.2.63}/vendor/llama.cpp/.github/workflows/build.yml

@@ -32,6 +32,8 @@ jobs:
       - name: Clone
         id: checkout
         uses: actions/checkout@v4
+        with:
+          fetch-depth: 0
 
       - name: Dependencies
         id: depends
@@ -88,6 +90,8 @@ jobs:
       - name: Clone
         id: checkout
         uses: actions/checkout@v4
+        with:
+          fetch-depth: 0
 
       - name: Dependencies
         id: depends
@@ -206,6 +210,8 @@ jobs:
       - name: Clone
         id: checkout
         uses: actions/checkout@v4
+        with:
+          fetch-depth: 0
 
       - name: Dependencies
         id: depends
@@ -238,6 +244,33 @@ jobs:
           ./bin/convert-llama2c-to-ggml --copy-vocab-from-model ./tok512.bin --llama2c-model stories260K.bin --llama2c-output-model stories260K.gguf
           ./bin/main -m stories260K.gguf -p "One day, Lily met a Shoggoth" -n 500 -c 256
 
+      - name: Determine tag name
+        id: tag
+        shell: bash
+        run: |
+          BUILD_NUMBER="$(git rev-list --count HEAD)"
+          SHORT_HASH="$(git rev-parse --short=7 HEAD)"
+          if [[ "${{ env.BRANCH_NAME }}" == "master" ]]; then
+            echo "name=b${BUILD_NUMBER}" >> $GITHUB_OUTPUT
+          else
+            SAFE_NAME=$(echo "${{ env.BRANCH_NAME }}" | tr '/' '-')
+            echo "name=${SAFE_NAME}-b${BUILD_NUMBER}-${SHORT_HASH}" >> $GITHUB_OUTPUT
+          fi
+
+      - name: Pack artifacts
+        id: pack_artifacts
+        if: ${{ ( github.event_name == 'push' && github.ref == 'refs/heads/master' ) || github.event.inputs.create_release == 'true' }}
+        run: |
+          cp LICENSE ./build/bin/
+          zip -r llama-${{ steps.tag.outputs.name }}-bin-ubuntu-x64.zip ./build/bin/*
+
+      - name: Upload artifacts
+        if: ${{ ( github.event_name == 'push' && github.ref == 'refs/heads/master' ) || github.event.inputs.create_release == 'true' }}
+        uses: actions/upload-artifact@v4
+        with:
+          path: llama-${{ steps.tag.outputs.name }}-bin-ubuntu-x64.zip
+          name: llama-bin-ubuntu-x64.zip
+
 #  ubuntu-latest-cmake-sanitizer:
 #    runs-on: ubuntu-latest
 #

{llama_cpp_python-0.2.62 → llama_cpp_python-0.2.63}/vendor/llama.cpp/README-sycl.md

@@ -229,12 +229,11 @@ source /opt/intel/oneapi/setvars.sh
 # Build LLAMA with MKL BLAS acceleration for intel GPU
 mkdir -p build && cd build
 
-# Option 1: Use FP16 for better performance in long-prompt  inference
-cmake --build .. -DLLAMA_SYCL=ON -DCMAKE_C_COMPILER=icx -DCMAKE_CXX_COMPILER=icpx -DLLAMA_SYCL_F16=ON
-# Or without "--build", run "make" next
+# Option 1: Use FP16 for better performance in long-prompt inference
+#cmake .. -DLLAMA_SYCL=ON -DCMAKE_C_COMPILER=icx -DCMAKE_CXX_COMPILER=icpx -DLLAMA_SYCL_F16=ON
 
 # Option 2: Use FP32 by default
-cmake --build .. -DLLAMA_SYCL=ON -DCMAKE_C_COMPILER=icx -DCMAKE_CXX_COMPILER=icpx
+cmake .. -DLLAMA_SYCL=ON -DCMAKE_C_COMPILER=icx -DCMAKE_CXX_COMPILER=icpx
 
 #build all binary
 cmake --build . --config Release -j -v
@@ -252,10 +251,10 @@ export CPLUS_INCLUDE_DIR=/path/to/oneMKL/include:$CPLUS_INCLUDE_DIR
 mkdir -p build && cd build
 
 # Option 1: Use FP16 for better performance in long-prompt  inference
-cmake --build .. -DLLAMA_SYCL=ON -DLLAMA_SYCL_TARGET=NVIDIA -DCMAKE_C_COMPILER=icx -DCMAKE_CXX_COMPILER=icpx -DLLAMA_SYCL_F16=ON
+cmake .. -DLLAMA_SYCL=ON -DLLAMA_SYCL_TARGET=NVIDIA -DCMAKE_C_COMPILER=icx -DCMAKE_CXX_COMPILER=icpx -DLLAMA_SYCL_F16=ON
 
 # Option 2: Use FP32 by default
-cmake --build .. -DLLAMA_SYCL=ON -DLLAMA_SYCL_TARGET=NVIDIA -DCMAKE_C_COMPILER=icx -DCMAKE_CXX_COMPILER=icpx
+cmake .. -DLLAMA_SYCL=ON -DLLAMA_SYCL_TARGET=NVIDIA -DCMAKE_C_COMPILER=icx -DCMAKE_CXX_COMPILER=icpx
 
 #build all binary
 cmake --build . --config Release -j -v

{llama_cpp_python-0.2.62 → llama_cpp_python-0.2.63}/vendor/llama.cpp/README.md

@@ -122,6 +122,7 @@ Typically finetunes of the base models below are supported as well.
 - [x] [Command-R](https://huggingface.co/CohereForAI/c4ai-command-r-v01)
 - [x] [SEA-LION](https://huggingface.co/models?search=sea-lion)
 - [x] [GritLM-7B](https://huggingface.co/GritLM/GritLM-7B) + [GritLM-8x7B](https://huggingface.co/GritLM/GritLM-8x7B)
+- [x] [OLMo](https://allenai.org/olmo)
 
 (instructions for supporting more models: [HOWTO-add-model.md](./docs/HOWTO-add-model.md))
 

{llama_cpp_python-0.2.62 → llama_cpp_python-0.2.63}/vendor/llama.cpp/common/build-info.cpp

@@ -1,4 +1,4 @@
 int LLAMA_BUILD_NUMBER = 1;
-char const *LLAMA_COMMIT = "3b8f1ec";
+char const *LLAMA_COMMIT = "0e4802b";
 char const *LLAMA_COMPILER = "cc (Ubuntu 11.4.0-1ubuntu1~22.04) 11.4.0";
 char const *LLAMA_BUILD_TARGET = "x86_64-linux-gnu";

{llama_cpp_python-0.2.62 → llama_cpp_python-0.2.63}/vendor/llama.cpp/convert-hf-to-gguf.py

@@ -2636,6 +2636,66 @@ class CommandR2Model(Model):
         self.gguf_writer.add_rope_scaling_type(gguf.RopeScalingType.NONE)
 
 
+@Model.register("OlmoForCausalLM")
+@Model.register("OLMoForCausalLM")
+class OlmoModel(Model):
+    model_arch = gguf.MODEL_ARCH.OLMO
+
+    def set_gguf_parameters(self):
+        super().set_gguf_parameters()
+        self.gguf_writer.add_layer_norm_eps(1e-5)
+        if "clip_qkv" in self.hparams is not None:
+            self.gguf_writer.add_clamp_kqv(self.hparams["clip_qkv"])
+
+    # Same as super class, but permuting q_proj, k_proj
+    # Copied from: LlamaModel
+    def write_tensors(self):
+        block_count = self.hparams.get("n_layers", self.hparams.get("num_hidden_layers", self.hparams.get("n_layer")))
+        tensor_map = gguf.get_tensor_name_map(self.model_arch, block_count)
+        n_head = self.hparams.get("num_attention_heads")
+        n_kv_head = self.hparams.get("num_key_value_heads")
+        for name, data_torch in self.get_tensors():
+            old_dtype = data_torch.dtype
+
+            # convert any unsupported data types to float32
+            if data_torch.dtype not in (torch.float16, torch.float32):
+                data_torch = data_torch.to(torch.float32)
+
+            data = data_torch.numpy()
+
+            if name.endswith("q_proj.weight"):
+                data = permute(data, n_head, n_head)
+            if name.endswith("k_proj.weight"):
+                data = permute(data, n_head, n_kv_head)
+
+            data = data.squeeze()
+
+            # map tensor names
+            new_name = tensor_map.get_name(name, try_suffixes=(".weight", ".bias"))
+            if new_name is None:
+                print(f"Can not map tensor {name!r}")
+                sys.exit()
+
+            n_dims = len(data.shape)
+            data_dtype = data.dtype
+
+            # if f32 desired, convert any float16 to float32
+            if self.ftype == 0 and data_dtype == np.float16:
+                data = data.astype(np.float32)
+
+            # 1d tensors need to be converted to float32
+            if self.ftype == 1 and data_dtype == np.float16 and n_dims == 1:
+                data = data.astype(np.float32)
+
+            # if f16 desired, convert any float32 2-dim weight tensors to float16
+            if self.ftype == 1 and data_dtype == np.float32 and n_dims == 2:
+                data = data.astype(np.float16)
+
+            print(f"{new_name}, n_dims = {n_dims}, {old_dtype} --> {data.dtype}")
+
+            self.gguf_writer.add_tensor(new_name, data)
+
+
 ###### CONVERSION LOGIC ######
 
 

{llama_cpp_python-0.2.62 → llama_cpp_python-0.2.63}/vendor/llama.cpp/examples/imatrix/imatrix.cpp

@@ -44,7 +44,7 @@ private:
     std::mutex                             m_mutex;
     int                                    m_last_call = 0;
     std::vector<float>                     m_src1_data;
-    std::vector<int>                       m_ids; // the expert ids from ggml_mul_mat_id
+    std::vector<char>                      m_ids; // the expert ids from ggml_mul_mat_id
                                                   //
     void save_imatrix(const char * file_name) const;
     void keep_imatrix(int ncall) const;
@@ -81,6 +81,7 @@ bool IMatrixCollector::collect_imatrix(struct ggml_tensor * t, bool ask, void *
     if (ask) {
         if (t->op == GGML_OP_MUL_MAT_ID) return true; // collect all indirect matrix multiplications
         if (t->op != GGML_OP_MUL_MAT) return false;
+        // why are small batches ignored (<16 tokens)?
         if (src1->ne[1] < 16 || src1->type != GGML_TYPE_F32) return false;
         if (!(wname.substr(0, 4) == "blk." || (m_params.collect_output_weight && wname == "output.weight"))) return false;
         return true;
@@ -101,14 +102,19 @@ bool IMatrixCollector::collect_imatrix(struct ggml_tensor * t, bool ask, void *
     // this has been adapted to the new format of storing merged experts in a single 3d tensor
     // ref: https://github.com/ggerganov/llama.cpp/pull/6387
     if (t->op == GGML_OP_MUL_MAT_ID) {
-        const int idx  = ((int32_t *) t->op_params)[0];
+        //   ids  -> [n_experts_used, n_tokens]
+        //   src1 -> [cols, n_expert_used, n_tokens]
         const ggml_tensor * ids = t->src[2];
         const int n_as = src0->ne[2];
+        const int n_ids = ids->ne[0];
 
         // the top-k selected expert ids are stored in the ids tensor
         // for simplicity, always copy ids to host, because it is small
-        GGML_ASSERT(ids->ne[1] == src1->ne[1]);
-        m_ids.resize(ggml_nbytes(ids)/sizeof(int));
+        // take into account that ids is not contiguous!
+
+        GGML_ASSERT(ids->ne[1] == src1->ne[2]);
+
+        m_ids.resize(ggml_nbytes(ids));
         ggml_backend_tensor_get(ids, m_ids.data(), 0, ggml_nbytes(ids));
 
         auto & e = m_stats[wname];
@@ -118,26 +124,35 @@ bool IMatrixCollector::collect_imatrix(struct ggml_tensor * t, bool ask, void *
         //       using the following line, we can correct for that if needed by replacing the line above with:
         //if (idx == t->src[0]->ne[0] - 1) ++e.ncall;
 
+        if (e.values.empty()) {
+            e.values.resize(src1->ne[0]*n_as, 0);
+        }
+        else if (e.values.size() != (size_t)src1->ne[0]*n_as) {
+            fprintf(stderr, "Oops: inconsistent size for %s (%d vs %d)\n", wname.c_str(), (int)e.values.size(), (int)src1->ne[0]*n_as);
+            exit(1); //GGML_ASSERT(false);
+        }
+        if (m_params.verbosity > 1) {
+            printf("%s[%d]: %32s, %s, %5d x %5d, %d\n", __func__, m_last_call, wname.c_str(), ggml_op_name(t->op), (int)src1->ne[0], (int)src1->ne[2], (int)src1->type);
+        }
         // loop over all possible experts, regardless if they are used or not in the batch
         for (int ex = 0; ex < n_as; ++ex) {
             size_t e_start = ex*src1->ne[0];
-            if (e.values.empty()) {
-                e.values.resize(src1->ne[0]*n_as, 0);
-            }
-            else if (e.values.size() != (size_t)src1->ne[0]*n_as) {
-                fprintf(stderr, "Oops: inconsistent size for %s (%d vs %d)\n", wname.c_str(), (int)e.values.size(), (int)src1->ne[0]*n_as);
-                exit(1); //GGML_ASSERT(false);
-            }
-            if (m_params.verbosity > 1) {
-                printf("%s[%d]: %32s, %s, %5d x %5d, %d\n", __func__, m_last_call, wname.c_str(), ggml_op_name(t->op), (int)src1->ne[0], (int)src1->ne[1], (int)src1->type);
-            }
-            for (int row = 0; row < (int)src1->ne[1]; ++row) {
-                const int excur = m_ids[row*n_as + idx];
-                GGML_ASSERT(excur >= 0 && excur < n_as); // sanity check
-                if (excur != ex) continue;
-                const float * x = data + row * src1->ne[0];
-                for (int j = 0; j < (int)src1->ne[0]; ++j) {
-                    e.values[e_start + j] += x[j]*x[j];
+
+            for (int idx = 0; idx < n_ids; ++idx) {
+                for (int row = 0; row < (int)src1->ne[2]; ++row) {
+                    const int excur = *(const int32_t *) (m_ids.data() + row*ids->nb[1] + idx*ids->nb[0]);
+
+                    GGML_ASSERT(excur >= 0 && excur < n_as); // sanity check
+
+                    if (excur != ex) continue;
+
+                    const int64_t i11 = idx % src1->ne[1];
+                    const int64_t i12 = row;
+                    const float * x = (const float *)((const char *)data + i11*src1->nb[1] + i12*src1->nb[2]);
+
+                    for (int j = 0; j < (int)src1->ne[0]; ++j) {
+                        e.values[e_start + j] += x[j]*x[j];
+                    }
                 }
             }
             if (e.ncall > m_last_call) {